From 251a6476292c05a256b5b4a98ebf0c668b4a88d9 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Fri, 28 Feb 2025 10:48:02 -0500
Subject: [PATCH] WIP: experimentation w/ eager loading migrations for faster
 tenancy

---
 lib/data_layer.ex              | 48 +++++++++++++++++++-------------
 lib/migration_compile_cache.ex | 11 ++++++--
 lib/multitenancy.ex            | 38 +++++++++++++++++++------
 lib/repo.ex                    | 51 ++++++++++++++++++++++++++++------
 4 files changed, 110 insertions(+), 38 deletions(-)

diff --git a/lib/data_layer.ex b/lib/data_layer.ex
index b66cdfc7..ff646b3b 100644
--- a/lib/data_layer.ex
+++ b/lib/data_layer.ex
@@ -2102,27 +2102,37 @@ defmodule AshPostgres.DataLayer do
 
   @impl true
   def create(resource, changeset) do
-    changeset = %{
-      changeset
-      | data:
-          Map.update!(
-            changeset.data,
-            :__meta__,
-            &Map.put(&1, :source, table(resource, changeset))
-          )
-    }
+    start_time = :erlang.monotonic_time(:millisecond)
 
-    case bulk_create(resource, [changeset], %{
-           single?: true,
-           tenant: Map.get(changeset, :to_tenant, changeset.tenant),
-           action_select: changeset.action_select,
-           return_records?: true
-         }) do
-      {:ok, [result]} ->
-        {:ok, result}
+    try do
+      changeset = %{
+        changeset
+        | data:
+            Map.update!(
+              changeset.data,
+              :__meta__,
+              &Map.put(&1, :source, table(resource, changeset))
+            )
+      }
 
-      {:error, error} ->
-        {:error, error}
+      case bulk_create(resource, [changeset], %{
+             single?: true,
+             tenant: Map.get(changeset, :to_tenant, changeset.tenant),
+             action_select: changeset.action_select,
+             return_records?: true
+           }) do
+        {:ok, [result]} ->
+          {:ok, result}
+
+        {:error, error} ->
+          {:error, error}
+      end
+    after
+      end_time = :erlang.monotonic_time(:millisecond)
+
+      IO.puts(
+        "3Time taken: #{:erlang.convert_time_unit(end_time - start_time, :millisecond, :millisecond)} ms"
+      )
     end
   end
 
diff --git a/lib/migration_compile_cache.ex b/lib/migration_compile_cache.ex
index 802c21ad..1ff9c8dd 100644
--- a/lib/migration_compile_cache.ex
+++ b/lib/migration_compile_cache.ex
@@ -11,8 +11,15 @@ defmodule AshPostgres.MigrationCompileCache do
   once, which would error out.
   """
 
-  def start_link(opts \\ %{}) do
-    Agent.start_link(fn -> opts end, name: __MODULE__)
+  def start_link(state \\ %{}, opts \\ []) do
+    Agent.start_link(fn -> state end, opts)
+  end
+
+  def child_spec(opts) do
+    %{
+      id: __MODULE__,
+      start: {__MODULE__, :start_link, [%{}, opts]}
+    }
   end
 
   @doc """
diff --git a/lib/multitenancy.ex b/lib/multitenancy.ex
index 20e67c0c..a8499633 100644
--- a/lib/multitenancy.ex
+++ b/lib/multitenancy.ex
@@ -5,9 +5,18 @@ defmodule AshPostgres.MultiTenancy do
 
   @tenant_name_regex ~r/^[a-zA-Z0-9_-]+$/
   def create_tenant!(tenant_name, repo) do
+    start_time = :erlang.monotonic_time(:millisecond)
     Ecto.Adapters.SQL.query!(repo, "CREATE SCHEMA IF NOT EXISTS \"#{tenant_name}\"", [])
 
     migrate_tenant(tenant_name, repo)
+
+    end_time = :erlang.monotonic_time(:millisecond)
+
+    IO.puts(
+      "2Time taken: #{:erlang.convert_time_unit(end_time - start_time, :millisecond, :millisecond)} ms"
+    )
+
+    :ok
   end
 
   def migrate_tenant(tenant_name, repo, migrations_path \\ nil) do
@@ -21,13 +30,15 @@ defmodule AshPostgres.MultiTenancy do
       prefix: tenant_name
     )
 
-    [tenant_migrations_path, "**", "*.exs"]
-    |> Path.join()
-    |> Path.wildcard()
-    |> Enum.map(&extract_migration_info/1)
-    |> Enum.filter(& &1)
-    |> Enum.map(&load_migration!/1)
-    |> Enum.each(fn {version, mod} ->
+    migrations =
+      if repo.config()[:eager_load_tenant_migrations?] do
+        :persistent_term.get({repo, :tenant_migrations}, nil) ||
+          tenant_migrations(tenant_migrations_path)
+      else
+        tenant_migrations(tenant_migrations_path)
+      end
+
+    Enum.each(migrations, fn {version, mod} ->
       Ecto.Migration.Runner.run(
         repo,
         [],
@@ -42,6 +53,8 @@ defmodule AshPostgres.MultiTenancy do
 
       Ecto.Migration.SchemaMigration.up(repo, repo.config(), version, prefix: tenant_name)
     end)
+
+    :ok
   end
 
   # sobelow_skip ["SQL"]
@@ -56,6 +69,15 @@ defmodule AshPostgres.MultiTenancy do
     :ok
   end
 
+  def tenant_migrations(tenant_migrations_path) do
+    [tenant_migrations_path, "**", "*.exs"]
+    |> Path.join()
+    |> Path.wildcard()
+    |> Enum.map(&extract_migration_info/1)
+    |> Enum.filter(& &1)
+    |> Enum.map(&load_migration!/1)
+  end
+
   defp load_migration!({version, _, file}) when is_binary(file) do
     loaded_modules = file |> compile_file() |> Enum.map(&elem(&1, 0))
 
@@ -68,7 +90,7 @@ defmodule AshPostgres.MultiTenancy do
   end
 
   defp compile_file(file) do
-    AshPostgres.MigrationCompileCache.start_link()
+    AshPostgres.MigrationCompileCache.start_link(%{}, name: AshPostgres.MigrationCompileCache)
     AshPostgres.MigrationCompileCache.compile_file(file)
   end
 
diff --git a/lib/repo.ex b/lib/repo.ex
index 02d8c703..80becd39 100644
--- a/lib/repo.ex
+++ b/lib/repo.ex
@@ -126,7 +126,30 @@ defmodule AshPostgres.Repo do
       @before_compile AshPostgres.Repo.BeforeCompile
       require Logger
 
-      defoverridable insert: 2, insert: 1, insert!: 2, insert!: 1, transaction: 1, transaction: 2
+      defoverridable insert: 2,
+                     insert: 1,
+                     insert!: 2,
+                     insert!: 1,
+                     transaction: 1,
+                     transaction: 2,
+                     start_link: 1
+
+      def start_link(opts) do
+        config = config()
+
+        if config[:eager_load_tenant_migrations?] &&
+             !:persistent_term.get({__MODULE__, :tenant_migrations}, nil) do
+          config[:tenant_migrations_path]
+          |> Kernel.||(tenant_migrations_path())
+          |> Kernel.||(default_tenant_migration_path(config))
+          |> AshPostgres.MultiTenancy.tenant_migrations()
+          |> then(&:persistent_term.put({__MODULE__, :tenant_migrations}, &1))
+        else
+          opts
+        end
+
+        super(opts)
+      end
 
       def transaction(fun, opts \\ []) do
         super(fun, opts)
@@ -170,18 +193,28 @@ defmodule AshPostgres.Repo do
       end
 
       def init(type, config) do
-        new_config =
-          config
-          |> Keyword.put(:installed_extensions, installed_extensions())
-          |> Keyword.put(:tenant_migrations_path, tenant_migrations_path())
-          |> Keyword.put(:migrations_path, migrations_path())
-          |> Keyword.put(:default_prefix, default_prefix())
-
-        {:ok, new_config}
+        config
+        |> Keyword.put(:installed_extensions, installed_extensions())
+        |> Keyword.put(
+          :tenant_migrations_path,
+          tenant_migrations_path() || default_tenant_migration_path(config)
+        )
+        |> Keyword.put(:migrations_path, migrations_path())
+        |> Keyword.put(:default_prefix, default_prefix())
+        |> then(&{:ok, &1})
       end
 
       def on_transaction_begin(_reason), do: :ok
 
+      defp default_tenant_migration_path(options) do
+        repo_name = __MODULE__ |> Module.split() |> List.last() |> Macro.underscore()
+        otp_app = options[:otp_app]
+
+        :code.priv_dir(otp_app)
+        |> Path.join(repo_name)
+        |> Path.join("tenant_migrations")
+      end
+
       # copied from Ecto.Repo
       defp ash_postgres_prepare_opts(operation_name, []),
         do: default_options(operation_name)
