Add support for custom/composite PKs on adhoc queries

Author: renatomassaro

lib/feeb/db/query.ex

@@ -1,5 +1,6 @@
 defmodule Feeb.DB.Query do
   require Logger
+  alias Feeb.DB.Schema
   alias __MODULE__.Binding
 
   @initial_q {"", {[], []}, nil}
@@ -53,10 +54,14 @@ defmodule Feeb.DB.Query do
     adhoc_query_id
   end
 
-  def get_templated_query_id({context, domain, :__insert}, target_fields, meta) do
+  def get_templated_query_id(query_id, target_fields, meta \\ %{})
+
+  def get_templated_query_id({context, domain, :__insert} = query_id, target_fields, _meta) do
+    model = Schema.get_model_from_query_id(query_id)
+
     target_fields =
       if target_fields == :all do
-        meta.schema.__cols__()
+        model.__cols__()
       else
         raise "Not supported for now, add & test it once needed"
         target_fields
@@ -70,11 +75,14 @@ defmodule Feeb.DB.Query do
         real_query_id
 
       nil ->
-        compile_templated_query(:__insert, real_query_id, target_fields)
+        compile_templated_query(:__insert, real_query_id, target_fields, model)
     end
   end
 
-  def get_templated_query_id({context, domain, :__update}, target_fields, _meta) do
+  def get_templated_query_id({context, domain, :__update} = query_id, target_fields, _meta) do
+    model = Schema.get_model_from_query_id(query_id)
+
+    # TODO: Transparently include `updated_at` (if present in the model)
     target_fields = Enum.sort(target_fields)
 
     query_name_suffix =
@@ -92,22 +100,24 @@ defmodule Feeb.DB.Query do
         real_query_id
 
       nil ->
-        compile_templated_query(:__update, real_query_id, target_fields)
+        compile_templated_query(:__update, real_query_id, target_fields, model)
     end
   end
 
   def get_templated_query_id({_context, _domain, query_name} = query_id, target_fields, _meta)
       when query_name in [:__all, :__fetch, :__delete] do
+    model = Schema.get_model_from_query_id(query_id)
+
     case get(query_id) do
       {_, _, _} = _compiled_query ->
         query_id
 
       nil ->
-        compile_templated_query(query_name, query_id, target_fields)
+        compile_templated_query(query_name, query_id, target_fields, model)
     end
   end
 
-  defp compile_templated_query(:__insert, {_, domain, _} = query_id, target_fields) do
+  defp compile_templated_query(:__insert, {_, domain, _} = query_id, target_fields, _model) do
     columns_clause = target_fields |> Enum.join(", ")
 
     values_clause =
@@ -123,42 +133,48 @@ defmodule Feeb.DB.Query do
     query_id
   end
 
-  defp compile_templated_query(:__update, {_, domain, _} = query_id, target_fields) do
-    # Hard-coded for now, but we may need this to be dynamc in the future
-    primary_key_col = :id
+  defp compile_templated_query(:__update, {_, domain, _} = query_id, target_fields, model) do
+    primary_keys = model.__primary_keys__()
+    assert_adhoc_query!(primary_keys, query_id, model)
 
-    set_clause =
+    set_conditions =
       target_fields
       |> Enum.reduce([], fn field, acc ->
         ["#{field} = ?" | acc]
       end)
       |> Enum.reverse()
       |> Enum.join(", ")
 
-    sql = "UPDATE #{domain} SET #{set_clause} WHERE id = ?;"
-    adhoc_query = {sql, {[], target_fields ++ [primary_key_col]}, :update}
+    sql = "UPDATE #{domain} SET #{set_conditions} #{generate_where_clause(primary_keys)};"
+    adhoc_query = {sql, {[], target_fields ++ primary_keys}, :update}
     append_runtime_query(query_id, adhoc_query)
 
     query_id
   end
 
-  defp compile_templated_query(:__all, {_, domain, _} = query_id, _target_fields) do
+  defp compile_templated_query(:__all, {_, domain, _} = query_id, _target_fields, _model) do
     sql = "SELECT * FROM #{domain};"
     adhoc_query = {sql, {[:*], []}, :select}
     append_runtime_query(query_id, adhoc_query)
     query_id
   end
 
-  defp compile_templated_query(:__fetch, {_, domain, _} = query_id, _target_fields) do
-    sql = "SELECT * FROM #{domain} WHERE id = ?;"
-    adhoc_query = {sql, {[:*], [:id]}, :select}
+  defp compile_templated_query(:__fetch, {_, domain, _} = query_id, _target_fields, model) do
+    primary_keys = model.__primary_keys__()
+    assert_adhoc_query!(primary_keys, query_id, model)
+    sql = "SELECT * FROM #{domain} #{generate_where_clause(primary_keys)};"
+
+    adhoc_query = {sql, {[:*], primary_keys}, :select}
     append_runtime_query(query_id, adhoc_query)
     query_id
   end
 
-  defp compile_templated_query(:__delete, {_, domain, _} = query_id, _target_fields) do
-    sql = "DELETE FROM #{domain} WHERE id = ?;"
-    adhoc_query = {sql, {[], [:id]}, :delete}
+  defp compile_templated_query(:__delete, {_, domain, _} = query_id, _target_fields, model) do
+    primary_keys = model.__primary_keys__()
+    assert_adhoc_query!(primary_keys, query_id, model)
+    sql = "DELETE FROM #{domain} #{generate_where_clause(primary_keys)};"
+
+    adhoc_query = {sql, {[], primary_keys}, :delete}
     append_runtime_query(query_id, adhoc_query)
     query_id
   end
@@ -235,6 +251,24 @@ defmodule Feeb.DB.Query do
   defp get_returning_fields({_, _, operation}) when operation in [:update, :delete],
     do: "*"
 
+  defp generate_where_clause(primary_keys) when is_list(primary_keys) do
+    where_conditions =
+      primary_keys
+      |> Enum.reduce([], fn field, acc ->
+        ["#{field} = ?" | acc]
+      end)
+      |> Enum.reverse()
+      |> Enum.join(" AND ")
+
+    "WHERE #{where_conditions}"
+  end
+
+  defp assert_adhoc_query!(nil, query_id, model) do
+    raise("Can't generate adhoc query #{inspect(query_id)} because #{inspect(model)} has no PKs")
+  end
+
+  defp assert_adhoc_query!(_, _, _), do: :ok
+
   # Line-break
   defp handle_line(<<>>, qs, id, q) when not is_nil(id) do
     {sql, {fields_bindings, params_bindings}, query_type} = q

lib/feeb/db/repo.ex

@@ -345,22 +345,22 @@ defmodule Feeb.DB.Repo do
   defp create_schema_from_rows({_, :pragma, _}, _, rows), do: rows
 
   defp create_schema_from_rows(query_id, {_, {fields_bindings, _}, :select}, rows) do
-    model = get_model_from_query_id(query_id)
+    model = Schema.get_model_from_query_id(query_id)
     Enum.map(rows, fn row -> Schema.from_row(model, fields_bindings, row) end)
   end
 
   defp create_schema_from_rows(query_id, {_, {_, params_bindings}, :insert}, rows) do
-    model = get_model_from_query_id(query_id)
+    model = Schema.get_model_from_query_id(query_id)
     Enum.map(rows, fn row -> Schema.from_row(model, params_bindings, row) end)
   end
 
   defp create_schema_from_rows(query_id, {_, _, :update}, rows) do
-    model = get_model_from_query_id(query_id)
+    model = Schema.get_model_from_query_id(query_id)
     Enum.map(rows, fn row -> Schema.from_row(model, model.__cols__(), row) end)
   end
 
   defp create_schema_from_rows(query_id, {_, _, :delete}, rows) do
-    model = get_model_from_query_id(query_id)
+    model = Schema.get_model_from_query_id(query_id)
     Enum.map(rows, fn row -> Schema.from_row(model, model.__cols__(), row) end)
   end
 
@@ -399,10 +399,6 @@ defmodule Feeb.DB.Repo do
     |> trunc()
   end
 
-  defp get_model_from_query_id({context, domain, _}) do
-    :persistent_term.get({:db_table_models, {context, domain}})
-  end
-
   # TODO: These pragma functions will be removed once that gets turned into a hook
   def conditional_pragma_based_on_env(conn, :test, _), do: custom_pragma_for_test(conn)
   def conditional_pragma_based_on_env(conn, _, true), do: custom_pragma_for_test(conn)

lib/feeb/db/schema.ex

@@ -366,6 +366,9 @@ defmodule Feeb.DB.Schema do
   def get_private(%{__private__: private}, k), do: private[k]
   def get_private!(%{__private__: private}, k), do: Map.fetch!(private, k)
 
+  def get_model_from_query_id({context, table, _}),
+    do: :persistent_term.get({:db_table_models, {context, table}})
+
   defp add_missing_values(struct, f, f), do: struct
 
   defp add_missing_values(struct, all_fields, added_fields) do

test/db/db_test.exs

@@ -2,7 +2,7 @@ defmodule Feeb.DBTest do
   use Test.Feeb.DBCase, async: true
   alias Feeb.DB, as: DB
   alias Feeb.DB.LocalState
-  alias Sample.{AllTypes, CustomTypes, Friend, Post}
+  alias Sample.{AllTypes, CustomTypes, Friend, OrderItems, Post}
   alias Sample.Types.TypedID
 
   @context :test
@@ -406,6 +406,29 @@ defmodule Feeb.DBTest do
       assert nil == DB.one({:friends, :get_by_id}, [0])
     end
 
+    test ":fetch templated query works", %{shard_id: shard_id} do
+      DB.begin(@context, shard_id, :write)
+
+      assert %{id: 1, name: "Phoebe"} = DB.one({:friends, :fetch}, [1])
+      assert nil == DB.one({:friends, :fetch}, [500])
+    end
+
+    test ":fetch templated query works on schema with composite PKs", %{shard_id: shard_id} do
+      DB.begin(@context, shard_id, :write)
+
+      %{order_id: 1, product_id: 2, quantity: 10, price: 50}
+      |> OrderItems.new()
+      |> DB.insert!()
+
+      order_item = DB.one({:order_items, :fetch}, [1, 2])
+      assert order_item.order_id == 1
+      assert order_item.product_id == 2
+      assert order_item.quantity == 10
+      assert order_item.price == 50
+
+      assert nil == DB.one({:order_items, :fetch}, [2, 1])
+    end
+
     test "supports the :format flag", %{shard_id: shard_id} do
       DB.begin(@context, shard_id, :write)
 
@@ -560,6 +583,20 @@ defmodule Feeb.DBTest do
       refute post.updated_at == new_post.updated_at
     end
 
+    test "updates the struct on schema with composite PKs", %{shard_id: shard_id} do
+      DB.begin(@context, shard_id, :write)
+
+      order_item =
+        %{order_id: 1, product_id: 2, quantity: 10, price: 50}
+        |> OrderItems.new()
+        |> DB.insert!()
+
+      assert {:ok, %{order_id: 1, product_id: 2, quantity: 20}} =
+               order_item
+               |> OrderItems.update(%{quantity: 20})
+               |> DB.update()
+    end
+
     test "update attempt that failed to find a matching entry", %{shard_id: shard_id} do
       DB.begin(@context, shard_id, :write)
 
@@ -655,6 +692,20 @@ defmodule Feeb.DBTest do
       refute DB.one({:friends, :get_by_id}, [1])
     end
 
+    test "deletes the struct (schema with composite PKs)", %{shard_id: shard_id} do
+      DB.begin(@context, shard_id, :write)
+
+      order_item =
+        %{order_id: 1, product_id: 2, quantity: 10, price: 50}
+        |> OrderItems.new()
+        |> DB.insert!()
+
+      assert {:ok, _} = DB.delete(order_item)
+
+      # OrderItem has been deleted
+      assert [] == DB.all(OrderItems)
+    end
+
     test "delete attempt that failed to find a matching entry", %{shard_id: shard_id} do
       DB.begin(@context, shard_id, :write)
       friend = DB.one({:friends, :get_by_id}, [1])