Merge branch 'align-schema-ordering-for-jh' into 'master'

Align the ordering of partitioned tables

See merge request gitlab-org/gitlab!72697

config/initializers/postgres_partitioning.rb

@@ -10,6 +10,29 @@ if Gitlab.ee?
     IncidentManagement::PendingEscalations::Alert,
     IncidentManagement::PendingEscalations::Issue
   ])
+else
+  Gitlab::Database::Partitioning.register_tables([
+    {
+      table_name: 'incident_management_pending_alert_escalations',
+      partitioned_column: :process_at, strategy: :monthly
+    },
+    {
+      table_name: 'incident_management_pending_issue_escalations',
+      partitioned_column: :process_at, strategy: :monthly
+    }
+  ])
+end
+
+# The following tables are already defined as models
+unless Gitlab.jh?
+  Gitlab::Database::Partitioning.register_tables([
+    # This should be synchronized with the following model:
+    # https://gitlab.com/gitlab-jh/gitlab/-/blob/main-jh/jh/app/models/phone/verification_code.rb
+    {
+      table_name: 'verification_codes',
+      partitioned_column: :created_at, strategy: :monthly
+    }
+  ])
 end
 
 begin

db/structure.sql

@@ -135,6 +135,19 @@ CREATE TABLE incident_management_pending_issue_escalations (
 )
 PARTITION BY RANGE (process_at);
 
+CREATE TABLE verification_codes (
+    created_at timestamp with time zone DEFAULT now() NOT NULL,
+    visitor_id_code text NOT NULL,
+    code text NOT NULL,
+    phone text NOT NULL,
+    CONSTRAINT check_9b84e6aaff CHECK ((char_length(code) <= 8)),
+    CONSTRAINT check_ccc542256b CHECK ((char_length(visitor_id_code) <= 64)),
+    CONSTRAINT check_f5684c195b CHECK ((char_length(phone) <= 32))
+)
+PARTITION BY RANGE (created_at);
+
+COMMENT ON TABLE verification_codes IS 'JiHu-specific table';
+
 CREATE TABLE web_hook_logs (
     id bigint NOT NULL,
     web_hook_id integer NOT NULL,
@@ -20185,19 +20198,6 @@ CREATE SEQUENCE users_statistics_id_seq
 
 ALTER SEQUENCE users_statistics_id_seq OWNED BY users_statistics.id;
 
-CREATE TABLE verification_codes (
-    created_at timestamp with time zone DEFAULT now() NOT NULL,
-    visitor_id_code text NOT NULL,
-    code text NOT NULL,
-    phone text NOT NULL,
-    CONSTRAINT check_9b84e6aaff CHECK ((char_length(code) <= 8)),
-    CONSTRAINT check_ccc542256b CHECK ((char_length(visitor_id_code) <= 64)),
-    CONSTRAINT check_f5684c195b CHECK ((char_length(phone) <= 32))
-)
-PARTITION BY RANGE (created_at);
-
-COMMENT ON TABLE verification_codes IS 'JiHu-specific table';
-
 CREATE TABLE vulnerabilities (
     id bigint NOT NULL,
     milestone_id bigint,

lib/gitlab/database/partitioning.rb

@@ -3,15 +3,31 @@
 module Gitlab
   module Database
     module Partitioning
-      def self.register_models(models)
+      class TableWithoutModel
+        include PartitionedTable::ClassMethods
+
+        attr_reader :table_name
+
+        def initialize(table_name:, partitioned_column:, strategy:)
+          @table_name = table_name
+          partitioned_by(partitioned_column, strategy: strategy)
+        end
+
+        def connection
+          Gitlab::Database::SharedModel.connection
+        end
+      end
+
+      class << self
+        def register_models(models)
           registered_models.merge(models)
         end
 
-      def self.registered_models
-        @registered_models ||= Set.new
+        def register_tables(tables)
+          registered_tables.merge(tables)
         end
 
-      def self.sync_partitions(models_to_sync = registered_models)
+        def sync_partitions(models_to_sync = registered_for_sync)
           Gitlab::AppLogger.info(message: 'Syncing dynamic postgres partitions')
 
           Gitlab::Database::EachDatabase.each_model_connection(models_to_sync) do |model|
@@ -21,7 +37,7 @@ module Gitlab
           Gitlab::AppLogger.info(message: 'Finished sync of dynamic postgres partitions')
         end
 
-      def self.report_metrics(models_to_monitor = registered_models)
+        def report_metrics(models_to_monitor = registered_models)
           partition_monitoring = PartitionMonitoring.new
 
           Gitlab::Database::EachDatabase.each_model_connection(models_to_monitor) do |model|
@@ -29,7 +45,7 @@ module Gitlab
           end
         end
 
-      def self.drop_detached_partitions
+        def drop_detached_partitions
           Gitlab::AppLogger.info(message: 'Dropping detached postgres partitions')
 
           Gitlab::Database::EachDatabase.each_database_connection do
@@ -38,6 +54,23 @@ module Gitlab
 
           Gitlab::AppLogger.info(message: 'Finished dropping detached postgres partitions')
         end
+
+        private
+
+        def registered_models
+          @registered_models ||= Set.new
+        end
+
+        def registered_tables
+          @registered_tables ||= Set.new
+        end
+
+        def registered_for_sync
+          registered_models + registered_tables.map do |table|
+            TableWithoutModel.new(**table)
+          end
+        end
+      end
     end
   end
 end

spec/lib/gitlab/database/partitioning_spec.rb

@@ -45,8 +45,13 @@ RSpec.describe Gitlab::Database::Partitioning do
       let(:model) { double('model') }
 
       it 'manages partitions for each registered model' do
+        registered_for_sync = described_class.__send__(:registered_for_sync)
+
+        allow(described_class).to receive(:registered_for_sync)
+          .and_return(registered_for_sync)
+
         expect(Gitlab::Database::EachDatabase).to receive(:each_model_connection)
-          .with(described_class.registered_models)
+          .with(registered_for_sync)
           .and_yield(model)
 
         expect(partition_manager_class).to receive(:new).with(model).and_return(partition_manager)
@@ -71,7 +76,7 @@ RSpec.describe Gitlab::Database::Partitioning do
         end
 
         expect(Gitlab::Database::EachDatabase).to receive(:each_model_connection)
-          .with(described_class.registered_models)
+          .with(described_class.__send__(:registered_models))
           .and_yield(model1)
           .and_yield(model2)
 
@@ -123,7 +128,7 @@ RSpec.describe Gitlab::Database::Partitioning do
 
   context 'ensure that the registered models have partitioning strategy' do
     it 'fails when partitioning_strategy is not specified for the model' do
-      expect(described_class.registered_models).to all(respond_to(:partitioning_strategy))
+      expect(described_class.__send__(:registered_models)).to all(respond_to(:partitioning_strategy))
     end
   end
 end