diff --git a/README.md b/README.md index e7708db..cff76ad 100644 --- a/README.md +++ b/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. diff --git a/lib/code0/zero_track/database/migration_helpers/table_partitioning.rb b/lib/code0/zero_track/database/migration_helpers/table_partitioning.rb index ba42bf9..4de8465 100644 --- a/lib/code0/zero_track/database/migration_helpers/table_partitioning.rb +++ b/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