1 files changed, 96 insertions, 0 deletions
diff --git a/spec/support/helpers/database/partitioning_helpers.rb b/spec/support/helpers/database/partitioning_helpers.rb
new file mode 100644
index 00000000000..80b31fe0603
--- /dev/null
+++ b/spec/support/helpers/database/partitioning_helpers.rb
@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module Database
+  module PartitioningHelpers
+    def expect_table_partitioned_by(table, columns, part_type: :range)
+      columns_with_part_type = columns.map { |c| [part_type.to_s, c] }
+      actual_columns = find_partitioned_columns(table)
+
+      expect(columns_with_part_type).to match_array(actual_columns)
+    end
+
+    def expect_range_partition_of(partition_name, table_name, min_value, max_value)
+      definition = find_partition_definition(partition_name, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)
+
+      expect(definition).not_to be_nil
+      expect(definition['base_table']).to eq(table_name.to_s)
+      expect(definition['condition']).to eq("FOR VALUES FROM (#{min_value}) TO (#{max_value})")
+    end
+
+    def expect_total_partitions(table_name, count, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)
+      partitions = find_partitions(table_name, schema: schema)
+
+      expect(partitions.size).to eq(count)
+    end
+
+    def expect_range_partitions_for(table_name, partitions)
+      partitions.each do |suffix, (min_value, max_value)|
+        partition_name = "#{table_name}_#{suffix}"
+        expect_range_partition_of(partition_name, table_name, min_value, max_value)
+      end
+
+      expect_total_partitions(table_name, partitions.size, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)
+    end
+
+    def expect_hash_partition_of(partition_name, table_name, modulus, remainder)
+      definition = find_partition_definition(partition_name, schema: Gitlab::Database::STATIC_PARTITIONS_SCHEMA)
+
+      expect(definition).not_to be_nil
+      expect(definition['base_table']).to eq(table_name.to_s)
+      expect(definition['condition']).to eq("FOR VALUES WITH (modulus #{modulus}, remainder #{remainder})")
+    end
+
+    private
+
+    def find_partitioned_columns(table)
+      connection.select_rows(<<~SQL)
+        select
+          case partstrat
+          when 'l' then 'list'
+          when 'r' then 'range'
+          when 'h' then 'hash'
+          end as partstrat,
+          cols.column_name
+        from (
+          select partrelid, partstrat, unnest(partattrs) as col_pos
+          from pg_partitioned_table
+        ) pg_part
+        inner join pg_class
+        on pg_part.partrelid = pg_class.oid
+        inner join information_schema.columns cols
+        on cols.table_name = pg_class.relname
+        and cols.ordinal_position = pg_part.col_pos
+        where pg_class.relname = '#{table}';
+      SQL
+    end
+
+    def find_partition_definition(partition, schema: )
+      connection.select_one(<<~SQL)
+        select
+          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 = '#{schema}'
+          and pg_class.relname = '#{partition}'
+          and pg_class.relispartition
+      SQL
+    end
+
+    def find_partitions(partition, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)
+      connection.select_rows(<<~SQL)
+        select
+          pg_class.relname
+        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 = '#{schema}'
+          and parent_class.relname = '#{partition}'
+          and pg_class.relispartition
+      SQL
+    end
+  end
+end