Add latest changes from gitlab-org/gitlab@master

1 files changed, 12 insertions, 32 deletions

diff --git a/spec/lib/gitlab/database/partitioning_spec.rb b/spec/lib/gitlab/database/partitioning_spec.rb
index 4c0fde46b2f..4aa9d5f6df0 100644
--- a/spec/lib/gitlab/database/partitioning_spec.rb
+++ b/spec/lib/gitlab/database/partitioning_spec.rb
@@ -2,11 +2,11 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::Partitioning do
+RSpec.describe Gitlab::Database::Partitioning, feature_category: :database do
   include Database::PartitioningHelpers
   include Database::TableSchemaHelpers
 
-  let(:connection) { ApplicationRecord.connection }
+  let(:main_connection) { ApplicationRecord.connection }
 
   around do |example|
     previously_registered_models = described_class.registered_models.dup
@@ -84,7 +84,7 @@ RSpec.describe Gitlab::Database::Partitioning do
 
     before do
       table_names.each do |table_name|
-        connection.execute(<<~SQL)
+        execute_on_each_database(<<~SQL)
           CREATE TABLE #{table_name} (
             id serial not null,
             created_at timestamptz not null,
@@ -101,32 +101,12 @@ RSpec.describe Gitlab::Database::Partitioning do
     end
 
     context 'with multiple databases' do
-      before do
-        table_names.each do |table_name|
-          ci_connection.execute("DROP TABLE IF EXISTS #{table_name}")
-
-          ci_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
-
-      after do
-        table_names.each do |table_name|
-          ci_connection.execute("DROP TABLE IF EXISTS #{table_name}")
-        end
-      end
-
       it 'creates partitions in each database' do
-        skip_if_multiple_databases_not_setup(:ci)
+        skip_if_shared_database(:ci)
 
         expect { described_class.sync_partitions(models) }
-          .to change { find_partitions(table_names.first, conn: connection).size }.from(0)
-          .and change { find_partitions(table_names.last, conn: connection).size }.from(0)
+          .to change { find_partitions(table_names.first, conn: main_connection).size }.from(0)
+          .and change { find_partitions(table_names.last, conn: main_connection).size }.from(0)
           .and change { find_partitions(table_names.first, conn: ci_connection).size }.from(0)
           .and change { find_partitions(table_names.last, conn: ci_connection).size }.from(0)
       end
@@ -161,10 +141,12 @@ RSpec.describe Gitlab::Database::Partitioning do
       end
 
       before do
+        skip_if_shared_database(:ci)
+
         (table_names + ['partitioning_test3']).each do |table_name|
-          ci_connection.execute("DROP TABLE IF EXISTS #{table_name}")
+          execute_on_each_database("DROP TABLE IF EXISTS #{table_name}")
 
-          ci_connection.execute(<<~SQL)
+          execute_on_each_database(<<~SQL)
           CREATE TABLE #{table_name} (
             id serial not null,
             created_at timestamptz not null,
@@ -181,14 +163,12 @@ RSpec.describe Gitlab::Database::Partitioning do
       end
 
       it 'manages partitions for models for the given database', :aggregate_failures do
-        skip_if_multiple_databases_not_setup(:ci)
-
         expect { described_class.sync_partitions([models.first, ci_model], only_on: 'ci') }
           .to change { find_partitions(ci_model.table_name, conn: ci_connection).size }.from(0)
 
-        expect(find_partitions(models.first.table_name).size).to eq(0)
+        expect(find_partitions(models.first.table_name, conn: main_connection).size).to eq(0)
         expect(find_partitions(models.first.table_name, conn: ci_connection).size).to eq(0)
-        expect(find_partitions(ci_model.table_name).size).to eq(0)
+        expect(find_partitions(ci_model.table_name, conn: main_connection).size).to eq(0)
       end
     end
   end