Add view-backed db model for postgres partitions

Add a new ActiveRecord model that can be used to query information about
database partitions. The model is backed by a database view, and should
reduce the number of ad-hoc queries to the Postgres system catalogs.

changelogs/unreleased/241267-add-postgres-partition-model.yml

+---
+title: Add database view for partitions
+merge_request: 45592
+author:
+type: other

db/migrate/20201019172704_add_partitions_view.rb

+# frozen_string_literal: true
+
+class AddPartitionsView < ActiveRecord::Migration[6.0]
+  DOWNTIME = false
+
+  def up
+    execute(<<~SQL)
+      CREATE OR REPLACE VIEW postgres_partitions AS
+      SELECT
+        pg_namespace.nspname::text || '.'::text || pg_class.relname::text AS identifier,
+        pg_class.oid AS oid,
+        pg_namespace.nspname AS schema,
+        pg_class.relname AS name,
+        parent_namespace.nspname::text || '.'::text || parent_class.relname::text AS parent_identifier,
+        pg_get_expr(pg_class.relpartbound, pg_inherits.inhrelid) AS condition
+      FROM pg_class
+      INNER JOIN pg_namespace
+      ON pg_namespace.oid = pg_class.relnamespace
+      INNER JOIN pg_inherits
+      ON pg_class.oid = pg_inherits.inhrelid
+      INNER JOIN pg_class parent_class
+      ON pg_inherits.inhparent = parent_class.oid
+      INNER JOIN pg_namespace parent_namespace
+      ON parent_class.relnamespace = parent_namespace.oid
+      WHERE pg_class.relispartition
+      AND pg_namespace.nspname IN (
+        current_schema(),
+        'gitlab_partitions_dynamic',
+        'gitlab_partitions_static'
+      )
+    SQL
+  end
+
+  def down
+    execute(<<~SQL)
+      DROP VIEW IF EXISTS postgres_partitions
+    SQL
+  end
+end

db/schema_migrations/20201019172704

+2a426e1a7cc6283e777667a1b4a6987f011c6cfd189ec702abc55d13a8499eca
\ No newline at end of file

db/structure.sql

@@ -14673,6 +14673,20 @@ CREATE VIEW postgres_partitioned_tables AS
             ELSE NULL::text
         END;
 
+CREATE VIEW postgres_partitions AS
+ SELECT (((pg_namespace.nspname)::text || '.'::text) || (pg_class.relname)::text) AS identifier,
+    pg_class.oid,
+    pg_namespace.nspname AS schema,
+    pg_class.relname AS name,
+    (((parent_namespace.nspname)::text || '.'::text) || (parent_class.relname)::text) AS parent_identifier,
+    pg_get_expr(pg_class.relpartbound, pg_inherits.inhrelid) AS condition
+   FROM ((((pg_class
+     JOIN pg_namespace ON ((pg_namespace.oid = pg_class.relnamespace)))
+     JOIN pg_inherits ON ((pg_class.oid = pg_inherits.inhrelid)))
+     JOIN pg_class parent_class ON ((pg_inherits.inhparent = parent_class.oid)))
+     JOIN pg_namespace parent_namespace ON ((parent_class.relnamespace = parent_namespace.oid)))
+  WHERE (pg_class.relispartition AND (pg_namespace.nspname = ANY (ARRAY["current_schema"(), 'gitlab_partitions_dynamic'::name, 'gitlab_partitions_static'::name])));
+
 CREATE TABLE postgres_reindex_actions (
     id bigint NOT NULL,
     action_start timestamp with time zone NOT NULL,

lib/gitlab/database/partitioning/monthly_strategy.rb

@@ -17,23 +17,8 @@ module Gitlab
         end
 
         def current_partitions
-          result = connection.select_all(<<~SQL)
-            select
-              pg_class.relname,
-              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 = #{connection.quote(Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)}
-              and parent_class.relname = #{connection.quote(table_name)}
-              and pg_class.relispartition
-            order by pg_class.relname
-          SQL
-
-          result.map do |record|
-            TimePartition.from_sql(table_name, record['relname'], record['condition'])
+          Gitlab::Database::PostgresPartition.for_parent_table(table_name).map do |partition|
+            TimePartition.from_sql(table_name, partition.name, partition.condition)
           end
         end
 

lib/gitlab/database/postgres_partition.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class PostgresPartition < ActiveRecord::Base
+      self.primary_key = :identifier
+
+      belongs_to :postgres_partitioned_table, foreign_key: 'parent_identifier', primary_key: 'identifier'
+
+      scope :by_identifier, ->(identifier) do
+        raise ArgumentError, "Partition name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
+
+        find(identifier)
+      end
+
+      scope :for_parent_table, ->(name) { where("parent_identifier = concat(current_schema(), '.', ?)", name).order(:name) }
+
+      def to_s
+        name
+      end
+    end
+  end
+end

lib/gitlab/database/postgres_partitioned_table.rb

@@ -7,6 +7,8 @@ module Gitlab
 
       self.primary_key = :identifier
 
+      has_many :postgres_partitions, foreign_key: 'parent_identifier', primary_key: 'identifier'
+
       scope :by_identifier, ->(identifier) do
         raise ArgumentError, "Table name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
 

spec/lib/gitlab/database/postgres_partition_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::PostgresPartition, type: :model do
+  let(:schema) { 'gitlab_partitions_dynamic' }
+  let(:name) { '_test_partition_01' }
+  let(:identifier) { "#{schema}.#{name}" }
+
+  before do
+    ActiveRecord::Base.connection.execute(<<~SQL)
+      CREATE TABLE public._test_partitioned_table (
+        id serial NOT NULL,
+        created_at timestamptz NOT NULL,
+        PRIMARY KEY (id, created_at)
+      ) PARTITION BY RANGE(created_at);
+
+      CREATE TABLE #{identifier} PARTITION OF public._test_partitioned_table
+      FOR VALUES FROM ('2020-01-01') to ('2020-02-01');
+    SQL
+  end
+
+  def find(identifier)
+    described_class.by_identifier(identifier)
+  end
+
+  describe 'associations' do
+    it { is_expected.to belong_to(:postgres_partitioned_table).with_primary_key('identifier').with_foreign_key('parent_identifier') }
+  end
+
+  it_behaves_like 'a postgres model'
+
+  describe '.for_parent_table' do
+    let(:second_name) { '_test_partition_02' }
+
+    before do
+      ActiveRecord::Base.connection.execute(<<~SQL)
+        CREATE TABLE #{schema}.#{second_name} PARTITION OF public._test_partitioned_table
+        FOR VALUES FROM ('2020-02-01') to ('2020-03-01');
+
+        CREATE TABLE #{schema}._test_other_table (
+          id serial NOT NULL,
+          created_at timestamptz NOT NULL,
+          PRIMARY KEY (id, created_at)
+        ) PARTITION BY RANGE(created_at);
+
+        CREATE TABLE #{schema}._test_other_partition_01 PARTITION OF #{schema}._test_other_table
+        FOR VALUES FROM ('2020-01-01') to ('2020-02-01');
+      SQL
+    end
+
+    it 'returns partitions for the parent table in the current schema' do
+      partitions = described_class.for_parent_table('_test_partitioned_table')
+
+      expect(partitions.count).to eq(2)
+      expect(partitions.pluck(:name)).to eq([name, second_name])
+    end
+
+    it 'does not return partitions for tables not in the current schema' do
+      expect(described_class.for_parent_table('_test_other_table').count).to eq(0)
+    end
+  end
+
+  describe '#parent_identifier' do
+    it 'returns the parent table identifier' do
+      expect(find(identifier).parent_identifier).to eq('public._test_partitioned_table')
+    end
+  end
+
+  describe '#condition' do
+    it 'returns the condition for the partitioned values' do
+      expect(find(identifier).condition).to eq("FOR VALUES FROM ('2020-01-01 00:00:00+00') TO ('2020-02-01 00:00:00+00')")
+    end
+  end
+end

spec/lib/gitlab/database/postgres_partitioned_table_spec.rb

@@ -33,6 +33,10 @@ RSpec.describe Gitlab::Database::PostgresPartitionedTable, type: :model do
     described_class.by_identifier(identifier)
   end
 
+  describe 'associations' do
+    it { is_expected.to have_many(:postgres_partitions).with_primary_key('identifier').with_foreign_key('parent_identifier') }
+  end
+
   it_behaves_like 'a postgres model'
 
   describe '#dynamic?' do