Add view-backed db model for partitioned tables

Add a new ActiveRecord model that can be used to query information about
partitioned tables. 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-partitioned-table-model.yml

+---
+title: Add database view for partitioned tables
+merge_request: 45591
+author:
+type: other

db/migrate/20201019161924_add_partitioned_table_view.rb

+# frozen_string_literal: true
+
+class AddPartitionedTableView < ActiveRecord::Migration[6.0]
+  DOWNTIME = false
+
+  def up
+    execute(<<~SQL)
+      CREATE OR REPLACE VIEW postgres_partitioned_tables 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,
+        CASE partitioned_tables.partstrat
+        WHEN 'l' THEN 'list'
+        WHEN 'r' THEN 'range'
+        WHEN 'h' THEN 'hash'
+        END as strategy,
+        array_agg(pg_attribute.attname) as key_columns
+        FROM (
+          SELECT
+            partrelid,
+            partstrat,
+            unnest(partattrs) as column_position
+          FROM pg_partitioned_table
+        ) partitioned_tables
+        INNER JOIN pg_class
+        ON partitioned_tables.partrelid = pg_class.oid
+        INNER JOIN pg_namespace
+        ON pg_class.relnamespace = pg_namespace.oid
+        INNER JOIN pg_attribute
+        ON pg_attribute.attrelid = pg_class.oid
+        AND pg_attribute.attnum = partitioned_tables.column_position
+        WHERE pg_namespace.nspname = current_schema()
+        GROUP BY identifier, pg_class.oid, schema, name, strategy;
+    SQL
+  end
+
+  def down
+    execute(<<~SQL)
+      DROP VIEW IF EXISTS postgres_partitioned_tables
+    SQL
+  end
+end

db/schema_migrations/20201019161924

+47aba29a35e24113c9f198c731ba95597a2a6cd5d16b01a958644ce3e1a96170
\ No newline at end of file

db/structure.sql

@@ -14645,6 +14645,34 @@ CREATE VIEW postgres_indexes AS
      JOIN pg_indexes ON ((pg_class.relname = pg_indexes.indexname)))
   WHERE ((pg_namespace.nspname <> 'pg_catalog'::name) AND (pg_namespace.nspname = ANY (ARRAY["current_schema"(), 'gitlab_partitions_dynamic'::name, 'gitlab_partitions_static'::name])));
 
+CREATE VIEW postgres_partitioned_tables 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,
+        CASE partitioned_tables.partstrat
+            WHEN 'l'::"char" THEN 'list'::text
+            WHEN 'r'::"char" THEN 'range'::text
+            WHEN 'h'::"char" THEN 'hash'::text
+            ELSE NULL::text
+        END AS strategy,
+    array_agg(pg_attribute.attname) AS key_columns
+   FROM (((( SELECT pg_partitioned_table.partrelid,
+            pg_partitioned_table.partstrat,
+            unnest(pg_partitioned_table.partattrs) AS column_position
+           FROM pg_partitioned_table) partitioned_tables
+     JOIN pg_class ON ((partitioned_tables.partrelid = pg_class.oid)))
+     JOIN pg_namespace ON ((pg_class.relnamespace = pg_namespace.oid)))
+     JOIN pg_attribute ON (((pg_attribute.attrelid = pg_class.oid) AND (pg_attribute.attnum = partitioned_tables.column_position))))
+  WHERE (pg_namespace.nspname = "current_schema"())
+  GROUP BY (((pg_namespace.nspname)::text || '.'::text) || (pg_class.relname)::text), pg_class.oid, pg_namespace.nspname, pg_class.relname,
+        CASE partitioned_tables.partstrat
+            WHEN 'l'::"char" THEN 'list'::text
+            WHEN 'r'::"char" THEN 'range'::text
+            WHEN 'h'::"char" THEN 'hash'::text
+            ELSE NULL::text
+        END;
+
 CREATE TABLE postgres_reindex_actions (
     id bigint NOT NULL,
     action_start timestamp with time zone NOT NULL,

lib/gitlab/database/postgres_partitioned_table.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class PostgresPartitionedTable < ActiveRecord::Base
+      DYNAMIC_PARTITION_STRATEGIES = %w[range list].freeze
+
+      self.primary_key = :identifier
+
+      scope :by_identifier, ->(identifier) do
+        raise ArgumentError, "Table name is not fully qualified with a schema: #{identifier}" unless identifier =~ /^\w+\.\w+$/
+
+        find(identifier)
+      end
+
+      def dynamic?
+        DYNAMIC_PARTITION_STRATEGIES.include?(strategy)
+      end
+
+      def static?
+        !dynamic?
+      end
+
+      def to_s
+        name
+      end
+    end
+  end
+end

spec/lib/gitlab/database/postgres_index_spec.rb

@@ -3,9 +3,13 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::PostgresIndex do
+  let(:schema) { 'public' }
+  let(:name) { 'foo_idx' }
+  let(:identifier) { "#{schema}.#{name}" }
+
   before do
     ActiveRecord::Base.connection.execute(<<~SQL)
-      CREATE INDEX foo_idx ON public.users (name);
+      CREATE INDEX #{name} ON public.users (name);
       CREATE UNIQUE INDEX bar_key ON public.users (id);
 
       CREATE TABLE example_table (id serial primary key);
@@ -16,19 +20,7 @@ RSpec.describe Gitlab::Database::PostgresIndex do
     described_class.by_identifier(name)
   end
 
-  describe '.by_identifier' do
-    it 'finds the index' do
-      expect(find('public.foo_idx')).to be_a(Gitlab::Database::PostgresIndex)
-    end
-
-    it 'raises an error if not found' do
-      expect { find('public.idontexist') }.to raise_error(ActiveRecord::RecordNotFound)
-    end
-
-    it 'raises ArgumentError if given a non-fully qualified index name' do
-      expect { find('foo') }.to raise_error(ArgumentError, /not fully qualified/)
-    end
-  end
+  it_behaves_like 'a postgres model'
 
   describe '.regular' do
     it 'only non-unique indexes' do
@@ -76,7 +68,7 @@ RSpec.describe Gitlab::Database::PostgresIndex do
 
   describe '#valid_index?' do
     it 'returns true if the index is invalid' do
-      expect(find('public.foo_idx')).to be_valid_index
+      expect(find(identifier)).to be_valid_index
     end
 
     it 'returns false if the index is marked as invalid' do
@@ -86,31 +78,13 @@ RSpec.describe Gitlab::Database::PostgresIndex do
         WHERE pg_class.relname = 'foo_idx' AND pg_index.indexrelid = pg_class.oid
       SQL
 
-      expect(find('public.foo_idx')).not_to be_valid_index
-    end
-  end
-
-  describe '#to_s' do
-    it 'returns the index name' do
-      expect(find('public.foo_idx').to_s).to eq('foo_idx')
-    end
-  end
-
-  describe '#name' do
-    it 'returns the name' do
-      expect(find('public.foo_idx').name).to eq('foo_idx')
-    end
-  end
-
-  describe '#schema' do
-    it 'returns the index schema' do
-      expect(find('public.foo_idx').schema).to eq('public')
+      expect(find(identifier)).not_to be_valid_index
     end
   end
 
   describe '#definition' do
     it 'returns the index definition' do
-      expect(find('public.foo_idx').definition).to eq('CREATE INDEX foo_idx ON public.users USING btree (name)')
+      expect(find(identifier).definition).to eq('CREATE INDEX foo_idx ON public.users USING btree (name)')
     end
   end
 end

spec/lib/gitlab/database/postgres_partitioned_table_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::PostgresPartitionedTable, type: :model do
+  let(:schema) { 'public' }
+  let(:name) { 'foo_range' }
+  let(:identifier) { "#{schema}.#{name}" }
+
+  before do
+    ActiveRecord::Base.connection.execute(<<~SQL)
+      CREATE TABLE #{identifier} (
+        id serial NOT NULL,
+        created_at timestamptz NOT NULL,
+        PRIMARY KEY (id, created_at)
+      ) PARTITION BY RANGE(created_at);
+
+      CREATE TABLE public.foo_list (
+        id serial NOT NULL,
+        row_type text NOT NULL,
+        PRIMARY KEY (id, row_type)
+      ) PARTITION BY LIST(row_type);
+
+      CREATE TABLE public.foo_hash (
+        id serial NOT NULL,
+        row_value int NOT NULL,
+        PRIMARY KEY (id, row_value)
+      ) PARTITION BY HASH (row_value);
+    SQL
+  end
+
+  def find(identifier)
+    described_class.by_identifier(identifier)
+  end
+
+  it_behaves_like 'a postgres model'
+
+  describe '#dynamic?' do
+    it 'returns true for tables partitioned by range' do
+      expect(find('public.foo_range')).to be_dynamic
+    end
+
+    it 'returns true for tables partitioned by list' do
+      expect(find('public.foo_list')).to be_dynamic
+    end
+
+    it 'returns false for tables partitioned by hash' do
+      expect(find('public.foo_hash')).not_to be_dynamic
+    end
+  end
+
+  describe '#static?' do
+    it 'returns false for tables partitioned by range' do
+      expect(find('public.foo_range')).not_to be_static
+    end
+
+    it 'returns false for tables partitioned by list' do
+      expect(find('public.foo_list')).not_to be_static
+    end
+
+    it 'returns true for tables partitioned by hash' do
+      expect(find('public.foo_hash')).to be_static
+    end
+  end
+
+  describe '#strategy' do
+    it 'returns the partitioning strategy' do
+      expect(find(identifier).strategy).to eq('range')
+    end
+  end
+
+  describe '#key_columns' do
+    it 'returns the partitioning key columns' do
+      expect(find(identifier).key_columns).to match_array(['created_at'])
+    end
+  end
+end

spec/support/shared_examples/lib/gitlab/database/postgres_model_shared_examples.rb

+# frozen_string_literal: true
+
+RSpec.shared_examples 'a postgres model' do
+  describe '.by_identifier' do
+    it "finds the #{described_class}" do
+      expect(find(identifier)).to be_a(described_class)
+    end
+
+    it 'raises an error if not found' do
+      expect { find('public.idontexist') }.to raise_error(ActiveRecord::RecordNotFound)
+    end
+
+    it 'raises ArgumentError if given a non-fully qualified identifier' do
+      expect { find('foo') }.to raise_error(ArgumentError, /not fully qualified/)
+    end
+  end
+
+  describe '#to_s' do
+    it 'returns the name' do
+      expect(find(identifier).to_s).to eq(name)
+    end
+  end
+
+  describe '#schema' do
+    it 'returns the schema' do
+      expect(find(identifier).schema).to eq(schema)
+    end
+  end
+
+  describe '#name' do
+    it 'returns the name' do
+      expect(find(identifier).name).to eq(name)
+    end
+  end
+end