Repository: elixir-ecto/ecto (elixir)

Issue:
Set schema prefix on insert
Referenced from: https://github.com/phoenixframework/phoenix_ecto/issues/82

### Environment

Elixir: 1.4.2
PostgreSQL 9.6.2
Ecto: 2.1.4
Phoenix:  1.2.3
postgrex 0.13.2
OS: macOS Sierra 10.12.4

### Current behavior

Here is a factory module based off of what is in my app:

    defmodule MyApp.Factory do
      def build(:tenant) do 
         %MyApp.Tenant{name: "tenant"}
      end

      def build(:user) do
        %Scalpel.User{username: "example-#{System.unique_integer()}"}
      end
      
      def build(factory_name, attributes) do
        factory_name |> build() |> struct(attributes)
      end
      
      def insert!(:tenant) do 
        tenant = Repo.insert!(build(:tenant))
        # Runs all of the tenant specific migrations
        MyApp.TenantActions.create_tenant(Repo, tenant)
        tenant
      end

      def insert!(factory_name, attributes \\ [], prefix \\ "public") do
        Repo.insert! build(factory_name, attributes), prefix: prefix
      end
    end

    tenant = insert!(:tenant)
    user = insert(:user, [], build_prefix(tenant))

    let(:response) do 
      build_conn |> assign(:current_user, user)
    end

  In a controller we have a plug function that loads the user's profile

    plug :load_profile, "before show" when action in [:show]

    #  We get a failure here in our test because the profile table only 
    #  exists under the schema of the tenant.
    def load_profile(%{current_user: user }) do 
       user |> Ecto.assoc(:profile) |> Repo.one()
    end

    def load_profile(conn, _opts) do 
       cond do 
          conn.assigns[:profile] ->
               conn
          profile = load_profile(conn.assigns) -> 
            conn |> assign(:profile, profile)
          true -> 
              conn |> redirect(to: page_path(conn, :index)) |> halt()
       end 
    end

### Expected Behavior

My expectation is that if I have inserted the user with the tenant prefix, it should update the record with the schema it has been given. 



Pull request:
Set schema prefix on insert, delete, and update
Closes #2051  

Patch:
diff --git a/lib/ecto/adapters/mysql/connection.ex b/lib/ecto/adapters/mysql/connection.ex
index 99719673..363189ed 100644
--- a/lib/ecto/adapters/mysql/connection.ex
+++ b/lib/ecto/adapters/mysql/connection.ex
@@ -98,9 +98,14 @@ if Code.ensure_loaded?(Mariaex) do
       sources = create_names(query)
       {from, name} = get_source(query, sources, 0, from)
 
+      fields = if prefix do
+        update_fields(:on_conflict, query, sources)
+      else
+        update_fields(:update, query, sources)
+      end
+
       join   = join(query, sources)
       prefix = prefix || ["UPDATE ", from, " AS ", name, join, " SET "]
-      fields = update_fields(query, sources)
       where  = where(query, sources)
 
       IO.iodata_to_binary([prefix, fields | where])
@@ -217,23 +222,32 @@ if Code.ensure_loaded?(Mariaex) do
       [" FROM ", from, " AS " | name]
     end
 
-    defp update_fields(%Query{updates: updates} = query, sources) do
-      for(%{expr: expr} <- updates,
-          {op, kw} <- expr,
-          {key, value} <- kw,
-          do: update_op(op, key, value, sources, query)) |> Enum.intersperse(", ")
+    defp update_fields(type, %Query{updates: updates} = query, sources) do
+     fields = for(%{expr: expr} <- updates,
+                   {op, kw} <- expr,
+                   {key, value} <- kw,
+                   do: update_op(op, update_key(type, key, query, sources), value, sources, query))
+      Enum.intersperse(fields, ", ")
+    end
+
+    defp update_key(:update, key, %Query{from: from} = query, sources) do
+      {_from, name} = get_source(query, sources, 0, from)
+
+      [name, ?. | quote_name(key)]
+    end
+    defp update_key(:on_conflict, key, _query, _sources) do
+      quote_name(key)
     end
 
-    defp update_op(:set, key, value, sources, query) do
-      [quote_name(key), " = " | expr(value, sources, query)]
+    defp update_op(:set, quoted_key, value, sources, query) do
+      [quoted_key, " = " | expr(value, sources, query)]
     end
 
-    defp update_op(:inc, key, value, sources, query) do
-      quoted = quote_name(key)
-      [quoted, " = ", quoted, " + " | expr(value, sources, query)]
+    defp update_op(:inc, quoted_key, value, sources, query) do
+      [quoted_key, " = ", quoted_key, " + " | expr(value, sources, query)]
     end
 
-    defp update_op(command, _key, _value, _sources, query) do
+    defp update_op(command, _quoted_key, _value, _sources, query) do
       error!(query, "Unknown update operation #{inspect command} for MySQL")
     end
 
diff --git a/lib/ecto/query/planner.ex b/lib/ecto/query/planner.ex
index 6b255a42..4db5a0bd 100644
--- a/lib/ecto/query/planner.ex
+++ b/lib/ecto/query/planner.ex
@@ -236,9 +236,10 @@ defmodule Ecto.Query.Planner do
   defp prepare_source(query, %Ecto.SubQuery{query: inner_query} = subquery, adapter) do
     try do
       {inner_query, params, key} = prepare(inner_query, :all, adapter, 0)
-      inner_query = inner_query |> returning(true) |> normalize_select()
-      %{subquery | query: inner_query, params: params, cache: key,
-                   fields: subquery_fields(inner_query, inner_query.select)}
+      %{select: select} = inner_query = inner_query |> returning(true) |> normalize_select()
+      {select, fields} = subquery_fields(inner_query, select)
+      %{subquery | query: %{inner_query | select: select},
+                   params: params, cache: key, fields: fields}
     rescue
       e -> raise Ecto.SubQueryError, query: query, exception: e
     end
@@ -255,29 +256,55 @@ defmodule Ecto.Query.Planner do
       when assocs != [] or preloads != [] do
     error!(query, "cannot preload associations in subquery")
   end
-  defp subquery_fields(query, %{expr: {:%{}, _, pairs} = expr}) do
+  defp subquery_fields(query, %{expr: {:%{}, meta, [{:|, _, [{:&, [], [ix]}, pairs]}]} = expr,
+                                fields: [{:&, _, [ix, [_ | _] = fields, _]} | _]} = select) do
+    # In case of map updates, we need to remove duplicated fields
+    # at query time because we use the field names as aliases and
+    # duplicate aliases will lead to invalid queries.
+    map_fields = subquery_map_fields(query, expr, pairs)
+    map_keys = Keyword.keys(map_fields)
+    invalid_keys = map_keys -- fields
+
+    if invalid_keys != [] do
+      error!(query, "invalid key `#{inspect hd(invalid_keys)}` on map update in subquery")
+    end
+
+    fields = subquery_selector_fields(fields -- map_keys, ix) ++ map_fields
+    {%{select | expr: {:%{}, meta, fields}, fields: fields}, fields}
+  end
+  defp subquery_fields(query, %{expr: {:%{}, _, pairs} = expr} = select) do
+    {select, subquery_map_fields(query, expr, pairs)}
+  end
+  defp subquery_fields(_query, %{expr: {:&, _, [ix]},
+                                 fields: [{:&, _, [ix, [_ | _] = fields, _]}]} = select) do
+    {select, subquery_selector_fields(fields, ix)}
+  end
+  defp subquery_fields(_query, %{expr: {{:., _, [{:&, _, [ix]}, field]}, _, []},
+                                 fields: [{{:., _, [{:&, _, [ix]}, field]}, _, []}]} = select) do
+    {select, [{field, {{:., [], [{:&, [], [ix]}, field]}, [], []}}]}
+  end
+  defp subquery_fields(query, %{expr: expr}) do
+    error!(query, "subquery must select a source (t), a field (t.field) or a map, got: `#{Macro.to_string(expr)}`")
+  end
+
+  defp subquery_selector_fields(fields, ix) do
+    Enum.map(fields, fn field ->
+      {field, {{:., [], [{:&, [], [ix]}, field]}, [], []}}
+    end)
+  end
+
+  defp subquery_map_fields(query, expr, pairs) do
     Enum.map(pairs, fn
       {key, _} when not is_atom(key) ->
         error!(query, "only atom keys are allowed when selecting a map in subquery, got: `#{Macro.to_string(expr)}`")
-      {key, expr} ->
-        if valid_subquery_value?(expr) do
-          {key, expr}
+      {key, value} ->
+        if valid_subquery_value?(value) do
+          {key, value}
         else
           error!(query, "containers (maps, lists, tuples) or sources (t) are not allowed when selecting a map in subquery, got: `#{Macro.to_string(expr)}`")
         end
     end)
   end
-  defp subquery_fields(_query, %{fields: [{:&, _, [ix, [_|_] = fields, _]}]}) do
-    Enum.map(fields, fn field ->
-      {field, {{:., [], [{:&, [], [ix]}, field]}, [], []}}
-    end)
-  end
-  defp subquery_fields(_query, %{fields: [{{:., _, [{:&, _, [ix]}, field]}, _, []}]}) do
-    [{field, {{:., [], [{:&, [], [ix]}, field]}, [], []}}]
-  end
-  defp subquery_fields(query, %{expr: expr}) do
-    error!(query, "subquery must select a source (t), a field (t.field) or a map, got: `#{Macro.to_string(expr)}`")
-  end
 
   defp valid_subquery_value?({_, _}), do: false
   defp valid_subquery_value?(args) when is_list(args), do: false
@@ -847,7 +874,7 @@ defmodule Ecto.Query.Planner do
   defp collect_fields({:{}, _, elems}, query, take, from),
     do: collect_fields(elems, query, take, from)
   defp collect_fields({:%{}, _, [{:|, _, [data, pairs]}]}, query, take, from),
-    do: collect_fields([data|pairs], query, take, from)
+    do: collect_fields([data | pairs], query, take, from)
   defp collect_fields({:%{}, _, pairs}, query, take, from),
     do: collect_fields(pairs, query, take, from)
   defp collect_fields({:%, _, [_name, expr]}, query, take, from),
diff --git a/lib/ecto/repo/schema.ex b/lib/ecto/repo/schema.ex
index 5ca79643..c7b8eb1e 100644
--- a/lib/ecto/repo/schema.ex
+++ b/lib/ecto/repo/schema.ex
@@ -199,7 +199,7 @@ defmodule Ecto.Repo.Schema do
         case apply(changeset, adapter, :insert, args) do
           {:ok, values} ->
             changeset
-            |> load_changes(:loaded, extra ++ values, autogen, adapter)
+            |> load_changes(:loaded, extra ++ values, autogen, adapter, opts)
             |> process_children(children, user_changeset, opts)
           {:error, _} = error ->
             error
@@ -270,7 +270,7 @@ defmodule Ecto.Repo.Schema do
           case apply(changeset, adapter, action, args) do
             {:ok, values} ->
               changeset
-              |> load_changes(:loaded, values, autogen, adapter)
+              |> load_changes(:loaded, values, autogen, adapter, opts)
               |> process_children(children, user_changeset, opts)
             {:error, _} = error ->
               error
@@ -352,7 +352,7 @@ defmodule Ecto.Repo.Schema do
       args = [repo, metadata(struct, opts), filters, opts]
       case apply(changeset, adapter, :delete, args) do
         {:ok, values} ->
-          {:ok, load_changes(changeset, :deleted, values, [], adapter).data}
+          {:ok, load_changes(changeset, :deleted, values, [], adapter, opts).data}
         {:error, _} = error ->
           error
         {:invalid, constraints} ->
@@ -541,13 +541,14 @@ defmodule Ecto.Repo.Schema do
     end
   end
 
-  defp load_changes(%{types: types, changes: changes} = changeset, state, values, autogen, adapter) do
+  defp load_changes(%{types: types, changes: changes} = changeset, state, values, autogen, adapter, opts) do
     # It is ok to use types from changeset because we have
     # already filtered the results to be only about fields.
     data =
       changeset.data
       |> merge_changes(changes)
       |> merge_autogen(autogen)
+      |> apply_metadata(opts)
       |> load_each(values, types, adapter)
     data = put_in(data.__meta__.state, state)
     Map.put(changeset, :data, data)
@@ -617,6 +618,11 @@ defmodule Ecto.Repo.Schema do
     end
   end
 
+  defp apply_metadata(struct, opts) do
+    prefix = Ecto.get_meta(struct, :prefix)
+    Ecto.put_meta(struct, prefix: Keyword.get(opts, :prefix, prefix))
+  end
+
   defp process_children(changeset, assocs, user_changeset, opts) do
     case Ecto.Association.on_repo_change(changeset, assocs, opts) do
       {:ok, struct} -> {:ok, struct}


Test patch:
diff --git a/integration_test/sql/subquery.exs b/integration_test/sql/subquery.exs
index 1a3e32b6..254fbfaf 100644
--- a/integration_test/sql/subquery.exs
+++ b/integration_test/sql/subquery.exs
@@ -18,7 +18,7 @@ defmodule Ecto.Integration.SubQueryTest do
            TestRepo.all(from p in subquery(query), select: p)
   end
 
-  test "from: subqueries with select expression" do
+  test "from: subqueries with map and select expression" do
     TestRepo.insert!(%Post{text: "hello", public: true})
 
     query = from p in Post, select: %{text: p.text, pub: not p.public}
@@ -32,6 +32,20 @@ defmodule Ecto.Integration.SubQueryTest do
            TestRepo.all(from p in subquery(query), select: {p, p.pub})
   end
 
+  test "from: subqueries with map update and select expression" do
+    TestRepo.insert!(%Post{text: "hello", public: true})
+
+    query = from p in Post, select: %{p | public: not p.public}
+    assert ["hello"] =
+           TestRepo.all(from p in subquery(query), select: p.text)
+    assert [%{text: "hello", public: false}] =
+           TestRepo.all(from p in subquery(query), select: p)
+    assert [{"hello", %{text: "hello", public: false}}] =
+           TestRepo.all(from p in subquery(query), select: {p.text, p})
+    assert [{%{text: "hello", public: false}, false}] =
+           TestRepo.all(from p in subquery(query), select: {p, p.public})
+  end
+
   test "from: subqueries with aggregates" do
     TestRepo.insert!(%Post{visits: 10})
     TestRepo.insert!(%Post{visits: 11})
diff --git a/test/ecto/adapters/mysql_test.exs b/test/ecto/adapters/mysql_test.exs
index 11c0eed5..02f790b7 100644
--- a/test/ecto/adapters/mysql_test.exs
+++ b/test/ecto/adapters/mysql_test.exs
@@ -313,36 +313,36 @@ defmodule Ecto.Adapters.MySQLTest do
   test "update all" do
     query = from(m in Schema, update: [set: [x: 0]]) |> normalize(:update_all)
     assert SQL.update_all(query) ==
-           ~s{UPDATE `schema` AS s0 SET `x` = 0}
+           ~s{UPDATE `schema` AS s0 SET s0.`x` = 0}
 
     query = from(m in Schema, update: [set: [x: 0], inc: [y: 1, z: -3]]) |> normalize(:update_all)
     assert SQL.update_all(query) ==
-           ~s{UPDATE `schema` AS s0 SET `x` = 0, `y` = `y` + 1, `z` = `z` + -3}
+           ~s{UPDATE `schema` AS s0 SET s0.`x` = 0, s0.`y` = s0.`y` + 1, s0.`z` = s0.`z` + -3}
 
     query = from(e in Schema, where: e.x == 123, update: [set: [x: 0]]) |> normalize(:update_all)
     assert SQL.update_all(query) ==
-           ~s{UPDATE `schema` AS s0 SET `x` = 0 WHERE (s0.`x` = 123)}
+           ~s{UPDATE `schema` AS s0 SET s0.`x` = 0 WHERE (s0.`x` = 123)}
 
     query = from(m in Schema, update: [set: [x: ^0]]) |> normalize(:update_all)
     assert SQL.update_all(query) ==
-           ~s{UPDATE `schema` AS s0 SET `x` = ?}
+           ~s{UPDATE `schema` AS s0 SET s0.`x` = ?}
 
     query = Schema |> join(:inner, [p], q in Schema2, p.x == q.z)
                   |> update([_], set: [x: 0]) |> normalize(:update_all)
     assert SQL.update_all(query) ==
-           ~s{UPDATE `schema` AS s0 INNER JOIN `schema2` AS s1 ON s0.`x` = s1.`z` SET `x` = 0}
+           ~s{UPDATE `schema` AS s0 INNER JOIN `schema2` AS s1 ON s0.`x` = s1.`z` SET s0.`x` = 0}
 
     query = from(e in Schema, where: e.x == 123, update: [set: [x: 0]],
                              join: q in Schema2, on: e.x == q.z) |> normalize(:update_all)
     assert SQL.update_all(query) ==
            ~s{UPDATE `schema` AS s0 INNER JOIN `schema2` AS s1 ON s0.`x` = s1.`z` } <>
-           ~s{SET `x` = 0 WHERE (s0.`x` = 123)}
+           ~s{SET s0.`x` = 0 WHERE (s0.`x` = 123)}
   end
 
   test "update all with prefix" do
     query = from(m in Schema, update: [set: [x: 0]]) |> normalize(:update_all)
     assert SQL.update_all(%{query | prefix: "prefix"}) ==
-           ~s{UPDATE `prefix`.`schema` AS s0 SET `x` = 0}
+           ~s{UPDATE `prefix`.`schema` AS s0 SET s0.`x` = 0}
   end
 
   test "delete all" do
diff --git a/test/ecto/query/planner_test.exs b/test/ecto/query/planner_test.exs
index 49f4d7c0..19c7942b 100644
--- a/test/ecto/query/planner_test.exs
+++ b/test/ecto/query/planner_test.exs
@@ -296,31 +296,54 @@ defmodule Ecto.Query.PlannerTest do
     end
   end
 
-  test "prepare: subqueries validates select fields" do
-    query = prepare(from(subquery(Post), [])) |> elem(0)
-    assert [{:id, {{:., [], [{:&, [], [0]}, :id]}, [], []}},
-            {:title, {{:., [], [{:&, [], [0]}, :title]}, [], []}} | _] = query.from.fields
+  describe "prepare: subqueries select" do
+    test "supports implicit select" do
+      query = prepare(from(subquery(Post), [])) |> elem(0)
+      assert [{:id, {{:., [], [{:&, [], [0]}, :id]}, [], []}},
+              {:title, {{:., [], [{:&, [], [0]}, :title]}, [], []}} | _] = query.from.fields
+    end
 
-    query = from p in "posts", select: p.code
-    query = prepare(from(subquery(query), [])) |> elem(0)
-    assert [code: {{:., [], [{:&, [], [0]}, :code]}, [], []}] = query.from.fields
+    test "supports field selector" do
+      query = from p in "posts", select: p.code
+      query = prepare(from(subquery(query), [])) |> elem(0)
+      assert [code: {{:., [], [{:&, [], [0]}, :code]}, [], []}] = query.from.fields
 
-    query = from p in Post, select: p.code
-    query = prepare(from(subquery(query), [])) |> elem(0)
-    assert [code: {{:., [], [{:&, [], [0]}, :code]}, [], []}] = query.from.fields
+      query = from p in Post, select: p.code
+      query = prepare(from(subquery(query), [])) |> elem(0)
+      assert [code: {{:., [], [{:&, [], [0]}, :code]}, [], []}] = query.from.fields
+    end
 
-    query = from p in Post, join: c in assoc(p, :comments), select: %{c: p.code}
-    query = prepare(from(subquery(query), [])) |> elem(0)
-    assert [c: {{:., [], [{:&, [], [0]}, :code]}, [], []}] = query.from.fields
+    test "supports maps" do
+      query = from p in Post, select: %{code: p.code}
+      query = prepare(from(subquery(query), [])) |> elem(0)
+      assert [code: {{:., [], [{:&, [], [0]}, :code]}, [], []}] = query.from.fields
+    end
 
-    query = from p in Post, select: fragment("? + ?", p.id, p.id)
-    assert_raise Ecto.SubQueryError, ~r/subquery must select a source \(t\), a field \(t\.field\) or a map/, fn ->
-      prepare(from(subquery(query), []))
+    test "supports update in maps" do
+      query = from p in Post, select: %{p | code: p.title}
+      query = prepare(from(subquery(query), [])) |> elem(0)
+      assert [{:id, {{:., [], [{:&, [], [0]}, :id]}, [], []}},
+              {:title, {{:., [], [{:&, [], [0]}, :title]}, [], []}} | _] = query.from.fields
+      assert {:code, {{:., [], [{:&, [], [0]}, :title]}, [], []}} = List.last(query.from.fields)
+
+      query = from p in Post, select: %{p | unknown: p.title}
+      assert_raise Ecto.SubQueryError, ~r/invalid key `:unknown` on map update in subquery/, fn ->
+        prepare(from(subquery(query), []))
+      end
     end
 
-    query = from p in Post, select: %{p.id => p.title}
-    assert_raise Ecto.SubQueryError, ~r/only atom keys are allowed/, fn ->
-      prepare(from(subquery(query), []))
+    test "requires atom keys for maps" do
+      query = from p in Post, select: %{p.id => p.title}
+      assert_raise Ecto.SubQueryError, ~r/only atom keys are allowed/, fn ->
+        prepare(from(subquery(query), []))
+      end
+    end
+
+    test "raises on custom expressions" do
+      query = from p in Post, select: fragment("? + ?", p.id, p.id)
+      assert_raise Ecto.SubQueryError, ~r/subquery must select a source \(t\), a field \(t\.field\) or a map/, fn ->
+        prepare(from(subquery(query), []))
+      end
     end
   end
 
diff --git a/test/ecto/repo_test.exs b/test/ecto/repo_test.exs
index e3d23580..1513ebf3 100644
--- a/test/ecto/repo_test.exs
+++ b/test/ecto/repo_test.exs
@@ -171,6 +171,23 @@ defmodule Ecto.RepoTest do
     assert {:ok, %MySchema{}} = TestRepo.delete(valid)
   end
 
+  test "insert, update, insert_or_update and delete sets schema prefix" do
+    valid = Ecto.Changeset.cast(%MySchema{id: 1}, %{}, [])
+    assert {:ok, schema} = TestRepo.insert(valid, prefix: "public")
+    {schema_prefix, _} = schema.__meta__.source
+    assert schema_prefix == "public"
+
+    valid = Ecto.Changeset.cast(%MySchema{id: 1}, %{x: "foo"}, [:x])
+    assert {:ok, schema} = TestRepo.update(valid, prefix: "public")
+    {schema_prefix, _} = schema.__meta__.source
+    assert schema_prefix == "public"
+
+    valid = Ecto.Changeset.cast(%MySchema{id: 1}, %{}, [])
+    assert {:ok, schema} = TestRepo.delete(valid, prefix: "public")
+    {schema_prefix, _} = schema.__meta__.source
+    assert schema_prefix == "public"
+  end
+
   test "insert, update, insert_or_update and delete errors on invalid changeset" do
     invalid = %Ecto.Changeset{valid?: false, data: %MySchema{}}