Partitioning: Initial suggestion

.gitignore

@@ -24,4 +24,5 @@ ash_postgres-*.tar
 
 test_migration_path
 test_snapshots_path
+test_tenant_migration_path
 

lib/data_layer.ex

@@ -252,6 +252,31 @@ defmodule AshPostgres.DataLayer do
     ]
   }
 
+  @partitioning %Spark.Dsl.Section{
+    name: :partitioning,
+    describe: """
+    A section for configuring the initial partitioning of the table
+    """,
+    examples: [
+      """
+      partitioning do
+        method :list
+        attribute :post
+      end
+      """
+    ],
+    schema: [
+      method: [
+        type: {:one_of, [:range, :list, :hash]},
+        doc: "Specifying what partitioning method to use"
+      ],
+      attribute: [
+        type: :atom,
+        doc: "The attribute to partition on"
+      ]
+    ]
+  }
+
   @postgres %Spark.Dsl.Section{
     name: :postgres,
     describe: """
@@ -262,7 +287,8 @@ defmodule AshPostgres.DataLayer do
       @custom_statements,
       @manage_tenant,
       @references,
-      @check_constraints
+      @check_constraints,
+      @partitioning
     ],
     modules: [
       :repo

lib/data_layer/info.ex

@@ -222,4 +222,14 @@ defmodule AshPostgres.DataLayer.Info do
   def manage_tenant_update?(resource) do
     Extension.get_opt(resource, [:postgres, :manage_tenant], :update?, false)
   end
+
+  @doc "Partitioning method"
+  def partitioning_method(resource) do
+    Extension.get_opt(resource, [:postgres, :partitioning], :method, nil)
+  end
+
+  @doc "Partitioning attribute"
+  def partitioning_attribute(resource) do
+    Extension.get_opt(resource, [:postgres, :partitioning], :attribute, nil)
+  end
 end

lib/migration_generator/migration_generator.ex

@@ -1192,14 +1192,25 @@ defmodule AshPostgres.MigrationGenerator do
 
   defp group_into_phases(
          [
-           %Operation.CreateTable{table: table, schema: schema, multitenancy: multitenancy} | rest
+           %Operation.CreateTable{
+             table: table,
+             schema: schema,
+             multitenancy: multitenancy,
+             partitioning: partitioning
+           }
+           | rest
          ],
          nil,
          acc
        ) do
     group_into_phases(
       rest,
-      %Phase.Create{table: table, schema: schema, multitenancy: multitenancy},
+      %Phase.Create{
+        table: table,
+        schema: schema,
+        multitenancy: multitenancy,
+        partitioning: partitioning
+      },
       acc
     )
   end
@@ -1801,7 +1812,8 @@ defmodule AshPostgres.MigrationGenerator do
         table: snapshot.table,
         schema: snapshot.schema,
         multitenancy: snapshot.multitenancy,
-        old_multitenancy: empty_snapshot.multitenancy
+        old_multitenancy: empty_snapshot.multitenancy,
+        partitioning: snapshot.partitioning
       }
       | acc
     ])
@@ -2851,7 +2863,8 @@ defmodule AshPostgres.MigrationGenerator do
       repo: AshPostgres.DataLayer.Info.repo(resource, :mutate),
       multitenancy: multitenancy(resource),
       base_filter: AshPostgres.DataLayer.Info.base_filter_sql(resource),
-      has_create_action: has_create_action?(resource)
+      has_create_action: has_create_action?(resource),
+      partitioning: partitioning(resource)
     }
 
     hash =
@@ -2926,6 +2939,20 @@ defmodule AshPostgres.MigrationGenerator do
     end)
   end
 
+  defp partitioning(resource) do
+    method = AshPostgres.DataLayer.Info.partitioning_method(resource)
+    attribute = AshPostgres.DataLayer.Info.partitioning_attribute(resource)
+
+    if not is_nil(method) and not is_nil(attribute) do
+      %{
+        method: method,
+        attribute: attribute
+      }
+    else
+      nil
+    end
+  end
+
   defp multitenancy(resource) do
     strategy = Ash.Resource.Info.multitenancy_strategy(resource)
     attribute = Ash.Resource.Info.multitenancy_attribute(resource)

lib/migration_generator/operation.ex

@@ -131,7 +131,7 @@ defmodule AshPostgres.MigrationGenerator.Operation do
 
   defmodule CreateTable do
     @moduledoc false
-    defstruct [:table, :schema, :multitenancy, :old_multitenancy]
+    defstruct [:table, :schema, :multitenancy, :old_multitenancy, :partitioning]
   end
 
   defmodule AddAttribute do
@@ -1055,17 +1055,24 @@ defmodule AshPostgres.MigrationGenerator.Operation do
     @moduledoc false
     defstruct [:schema, :table, :keys, no_phase: true]
 
-    def up(%{schema: schema, table: table, keys: keys}) do
+    def up(%{schema: schema, table: table, keys: keys, multitenancy: multitenancy}) do
       keys = Enum.join(keys, ", ")
 
-      if schema do
-        """
-        execute("ALTER TABLE \\\"#{schema}.#{table}\\\" ADD PRIMARY KEY (#{keys})")
-        """
-      else
-        """
-        execute("ALTER TABLE \\\"#{table}\\\" ADD PRIMARY KEY (#{keys})")
-        """
+      cond do
+        multitenancy.strategy == :context ->
+          """
+          execute("ALTER TABLE \\\"\#{prefix()}\\\".\\\"#{table}\\\" ADD PRIMARY KEY (#{keys})")
+          """
+
+        schema ->
+          """
+          execute("ALTER TABLE \\\"#{schema}.#{table}\\\" ADD PRIMARY KEY (#{keys})")
+          """
+
+        true ->
+          """
+          execute("ALTER TABLE \\\"#{table}\\\" ADD PRIMARY KEY (#{keys})")
+          """
       end
     end
 
@@ -1082,33 +1089,54 @@ defmodule AshPostgres.MigrationGenerator.Operation do
       ""
     end
 
-    def down(%{schema: schema, table: table, remove_old?: remove_old?, keys: keys}) do
+    def down(%{
+          schema: schema,
+          table: table,
+          remove_old?: remove_old?,
+          keys: keys,
+          multitenancy: multitenancy
+        }) do
       keys = Enum.join(keys, ", ")
 
-      if schema do
-        remove_old =
-          if remove_old? do
-            """
-            execute("ALTER TABLE \\\"#{schema}.#{table}\\\" DROP constraint #{table}_pkey")
-            """
-          end
+      cond do
+        multitenancy.strategy == :context ->
+          remove_old =
+            if remove_old? do
+              """
+              execute("ALTER TABLE \\\"\#{prefix()}\\\".\\\"#{table}\\\" DROP constraint #{table}_pkey")
+              """
+            end
 
-        """
-        #{remove_old}
-        execute("ALTER TABLE \\\"#{schema}.#{table}\\\" ADD PRIMARY KEY (#{keys})")
-        """
-      else
-        remove_old =
-          if remove_old? do
-            """
-            execute("ALTER TABLE \\\"#{table}\\\" DROP constraint #{table}_pkey")
-            """
-          end
+          """
+          #{remove_old}
+          execute("ALTER TABLE \\\"\#{prefix()}\\\".\\\"#{table}\\\" ADD PRIMARY KEY (#{keys})")
+          """
+
+        not is_nil(schema) ->
+          remove_old =
+            if remove_old? do
+              """
+              execute("ALTER TABLE \\\"#{schema}.#{table}\\\" DROP constraint #{table}_pkey")
+              """
+            end
+
+          """
+          #{remove_old}
+          execute("ALTER TABLE \\\"#{schema}.#{table}\\\" ADD PRIMARY KEY (#{keys})")
+          """
+
+        true ->
+          remove_old =
+            if remove_old? do
+              """
+              execute("ALTER TABLE \\\"#{table}\\\" DROP constraint #{table}_pkey")
+              """
+            end
 
-        """
-        #{remove_old}
-        execute("ALTER TABLE \\\"#{table}\\\" ADD PRIMARY KEY (#{keys})")
-        """
+          """
+          #{remove_old}
+          execute("ALTER TABLE \\\"#{table}\\\" ADD PRIMARY KEY (#{keys})")
+          """
       end
     end
   end
@@ -1117,11 +1145,16 @@ defmodule AshPostgres.MigrationGenerator.Operation do
     @moduledoc false
     defstruct [:schema, :table, no_phase: true]
 
-    def up(%{schema: schema, table: table}) do
-      if schema do
-        "drop constraint(#{inspect(table)}, \"#{table}_pkey\", prefix: \"#{schema}\")"
-      else
-        "drop constraint(#{inspect(table)}, \"#{table}_pkey\")"
+    def up(%{schema: schema, table: table, multitenancy: multitenancy}) do
+      cond do
+        multitenancy.strategy == :context ->
+          "drop constraint(#{inspect(table)}, \"#{table}_pkey\", prefix: prefix())"
+
+        schema ->
+          "drop constraint(#{inspect(table)}, \"#{table}_pkey\", prefix: \"#{schema}\")"
+
+        true ->
+          "drop constraint(#{inspect(table)}, \"#{table}_pkey\")"
       end
     end
 
@@ -1138,7 +1171,7 @@ defmodule AshPostgres.MigrationGenerator.Operation do
       ""
     end
 
-    def down(%{schema: schema, table: table, commented?: commented?}) do
+    def down(%{schema: schema, table: table, commented?: commented?, multitenancy: multitenancy}) do
       comment =
         if commented? do
           """
@@ -1149,10 +1182,15 @@ defmodule AshPostgres.MigrationGenerator.Operation do
           ""
         end
 
-      if schema do
-        "#{comment}drop constraint(#{inspect(table)}, \"#{table}_pkey\", prefix: \"#{schema}\")"
-      else
-        "#{comment}drop constraint(#{inspect(table)}, \"#{table}_pkey\")"
+      cond do
+        multitenancy.strategy == :context ->
+          "#{comment}drop constraint(#{inspect(table)}, \"#{table}_pkey\", prefix: prefix())"
+
+        schema ->
+          "#{comment}drop constraint(#{inspect(table)}, \"#{table}_pkey\", prefix: \"#{schema}\")"
+
+        true ->
+          "#{comment}drop constraint(#{inspect(table)}, \"#{table}_pkey\")"
       end
     end
   end

lib/migration_generator/phase.ex

@@ -3,24 +3,34 @@ defmodule AshPostgres.MigrationGenerator.Phase do
 
   defmodule Create do
     @moduledoc false
-    defstruct [:table, :schema, :multitenancy, operations: [], commented?: false]
+    defstruct [
+      :table,
+      :schema,
+      :multitenancy,
+      partitioning: nil,
+      operations: [],
+      commented?: false
+    ]
 
     import AshPostgres.MigrationGenerator.Operation.Helper, only: [as_atom: 1]
 
-    def up(%{schema: schema, table: table, operations: operations, multitenancy: multitenancy}) do
+    def up(%{
+          schema: schema,
+          table: table,
+          operations: operations,
+          multitenancy: multitenancy,
+          partitioning: partitioning
+        }) do
       if multitenancy.strategy == :context do
-        "create table(:#{as_atom(table)}, primary_key: false, prefix: prefix()) do\n" <>
+        arguments = arguments([prefix(true), options(partitioning: partitioning)])
+
+        "create table(:#{as_atom(table)}, primary_key: false#{arguments}) do\n" <>
           Enum.map_join(operations, "\n", fn operation -> operation.__struct__.up(operation) end) <>
           "\nend"
       else
-        opts =
-          if schema do
-            ", prefix: \"#{schema}\""
-          else
-            ""
-          end
+        arguments = arguments([prefix(schema), options(partitioning: partitioning)])
 
-        "create table(:#{as_atom(table)}, primary_key: false#{opts}) do\n" <>
+        "create table(:#{as_atom(table)}, primary_key: false#{arguments}) do\n" <>
           Enum.map_join(operations, "\n", fn operation -> operation.__struct__.up(operation) end) <>
           "\nend"
       end
@@ -40,6 +50,28 @@ defmodule AshPostgres.MigrationGenerator.Phase do
         "drop table(:#{as_atom(table)}#{opts})"
       end
     end
+
+    def arguments([nil, nil]), do: ""
+    def arguments(arguments), do: ", " <> Enum.join(Enum.reject(arguments, &is_nil(&1)), ",")
+
+    def prefix(true), do: "prefix: prefix()"
+    def prefix(schema) when is_binary(schema) and schema != "", do: "prefix: \"#{schema}\""
+    def prefix(_), do: nil
+
+    def options(_options, _acc \\ [])
+    def options([], []), do: nil
+    def options([], acc), do: "options: \"#{Enum.join(acc, " ")}\""
+
+    def options([{:partitioning, %{method: method, attribute: attribute}} | rest], acc) do
+      option = "PARTITION BY #{String.upcase(Atom.to_string(method))} (#{attribute})"
+
+      rest
+      |> options(acc ++ [option])
+    end
+
+    def options([_ | rest], acc) do
+      options(rest, acc)
+    end
   end
 
   defmodule Alter do