1 files changed, 65 insertions, 24 deletions

diff --git a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
index dca4548a0a3..8027990a546 100644
--- a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
@@ -21,20 +21,11 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
 
     let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table, connection: connection) }
     let(:connection) { ActiveRecord::Base.connection }
-    let(:table) { "issues" }
+    let(:table) { "my_model_example_table" }
     let(:partitioning_strategy) do
       double(missing_partitions: partitions, extra_partitions: [], after_adding_partitions: nil)
     end
 
-    before do
-      allow(connection).to receive(:table_exists?).and_call_original
-      allow(connection).to receive(:table_exists?).with(table).and_return(true)
-      allow(connection).to receive(:execute).and_call_original
-      expect(partitioning_strategy).to receive(:validate_and_fix)
-
-      stub_exclusive_lease(described_class::MANAGEMENT_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"),
@@ -42,19 +33,63 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
       ]
     end
 
-    it 'creates the partition' do
-      expect(connection).to receive(:execute).with("LOCK TABLE \"#{table}\" IN ACCESS EXCLUSIVE MODE")
-      expect(connection).to receive(:execute).with(partitions.first.to_sql)
-      expect(connection).to receive(:execute).with(partitions.second.to_sql)
+    context 'when the given table is partitioned' do
+      before do
+        create_partitioned_table(connection, table)
 
-      sync_partitions
+        allow(connection).to receive(:table_exists?).and_call_original
+        allow(connection).to receive(:table_exists?).with(table).and_return(true)
+        allow(connection).to receive(:execute).and_call_original
+        expect(partitioning_strategy).to receive(:validate_and_fix)
+
+        stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
+      end
+
+      it 'creates the partition' do
+        expect(connection).to receive(:execute).with("LOCK TABLE \"#{table}\" IN ACCESS EXCLUSIVE MODE")
+        expect(connection).to receive(:execute).with(partitions.first.to_sql)
+        expect(connection).to receive(:execute).with(partitions.second.to_sql)
+
+        sync_partitions
+      end
+
+      context 'with eplicitly provided connection' do
+        let(:connection) { Ci::ApplicationRecord.connection }
+
+        it 'uses the explicitly provided connection when any' do
+          skip_if_multiple_databases_not_setup
+
+          expect(connection).to receive(:execute).with("LOCK TABLE \"#{table}\" IN ACCESS EXCLUSIVE MODE")
+          expect(connection).to receive(:execute).with(partitions.first.to_sql)
+          expect(connection).to receive(:execute).with(partitions.second.to_sql)
+
+          described_class.new(model, connection: connection).sync_partitions
+        end
+      end
+
+      context 'when an error occurs during partition management' do
+        it 'does not raise an error' do
+          expect(partitioning_strategy).to receive(:missing_partitions).and_raise('this should never happen (tm)')
+
+          expect { sync_partitions }.not_to raise_error
+        end
+      end
     end
 
-    context 'when an error occurs during partition management' do
-      it 'does not raise an error' do
-        expect(partitioning_strategy).to receive(:missing_partitions).and_raise('this should never happen (tm)')
+    context 'when the table is not partitioned' do
+      let(:table) { 'this_does_not_need_to_be_real_table' }
+
+      it 'does not try creating the partitions' do
+        expect(connection).not_to receive(:execute).with("LOCK TABLE \"#{table}\" IN ACCESS EXCLUSIVE MODE")
+        expect(Gitlab::AppLogger).to receive(:warn).with(
+          {
+            message: 'Skipping synching partitions',
+            table_name: table,
+            connection_name: 'main'
+          }
+        )
 
-        expect { sync_partitions }.not_to raise_error
+        sync_partitions
       end
     end
   end
@@ -74,11 +109,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     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
+      create_partitioned_table(connection, 'my_model_example_table')
     end
 
     it 'creates partitions' do
@@ -98,6 +129,8 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     end
 
     before do
+      create_partitioned_table(connection, table)
+
       allow(connection).to receive(:table_exists?).and_call_original
       allow(connection).to receive(:table_exists?).with(table).and_return(true)
       expect(partitioning_strategy).to receive(:validate_and_fix)
@@ -260,4 +293,12 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
       expect { described_class.new(my_model).sync_partitions }.to change { has_partition(my_model, 2.months.ago.beginning_of_month) }.from(true).to(false).and(change { num_partitions(my_model) }.by(0))
     end
   end
+
+  def create_partitioned_table(connection, table)
+    connection.execute(<<~SQL)
+      CREATE TABLE #{table}
+      (id serial not null, created_at timestamptz not null, primary key (id, created_at))
+      PARTITION BY RANGE (created_at);
+    SQL
+  end
 end