CubDB Query helpers

scratch.exs

Mix.install([
  {:cubdb, "~> 2.0"}
])

{:ok, _} = CubDB.start_link(data_dir: "./cubdb", name: DB)

defmodule CubQuery do
  @moduledoc """
  Key and query utilities for CubDB
  """

  @doc false
  defmacro __using__([]) do
    quote do
      @before_compile CubQuery
      import CubQuery, only: [defkey: 2]
    end
  end

  @doc false
  defmacro __before_compile__(_env) do
    quote do
      def __key__(tag, _attrs) do
        raise ArgumentError, "Unknown key: #{inspect(tag)}. Did you define it using `defkey`?"
      end

      def __key_range__(tag, _attrs) do
        raise ArgumentError, "Unknown key: #{inspect(tag)}. Did you define it using `defkey`?"
      end

      def __key_fields__(tag) do
        raise ArgumentError, "Unknown key: #{inspect(tag)}. Did you define it using `defkey`?"
      end

      def key(tag, attrs \\ []) do
        __key__(tag, attrs)
      end

      def key_range(tag, attrs \\ []) do
        __key_range__(tag, attrs)
      end

      def query(db, tag, attrs \\ [], select \\ &CubDB.select/2) do
        CubQuery.query(db, __MODULE__, tag, attrs, select)
      end
    end
  end

  @doc """
  Define a key to be used in CubDB

  ## Example

      defmodule MyModule do
        use CubQuery

        defkey :thing, :thing_range, category: nil, id: ""
      end

      iex> MyModule.key(:thing, id: "A")
      {:thing, {:foo}, {"A"}}

      iex> range = MyModule.key_range(:thing, category: :foo)
      [min_key: {:thing, {:foo}, {}},
       max_key: {:thing, {:foo}, {0, 0}}]

      iex> CubDB.select(db, range)
      ...> |> Enum.to_list()
      [{{:thing, {:foo}, {"A"}}, ...},
       {{:thing, {:foo}, {"B"}}, ...}]
  """
  defmacro defkey(tag, attrs \\ []) do
    defaults = normalize_attributes(attrs)
    fields = Keyword.keys(defaults)

    quote bind_quoted: [tag: tag, defaults: defaults, fields: fields] do
      def __key_fields__(unquote(tag)) do
        unquote(fields)
      end

      def __key__(unquote(tag), attrs) do
        CubQuery.__key__(
          unquote(tag),
          unquote(defaults),
          attrs
        )
      end

      def __key_range__(unquote(tag), attrs) do
        CubQuery.__key_range__(
          unquote(tag),
          unquote(defaults),
          attrs
        )
      end
    end
  end

  defp normalize_attributes(attrs) do
    Enum.map(attrs, fn
      {key, value} when is_atom(key) -> {key, value}
      key when is_atom(key) -> {key, nil}
      invalid -> raise ArgumentError, "invalid key attribute: #{inspect(invalid)}"
    end)
  end

  @doc false
  def __key__(tag, defaults, attrs) do
    :ok = validate_key_attrs!(attrs, defaults)

    defaults
    |> Enum.map(fn {k, default} ->
      case Keyword.fetch(attrs, k) do
        {:ok, v} -> {v}
        :error -> {default}
      end
    end)
    |> List.to_tuple()
    |> Tuple.insert_at(0, tag)
  end

  defp validate_key_attrs!(attrs, defaults) do
    _ = Keyword.validate!(attrs, defaults)

    :ok
  end

  @doc false
  def __key_range__(tag, defaults, attrs) do
    :ok = validate_range_attrs!(attrs, defaults, strict_ordering: false)

    [
      min_key: range_key(tag, defaults, attrs, {}),
      max_key: range_key(tag, defaults, attrs, {nil, nil})
    ]
  end

  defp range_key(tag, defaults, attrs, value_when_missing) do
    {values, _} =
      Enum.map_reduce(defaults, false, fn {k, _}, seen_missing? ->
        case {Keyword.fetch(attrs, k), seen_missing?} do
          {_, true} -> {nil, true}
          {:error, _} -> {value_when_missing, true}
          {{:ok, v}, _} -> {{v}, false}
        end
      end)

    List.to_tuple([tag | values])
  end

  defp validate_range_attrs!(attrs, defaults, strict_ordering: strict_ordering?) do
    _ = Keyword.validate!(attrs, defaults)

    if strict_ordering? do
      validate_strict_order!(attrs, defaults)
    else
      :ok
    end
  end

  defp validate_strict_order!(attrs, defaults) do
    missing =
      defaults
      |> Enum.reverse()
      |> Stream.map(fn {k, _} ->
        {k, Keyword.fetch(attrs, k)}
      end)
      |> Stream.drop_while(fn kv ->
        match?({_k, :error}, kv)
      end)
      |> Stream.filter(fn kv ->
        match?({_k, :error}, kv)
      end)
      |> Stream.map(&elem(&1, 0))
      |> Enum.reverse()

    if missing == [] do
      :ok
    else
      raise ArgumentError, "missing fields: #{inspect(missing)}"
    end
  end

  @doc """
  Issue a CubDB query for key-value pairs matching the given key attributes.

  ## Examples

      iex> CubQuery.query(db, MyKeyModule, :my_key, category: :foo)
      [...]

      iex> CubDB.with_snapshot(fn snapshot ->
      ...>   CubQuery.query(db, MyKeyModule, :my_key, [category: :foo], &CubDB.Tx.select/2)
      ...> end)
      [...]

  """
  def query(db, key_module, tag, attrs \\ [], select \\ &CubDB.select/2) do
    fields = key_module.__key_fields__(tag)

    db
    |> select.(key_module.key_range(tag, attrs))
    |> Stream.filter(fn {key, value} -> key_match?(key, attrs, fields) end)
  end

  defp key_match?(key, attrs, fields) do
    values = key |> Tuple.to_list() |> tl()

    fields
    |> Enum.zip(values)
    |> Enum.all?(fn {attr, {value}} ->
      case Keyword.fetch(attrs, attr) do
        :error -> true
        {:ok, ^value} -> true
        _ -> false
      end
    end)
  end
end

defmodule Example do
  use CubQuery

  defkey(:a, category: nil, value: nil)
  defkey(:b, category: nil, subcategory: nil, value: nil)

  def init do
    as =
      for cat <- 1..2,
          val <- 1..5 do
        cat = :"category#{cat}"
        val = val |> to_string() |> String.duplicate(val)
        {key(:a, category: cat, value: val), nil}
      end

    bs =
      for cat <- 1..2,
          sub <- 1..3,
          val <- 1..5 do
        cat = :"category#{cat}"
        sub = :"subcategory#{sub}"
        val = val |> to_string() |> String.duplicate(val)
        {key(:b, category: cat, subcategory: sub, value: val), nil}
      end

    :ok = CubDB.clear(DB)
    :ok = CubDB.put_multi(DB, as ++ bs)
  end
end

Example.init()

Example.query(DB, :a) |> Enum.to_list() |> IO.inspect(label: "all a records")

Example.query(DB, :a, category: :category1)
|> Enum.to_list()
|> IO.inspect(label: "a records in category1")

Example.query(DB, :a, value: "22")
|> Enum.to_list()
|> IO.inspect(label: "a records with value 22")

Example.query(DB, :b, subcategory: :subcategory2)
|> Enum.to_list()
|> IO.inspect(label: "b records in subcategory2")

Example.query(DB, :b, category: :category1, value: "4444")
|> Enum.to_list()
|> IO.inspect(label: "b records in category1 with value 4444")