1 files changed, 93 insertions, 18 deletions

diff --git a/spec/lib/gitlab/database/partitioning_spec.rb b/spec/lib/gitlab/database/partitioning_spec.rb
index 486af9413e8..6841af8f69d 100644
--- a/spec/lib/gitlab/database/partitioning_spec.rb
+++ b/spec/lib/gitlab/database/partitioning_spec.rb
@@ -3,46 +3,121 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Partitioning do
+  include Database::PartitioningHelpers
+  include Database::TableSchemaHelpers
+
+  let(:connection) { ApplicationRecord.connection }
+
   describe '.sync_partitions' do
-    let(:partition_manager_class) { described_class::MultiDatabasePartitionManager }
-    let(:partition_manager) { double('partition manager') }
+    let(:table_names) { %w[partitioning_test1 partitioning_test2] }
+    let(:models) do
+      table_names.map do |table_name|
+        Class.new(ApplicationRecord) do
+          include PartitionedTable
+
+          self.table_name = table_name
+          partitioned_by :created_at, strategy: :monthly
+        end
+      end
+    end
+
+    before do
+      table_names.each do |table_name|
+        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);
+        SQL
+      end
+    end
+
+    it 'manages partitions for each given model' do
+      expect { described_class.sync_partitions(models)}
+        .to change { find_partitions(table_names.first).size }.from(0)
+        .and change { find_partitions(table_names.last).size }.from(0)
+    end
 
     context 'when no partitioned models are given' do
-      it 'calls the partition manager with the registered models' do
-        expect(partition_manager_class).to receive(:new)
+      let(:partition_manager_class) { described_class::PartitionManager }
+      let(:partition_manager) { double('partition manager') }
+      let(:model) { double('model') }
+
+      it 'manages partitions for each registered model' do
+        expect(Gitlab::Database::EachDatabase).to receive(:each_model_connection)
           .with(described_class.registered_models)
-          .and_return(partition_manager)
+          .and_yield(model)
 
+        expect(partition_manager_class).to receive(:new).with(model).and_return(partition_manager)
         expect(partition_manager).to receive(:sync_partitions)
 
         described_class.sync_partitions
       end
     end
+  end
 
-    context 'when partitioned models are given' do
-      it 'calls the partition manager with the given models' do
-        models = ['my special model']
+  describe '.report_metrics' do
+    let(:model1) { double('model') }
+    let(:model2) { double('model') }
 
-        expect(partition_manager_class).to receive(:new)
-          .with(models)
-          .and_return(partition_manager)
+    let(:partition_monitoring_class) { described_class::PartitionMonitoring }
 
-        expect(partition_manager).to receive(:sync_partitions)
+    context 'when no partitioned models are given' do
+      it 'reports metrics for each registered model' do
+        expect_next_instance_of(partition_monitoring_class) do |partition_monitor|
+          expect(partition_monitor).to receive(:report_metrics_for_model).with(model1)
+          expect(partition_monitor).to receive(:report_metrics_for_model).with(model2)
+        end
+
+        expect(Gitlab::Database::EachDatabase).to receive(:each_model_connection)
+          .with(described_class.registered_models)
+          .and_yield(model1)
+          .and_yield(model2)
+
+        described_class.report_metrics
+      end
+    end
+
+    context 'when partitioned models are given' do
+      it 'reports metrics for each given model' do
+        expect_next_instance_of(partition_monitoring_class) do |partition_monitor|
+          expect(partition_monitor).to receive(:report_metrics_for_model).with(model1)
+          expect(partition_monitor).to receive(:report_metrics_for_model).with(model2)
+        end
+
+        expect(Gitlab::Database::EachDatabase).to receive(:each_model_connection)
+          .with([model1, model2])
+          .and_yield(model1)
+          .and_yield(model2)
 
-        described_class.sync_partitions(models)
+        described_class.report_metrics([model1, model2])
       end
     end
   end
 
   describe '.drop_detached_partitions' do
-    let(:partition_dropper_class) { described_class::MultiDatabasePartitionDropper }
+    let(:table_names) { %w[detached_test_partition1 detached_test_partition2] }
 
-    it 'delegates to the partition dropper' do
-      expect_next_instance_of(partition_dropper_class) do |partition_dropper|
-        expect(partition_dropper).to receive(:drop_detached_partitions)
+    before do
+      table_names.each do |table_name|
+        connection.create_table("#{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.#{table_name}")
+
+        Postgresql::DetachedPartition.create!(table_name: table_name, drop_after: 1.year.ago)
       end
+    end
+
+    it 'drops detached partitions for each database' do
+      expect(Gitlab::Database::EachDatabase).to receive(:each_database_connection).and_yield
+
+      expect { described_class.drop_detached_partitions }
+        .to change { Postgresql::DetachedPartition.count }.from(2).to(0)
+        .and change { table_exists?(table_names.first) }.from(true).to(false)
+        .and change { table_exists?(table_names.last) }.from(true).to(false)
+    end
 
-      described_class.drop_detached_partitions
+    def table_exists?(table_name)
+      table_oid(table_name).present?
     end
   end