Update specs to better test index behavior in db

Update the specs to more thoroughly verify that the creation and removal
of indexes on partitioned works as expected across Postgres versions.

spec/lib/gitlab/database/partitioning_migration_helpers/index_helpers_spec.rb

@@ -3,42 +3,45 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
+  include TableSchemaHelpers
+
   let(:migration) do
     ActiveRecord::Migration.new.extend(described_class)
   end
 
   let(:table_name) { '_test_partitioned_table' }
+  let(:column_name) { 'created_at' }
   let(:index_name) { '_test_partitioning_index_name' }
-  let(:partitioned_table) { double(name: 'partitioned_table') }
+  let(:partition_schema) { 'gitlab_partitions_dynamic' }
+  let(:partition1_identifier) { "#{partition_schema}.#{table_name}_202001" }
+  let(:partition2_identifier) { "#{partition_schema}.#{table_name}_202002" }
+  let(:partition1_index) { "index_#{table_name}_202001_#{column_name}" }
+  let(:partition2_index) { "index_#{table_name}_202002_#{column_name}" }
 
   before do
     allow(migration).to receive(:puts)
 
-    allow(Gitlab::Database::PostgresPartitionedTable).to receive(:find_by_name_in_current_schema).with(table_name)
-      .and_return(partitioned_table)
-  end
-
-  describe '#add_concurrent_partitioned_index' do
-    let(:column_name) { '_test_column_name' }
-
-    let(:partition_name_1) { "#{table_name}_202001" }
-    let(:partition_name_2) { "#{table_name}_202002" }
-    let(:partition_name_3) { "#{table_name}_202003" }
+    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);
 
-    let(:index_name_1) { '_test_index_name_1' }
-    let(:index_name_2) { '_test_index_name_2' }
-    let(:index_name_3) { '_test_index_name_3' }
+      CREATE TABLE #{partition1_identifier} PARTITION OF #{table_name}
+      FOR VALUES FROM ('2020-01-01') TO ('2020-02-01');
 
-    let(:partition_1) { double(identifier: partition_name_1) }
-    let(:partition_2) { double(identifier: partition_name_2) }
-    let(:partition_3) { double(identifier: partition_name_3) }
-    let(:current_partitions) { [partition_1, partition_2, partition_3] }
+      CREATE TABLE #{partition2_identifier} PARTITION OF #{table_name}
+      FOR VALUES FROM ('2020-02-01') TO ('2020-03-01');
+    SQL
+  end
 
+  describe '#add_concurrent_partitioned_index' do
     before do
-      allow(partitioned_table).to receive(:postgres_partitions).and_return(current_partitions)
-
-      allow(migration).to receive(:generated_index_name).and_return(index_name_1, index_name_2, index_name_3)
       allow(migration).to receive(:index_name_exists?).with(table_name, index_name).and_return(false)
+
+      allow(migration).to receive(:generated_index_name).and_return(partition1_index, partition2_index)
+
       allow(migration).to receive(:with_lock_retries).and_yield
     end
 
@@ -46,17 +49,22 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
       it 'creates the index on each partition, and the parent table', :aggregate_failures do
         expect(migration).to receive(:index_name_exists?).with(table_name, index_name).and_return(false)
 
-        expect(migration).to receive(:add_concurrent_index)
-          .with(partition_name_1, column_name, name: index_name_1).ordered
-        expect(migration).to receive(:add_concurrent_index)
-          .with(partition_name_2, column_name, name: index_name_2).ordered
-        expect(migration).to receive(:add_concurrent_index)
-          .with(partition_name_3, column_name, name: index_name_3).ordered
+        expect_add_concurrent_index_and_call_original(partition1_identifier, column_name, partition1_index)
+        expect_add_concurrent_index_and_call_original(partition2_identifier, column_name, partition2_index)
 
-        expect(migration).to receive(:with_lock_retries).and_yield.ordered
-        expect(migration).to receive(:add_index).with(table_name, column_name, name: index_name).ordered
+        expect(migration).to receive(:with_lock_retries).ordered.and_yield
+        expect(migration).to receive(:add_index).with(table_name, column_name, name: index_name).ordered.and_call_original
 
         migration.add_concurrent_partitioned_index(table_name, column_name, name: index_name)
+
+        expect_index_to_exist(partition1_index, schema: partition_schema)
+        expect_index_to_exist(partition2_index, schema: partition_schema)
+        expect_index_to_exist(index_name)
+      end
+
+      def expect_add_concurrent_index_and_call_original(table, column, index)
+        expect(migration).to receive(:add_concurrent_index).ordered.with(table, column, name: index)
+          .and_wrap_original { |_, table, column, options| connection.add_index(table, column, options) }
       end
     end
 
@@ -73,11 +81,15 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
     end
 
     context 'when additional index options are given' do
-      let(:current_partitions) { [partition_1] }
+      before do
+        connection.execute(<<~SQL)
+          DROP TABLE #{partition2_identifier}
+        SQL
+      end
 
       it 'forwards them to the index helper methods', :aggregate_failures do
         expect(migration).to receive(:add_concurrent_index)
-          .with(partition_name_1, column_name, name: index_name_1, where: 'x > 0', unique: true)
+          .with(partition1_identifier, column_name, name: partition1_index, where: 'x > 0', unique: true)
 
         expect(migration).to receive(:add_index)
           .with(table_name, column_name, name: index_name, where: 'x > 0', unique: true)
@@ -90,6 +102,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
     context 'when a name argument for the index is not given' do
       it 'raises an error', :aggregate_failures do
         expect(migration).not_to receive(:add_concurrent_index)
+        expect(migration).not_to receive(:with_lock_retries)
         expect(migration).not_to receive(:add_index)
 
         expect do
@@ -106,6 +119,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
 
       it 'raises an error', :aggregate_failures do
         expect(migration).not_to receive(:add_concurrent_index)
+        expect(migration).not_to receive(:with_lock_retries)
         expect(migration).not_to receive(:add_index)
 
         expect do
@@ -116,19 +130,29 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
   end
 
   describe '#remove_concurrent_partitioned_index_by_name' do
-    before do
-      allow(migration).to receive(:index_name_exists?).with(table_name, index_name).and_return(true)
-      allow(migration).to receive(:with_lock_retries).and_yield
-    end
-
     context 'when the index exists' do
-      it 'drops the index on the parent table', :aggregate_failures do
-        expect(migration).to receive(:index_name_exists?).with(table_name, index_name).and_return(true)
+      before do
+        connection.execute(<<~SQL)
+          CREATE INDEX #{partition1_index} ON #{partition1_identifier} (#{column_name});
+          CREATE INDEX #{partition2_index} ON #{partition2_identifier} (#{column_name});
+
+          CREATE INDEX #{index_name} ON #{table_name} (#{column_name});
+        SQL
+      end
 
-        expect(migration).to receive(:with_lock_retries).and_yield.ordered
-        expect(migration).to receive(:remove_index).with(table_name, name: index_name).ordered
+      it 'drops the index on the parent table, cascading to all partitions', :aggregate_failures do
+        expect_index_to_exist(partition1_index, schema: partition_schema)
+        expect_index_to_exist(partition2_index, schema: partition_schema)
+        expect_index_to_exist(index_name)
+
+        expect(migration).to receive(:with_lock_retries).ordered.and_yield
+        expect(migration).to receive(:remove_index).with(table_name, name: index_name).ordered.and_call_original
 
         migration.remove_concurrent_partitioned_index_by_name(table_name, index_name)
+
+        expect_index_not_to_exist(partition1_index, schema: partition_schema)
+        expect_index_not_to_exist(partition2_index, schema: partition_schema)
+        expect_index_not_to_exist(index_name)
       end
     end
 
@@ -150,6 +174,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
       end
 
       it 'raises an error', :aggregate_failures do
+        expect(migration).not_to receive(:with_lock_retries)
         expect(migration).not_to receive(:remove_index)
 
         expect do

spec/support/helpers/table_schema_helpers.rb

@@ -16,6 +16,14 @@ module TableSchemaHelpers
     expect(table_oid(replacement_table)).to be_nil
   end
 
+  def expect_index_to_exist(name, schema: nil)
+    expect(index_exists_by_name(name, schema: schema)).to eq(true)
+  end
+
+  def expect_index_not_to_exist(name, schema: nil)
+    expect(index_exists_by_name(name, schema: schema)).to be_nil
+  end
+
   def table_oid(name)
     connection.select_value(<<~SQL)
       SELECT oid
@@ -76,4 +84,19 @@ module TableSchemaHelpers
         AND contype = 'p'
     SQL
   end
+
+  def index_exists_by_name(index, schema: nil)
+    schema = schema ? "'#{schema}'" : 'current_schema'
+
+    connection.select_value(<<~SQL)
+      SELECT true
+      FROM pg_catalog.pg_index i
+      INNER JOIN pg_catalog.pg_class c
+        ON c.oid = i.indexrelid
+      INNER JOIN pg_catalog.pg_namespace n
+        ON c.relnamespace = n.oid
+      WHERE c.relname = '#{index}'
+        AND n.nspname = #{schema}
+    SQL
+  end
 end