improvement: Add immutable version of ash_raise_error function to support extensions like Citus (#620)

Author: stevebrambilla

lib/extensions/immutable_raise_error.ex

@@ -0,0 +1,74 @@
+defmodule AshPostgres.Extensions.ImmutableRaiseError do
+  @moduledoc """
+  An extension that installs an immutable version of ash_raise_error.
+
+  This can be used to improve compatibility with Postgres sharding extensions like Citus,
+  which requires functions used in CASE or COALESCE expressions to be immutable.
+
+  The new `ash_raise_error_immutable` functions add an additional row-dependent argument to ensure
+  the planner doesn't constant-fold error expressions.
+
+  To install, add this module to your repo's `installed_extensions` list:
+
+  ```elixir
+  def installed_extensions do
+    ["ash-functions", AshPostgres.Extensions.ImmutableRaiseError]
+  end
+  ```
+
+  And run `mix ash_postgres.generate_migrations` to generate the migrations.
+
+  Once installed, you can control whether the immutable function is used by adding this to your
+  repo:
+
+  ```elixir
+  def immutable_expr_error?, do: true
+  ```
+  """
+
+  use AshPostgres.CustomExtension, name: "immutable_raise_error", latest_version: 1
+
+  @impl true
+  def install(0) do
+    ash_raise_error_immutable()
+  end
+
+  @impl true
+  def uninstall(_version) do
+    "execute(\"DROP FUNCTION IF EXISTS ash_raise_error_immutable(jsonb, ANYCOMPATIBLE), ash_raise_error_immutable(jsonb, ANYELEMENT, ANYCOMPATIBLE)\")"
+  end
+
+  defp ash_raise_error_immutable do
+    """
+    execute(\"\"\"
+    CREATE OR REPLACE FUNCTION ash_raise_error_immutable(json_data jsonb, token ANYCOMPATIBLE)
+    RETURNS BOOLEAN AS $$
+    BEGIN
+        -- Raise an error with the provided JSON data.
+        -- The JSON object is converted to text for inclusion in the error message.
+        -- 'token' is intentionally ignored; its presence makes the call non-constant at the call site.
+        RAISE EXCEPTION 'ash_error: %', json_data::text;
+        RETURN NULL;
+    END;
+    $$ LANGUAGE plpgsql
+    IMMUTABLE
+    SET search_path = '';
+    \"\"\")
+
+    execute(\"\"\"
+    CREATE OR REPLACE FUNCTION ash_raise_error_immutable(json_data jsonb, type_signal ANYELEMENT, token ANYCOMPATIBLE)
+    RETURNS ANYELEMENT AS $$
+    BEGIN
+        -- Raise an error with the provided JSON data.
+        -- The JSON object is converted to text for inclusion in the error message.
+        -- 'token' is intentionally ignored; its presence makes the call non-constant at the call site.
+        RAISE EXCEPTION 'ash_error: %', json_data::text;
+        RETURN NULL;
+    END;
+    $$ LANGUAGE plpgsql
+    IMMUTABLE
+    SET search_path = '';
+    \"\"\")
+    """
+  end
+end

lib/repo.ex

@@ -113,6 +113,13 @@ defmodule AshPostgres.Repo do
   @doc "Disable expression errors for this repo"
   @callback disable_expr_error?() :: boolean
 
+  @doc """
+  Opt-in to using immutable versions of the expression error functions.
+
+  Requires the `AshPostgres.Extensions.ImmutableRaiseError` extension.
+  """
+  @callback immutable_expr_error?() :: boolean
+
   defmacro __using__(opts) do
     quote bind_quoted: [opts: opts] do
       if Keyword.get(opts, :define_ecto_repo?, true) do
@@ -145,6 +152,7 @@ defmodule AshPostgres.Repo do
       def drop?, do: true
       def disable_atomic_actions?, do: false
       def disable_expr_error?, do: false
+      def immutable_expr_error?, do: false
 
       # default to false in 4.0
       def prefer_transaction?, do: true
@@ -315,7 +323,8 @@ defmodule AshPostgres.Repo do
                      create?: 0,
                      drop?: 0,
                      disable_atomic_actions?: 0,
-                     disable_expr_error?: 0
+                     disable_expr_error?: 0,
+                     immutable_expr_error?: 0
 
       # We do this switch because `!@warn_on_missing_ash_functions` in the function body triggers
       # a dialyzer error

lib/sql_implementation.ex

@@ -334,4 +334,9 @@ defmodule AshPostgres.SqlImplementation do
 
     {types, new_returns || returns}
   end
+
+  @impl true
+  def immutable_errors?(repo) do
+    repo.immutable_expr_error?()
+  end
 end

mix.exs

@@ -168,7 +168,9 @@ defmodule AshPostgres.MixProject do
     [
       {:ash, ash_version("~> 3.5 and >= 3.5.35")},
       {:spark, "~> 2.3 and >= 2.3.4"},
-      {:ash_sql, ash_sql_version("~> 0.2 and >= 0.2.90")},
+      # TODO: bump to next ash_sql release
+      # {:ash_sql, ash_sql_version("~> 0.2 and >= 0.2.90")},
+      {:ash_sql, git: "https://github.com/ash-project/ash_sql.git"},
       {:igniter, "~> 0.6 and >= 0.6.14", optional: true},
       {:ecto_sql, "~> 3.13"},
       {:ecto, "~> 3.13"},

mix.lock

@@ -1,6 +1,6 @@
 %{
   "ash": {:hex, :ash, "3.5.43", "222f9a8ac26ad3b029f8e69306cc83091c992d858b4538af12e33a148f301cab", [:mix], [{:decimal, "~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:ecto, "~> 3.7", [hex: :ecto, repo: "hexpm", optional: false]}, {:ets, "~> 0.8", [hex: :ets, repo: "hexpm", optional: false]}, {:igniter, ">= 0.6.29 and < 1.0.0-0", [hex: :igniter, repo: "hexpm", optional: true]}, {:jason, ">= 1.0.0", [hex: :jason, repo: "hexpm", optional: false]}, {:picosat_elixir, "~> 0.2", [hex: :picosat_elixir, repo: "hexpm", optional: true]}, {:plug, ">= 0.0.0", [hex: :plug, repo: "hexpm", optional: true]}, {:reactor, "~> 0.11", [hex: :reactor, repo: "hexpm", optional: false]}, {:simple_sat, ">= 0.1.1 and < 1.0.0-0", [hex: :simple_sat, repo: "hexpm", optional: true]}, {:spark, ">= 2.3.3 and < 3.0.0-0", [hex: :spark, repo: "hexpm", optional: false]}, {:splode, ">= 0.2.6 and < 1.0.0-0", [hex: :splode, repo: "hexpm", optional: false]}, {:stream_data, "~> 1.0", [hex: :stream_data, repo: "hexpm", optional: false]}, {:telemetry, "~> 1.1", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "48b2aa274c524f5b968c563dd56aec8f9b278c529c8aa46e6fe0ca564c26cc1c"},
-  "ash_sql": {:hex, :ash_sql, "0.3.0", "2c43ddcc8c7fb51dc25ba3bca965d8b68e7aaecb290cabfed3cf213965aca937", [:mix], [{:ash, ">= 3.5.43 and < 4.0.0-0", [hex: :ash, repo: "hexpm", optional: false]}, {:ecto, "~> 3.9", [hex: :ecto, repo: "hexpm", optional: false]}, {:ecto_sql, "~> 3.9", [hex: :ecto_sql, repo: "hexpm", optional: false]}], "hexpm", "ef03759d6a1d4cb189fcadbd183cf047f0565d7d88ea8612d85c9e6e724835e7"},
+  "ash_sql": {:git, "https://github.com/ash-project/ash_sql.git", "65854408e7ce129f78fabafb0a4393f0142da6a6", []},
   "benchee": {:hex, :benchee, "1.4.0", "9f1f96a30ac80bab94faad644b39a9031d5632e517416a8ab0a6b0ac4df124ce", [:mix], [{:deep_merge, "~> 1.0", [hex: :deep_merge, repo: "hexpm", optional: false]}, {:statistex, "~> 1.0", [hex: :statistex, repo: "hexpm", optional: false]}, {:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: true]}], "hexpm", "299cd10dd8ce51c9ea3ddb74bb150f93d25e968f93e4c1fa31698a8e4fa5d715"},
   "bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"},
   "credo": {:hex, :credo, "1.7.12", "9e3c20463de4b5f3f23721527fcaf16722ec815e70ff6c60b86412c695d426c1", [:mix], [{:bunt, "~> 0.2.1 or ~> 1.0", [hex: :bunt, repo: "hexpm", optional: false]}, {:file_system, "~> 0.2 or ~> 1.0", [hex: :file_system, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "8493d45c656c5427d9c729235b99d498bd133421f3e0a683e5c1b561471291e5"},