Implement table partition management

README.md

@@ -65,4 +65,131 @@ can be filled with the correct entries when the schema is loaded from the schema
 
 This approach is prone to git conflicts, so you can switch to a file based persistence
 with `config.zero_track.active_record.schema_migrations = true`. Instead of an `INSERT INTO` in
-the `db/structure.sql`, this mode creates files in the `db/schema_migrations` directory.
+the `db/structure.sql`, this mode creates files in the `db/schema_migrations` directory.
+
+### Table Partitioning
+
+PostgreSQL supports declarative table partitioning. The partition manager automates the
+management of partitions without manual operations or extensions on the PostgreSQL server.
+
+#### Configuration
+
+```ruby
+config.zero_track.db_partitioning.dynamic_partition_schema = 'partitions_dynamic' # default
+config.zero_track.db_partitioning.base_ar_class = 'ActiveRecord::Base' # default
+```
+
+- `dynamic_partition_schema`: The PostgreSQL schema where dynamic partitions are stored.
+- `base_ar_class`: The ActiveRecord base class used for the internal partitioning models.
+
+#### Migration Helpers
+
+Include the migration helpers by inheriting from `Code0::ZeroTrack::Database::Migration[1.0]` (or the
+appropriate version). The following methods become available:
+
+`create_partition_by_date_table(table_name, partition_column:, **options, &block)` creates a table
+partitioned by range on the given column. It automatically sets up a composite primary key
+of `(id, partition_column)`.
+
+`create_dynamic_partition_schema` / `drop_dynamic_partition_schema` creates or drops the schema
+used for storing dynamic partitions.
+
+`create_partitioning_views` / `drop_partitioning_views` creates or drops the PostgreSQL views
+(`postgres_partitioned_tables`, `postgres_partitions`, `postgres_detached_partitions`) that the
+partition manager uses to inspect existing partitions.
+
+Example migration:
+
+```ruby
+class CreatePartitionedEvents < Code0::ZeroTrack::Database::Migration[1.0]
+  def change
+    create_dynamic_partition_schema
+    create_partitioning_views
+
+    create_partition_by_date_table :events, partition_column: :created_at do |t|
+      t.text :name, null: false
+      t.timestamps_with_timezone null: false
+    end
+  end
+end
+```
+
+The schema and views only need to be created once before creating the first partitioned table.
+
+Tables don't necessarily have to be created with the provided helper. The partition manager will
+work as long as the model is correctly configured.
+
+#### Defining a Partitioned Model
+
+Include `Code0::ZeroTrack::Database::Partitioning::PartitionedTable` in your model and declare
+the partitioning strategy:
+
+```ruby
+class Event < ApplicationRecord
+  include Code0::ZeroTrack::Database::Partitioning::PartitionedTable
+
+  partition_by :created_at, strategy: :monthly, retain_for: 12.months
+end
+```
+
+Available strategies: `:daily` and `:monthly`.
+
+Options passed to `partition_by`:
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `strategy` | `:daily` or `:monthly` | *required* |
+| `headroom` | How far ahead to pre-create partitions | 30 days (daily) / 6 months (monthly) |
+| `retain_for` | How long to keep partitions before detaching (enables retention) | `nil` (disabled) |
+| `retain_detached_for` | How long to keep detached partitions before dropping | 7 days |
+
+#### Partition Manager
+
+Register models for automatic partition management:
+
+```ruby
+Code0::ZeroTrack::Database::Partitioning::PartitionManager.register_model(Event)
+Code0::ZeroTrack::Database::Partitioning::PartitionManager.register_model(EventDetail)
+```
+
+Then synchronize all registered models. The gem won't run this for you.
+Call it from a cron job, Sidekiq scheduler, or deploy script:
+
+```ruby
+Code0::ZeroTrack::Database::Partitioning::PartitionManager.sync_all_partitions!
+```
+
+Or manage a single model:
+
+```ruby
+manager = Code0::ZeroTrack::Database::Partitioning::PartitionManager.new(Event)
+manager.sync_partitions!
+```
+
+`sync_all_partitions!` first creates partitions for all registered models, then detaches and drops
+partitions for all models in reverse registration order.
+
+`sync_partitions!` performs three operations in order for a single model:
+1. Create and attach new partitions to cover the desired range (up to the configured headroom).
+2. Detach partitions that fall outside the desired range (when retention is enabled).
+3. Drop detached partitions that have been detached longer than `retain_detached_for`.
+
+If needed, the three phases can be called individually with `create_partitions!`, `detach_partitions!`
+and `drop_partitions!`. This only works on a partition manager for a specific model. There is no
+shortcut to run this on all registered models like the `sync_all_partitions!` method.
+
+Each table gets a PostgreSQL advisory lock, so concurrent calls won't conflict.
+
+When tables have foreign key relationships, registration order and retention configuration matter:
+
+- Register parent tables before child tables. `sync_all_partitions!` creates partitions in
+  registration order and detaches/drops them in reverse order. This way the parent tables
+  are created before the children and dropped after their children.
+- A child table's `retain_for` must be less than or equal to the parent table's `retain_for`.
+  If a child retains partitions longer than its parent, dropping the parent partition will
+  fail because the child's foreign key still references it.
+
+#### Schema Cleaner Integration
+
+Dynamic partition tables are automatically removed from `db/structure.sql` if the
+[schema cleaner](#configzero_trackactive_recordschema_cleaner) is enabled.

lib/code0/zero_track/database/migration.rb

@@ -11,6 +11,7 @@ class V1_0 < ::ActiveRecord::Migration[7.1]
           include Database::MigrationHelpers::IndexHelpers
           include Database::MigrationHelpers::RemoveColumnEnhancements
           include Database::MigrationHelpers::TableEnhancements
+          include Database::MigrationHelpers::TablePartitioning
         end
         # rubocop:enable Naming/ClassAndModuleCamelCase
 

lib/code0/zero_track/database/migration_helpers/table_partitioning.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+module Code0
+  module ZeroTrack
+    module Database
+      module MigrationHelpers
+        module TablePartitioning
+          def create_partition_by_date_table(table_name, partition_column:, **options, &block)
+            options[:options] = "PARTITION BY RANGE (#{quote_column_name(partition_column)})"
+            options[:id] = false
+
+            create_table(table_name, **options) do |t|
+              t.bigserial :id, null: false
+
+              block.call(t)
+            end
+
+            reversible do |dir|
+              dir.up do
+                execute <<~SQL.squish
+                  ALTER TABLE #{quote_table_name(table_name)}
+                  ADD PRIMARY KEY (#{quote_column_name(:id)}, #{quote_column_name(partition_column)})
+                SQL
+              end
+            end
+          end
+
+          def create_dynamic_partition_schema
+            schema = quote_table_name(Rails.application.config.zero_track.db_partitioning.dynamic_partition_schema)
+            execute "CREATE SCHEMA #{schema}"
+          end
+
+          def drop_dynamic_partition_schema
+            schema = quote_table_name(Rails.application.config.zero_track.db_partitioning.dynamic_partition_schema)
+            execute "DROP SCHEMA #{schema}"
+          end
+
+          def create_partitioning_views
+            dynamic_schema = Rails.application.config.zero_track.db_partitioning.dynamic_partition_schema
+
+            execute <<-SQL.squish
+              CREATE OR REPLACE VIEW postgres_partitioned_tables AS
+              SELECT c.oid::regclass::text AS identifier,
+                     c.oid,
+                     n.nspname AS schema,
+                     c.relname AS name,
+                     CASE p.partstrat
+                       WHEN 'l' THEN 'list'
+                       WHEN 'r' THEN 'range'
+                       WHEN 'h' THEN 'hash'
+                     END AS strategy,
+                     pg_get_partkeydef(c.oid) AS partition_key
+              FROM pg_partitioned_table p
+              JOIN pg_class c ON c.oid = p.partrelid
+              JOIN pg_namespace n ON n.oid = c.relnamespace
+              WHERE n.nspname = current_schema();
+            SQL
+
+            execute <<-SQL.squish
+              CREATE OR REPLACE VIEW postgres_partitions AS
+              SELECT c.oid::regclass::text AS identifier,
+                     c.oid,
+                     n.nspname AS schema,
+                     c.relname AS name,
+                     i.inhparent::regclass::text AS parent_identifier,
+                     pg_get_expr(c.relpartbound, c.oid) AS condition,
+                     obj_description(c.oid) AS comment,
+                     i.inhrelid IS NOT NULL AS attached
+              FROM pg_class c
+              LEFT JOIN pg_inherits i ON c.oid = i.inhrelid
+              JOIN pg_namespace n ON n.oid = c.relnamespace
+              WHERE c.relispartition
+                AND c.relkind = 'r'
+                AND n.nspname IN (current_schema(), #{quote(dynamic_schema)});
+            SQL
+
+            execute <<-SQL.squish
+              CREATE OR REPLACE VIEW postgres_detached_partitions AS
+              SELECT c.oid::regclass::text AS identifier,
+                     c.oid,
+                     n.nspname AS schema,
+                     c.relname AS name,
+                     obj_description(c.oid)::jsonb ->> 'table' AS parent_identifier,
+                     (obj_description(c.oid)::jsonb ->> 'detached_at')::timestamptz AS detached_at
+              FROM pg_class c
+              JOIN pg_namespace n ON n.oid = c.relnamespace
+              WHERE c.relkind = 'r'
+                AND n.nspname = #{quote(dynamic_schema)}
+                AND NOT EXISTS (
+                  SELECT 1 FROM pg_inherits WHERE inhrelid = c.oid
+                )
+                AND obj_description(c.oid)::jsonb ? 'table'
+                AND obj_description(c.oid)::jsonb ? 'detached_at';
+            SQL
+          end
+
+          def drop_partitioning_views
+            execute 'DROP VIEW IF EXISTS postgres_detached_partitions'
+            execute 'DROP VIEW IF EXISTS postgres_partitions'
+            execute 'DROP VIEW IF EXISTS postgres_partitioned_tables'
+          end
+        end
+      end
+    end
+  end
+end

lib/code0/zero_track/database/partitioning/partition_manager.rb

@@ -0,0 +1,148 @@
+# frozen_string_literal: true
+
+require 'zlib'
+
+module Code0
+  module ZeroTrack
+    module Database
+      module Partitioning
+        class PartitionManager
+          include Loggable
+
+          cattr_accessor :models
+          self.models = []
+
+          def self.register_model(clazz)
+            models << clazz
+          end
+
+          def self.reset_registered_models!
+            models.clear
+          end
+
+          def self.sync_all_partitions!
+            models.each do |model|
+              new(model).create_partitions!
+            end
+
+            models.reverse_each do |model|
+              manager = new(model)
+              manager.detach_partitions!
+              manager.drop_partitions!
+            end
+          end
+
+          attr_reader :model
+
+          def initialize(model)
+            if model.try(:partitioning_strategy).nil?
+              raise ArgumentError, "Model #{model} not configured for partitioning"
+            end
+
+            @model = model
+          end
+
+          def sync_partitions!
+            create_partitions!
+
+            detach_partitions!
+
+            drop_partitions!
+          end
+
+          def create_partitions!
+            with_lock do |connection|
+              model.partitioning_strategy.partitions_to_create.each do |partition|
+                create_partition(partition, connection)
+                attach_partition(partition, connection)
+              end
+            end
+          end
+
+          def detach_partitions!
+            with_lock do |connection|
+              model.partitioning_strategy.partitions_to_detach.each do |partition|
+                detach_partition(partition, connection)
+              end
+            end
+          end
+
+          def drop_partitions!
+            with_lock do |connection|
+              model.partitioning_strategy.partitions_to_drop.each do |detached_partition|
+                drop_partition(detached_partition, connection)
+              end
+            end
+          end
+
+          private
+
+          def create_partition(partition, connection)
+            connection.execute(partition.to_create_sql(connection))
+            logger.info(
+              message: 'Created new partition',
+              table_name: partition.model.table_name,
+              partition_name: partition.partition_name
+            )
+          end
+
+          def attach_partition(partition, connection)
+            connection.execute(partition.to_attach_sql(connection))
+            logger.info(
+              message: 'Attached partition',
+              table_name: partition.model.table_name,
+              partition_name: partition.partition_name
+            )
+          end
+
+          def detach_partition(partition, connection)
+            connection.execute(partition.to_detach_sql(connection))
+
+            partition_comment = connection.quote({ table: model.table_name, detached_at: Time.current.iso8601 }.to_json)
+            fully_qualified_partition = partition.fully_qualified_partition(connection)
+            connection.execute("COMMENT ON TABLE #{fully_qualified_partition} IS #{partition_comment}")
+
+            logger.info(
+              message: 'Detached partition',
+              table_name: partition.model.table_name,
+              partition_name: partition.partition_name
+            )
+          end
+
+          def drop_partition(detached_partition, connection)
+            schema_name = connection.quote_table_name(detached_partition.schema)
+            partition_name = connection.quote_table_name(detached_partition.name)
+            qualified_name = "#{schema_name}.#{partition_name}"
+            connection.execute("DROP TABLE #{qualified_name}")
+
+            logger.info(
+              message: 'Dropped partition',
+              table_name: detached_partition.parent_identifier,
+              partition_name: detached_partition.name
+            )
+          end
+
+          def with_lock
+            lock_key = lock_key_for(model.table_name)
+
+            with_connection do |connection|
+              connection.transaction do
+                connection.execute("SELECT pg_advisory_xact_lock(#{lock_key})")
+                yield connection
+              end
+            end
+          end
+
+          def lock_key_for(table_name)
+            namespace = 'zero_track:partition_sync'
+            Zlib.crc32("#{namespace}:#{table_name}")
+          end
+
+          def with_connection(&block)
+            model.with_connection(&block)
+          end
+        end
+      end
+    end
+  end
+end