improvement: dynamically select and allow setting a repo

lib/data_layer.ex

@@ -2685,6 +2685,7 @@ defmodule AshPostgres.DataLayer do
   @impl true
   def update(resource, changeset) do
     source = resolve_source(resource, changeset)
+
     modifying =
       Map.keys(changeset.attributes) ++
         Keyword.keys(changeset.atomics) ++ Ash.Resource.Info.primary_key(resource)
@@ -2945,7 +2946,7 @@ defmodule AshPostgres.DataLayer do
     end
   end
 
-  def install(igniter, module, Ash.Resource, _path, _argv) do
+  def install(igniter, module, Ash.Resource, _path, argv) do
     table_name =
       module
       |> Module.split()
@@ -2953,7 +2954,16 @@ defmodule AshPostgres.DataLayer do
       |> Macro.underscore()
       |> Inflex.pluralize()
 
-    repo = Igniter.Code.Module.module_name("Repo")
+    {options, _, _} = OptionParser.parse(argv, switches: [repo: :string])
+
+    repo =
+      case options[:repo] do
+        nil ->
+          Igniter.Code.Module.module_name("Repo")
+
+        repo ->
+          Igniter.Code.Module.parse(repo)
+      end
 
     igniter
     |> Spark.Igniter.set_option(module, [:postgres, :table], table_name)

lib/igniter.ex

@@ -0,0 +1,84 @@
+defmodule AshPostgres.Igniter do
+  @moduledoc "Codemods and utilities for working with AshPostgres & Igniter"
+
+  @doc false
+  def default_repo_contents(otp_app) do
+    """
+    use AshPostgres.Repo, otp_app: #{inspect(otp_app)}
+
+    def installed_extensions do
+      # Add extensions here, and the migration generator will install them.
+      ["ash-functions"]
+    end
+    """
+  end
+
+  def add_postgres_extension(igniter, repo_name, extension) do
+    Igniter.Code.Module.find_and_update_module!(igniter, repo_name, fn zipper ->
+      case Igniter.Code.Function.move_to_def(zipper, :installed_extensions, 0) do
+        {:ok, zipper} ->
+          case Igniter.Code.List.append_new_to_list(zipper, extension) do
+            {:ok, zipper} ->
+              {:ok, zipper}
+
+            _ ->
+              {:warning,
+               "Could not add installed extension #{inspect(extension)} to #{inspect(repo_name)}.installed_extensions/0"}
+          end
+
+        _ ->
+          zipper = Sourceror.Zipper.rightmost(zipper)
+
+          code = """
+          def installed_extensions do
+            [#{inspect(extension)}]
+          end
+          """
+
+          {:ok, Igniter.Code.Common.add_code(zipper, code)}
+      end
+    end)
+  end
+
+  def select_repo(igniter, opts \\ []) do
+    label = Keyword.get(opts, :label, "Which repo should be used?")
+    generate = Keyword.get(opts, :generate?, false)
+
+    case list_repos(igniter) do
+      {igniter, []} ->
+        if generate do
+          repo = Igniter.Code.Module.module_name("Repo")
+          otp_app = Igniter.Project.Application.app_name()
+
+          igniter =
+            Igniter.Code.Module.create_module(igniter, repo, default_repo_contents(otp_app))
+
+          {igniter, repo}
+        else
+          {igniter, nil}
+        end
+
+      {igniter, [repo]} ->
+        {igniter, repo}
+
+      {igniter, repos} ->
+        {igniter, Owl.IO.select(repos, label: label, render_as: &inspect/1)}
+    end
+  end
+
+  def list_repos(igniter) do
+    Igniter.Code.Module.find_all_matching_modules(igniter, fn _mod, zipper ->
+      move_to_repo_use(zipper) != :error
+    end)
+  end
+
+  defp move_to_repo_use(zipper) do
+    Igniter.Code.Function.move_to_function_call(zipper, :use, 2, fn zipper ->
+      Igniter.Code.Function.argument_equals?(
+        zipper,
+        0,
+        AshPostgres.Repo
+      )
+    end)
+  end
+end

lib/mix/tasks/ash_postgres.install.ex

@@ -250,20 +250,10 @@ defmodule Mix.Tasks.AshPostgres.Install do
   end
 
   defp setup_repo_module(igniter, otp_app, repo) do
-    default_repo_contents =
-      """
-      use AshPostgres.Repo, otp_app: #{inspect(otp_app)}
-
-      def installed_extensions do
-        # Add extensions here, and the migration generator will install them.
-        ["ash-functions"]
-      end
-      """
-
     Igniter.Code.Module.find_and_update_or_create_module(
       igniter,
       repo,
-      default_repo_contents,
+      AshPostgres.Igniter.default_repo_contents(otp_app),
       fn zipper ->
         zipper
         |> set_otp_app(otp_app)