Support custom field selection on user-defined select queries

Author: renatomassaro

lib/feeb/db.ex

@@ -107,7 +107,9 @@ defmodule Feeb.DB do
   def one({domain, :fetch}, value, opts), do: one({domain, :fetch}, [value], opts)
 
   def one({domain, query_name}, bindings, opts) when is_list(bindings) do
-    one({get_context!(), domain, query_name}, bindings, opts)
+    {get_context!(), domain, query_name}
+    |> get_query_id_for_select_query(opts)
+    |> one(bindings, opts)
   end
 
   def one({domain, query_name}, value, opts), do: one({domain, query_name}, [value], opts)
@@ -135,13 +137,15 @@ defmodule Feeb.DB do
     |> all([], opts)
   end
 
-  def all({domain, query_name}, bindings, opts) do
-    all({get_context!(), domain, query_name}, bindings, opts)
+  def all({domain, query_name}, bindings, opts) when is_list(bindings) do
+    {get_context!(), domain, query_name}
+    |> get_query_id_for_select_query(opts)
+    |> all(bindings, opts)
   end
 
-  def all({_, domain, query_name}, bindings, opts) do
-    bindings = if is_list(bindings), do: bindings, else: [bindings]
+  def all({domain, query_name}, value, opts), do: all({domain, query_name}, [value], opts)
 
+  def all({_, domain, query_name}, bindings, opts) do
     case GenServer.call(get_pid!(), {:query, :all, {domain, query_name}, bindings, opts}) do
       {:ok, rows} -> rows
       {:error, reason} -> raise reason
@@ -291,6 +295,18 @@ defmodule Feeb.DB do
     LocalState.get_current_context!().context
   end
 
+  defp get_query_id_for_select_query(original_query_id, []), do: original_query_id
+
+  defp get_query_id_for_select_query(original_query_id, opts) do
+    target_fields = opts[:select] || [:*]
+
+    if target_fields == [:*] do
+      original_query_id
+    else
+      Query.compile_adhoc_query(original_query_id, target_fields)
+    end
+  end
+
   defp get_bindings(query_id, struct) do
     {_, {_, params_bindings}, _} = Query.fetch!(query_id)
 

lib/feeb/db/query.ex

@@ -33,24 +33,32 @@ defmodule Feeb.DB.Query do
   end
 
   @doc """
-  Compiling an adhoc query is useful when you want to select custom fields
-  off of a "select *" query. It's like a subset of the original query
+  Compiling an adhoc query is useful when the user wants to select custom fields off of a `SELECT *`
+  query. It's essentially a subset of the original query, with specific fields being selected.
   """
   @spec compile_adhoc_query(term, term) :: no_return
-  def compile_adhoc_query({context, domain, query_name} = query_id, custom_fields) do
-    raise "Deprecated; consider implementing this feature as part of `get_templated_query_id/3`"
-    query_name = :"#{query_name}$#{Enum.join(custom_fields, "$")}"
-    adhoc_query_id = {context, domain, query_name}
+  def compile_adhoc_query({context, domain, query_name} = original_query_id, target_fields) do
+    model = Schema.get_model_from_query_id(original_query_id)
+    valid_fields = model.__cols__()
+    sorted_target_fields = Enum.sort(target_fields)
+
+    adhoc_query_name = :"#{query_name}$#{get_query_name_suffix(sorted_target_fields)}"
+    adhoc_query_id = {context, domain, adhoc_query_name}
 
-    {sql, {fields_bindings, params_bindings}, qt} = fetch!(query_id)
+    {sql, {fields_bindings, params_bindings}, qt} = fetch!(original_query_id)
 
     if fields_bindings != [:*] do
-      raise "#{inspect(query_id)}: Custom fields can only be used on 'SELECT *' queries"
+      raise "#{inspect(original_query_id)}: Custom selection can only be used on 'SELECT *' queries"
     end
 
-    # "Compile" new query
-    new_sql = String.replace(sql, "*", Enum.join(custom_fields, ", "))
-    adhoc_q = {new_sql, {custom_fields, params_bindings}, qt}
+    Enum.each(target_fields, fn field ->
+      if field not in valid_fields,
+        do: raise("Can't select #{inspect(field)}; not a valid field for #{model}")
+    end)
+
+    # "Compile" new query (replaces `SELECT *` with `SELECT <sorted_target_fields>`)
+    new_sql = String.replace(sql, "*", Enum.join(sorted_target_fields, ", "), global: false)
+    adhoc_q = {new_sql, {sorted_target_fields, params_bindings}, qt}
 
     append_runtime_query(adhoc_query_id, adhoc_q)
     adhoc_query_id

priv/test/queries/test/all_types.sql

@@ -7,6 +7,9 @@ select map_keys_atom from all_types;
 -- :get_map
 select map from all_types;
 
+-- :get_by_integer
+select * from all_types where integer = ?;
+
 -- :get_max_integer
 select max(integer) from all_types;
 

test/db/db_test.exs

@@ -437,6 +437,13 @@ defmodule Feeb.DBTest do
       assert nil == DB.one({:order_items, :fetch}, [2, 1])
     end
 
+    test "supports custom selection on user-defined queries", %{shard_id: shard_id} do
+      DB.begin(@context, shard_id, :read)
+
+      assert %{id: %NotLoaded{}, name: "Phoebe"} =
+               DB.one({:friends, :get_by_name}, "Phoebe", select: [:name])
+    end
+
     test "supports the :format flag", %{shard_id: shard_id} do
       DB.begin(@context, shard_id, :write)
 
@@ -536,6 +543,13 @@ defmodule Feeb.DBTest do
       assert [%{id: 1, title: "Post", body: %NotLoaded{}}] = DB.all(Post, [], select: [:id, :title])
     end
 
+    test "supports custom selection on user-defined queries", %{shard_id: shard_id} do
+      DB.begin(@context, shard_id, :read)
+
+      assert [%{id: %NotLoaded{}, name: "Phoebe"}] =
+               DB.all({:friends, :get_by_name}, "Phoebe", select: [:name])
+    end
+
     test "supports the :format flag", %{shard_id: shard_id} do
       DB.begin(@context, shard_id, :write)
 

test/db/query_test.exs

@@ -48,6 +48,48 @@ defmodule Feeb.DB.QueryTest do
     end
   end
 
+  describe "compile_adhoc_query/2" do
+    test "generates a subset of the original query" do
+      Query.compile(@all_types_path, {:test, :all_types})
+
+      original_query_id = {:test, :all_types, :get_by_integer}
+      query_id = Query.compile_adhoc_query(original_query_id, [:string, :atom, :uuid])
+
+      assert {sql, {target_fields, bindings}, query_type} = Query.fetch!(query_id, [])
+      assert query_type == :select
+      assert target_fields == [:atom, :string, :uuid]
+      assert bindings == [:integer]
+      assert sql == "select atom, string, uuid from all_types where integer = ?;"
+    end
+
+    test "raises on custom selection of a non-wildcard select query" do
+      Query.compile(@all_types_path, {:test, :all_types})
+
+      # This query is: `SELECT atom, integer FROM all_types;`
+      original_query_id = {:test, :all_types, :get_atom_and_integer}
+
+      %{message: error} =
+        assert_raise RuntimeError, fn ->
+          Query.compile_adhoc_query(original_query_id, [:string, :atom, :uuid])
+        end
+
+      assert error =~ "Custom selection can only be used on 'SELECT *' queries"
+    end
+
+    test "raises when trying to select a field that does not exist in the schema" do
+      Query.compile(@all_types_path, {:test, :all_types})
+
+      original_query_id = {:test, :all_types, :get_by_integer}
+
+      %{message: error} =
+        assert_raise RuntimeError, fn ->
+          Query.compile_adhoc_query(original_query_id, [:string, :i_dont_exist, :atom])
+        end
+
+      assert error =~ "Can't select :i_dont_exist; not a valid field for Elixir.Sample.AllTypes"
+    end
+  end
+
   describe "get_templated_query_id/3" do
     test ":__all" do
       Query.compile(@friends_path, {:test, :friends})