Merge branch 'ab/auto-create-partitions' into 'master'

Automatically create dynamic database partitions

See merge request gitlab-org/gitlab!35981

app/models/concerns/partitioned_table.rb

+# frozen_string_literal: true
+
+module PartitionedTable
+  extend ActiveSupport::Concern
+
+  class_methods do
+    attr_reader :partitioning_strategy
+
+    PARTITIONING_STRATEGIES = {
+      monthly: Gitlab::Database::Partitioning::MonthlyStrategy
+    }.freeze
+
+    def partitioned_by(partitioning_key, strategy:)
+      strategy_class = PARTITIONING_STRATEGIES[strategy.to_sym] || raise(ArgumentError, "Unknown partitioning strategy: #{strategy}")
+
+      @partitioning_strategy = strategy_class.new(self, partitioning_key)
+
+      Gitlab::Database::Partitioning::PartitionCreator.register(self)
+    end
+  end
+end

app/workers/all_queues.yml

@@ -235,6 +235,14 @@
   :weight: 1
   :idempotent:
   :tags: []
+- :name: cronjob:partition_creation
+  :feature_category: :database
+  :has_external_dependencies: 
+  :urgency: :low
+  :resource_boundary: :unknown
+  :weight: 1
+  :idempotent: true
+  :tags: []
 - :name: cronjob:personal_access_tokens_expiring
   :feature_category: :authentication_and_authorization
   :has_external_dependencies:

app/workers/partition_creation_worker.rb

+# frozen_string_literal: true
+
+class PartitionCreationWorker
+  include ApplicationWorker
+  include CronjobQueue # rubocop:disable Scalability/CronWorkerContext
+
+  feature_category :database
+  idempotent!
+
+  def perform
+    Gitlab::AppLogger.info("Checking state of dynamic postgres partitions")
+
+    Gitlab::Database::Partitioning::PartitionCreator.new.create_partitions
+  end
+end

config/initializers/1_settings.rb

@@ -503,6 +503,9 @@ Settings.cron_jobs['authorized_project_update_periodic_recalculate_worker']['job
 Settings.cron_jobs['update_container_registry_info_worker'] ||= Settingslogic.new({})
 Settings.cron_jobs['update_container_registry_info_worker']['cron'] ||= '0 0 * * *'
 Settings.cron_jobs['update_container_registry_info_worker']['job_class'] = 'UpdateContainerRegistryInfoWorker'
+Settings.cron_jobs['postgres_dynamic_partitions_creator'] ||= Settingslogic.new({})
+Settings.cron_jobs['postgres_dynamic_partitions_creator']['cron'] ||= '21 */6 * * *'
+Settings.cron_jobs['postgres_dynamic_partitions_creator']['job_class'] ||= 'PartitionCreationWorker'
 
 Gitlab.ee do
   Settings.cron_jobs['adjourned_group_deletion_worker'] ||= Settingslogic.new({})

config/initializers/postgres_partitioning.rb

+# frozen_string_literal: true
+
+# Make sure we have loaded partitioned models here
+# (even with eager loading disabled).
+
+begin
+  Gitlab::Database::Partitioning::PartitionCreator.new.create_partitions
+rescue ActiveRecord::ActiveRecordError, PG::Error
+  # ignore - happens when Rake tasks yet have to create a database, e.g. for testing
+end

lib/gitlab/database/partitioning/monthly_strategy.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Partitioning
+      class MonthlyStrategy
+        attr_reader :model, :partitioning_key
+
+        # We create this many partitions in the future
+        HEADROOM = 6.months
+
+        delegate :table_name, to: :model
+
+        def initialize(model, partitioning_key)
+          @model = model
+          @partitioning_key = partitioning_key
+        end
+
+        def current_partitions
+          result = connection.select_all(<<~SQL)
+            select
+              pg_class.relname,
+              parent_class.relname as base_table,
+              pg_get_expr(pg_class.relpartbound, inhrelid) as condition
+            from pg_class
+            inner join pg_inherits i on pg_class.oid = inhrelid
+            inner join pg_class parent_class on parent_class.oid = inhparent
+            inner join pg_namespace ON pg_namespace.oid = pg_class.relnamespace
+            where pg_namespace.nspname = #{connection.quote(Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)}
+              and parent_class.relname = #{connection.quote(table_name)}
+              and pg_class.relispartition
+            order by pg_class.relname
+          SQL
+
+          result.map do |record|
+            TimePartition.from_sql(table_name, record['relname'], record['condition'])
+          end
+        end
+
+        # Check the currently existing partitions and determine which ones are missing
+        def missing_partitions
+          desired_partitions - current_partitions
+        end
+
+        private
+
+        def desired_partitions
+          [].tap do |parts|
+            min_date, max_date = relevant_range
+
+            parts << partition_for(upper_bound: min_date)
+
+            while min_date < max_date
+              next_date = min_date.next_month
+
+              parts << partition_for(lower_bound: min_date, upper_bound: next_date)
+
+              min_date = next_date
+            end
+          end
+        end
+
+        # This determines the relevant time range for which we expect to have data
+        # (and therefore need to create partitions for).
+        #
+        # Note: We typically expect the first partition to be half-unbounded, i.e.
+        #       to start from MINVALUE to a specific date `x`. The range returned
+        #       does not include the range of the first, half-unbounded partition.
+        def relevant_range
+          if first_partition = current_partitions.min
+            # Case 1: First partition starts with MINVALUE, i.e. from is nil -> start with first real partition
+            # Case 2: Rather unexpectedly, first partition does not start with MINVALUE, i.e. from is not nil
+            #         In this case, use first partition beginning as a start
+            min_date = first_partition.from || first_partition.to
+          end
+
+          # In case we don't have a partition yet
+          min_date ||= Date.today
+          min_date = min_date.beginning_of_month
+
+          max_date = Date.today.end_of_month + HEADROOM
+
+          [min_date, max_date]
+        end
+
+        def partition_for(lower_bound: nil, upper_bound:)
+          TimePartition.new(table_name, lower_bound, upper_bound)
+        end
+
+        def connection
+          ActiveRecord::Base.connection
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/partitioning/partition_creator.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Partitioning
+      class PartitionCreator
+        def self.register(model)
+          raise ArgumentError, "Only models with a #partitioning_strategy can be registered." unless model.respond_to?(:partitioning_strategy)
+
+          models << model
+        end
+
+        def self.models
+          @models ||= Set.new
+        end
+
+        LEASE_TIMEOUT = 1.minute
+        LEASE_KEY = 'database_partition_creation_%s'
+
+        attr_reader :models
+
+        def initialize(models = self.class.models)
+          @models = models
+        end
+
+        def create_partitions
+          return unless Feature.enabled?(:postgres_dynamic_partition_creation, default_enabled: true)
+
+          models.each do |model|
+            # Double-checking before getting the lease:
+            # The prevailing situation is no missing partitions
+            next if missing_partitions(model).empty?
+
+            only_with_exclusive_lease(model) do
+              partitions_to_create = missing_partitions(model)
+
+              next if partitions_to_create.empty?
+
+              create(model, partitions_to_create)
+            end
+          rescue => e
+            Gitlab::AppLogger.error("Failed to create partition(s) for #{model.table_name}: #{e.class}: #{e.message}")
+          end
+        end
+
+        private
+
+        def missing_partitions(model)
+          return [] unless connection.table_exists?(model.table_name)
+
+          model.partitioning_strategy.missing_partitions
+        end
+
+        def only_with_exclusive_lease(model)
+          lease = Gitlab::ExclusiveLease.new(LEASE_KEY % model.table_name, timeout: LEASE_TIMEOUT)
+
+          yield if lease.try_obtain
+        ensure
+          lease&.cancel
+        end
+
+        def create(model, partitions)
+          connection.transaction do
+            with_lock_retries do
+              partitions.each do |partition|
+                connection.execute partition.to_sql
+
+                Gitlab::AppLogger.info("Created partition #{partition.partition_name} for table #{partition.table}")
+              end
+            end
+          end
+        end
+
+        def with_lock_retries(&block)
+          Gitlab::Database::WithLockRetries.new({
+            klass: self.class,
+            logger: Gitlab::AppLogger
+          }).run(&block)
+        end
+
+        def connection
+          ActiveRecord::Base.connection
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/partitioning/time_partition.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Partitioning
+      class TimePartition
+        include Comparable
+
+        def self.from_sql(table, partition_name, definition)
+          matches = definition.match(/FOR VALUES FROM \('?(?<from>.+)'?\) TO \('?(?<to>.+)'?\)/)
+
+          raise ArgumentError, "Unknown partition definition: #{definition}" unless matches
+
+          raise NotImplementedError, "Open-end time partitions with MAXVALUE are not supported yet" if matches[:to] == 'MAXVALUE'
+
+          from = matches[:from] == 'MINVALUE' ? nil : matches[:from]
+          to = matches[:to]
+
+          new(table, from, to, partition_name: partition_name)
+        end
+
+        attr_reader :table, :from, :to
+
+        def initialize(table, from, to, partition_name: nil)
+          @table = table.to_s
+          @from = date_or_nil(from)
+          @to = date_or_nil(to)
+          @partition_name = partition_name
+        end
+
+        def partition_name
+          return @partition_name if @partition_name
+
+          suffix = from&.strftime('%Y%m') || '000000'
+
+          "#{table}_#{suffix}"
+        end
+
+        def to_sql
+          from_sql = from ? conn.quote(from.strftime('%Y-%m-%d')) : 'MINVALUE'
+          to_sql = conn.quote(to.strftime('%Y-%m-%d'))
+
+          <<~SQL
+            CREATE TABLE IF NOT EXISTS #{fully_qualified_partition}
+            PARTITION OF #{conn.quote_table_name(table)}
+            FOR VALUES FROM (#{from_sql}) TO (#{to_sql})
+          SQL
+        end
+
+        def ==(other)
+          table == other.table && partition_name == other.partition_name && from == other.from && to == other.to
+        end
+        alias_method :eql?, :==
+
+        def hash
+          [table, partition_name, from, to].hash
+        end
+
+        def <=>(other)
+          return if table != other.table
+
+          partition_name <=> other.partition_name
+        end
+
+        private
+
+        def date_or_nil(obj)
+          return unless obj
+          return obj if obj.is_a?(Date)
+
+          Date.parse(obj)
+        end
+
+        def fully_qualified_partition
+          "%s.%s" % [conn.quote_table_name(Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA), conn.quote_table_name(partition_name)]
+        end
+
+        def conn
+          @conn ||= ActiveRecord::Base.connection
+        end
+      end
+    end
+  end
+end

lib/tasks/gitlab/db.rake

@@ -134,5 +134,37 @@ namespace :gitlab do
     Rake::Task['db:structure:load'].enhance do
       Rake::Task['gitlab:db:load_custom_structure'].invoke
     end
+
+    desc 'Create missing dynamic database partitions'
+    task :create_dynamic_partitions do
+      Gitlab::Database::Partitioning::PartitionCreator.new.create_partitions
+    end
+
+    # This is targeted towards deploys and upgrades of GitLab.
+    # Since we're running migrations already at this time,
+    # we also check and create partitions as needed here.
+    Rake::Task['db:migrate'].enhance do
+      Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
+    end
+
+    # When we load the database schema from db/structure.sql
+    # we don't have any dynamic partitions created. We don't really need to
+    # because application initializers/sidekiq take care of that, too.
+    # However, the presence of partitions for a table has influence on their
+    # position in db/structure.sql (which is topologically sorted).
+    #
+    # Other than that it's helpful to create partitions early when bootstrapping
+    # a new installation.
+    Rake::Task['db:structure:load'].enhance do
+      Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
+    end
+
+    # During testing, db:test:load restores the database schema from scratch
+    # which does not include dynamic partitions. We cannot rely on application
+    # initializers here as the application can continue to run while
+    # a rake task reloads the database schema.
+    Rake::Task['db:test:load'].enhance do
+      Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
+    end
   end
 end

spec/lib/gitlab/database/partitioning/monthly_strategy_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Partitioning::MonthlyStrategy do
+  describe '#current_partitions' do
+    subject { described_class.new(model, partitioning_key).current_partitions }
+
+    let(:model) { double('model', table_name: table_name) }
+    let(:partitioning_key) { double }
+    let(:table_name) { :partitioned_test }
+
+    before do
+      ActiveRecord::Base.connection.execute(<<~SQL)
+        CREATE TABLE #{table_name}
+          (id serial not null, created_at timestamptz not null, PRIMARY KEY (id, created_at))
+          PARTITION BY RANGE (created_at);
+
+        CREATE TABLE #{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.partitioned_test_000000
+        PARTITION OF #{table_name}
+        FOR VALUES FROM (MINVALUE) TO ('2020-05-01');
+
+        CREATE TABLE #{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.partitioned_test_202005
+        PARTITION OF #{table_name}
+        FOR VALUES FROM ('2020-05-01') TO ('2020-06-01');
+      SQL
+    end
+
+    it 'detects both partitions' do
+      expect(subject).to eq([
+        Gitlab::Database::Partitioning::TimePartition.new(table_name, nil, '2020-05-01', partition_name: 'partitioned_test_000000'),
+        Gitlab::Database::Partitioning::TimePartition.new(table_name, '2020-05-01', '2020-06-01', partition_name: 'partitioned_test_202005')
+    ])
+    end
+  end
+
+  describe '#missing_partitions' do
+    subject { described_class.new(model, partitioning_key).missing_partitions }
+
+    let(:model) do
+      Class.new(ActiveRecord::Base) do
+        self.table_name = 'partitioned_test'
+        self.primary_key = :id
+      end
+    end
+
+    let(:partitioning_key) { :created_at }
+
+    around do |example|
+      Timecop.freeze(Date.parse('2020-08-22')) { example.run }
+    end
+
+    context 'with existing partitions' do
+      before do
+        ActiveRecord::Base.connection.execute(<<~SQL)
+          CREATE TABLE #{model.table_name}
+            (id serial not null, created_at timestamptz not null, PRIMARY KEY (id, created_at))
+            PARTITION BY RANGE (created_at);
+
+          CREATE TABLE #{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.partitioned_test_000000
+          PARTITION OF #{model.table_name}
+          FOR VALUES FROM (MINVALUE) TO ('2020-05-01');
+
+          CREATE TABLE #{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.partitioned_test_202006
+          PARTITION OF #{model.table_name}
+          FOR VALUES FROM ('2020-06-01') TO ('2020-07-01');
+        SQL
+
+        # Insert some data, it doesn't make a difference
+        model.create!(created_at: Date.parse('2020-04-20'))
+        model.create!(created_at: Date.parse('2020-06-15'))
+      end
+
+      it 'detects the gap and the missing partition in May 2020' do
+        expect(subject).to include(Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-05-01', '2020-06-01'))
+      end
+
+      it 'detects the missing partitions at the end of the range and expects a partition for July 2020' do
+        expect(subject).to include(Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-07-01', '2020-08-01'))
+      end
+
+      it 'detects the missing partitions at the end of the range and expects a partition for August 2020' do
+        expect(subject).to include(Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-08-01', '2020-09-01'))
+      end
+
+      it 'creates partitions 6 months out from now (Sep 2020 through Feb 2021)' do
+        expect(subject).to include(
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-09-01', '2020-10-01'),
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-10-01', '2020-11-01'),
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-11-01', '2020-12-01'),
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-12-01', '2021-01-01'),
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2021-01-01', '2021-02-01'),
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2021-02-01', '2021-03-01')
+        )
+      end
+
+      it 'detects all missing partitions' do
+        expect(subject.size).to eq(9)
+      end
+    end
+
+    context 'without existing partitions' do
+      before do
+        ActiveRecord::Base.connection.execute(<<~SQL)
+          CREATE TABLE #{model.table_name}
+            (id serial not null, created_at timestamptz not null, PRIMARY KEY (id, created_at))
+            PARTITION BY RANGE (created_at);
+        SQL
+      end
+
+      it 'detects the missing catch-all partition at the beginning' do
+        expect(subject).to include(Gitlab::Database::Partitioning::TimePartition.new(model.table_name, nil, '2020-08-01'))
+      end
+
+      it 'detects the missing partition for today and expects a partition for August 2020' do
+        expect(subject).to include(Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-08-01', '2020-09-01'))
+      end
+
+      it 'creates partitions 6 months out from now (Sep 2020 through Feb 2021' do
+        expect(subject).to include(
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-09-01', '2020-10-01'),
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-10-01', '2020-11-01'),
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-11-01', '2020-12-01'),
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2020-12-01', '2021-01-01'),
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2021-01-01', '2021-02-01'),
+          Gitlab::Database::Partitioning::TimePartition.new(model.table_name, '2021-02-01', '2021-03-01')
+        )
+      end
+
+      it 'detects all missing partitions' do
+        expect(subject.size).to eq(8)
+      end
+    end
+
+    context 'with a regular partition but no catchall (MINVALUE, to) partition' do
+      before do
+        ActiveRecord::Base.connection.execute(<<~SQL)
+          CREATE TABLE #{model.table_name}
+            (id serial not null, created_at timestamptz not null, PRIMARY KEY (id, created_at))
+            PARTITION BY RANGE (created_at);
+
+            CREATE TABLE #{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.partitioned_test_202006
+            PARTITION OF #{model.table_name}
+            FOR VALUES FROM ('2020-06-01') TO ('2020-07-01');
+        SQL
+      end
+
+      it 'detects a missing catch-all partition to add before the existing partition' do
+        expect(subject).to include(Gitlab::Database::Partitioning::TimePartition.new(model.table_name, nil, '2020-06-01'))
+      end
+    end
+  end
+end

spec/lib/gitlab/database/partitioning/partition_creator_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Partitioning::PartitionCreator do
+  include PartitioningHelpers
+  include ExclusiveLeaseHelpers
+
+  describe '.register' do
+    let(:model) { double(partitioning_strategy: nil) }
+
+    it 'remembers registered models' do
+      expect { described_class.register(model) }.to change { described_class.models }.to include(model)
+    end
+  end
+
+  describe '#create_partitions (mocked)' do
+    subject { described_class.new(models).create_partitions }
+
+    let(:models) { [model] }
+    let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table) }
+    let(:partitioning_strategy) { double(missing_partitions: partitions) }
+    let(:table) { "some_table" }
+
+    before do
+      allow(ActiveRecord::Base.connection).to receive(:table_exists?).and_call_original
+      allow(ActiveRecord::Base.connection).to receive(:table_exists?).with(table).and_return(true)
+      allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original
+
+      stub_exclusive_lease(described_class::LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
+    end
+
+    let(:partitions) do
+      [
+        instance_double(Gitlab::Database::Partitioning::TimePartition, table: 'bar', partition_name: 'foo', to_sql: "SELECT 1"),
+        instance_double(Gitlab::Database::Partitioning::TimePartition, table: 'bar', partition_name: 'foo2', to_sql: "SELECT 2")
+      ]
+    end
+
+    it 'creates the partition' do
+      expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.first.to_sql)
+      expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.second.to_sql)
+
+      subject
+    end
+
+    context 'error handling with 2 models' do
+      let(:models) do
+        [
+          double(partitioning_strategy: strategy1, table_name: table),
+          double(partitioning_strategy: strategy2, table_name: table)
+        ]
+      end
+
+      let(:strategy1) { double('strategy1', missing_partitions: nil) }
+      let(:strategy2) { double('strategy2', missing_partitions: partitions) }
+
+      it 'still creates partitions for the second table' do
+        expect(strategy1).to receive(:missing_partitions).and_raise('this should never happen (tm)')
+        expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.first.to_sql)
+        expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.second.to_sql)
+
+        subject
+      end
+    end
+  end
+
+  describe '#create_partitions' do
+    subject { described_class.new([my_model]).create_partitions }
+
+    let(:connection) { ActiveRecord::Base.connection }
+    let(:my_model) do
+      Class.new(ApplicationRecord) do
+        include PartitionedTable
+
+        self.table_name = 'my_model_example_table'
+
+        partitioned_by :created_at, strategy: :monthly
+      end
+    end
+
+    before do
+      connection.execute(<<~SQL)
+        CREATE TABLE my_model_example_table
+        (id serial not null, created_at timestamptz not null, primary key (id, created_at))
+        PARTITION BY RANGE (created_at);
+      SQL
+    end
+
+    it 'creates partitions' do
+      expect { subject }.to change { find_partitions(my_model.table_name, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA).size }.from(0)
+
+      subject
+    end
+  end
+end

spec/lib/gitlab/database/partitioning/time_partition_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Partitioning::TimePartition do
+  describe '.from_sql' do
+    subject { described_class.from_sql(table, partition_name, definition) }
+
+    let(:table) { 'foo' }
+    let(:partition_name) { 'foo_bar' }
+    let(:definition) { 'FOR VALUES FROM (\'2020-04-01 00:00:00\') TO (\'2020-05-01 00:00:00\')' }
+
+    it 'uses specified table name' do
+      expect(subject.table).to eq(table)
+    end
+
+    it 'uses specified partition name' do
+      expect(subject.partition_name).to eq(partition_name)
+    end
+
+    it 'parses start date' do
+      expect(subject.from).to eq(Date.parse('2020-04-01'))
+    end
+
+    it 'parses end date' do
+      expect(subject.to).to eq(Date.parse('2020-05-01'))
+    end
+
+    context 'with MINVALUE as a start date' do
+      let(:definition) { 'FOR VALUES FROM (MINVALUE) TO (\'2020-05-01\')' }
+
+      it 'sets from to nil' do
+        expect(subject.from).to be_nil
+      end
+    end
+
+    context 'with MAXVALUE as an end date' do
+      let(:definition) { 'FOR VALUES FROM (\'2020-04-01\') TO (MAXVALUE)' }
+
+      it 'raises a NotImplementedError' do
+        expect { subject }.to raise_error(NotImplementedError)
+      end
+    end
+  end
+
+  describe '#partition_name' do
+    subject { described_class.new(table, from, to, partition_name: partition_name).partition_name }
+
+    let(:table) { 'foo' }
+    let(:from) { '2020-04-01 00:00:00' }
+    let(:to) { '2020-05-01 00:00:00' }
+    let(:partition_name) { nil }
+
+    it 'uses table as prefix' do
+      expect(subject).to start_with(table)
+    end
+
+    it 'uses Year-Month (from) as suffix' do
+      expect(subject).to end_with("_202004")
+    end
+
+    context 'without from date' do
+      let(:from) { nil }
+
+      it 'uses 000000 as suffix for first partition' do
+        expect(subject).to end_with("_000000")
+      end
+    end
+
+    context 'with partition name explicitly given' do
+      let(:partition_name) { "foo_bar" }
+
+      it 'uses given partition name' do
+        expect(subject).to eq(partition_name)
+      end
+    end
+  end
+
+  describe '#to_sql' do
+    subject { described_class.new(table, from, to).to_sql }
+
+    let(:table) { 'foo' }
+    let(:from) { '2020-04-01 00:00:00' }
+    let(:to) { '2020-05-01 00:00:00' }
+
+    it 'transforms to a CREATE TABLE statement' do
+      expect(subject).to eq(<<~SQL)
+        CREATE TABLE IF NOT EXISTS "#{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}"."foo_202004"
+        PARTITION OF "foo"
+        FOR VALUES FROM ('2020-04-01') TO ('2020-05-01')
+      SQL
+    end
+
+    context 'without from date' do
+      let(:from) { nil }
+
+      it 'uses MINVALUE instead' do
+        expect(subject).to eq(<<~SQL)
+          CREATE TABLE IF NOT EXISTS "#{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}"."foo_000000"
+          PARTITION OF "foo"
+          FOR VALUES FROM (MINVALUE) TO ('2020-05-01')
+        SQL
+      end
+    end
+  end
+
+  describe 'object equality - #eql' do
+    def expect_inequality(actual, other)
+      expect(actual.eql?(other)).to be_falsey
+      expect(actual).not_to eq(other)
+    end
+
+    def expect_equality(actual, other)
+      expect(actual).to eq(other)
+      expect(actual.eql?(other)).to be_truthy
+      expect(actual.hash).to eq(other.hash)
+    end
+
+    def make_new(table: 'foo', from: '2020-04-01 00:00:00', to: '2020-05-01 00:00:00', partition_name: 'foo_202004')
+      described_class.new(table, from, to, partition_name: partition_name)
+    end
+
+    it 'treats objects identical with identical attributes' do
+      expect_equality(make_new, make_new)
+    end
+
+    it 'different table leads to in-equality' do
+      expect_inequality(make_new, make_new(table: 'bar'))
+    end
+
+    it 'different from leads to in-equality' do
+      expect_inequality(make_new, make_new(from: '2020-05-01 00:00:00'))
+    end
+
+    it 'different to leads to in-equality' do
+      expect_inequality(make_new, make_new(to: '2020-06-01 00:00:00'))
+    end
+
+    it 'different partition_name leads to in-equality' do
+      expect_inequality(make_new, make_new(partition_name: 'different'))
+    end
+
+    it 'nil partition_name is ignored if auto-generated matches' do
+      expect_equality(make_new, make_new(partition_name: nil))
+    end
+  end
+
+  describe 'Comparable, #<=>' do
+    let(:table) { 'foo' }
+
+    it 'sorts by partition name, i.e. by month - MINVALUE partition first' do
+      partitions = [
+        described_class.new(table, '2020-04-01', '2020-05-01'),
+        described_class.new(table, '2020-02-01', '2020-03-01'),
+        described_class.new(table, nil, '2020-02-01'),
+        described_class.new(table, '2020-03-01', '2020-04-01')
+      ]
+
+      expect(partitions.sort).to eq([
+        described_class.new(table, nil, '2020-02-01'),
+        described_class.new(table, '2020-02-01', '2020-03-01'),
+        described_class.new(table, '2020-03-01', '2020-04-01'),
+        described_class.new(table, '2020-04-01', '2020-05-01')
+      ])
+    end
+
+    it 'returns nil for partitions of different tables' do
+      one = described_class.new('foo', '2020-02-01', '2020-03-01')
+      two = described_class.new('bar', '2020-02-01', '2020-03-01')
+
+      expect(one.<=>(two)).to be_nil
+    end
+  end
+end

spec/models/concerns/partitioned_table_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe PartitionedTable do
+  describe '.partitioned_by' do
+    subject { my_class.partitioned_by(key, strategy: :monthly) }
+
+    let(:key) { :foo }
+
+    let(:my_class) do
+      Class.new do
+        include PartitionedTable
+      end
+    end
+
+    it 'assigns the MonthlyStrategy as the partitioning strategy' do
+      subject
+
+      expect(my_class.partitioning_strategy).to be_a(Gitlab::Database::Partitioning::MonthlyStrategy)
+    end
+
+    it 'passes the partitioning key to the strategy instance' do
+      subject
+
+      expect(my_class.partitioning_strategy.partitioning_key).to eq(key)
+    end
+
+    it 'registers itself with the PartitionCreator' do
+      expect(Gitlab::Database::Partitioning::PartitionCreator).to receive(:register).with(my_class)
+
+      subject
+    end
+  end
+end

spec/workers/partition_creation_worker_spec.rb

+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe PartitionCreationWorker do
+  describe '#perform' do
+    let(:creator) { double(create_partitions: nil) }
+
+    before do
+      allow(Gitlab::Database::Partitioning::PartitionCreator).to receive(:new).and_return(creator)
+    end
+
+    it 'delegates to PartitionCreator' do
+      expect(creator).to receive(:create_partitions)
+
+      described_class.new.perform
+    end
+  end
+end