Support for composite/custom types in type parameters

Allow Exandra sets/maps/tuples to have composite/custom type parameters. The logic is based on [ecto's](https://github.com/elixir-ecto/ecto/blob/master/lib/ecto/schema.ex#L2443) so that the behavior is similar to standard `{:array, Custom}` declarations. This comes with the drawback that it's impossible to declare different parameters for elements of the same type at the same level, for example in `field :f, Exandra.Map, key: Exandra.Set, type: :integer, value: Exandra.Set, type: :string` both sets will be `#Set<:integer>`s, as only the first `type` will be considered.
vinniefranco · May 7, 2024 · f92d2c4 · f92d2c4
1 parent 6a1aa79
commit f92d2c4
Show file tree

Hide file tree

Showing 9 changed files with 148 additions and 23 deletions.
diff --git a/lib/exandra/map.ex b/lib/exandra/map.ex
@@ -1,12 +1,12 @@
 defmodule Exandra.Map do
   opts_schema = [
     key: [
-      type: :atom,
+      type: :any,
       required: true,
       doc: "The type of the keys in the map."
     ],
     value: [
-      type: :atom,
+      type: :any,
       required: true,
       doc: "The type of the values in the map."
     ],
@@ -37,6 +37,8 @@ defmodule Exandra.Map do
 
   use Ecto.ParameterizedType
 
+  alias Exandra.Types
+
   @opts_schema NimbleOptions.new!(opts_schema)
 
   # Made public for testing.
@@ -45,7 +47,16 @@ defmodule Exandra.Map do
 
   @impl Ecto.ParameterizedType
   def init(opts) do
+    {key, opts} = Keyword.pop_first(opts, :key)
+    {value, opts} = Keyword.pop_first(opts, :value)
+
+    key = Types.check_type!(__MODULE__, key, opts)
+    value = Types.check_type!(__MODULE__, value, opts)
+
     opts
+    |> Keyword.put(:key, key)
+    |> Keyword.put(:value, value)
+    |> Keyword.take(Keyword.keys(@opts_schema.schema))
     |> NimbleOptions.validate!(@opts_schema)
     |> Map.new()
   end

diff --git a/lib/exandra/set.ex b/lib/exandra/set.ex
@@ -1,7 +1,7 @@
 defmodule Exandra.Set do
   opts_schema = [
     type: [
-      type: :atom,
+      type: :any,
       required: true,
       doc: "The type of the elements in the set."
     ],
@@ -33,6 +33,8 @@ defmodule Exandra.Set do
 
   use Ecto.ParameterizedType
 
+  alias Exandra.Types
+
   @type t() :: MapSet.t()
 
   @opts_schema NimbleOptions.new!(opts_schema)
@@ -46,7 +48,12 @@ defmodule Exandra.Set do
 
   @impl Ecto.ParameterizedType
   def init(opts) do
+    {type, opts} = Keyword.pop_first(opts, :type)
+    checked_type = Types.check_type!(__MODULE__, type, opts)
+
     opts
+    |> Keyword.put(:type, checked_type)
+    |> Keyword.take(Keyword.keys(@opts_schema.schema))
     |> NimbleOptions.validate!(@opts_schema)
     |> Map.new()
   end

diff --git a/lib/exandra/tuple.ex b/lib/exandra/tuple.ex
@@ -1,7 +1,7 @@
 defmodule Exandra.Tuple do
   opts_schema = [
     types: [
-      type: {:list, :atom},
+      type: {:list, :any},
       required: true,
       doc: "The types of the elements in the tuple."
     ],
@@ -36,6 +36,8 @@ defmodule Exandra.Tuple do
 
   use Ecto.ParameterizedType
 
+  alias Exandra.Types
+
   @type t() :: Tuple.t()
 
   @opts_schema NimbleOptions.new!(opts_schema)
@@ -49,16 +51,19 @@ defmodule Exandra.Tuple do
 
   @impl Ecto.ParameterizedType
   def init(opts) do
-    opts =
-      opts
-      |> NimbleOptions.validate!(@opts_schema)
-      |> Map.new()
+    {types, opts} = Keyword.pop_first(opts, :types)
 
-    if opts.types == [] do
+    if types == [] do
       raise ArgumentError, "CQL tuples must have at least one element, got: []"
     end
 
+    types = types |> Enum.map(&Types.check_type!(__MODULE__, &1, opts))
+
     opts
+    |> Keyword.put(:types, types)
+    |> Keyword.take(Keyword.keys(@opts_schema.schema))
+    |> NimbleOptions.validate!(@opts_schema)
+    |> Map.new()
   end
 
   @impl Ecto.ParameterizedType

diff --git a/lib/exandra/types.ex b/lib/exandra/types.ex
@@ -43,4 +43,39 @@ defmodule Exandra.Types do
   end
 
   def for(_ecto_type, _opts), do: :error
+
+  @spec check_type!(module(), any(), keyword()) :: Ecto.Type.t()
+  def check_type!(name, type, opts) when is_atom(type) do
+    cond do
+      Ecto.Type.base?(type) -> type
+      Code.ensure_compiled(type) == {:module, type} -> check_parameterized(name, type, opts)
+      true -> raise ArgumentError, "#{name}: not a valid type parameter, got #{inspect(type)}"
+    end
+  end
+
+  def check_type!(name, {composite, inner}, opts) do
+    if Ecto.Type.composite?(composite) do
+      inner = check_type!(name, inner, opts)
+      {composite, inner}
+    else
+      raise ArgumentError, "#{name}: expected Ecto composite type, got: #{inspect(composite)}"
+    end
+  end
+
+  def check_type!(name, any, _opts),
+    do: raise(ArgumentError, "#{name}: unknown type parameter, got: #{inspect(any)}")
+
+  defp check_parameterized(name, type, opts) do
+    cond do
+      function_exported?(type, :type, 0) ->
+        type
+
+      function_exported?(type, :type, 1) ->
+        Ecto.ParameterizedType.init(type, opts)
+
+      true ->
+        raise ArgumentError,
+              "#{name}: expected Ecto.Type/Ecto.ParameterizedType, got: #{inspect(type)}"
+    end
+  end
 end
diff --git a/test/exandra/integration_test.exs b/test/exandra/integration_test.exs
@@ -24,6 +24,11 @@ defmodule Exandra.IntegrationTest do
         type: :my_complex,
         encoded_fields: [:meta]
 
+      field :my_composite, Exandra.Map,
+        key: Exandra.Tuple,
+        types: [:string, :integer],
+        value: :integer
+
       field :my_complex_udt, Exandra.UDT, type: :my_complex, encoded_fields: [:meta]
       field :my_list, {:array, :string}
       field :my_utc, :utc_datetime_usec
@@ -93,6 +98,7 @@ defmodule Exandra.IntegrationTest do
       my_complex_list_udt list<FROZEN<my_complex>>,
       my_complex_udt my_complex,
       my_tuple tuple<int, text>,
+      my_composite map<FROZEN<tuple<ascii, int>>, bigint>,
       my_embedded_udt my_embedded_type,
       my_list list<varchar>,
       my_utc timestamp,
@@ -274,6 +280,7 @@ defmodule Exandra.IntegrationTest do
         meta: %{"foo" => "bar", "baz" => %{"qux" => "quux"}},
         happened: ~U[2020-01-01T00:00:00Z]
       },
+      my_composite: %{{"foo", 1} => 1, {"bar", 8} => 4},
       my_bool: true,
       my_integer: 4,
       my_decimal: decimal1
@@ -296,6 +303,7 @@ defmodule Exandra.IntegrationTest do
         meta: %{"foo" => "bar", "baz" => %{"qux" => "quux"}},
         happened: ~U[2020-01-01T00:00:00Z]
       },
+      my_composite: %{{"baz", 0} => 2},
       my_bool: false,
       my_integer: 5,
       my_decimal: decimal2
@@ -313,6 +321,7 @@ defmodule Exandra.IntegrationTest do
       my_tuple: nil,
       my_complex_list_udt: nil,
       my_complex_udt: nil,
+      my_composite: nil,
       my_bool: nil,
       # my_integer is used for sorting.
       my_integer: 6,
@@ -347,6 +356,7 @@ defmodule Exandra.IntegrationTest do
                "meta" => %{"foo" => "bar", "baz" => %{"qux" => "quux"}},
                "happened" => ~U[2020-01-01T00:00:00.000Z]
              },
+             my_composite: %{{"bar", 8} => 4, {"foo", 1} => 1},
              my_bool: true,
              my_integer: 4,
              my_decimal: ^decimal1
@@ -373,6 +383,7 @@ defmodule Exandra.IntegrationTest do
                "meta" => %{"foo" => "bar", "baz" => %{"qux" => "quux"}},
                "happened" => ~U[2020-01-01T00:00:00.000Z]
              },
+             my_composite: %{{"baz", 0} => 2},
              my_bool: false,
              my_integer: 5,
              my_decimal: ^decimal2
@@ -390,6 +401,7 @@ defmodule Exandra.IntegrationTest do
              my_udt: %{},
              my_list_udt: nil,
              my_complex_list_udt: nil,
+             my_composite: %{},
              my_complex_udt: %{"meta" => %{}},
              my_bool: nil,
              my_integer: 6,

diff --git a/test/exandra/types/map_test.exs b/test/exandra/types/map_test.exs
@@ -10,7 +10,7 @@ defmodule Exandra.MapTest do
       field :my_int_key_map, Exandra.Map, key: :integer, value: :string
       field :my_int_value_map, Exandra.Map, key: :string, value: :integer
       field :my_int_map, Exandra.Map, key: :integer, value: :integer
-      field :my_atom_map, Exandra.Map, key: :atom, value: :integer
+      field :my_map_of_sets, Exandra.Map, key: Exandra.Set, type: :integer, value: :integer
     end
   end
 
@@ -19,12 +19,12 @@ defmodule Exandra.MapTest do
              :parameterized,
              Map,
              %{
-               field: :my_atom_map,
-               key: :atom,
+               field: :my_int_key_map,
+               key: :integer,
                schema: Schema,
-               value: :integer
+               value: :string
              }
-           } = Schema.__schema__(:type, :my_atom_map)
+           } = Schema.__schema__(:type, :my_int_key_map)
   end
 
   @p_dump_type {:parameterized, Map, Map.params(:dump)}
@@ -37,11 +37,11 @@ defmodule Exandra.MapTest do
     assert :self = Ecto.Type.embed_as(@p_self_type, :foo)
     assert :self = Ecto.Type.embed_as(@p_dump_type, :foo)
 
-    type = Schema.__schema__(:type, :my_atom_map)
-    assert {:ok, %{}} = Ecto.Type.load(type, :my_atom_map)
+    type = Schema.__schema__(:type, :my_int_key_map)
+    assert {:ok, %{}} = Ecto.Type.load(type, :my_int_key_map)
 
-    type = Schema.__schema__(:type, :my_int_map)
-    assert Ecto.Type.dump(type, %{1 => 3}) == {:ok, %{1 => 3}}
+    type = Schema.__schema__(:type, :my_int_key_map)
+    assert Ecto.Type.dump(type, %{1 => "a"}) == {:ok, %{1 => "a"}}
   end
 
   test "type/1" do

diff --git a/test/exandra/types/set_test.exs b/test/exandra/types/set_test.exs
@@ -7,7 +7,7 @@ defmodule Exandra.SetTest do
     use Ecto.Schema
 
     schema "my_schema" do
-      field :my_set, Set, type: :uuid
+      field :my_set, Set, type: :binary_id
     end
   end
 
@@ -17,7 +17,7 @@ defmodule Exandra.SetTest do
              Set,
              %{
                field: :my_set,
-               type: :uuid,
+               type: :binary_id,
                schema: Schema
              }
            } = Schema.__schema__(:type, :my_set)
@@ -64,6 +64,12 @@ defmodule Exandra.SetTest do
     assert {:ok, MapSet.new([1])} == Set.cast(1, %{type: :integer})
     assert {:ok, MapSet.new([1])} == Set.cast([1], %{type: :integer})
     assert :error = Set.cast(:asd, nil)
+
+    assert {:ok, MapSet.new([[1, 2, 3], [4, 5, 6]])} ==
+             Set.cast([[1, 2, 3], [4, 5, 6]], %{type: {:array, :integer}})
+
+    assert {:ok, MapSet.new([MapSet.new([1, 2, 3]), MapSet.new([4, 5, 6])])} ==
+             Set.cast([[1, 2, 3], [4, 5, 6]], %{type: {:parameterized, Set, %{type: :integer}}})
   end
 
   test "equal?/3" do

diff --git a/test/exandra/types/tuple_test.exs b/test/exandra/types/tuple_test.exs
@@ -8,9 +8,9 @@ defmodule Exandra.TupleTest do
              :parameterized,
              Tuple,
              %{
-               types: [:uuid, :integer, :string]
+               types: [:binary_id, :integer, :string]
              }
-           } = Ecto.ParameterizedType.init(Tuple, types: [:uuid, :integer, :string])
+           } = Ecto.ParameterizedType.init(Tuple, types: [:binary_id, :integer, :string])
   end
 
   @p_dump_type {:parameterized, Tuple, Tuple.params(:dump)}
@@ -23,7 +23,7 @@ defmodule Exandra.TupleTest do
     assert :self = Ecto.Type.embed_as(@p_self_type, :foo)
     assert :self = Ecto.Type.embed_as(@p_dump_type, :foo)
 
-    type = Ecto.ParameterizedType.init(Tuple, types: [:uuid, :integer, :string])
+    type = Ecto.ParameterizedType.init(Tuple, types: [:binary_id, :integer, :string])
     assert {:ok, nil} = Ecto.Type.load(type, :my_tuple)
 
     tuple = {Ecto.UUID.generate(), :rand.uniform(1000), :crypto.strong_rand_bytes(20)}
@@ -46,6 +46,17 @@ defmodule Exandra.TupleTest do
     assert {:ok, {1}} == Tuple.cast([1], %{types: [:integer]})
 
     assert {:ok, {1, "a"}} == Tuple.cast({1, "a"}, %{types: [:integer, :string]})
+
+    assert {:ok, {MapSet.new([1]), %{2 => 3}, [4], {5}, 6}} ==
+             Tuple.cast({[1], %{2 => 3}, [4], 5, 6}, %{
+               types: [
+                 {:parameterized, Exandra.Set, %{type: :integer}},
+                 {:parameterized, Exandra.Map, %{key: :integer, value: :integer}},
+                 {:array, :integer},
+                 {:parameterized, Exandra.Tuple, %{types: [:integer]}},
+                 :integer
+               ]
+             })
   end
 
   test "load/2" do

diff --git a/test/exandra/types_test.exs b/test/exandra/types_test.exs
@@ -30,4 +30,42 @@ defmodule Exandra.TypesTest do
     # With types that are modules that export the type/0 callback.
     assert Types.for(Exandra.Counter) == {:ok, "counter"}
   end
+
+  test "check_type!/3" do
+    assert Types.check_type!(nil, :string, []) == :string
+
+    assert Types.check_type!(nil, {:array, Exandra.Set}, type: :integer) ==
+             {:array, {:parameterized, Exandra.Set, %{type: :integer}}}
+
+    assert Types.check_type!(nil, Exandra.Set, type: :integer) ==
+             {:parameterized, Exandra.Set, %{type: :integer}}
+
+    assert Types.check_type!(nil, Exandra.Map, key: Exandra.Set, type: :integer, value: :integer) ==
+             {:parameterized, Exandra.Map,
+              %{key: {:parameterized, Exandra.Set, %{type: :integer}}, value: :integer}}
+
+    assert Types.check_type!(nil, Exandra.Set, type: Exandra.Set, type: :integer) ==
+             {:parameterized, Exandra.Set,
+              %{type: {:parameterized, Exandra.Set, %{type: :integer}}}}
+
+    assert Types.check_type!(nil, Exandra.Tuple,
+             types: [:integer, Exandra.Set, Exandra.Map],
+             type: :string,
+             key: :integer,
+             value: :binary_id
+           ) ==
+             {:parameterized, Exandra.Tuple,
+              %{
+                types: [
+                  :integer,
+                  {:parameterized, Exandra.Set, %{type: :string}},
+                  {:parameterized, Exandra.Map, %{key: :integer, value: :binary_id}}
+                ]
+              }}
+
+    assert_raise ArgumentError, fn -> Types.check_type!(nil, :nonexsisting, []) end
+    assert_raise ArgumentError, fn -> Types.check_type!(nil, {:array, :nonexisting}, []) end
+    assert_raise ArgumentError, fn -> Types.check_type!(nil, {:nonexisting, :string}, []) end
+    assert_raise ArgumentError, fn -> Types.check_type!(nil, Exandra.Types, []) end
+  end
 end