improvement: add ImmutableRaiseError extension

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,12 @@ defmodule AshPostgres.SqlImplementation do
 
     {types, new_returns || returns}
   end
+
+  # Conditionally apply @impl to keep compatibility across ash_sql versions
+  if {:immutable_errors?, 1} in AshSql.Implementation.behaviour_info(:callbacks) do
+    @impl true
+    def immutable_errors?(repo) do
+      repo.immutable_expr_error?()
+    end
+  end
 end