Merge pull request #14 from code0-tech/allow-non-serial-pk-on-partitioned-tables

Allow non-serial primary key on partitioned tables

Author: Taucher2003

README.md

@@ -87,9 +87,10 @@ config.zero_track.db_partitioning.base_ar_class = 'ActiveRecord::Base' # default
 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_partition_by_date_table(table_name, partition_column:, primary_key: nil, **options, &block)` creates a table
+partitioned by range on the given column. By default, it adds a `bigserial` id column and sets up a composite
+primary key of `(id, partition_column)`. If `primary_key` is provided (e.g. `primary_key: %i[date other_column]`),
+the id column is omitted and the given columns are used as the primary key instead.
 
 `create_dynamic_partition_schema` / `drop_dynamic_partition_schema` creates or drops the schema
 used for storing dynamic partitions.

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

@@ -5,21 +5,23 @@ module ZeroTrack
     module Database
       module MigrationHelpers
         module TablePartitioning
-          def create_partition_by_date_table(table_name, partition_column:, **options, &block)
+          def create_partition_by_date_table(table_name, partition_column:, primary_key: nil, **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
+              t.bigserial :id, null: false unless primary_key
 
               block.call(t)
             end
 
+            pk_columns = primary_key || [:id, partition_column]
+
             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)})
+                  ADD PRIMARY KEY (#{pk_columns.map { |c| quote_column_name(c) }.join(', ')})
                 SQL
               end
             end