Support non-string prefix (#4497)

lib/ecto/query.ex

@@ -882,7 +882,7 @@ defmodule Ecto.Query do
     subquery = wrap_in_subquery(query)
 
     case Keyword.fetch(opts, :prefix) do
-      {:ok, prefix} when is_binary(prefix) or is_nil(prefix) ->
+      {:ok, prefix} ->
         put_in(subquery.query.prefix, prefix)
 
       :error ->
@@ -909,12 +909,13 @@ defmodule Ecto.Query do
   @doc """
   Puts the given prefix in a query.
   """
-  def put_query_prefix(%Ecto.Query{} = query, prefix) when is_binary(prefix) do
+  def put_query_prefix(%Ecto.Query{} = query, prefix) do
     %{query | prefix: prefix}
   end
 
-  def put_query_prefix(other, prefix) when is_binary(prefix) do
-    other |> Ecto.Queryable.to_query() |> put_query_prefix(prefix)
+  def put_query_prefix(other, prefix) do
+    query = %Ecto.Query{} = Ecto.Queryable.to_query(other)
+    put_query_prefix(query, prefix)
   end
 
   @doc """

lib/ecto/query/builder/from.ex

@@ -77,15 +77,15 @@ defmodule Ecto.Query.Builder.From do
   If possible, it does all calculations at compile time to avoid
   runtime work.
   """
-  @spec build(Macro.t(), Macro.Env.t(), atom, {:ok, String.t | nil} | nil, hints) ::
+  @spec build(Macro.t(), Macro.Env.t(), atom, {:ok, Ecto.Schema.prefix | nil} | nil, hints) ::
           {Macro.t(), Keyword.t(), non_neg_integer | nil}
   def build(query, env, as, prefix, hints) do
     hints = Enum.map(hints, &hint!(&1))
 
     prefix = case prefix do
       nil -> nil
       {:ok, prefix} when is_binary(prefix) or is_nil(prefix) -> {:ok, prefix}
-      {:ok, {:^, _, [prefix]}} -> {:ok, quote(do: Ecto.Query.Builder.From.prefix!(unquote(prefix)))}
+      {:ok, {:^, _, [prefix]}} -> {:ok, prefix}
       {:ok, prefix} -> Builder.error!("`prefix` must be a compile time string or an interpolated value using ^, got: #{Macro.to_string(prefix)}")
     end
 
@@ -152,13 +152,6 @@ defmodule Ecto.Query.Builder.From do
     {:%, [], [Ecto.Query, {:%{}, [], query_fields}]}
   end
 
-  @doc """
-  Validates a prefix at runtime.
-  """
-  @spec prefix!(any) :: nil | String.t()
-  def prefix!(prefix) when is_binary(prefix) or is_nil(prefix), do: prefix
-  def prefix!(prefix), do: raise("`prefix` must be a string, got: #{inspect(prefix)}")
-
   @doc """
   Validates hints at compile time and runtime
   """
@@ -187,7 +180,7 @@ defmodule Ecto.Query.Builder.From do
   @doc """
   The callback applied by `build/2` to build the query.
   """
-  @spec apply(Ecto.Queryable.t(), non_neg_integer, Macro.t(), {:ok, String.t} | nil, hints) :: Ecto.Query.t()
+  @spec apply(Ecto.Queryable.t(), non_neg_integer, Macro.t(), {:ok, Ecto.Schema.prefix} | nil, hints) :: Ecto.Query.t()
   def apply(query, binds, as, prefix, hints) do
     query =
       query

lib/ecto/query/builder/join.ex

@@ -128,7 +128,7 @@ defmodule Ecto.Query.Builder.Join do
   If possible, it does all calculations at compile time to avoid
   runtime work.
   """
-  @spec build(Macro.t, atom, [Macro.t], Macro.t, Macro.t, Macro.t, atom, nil | {:ok, String.t | nil}, nil | String.t | [String.t], Macro.Env.t) ::
+  @spec build(Macro.t, atom, [Macro.t], Macro.t, Macro.t, Macro.t, atom, nil | {:ok, Ecto.Schema.prefix}, nil | String.t | [String.t], Macro.Env.t) ::
               {Macro.t, Keyword.t, non_neg_integer | nil}
   def build(query, qual, binding, expr, count_bind, on, as, prefix, maybe_hints, env) do
     {:ok, prefix} = prefix || {:ok, nil}
@@ -144,7 +144,7 @@ defmodule Ecto.Query.Builder.Join do
     prefix = case prefix do
       nil -> nil
       prefix when is_binary(prefix) -> prefix
-      {:^, _, [prefix]} -> quote(do: Ecto.Query.Builder.From.prefix!(unquote(prefix)))
+      {:^, _, [prefix]} -> prefix
       prefix -> Builder.error!("`prefix` must be a compile time string or an interpolated value using ^, got: #{Macro.to_string(prefix)}")
     end
 

lib/ecto/query/planner.ex

@@ -2311,11 +2311,7 @@ defmodule Ecto.Query.Planner do
   """
   def attach_prefix(%{prefix: nil} = query, opts) when is_list(opts) do
     case Keyword.fetch(opts, :prefix) do
-      {:ok, prefix} when is_binary(prefix) or is_nil(prefix) ->
-        %{query | prefix: prefix}
-
-      {:ok, prefix} when is_atom(prefix) ->
-        IO.warn("atom prefixes are deprecated. Please use a string instead.")
+      {:ok, prefix} ->
         %{query | prefix: prefix}
 
       :error ->

lib/ecto/schema.ex

@@ -478,7 +478,7 @@ defmodule Ecto.Schema do
   alias Ecto.Schema.Metadata
 
   @type source :: String.t()
-  @type prefix :: String.t() | nil
+  @type prefix :: any()
   @type schema :: %{optional(atom) => any, __struct__: atom, __meta__: Metadata.t()}
   @type embedded_schema :: %{optional(atom) => any, __struct__: atom}
   @type t :: schema | embedded_schema
@@ -662,7 +662,7 @@ defmodule Ecto.Schema do
           %{unquote_splicing(Macro.escape(@ecto_changeset_fields))}
         end
 
-        def __schema__(:prefix), do: unquote(prefix)
+        def __schema__(:prefix), do: unquote(Macro.escape(prefix))
         def __schema__(:source), do: unquote(source)
         def __schema__(:fields), do: unquote(Enum.map(fields, &elem(&1, 0)))
         def __schema__(:query_fields), do: unquote(Enum.map(query_fields, &elem(&1, 0)))
@@ -690,7 +690,7 @@ defmodule Ecto.Schema do
             %Ecto.Query{
               from: %Ecto.Query.FromExpr{
                 source: {unquote(source), __MODULE__},
-                prefix: unquote(prefix)
+                prefix: unquote(Macro.escape(prefix))
               }
             }
           end

test/ecto/query/planner_test.exs

@@ -817,6 +817,12 @@ defmodule Ecto.Query.PlannerTest do
       |> plan()
 
     assert query.sources == {{"comments", Comment, "global"}, {"comments", Comment, "local"}}
+
+    # Non-string schema prefix is supported
+    {query, _, _, _} =
+      from(c in Comment, join: Post, on: true) |> Map.put(:prefix, %{key: :global}) |> plan()
+
+    assert query.sources == {{"comments", Comment, %{key: :global}}, {"posts", Post, "my_prefix"}}
   end
 
   test "plan: combination queries" do

test/ecto/query_test.exs

@@ -242,6 +242,11 @@ defmodule Ecto.QueryTest do
       assert put_query_prefix(from("posts"), "hello").prefix == "hello"
       assert put_query_prefix(Schema, "hello").prefix == "hello"
     end
+
+    test "stores non-string prefix in query" do
+      assert put_query_prefix(from("posts"), %{key: :hello}).prefix == %{key: :hello}
+      assert put_query_prefix(Schema, %{key: :hello}).prefix == %{key: :hello}
+    end
   end
 
   describe "trailing bindings (...)" do
@@ -500,25 +505,16 @@ defmodule Ecto.QueryTest do
       assert hd(query.joins).prefix == join_prefix
     end
 
-    test "variables are validated at runtime" do
-      prefix = 123
-
-      assert_raise RuntimeError, ~r/`prefix` must be a string/, fn ->
-        from p in "posts", prefix: ^prefix
-      end
-
-      assert_raise RuntimeError, ~r/`prefix` must be a string/, fn ->
-        from p in "posts", join: "comments", on: true, prefix: ^prefix
-      end
-    end
-
     test "are supported and overridden from schemas" do
       query = from(Post)
       assert query.from.prefix == "another"
 
       query = from(Post, prefix: "hello")
       assert query.from.prefix == "hello"
 
+      query = from(Post, prefix: ^123)
+      assert query.from.prefix == 123
+
       query = from(Post, prefix: nil)
       assert query.from.prefix == nil
     end

test/ecto/repo_test.exs

@@ -81,6 +81,16 @@ defmodule Ecto.RepoTest do
     end
   end
 
+  defmodule MySchemaWithNonStringPrefix do
+    use Ecto.Schema
+
+    @schema_prefix %{key: :private}
+
+    schema "my_schema" do
+      field :x, :string
+    end
+  end
+
   defmodule MySchemaWithAssoc do
     use Ecto.Schema
 
@@ -318,6 +328,12 @@ defmodule Ecto.RepoTest do
       TestRepo.reload(struct_with_prefix)
       assert_received {:all, %{prefix: "another"}}
     end
+
+    test "supports non-string prefix" do
+      struct_with_prefix = put_meta(%MySchema{id: 2}, prefix: %{key: :another})
+      TestRepo.reload(struct_with_prefix)
+      assert_received {:all, %{prefix: %{key: :another}}}
+    end
   end
 
   defmodule DefaultOptionRepo do
@@ -1116,6 +1132,12 @@ defmodule Ecto.RepoTest do
 
       assert {:ok, schema} = TestRepo.delete(valid, prefix: "public")
       assert schema.__meta__.prefix == "public"
+
+      assert {:ok, schema} = TestRepo.insert(valid, prefix: %{key: :public})
+      assert schema.__meta__.prefix == %{key: :public}
+
+      assert {:ok, schema} = TestRepo.delete(valid, prefix: %{key: :public})
+      assert schema.__meta__.prefix == %{key: :public}
     end
 
     test "insert and delete prefix overrides schema_prefix with struct" do
@@ -1128,6 +1150,16 @@ defmodule Ecto.RepoTest do
       assert schema.__meta__.prefix == "public"
     end
 
+    test "insert and delete prefix overrides schema_prefix with struct when prefix is not a string" do
+      valid = %MySchemaWithNonStringPrefix{id: 1}
+
+      assert {:ok, schema} = TestRepo.insert(valid, prefix: %{key: :public})
+      assert schema.__meta__.prefix == %{key: :public}
+
+      assert {:ok, schema} = TestRepo.delete(valid, prefix: %{key: :public})
+      assert schema.__meta__.prefix == %{key: :public}
+    end
+
     test "insert, update, insert_or_update and delete sets schema prefix with changeset" do
       valid = Ecto.Changeset.cast(%MySchema{id: 1}, %{x: "foo"}, [:x])
 
@@ -1474,6 +1506,18 @@ defmodule Ecto.RepoTest do
 
     assert [schema] = TestRepo.all(MySchema, prefix: "public")
     assert schema.__meta__.prefix == "public"
+
+    assert schema = TestRepo.get(MySchema, 123, prefix: %{key: :public})
+    assert schema.__meta__.prefix == %{key: :public}
+
+    assert schema = TestRepo.get_by(MySchema, [id: 123], prefix: %{key: :public})
+    assert schema.__meta__.prefix == %{key: :public}
+
+    assert schema = TestRepo.one(MySchema, prefix: %{key: :public})
+    assert schema.__meta__.prefix == %{key: :public}
+
+    assert [schema] = TestRepo.all(MySchema, prefix: %{key: :public})
+    assert schema.__meta__.prefix == %{key: :public}
   end
 
   test "get, get_by, one and all ignores prefix if schema_prefix set" do
@@ -1488,6 +1532,18 @@ defmodule Ecto.RepoTest do
 
     assert [schema] = TestRepo.all(MySchemaWithPrefix, prefix: "public")
     assert schema.__meta__.prefix == "private"
+
+    assert schema = TestRepo.get(MySchemaWithNonStringPrefix, 123, prefix: %{key: :public})
+    assert schema.__meta__.prefix == %{key: :private}
+
+    assert schema = TestRepo.get_by(MySchemaWithNonStringPrefix, [id: 123], prefix: %{key: :public})
+    assert schema.__meta__.prefix == %{key: :private}
+
+    assert schema = TestRepo.one(MySchemaWithNonStringPrefix, prefix: %{key: :public})
+    assert schema.__meta__.prefix == %{key: :private}
+
+    assert [schema] = TestRepo.all(MySchemaWithNonStringPrefix, prefix: %{key: :public})
+    assert schema.__meta__.prefix == %{key: :private}
   end
 
   describe "changeset prepare" do

test/ecto/schema_test.exs

@@ -341,6 +341,15 @@ defmodule Ecto.SchemaTest do
     end
   end
 
+  defmodule SchemaWithNonStringPrefix do
+    use Ecto.Schema
+
+    @schema_prefix %{key: :tenant}
+    schema "company" do
+      field :name
+    end
+  end
+
   test "schema prefix metadata" do
     assert SchemaWithPrefix.__schema__(:source) == "company"
     assert SchemaWithPrefix.__schema__(:prefix) == "tenant"
@@ -366,6 +375,31 @@ defmodule Ecto.SchemaTest do
     assert query.from.prefix == "tenant"
   end
 
+  test "schema non-string prefix metadata" do
+    assert SchemaWithNonStringPrefix.__schema__(:source) == "company"
+    assert SchemaWithNonStringPrefix.__schema__(:prefix) == %{key: :tenant}
+    assert %SchemaWithNonStringPrefix{}.__meta__.source == "company"
+    assert %SchemaWithNonStringPrefix{}.__meta__.prefix == %{key: :tenant}
+  end
+
+  test "schema non-string prefix in queries from" do
+    import Ecto.Query
+
+    query = from(SchemaWithNonStringPrefix, select: 1)
+    assert query.from.prefix == %{key: :tenant}
+
+    query = from({"another_company", SchemaWithNonStringPrefix}, select: 1)
+    assert query.from.prefix == %{key: :tenant}
+
+    from = SchemaWithNonStringPrefix
+    query = from(from, select: 1)
+    assert query.from.prefix == %{key: :tenant}
+
+    from = {"another_company", SchemaWithNonStringPrefix}
+    query = from(from, select: 1)
+    assert query.from.prefix == %{key: :tenant}
+  end
+
   ## Schema context
   defmodule SchemaWithContext do
     use Ecto.Schema