Merge branch 'restrict-migrations' into 'master'

Add `restrict_gitlab_migration` to limit when a given migration is executed in context of decomposed databases

See merge request gitlab-org/gitlab!73756

lib/gitlab/database.rb

@@ -231,12 +231,26 @@ module Gitlab
       ::ActiveRecord::Base.configurations.configs_for(env_name: Rails.env).map(&:name)
     end
 
+    # This returns all matching schemas that a given connection can use
+    # Since the `ActiveRecord::Base` might change the connection (from main to ci)
+    # This does not look at literal connection names, but rather compares
+    # models that are holders for a given db_config_name
+    def self.gitlab_schemas_for_connection(connection)
+      connection_name = self.db_config_name(connection)
+      primary_model = self.database_base_models.fetch(connection_name)
+
+      self.schemas_to_base_models
+        .select { |_, models| models.include?(primary_model) }
+        .keys
+        .map!(&:to_sym)
+    end
+
     def self.db_config_for_connection(connection)
       return unless connection
 
-      # The LB connection proxy does not have a direct db_config
-      # that can be referenced
-      return if connection.is_a?(::Gitlab::Database::LoadBalancing::ConnectionProxy)
+      if connection.is_a?(::Gitlab::Database::LoadBalancing::ConnectionProxy)
+        return connection.load_balancer.configuration.primary_db_config
+      end
 
       # During application init we might receive `NullPool`
       return unless connection.respond_to?(:pool) &&

lib/gitlab/database/async_indexes/migration_helpers.rb

@@ -5,6 +5,8 @@ module Gitlab
     module AsyncIndexes
       module MigrationHelpers
         def unprepare_async_index(table_name, column_name, **options)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
           return unless async_index_creation_available?
 
           index_name = options[:name] || index_name(table_name, column_name)
@@ -15,6 +17,8 @@ module Gitlab
         end
 
         def unprepare_async_index_by_name(table_name, index_name, **options)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
           return unless async_index_creation_available?
 
           PostgresAsyncIndex.find_by(name: index_name).try do |async_index|
@@ -32,6 +36,8 @@ module Gitlab
         # If the requested index has already been created, it is not stored in the table for
         # asynchronous creation.
         def prepare_async_index(table_name, column_name, **options)
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas.require_ddl_mode!
+
           return unless async_index_creation_available?
 
           index_name = options[:name] || index_name(table_name, column_name)

lib/gitlab/database/gitlab_schema.rb

@@ -95,6 +95,10 @@ module Gitlab
       def self.tables_to_schema
         @tables_to_schema ||= YAML.load_file(Rails.root.join('lib/gitlab/database/gitlab_schemas.yml'))
       end
+
+      def self.schema_names
+        @schema_names ||= self.tables_to_schema.values.to_set
+      end
     end
   end
 end

lib/gitlab/database/load_balancing.rb

@@ -47,6 +47,8 @@ module Gitlab
       # Returns the role (primary/replica) of the database the connection is
       # connecting to.
       def self.db_role_for_connection(connection)
+        return ROLE_UNKNOWN if connection.is_a?(::Gitlab::Database::LoadBalancing::ConnectionProxy)
+
         db_config = Database.db_config_for_connection(connection)
         return ROLE_UNKNOWN unless db_config
 

lib/gitlab/database/load_balancing/configuration.rb

@@ -94,6 +94,10 @@ module Gitlab
           end
         end
 
+        def primary_db_config
+          primary_model_or_model_if_enabled.connection_db_config
+        end
+
         def replica_db_config
           @model.connection_db_config
         end

lib/gitlab/database/migration_helpers/restrict_gitlab_schema.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module MigrationHelpers
+      module RestrictGitlabSchema
+        extend ActiveSupport::Concern
+
+        MigrationSkippedError = Class.new(StandardError)
+
+        included do
+          class_attribute :allowed_gitlab_schemas
+        end
+
+        class_methods do
+          def restrict_gitlab_migration(gitlab_schema:)
+            unless Gitlab::Database::GitlabSchema.schema_names.include?(gitlab_schema)
+              raise "Unknown 'gitlab_schema: #{gitlab_schema}' specified. It needs to be one of: " \
+                "#{Gitlab::Database::GitlabSchema.schema_names.to_a}"
+            end
+
+            self.allowed_gitlab_schemas = [gitlab_schema]
+          end
+        end
+
+        def migrate(direction)
+          if unmatched_schemas.any?
+            # TODO: Today skipping migration would raise an exception.
+            # Ideally, skipped migration should be ignored (not loaded), or softly ignored.
+            # Read more in: https://gitlab.com/gitlab-org/gitlab/-/issues/355014
+            raise MigrationSkippedError, "Current migration is skipped since it modifies "\
+              "'#{self.class.allowed_gitlab_schemas}' which is outside of '#{allowed_schemas_for_connection}'"
+          end
+
+          Gitlab::Database::QueryAnalyzer.instance.within([validator_class]) do
+            validator_class.allowed_gitlab_schemas = self.allowed_gitlab_schemas
+
+            super
+          end
+        end
+
+        private
+
+        def validator_class
+          Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas
+        end
+
+        def unmatched_schemas
+          (self.allowed_gitlab_schemas || []) - allowed_schemas_for_connection
+        end
+
+        def allowed_schemas_for_connection
+          Gitlab::Database.gitlab_schemas_for_connection(connection)
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/query_analyzer.rb

@@ -30,13 +30,17 @@ module Gitlab
         end
       end
 
-      def within
+      def within(user_analyzers = nil)
         # Due to singleton nature of analyzers
         # only an outer invocation of the `.within`
         # is allowed to initialize them
-        return yield if already_within?
+        if already_within?
+          raise 'Query analyzers are already defined, cannot re-define them.' if user_analyzers
 
-        begin!
+          return yield
+        end
+
+        begin!(user_analyzers || all_analyzers)
 
         begin
           yield
@@ -61,21 +65,21 @@ module Gitlab
           next if analyzer.suppressed? && !analyzer.requires_tracking?(parsed)
 
           analyzer.analyze(parsed)
-        rescue StandardError, QueryAnalyzers::Base::QueryAnalyzerError => e
+        rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
           # We catch all standard errors to prevent validation errors to introduce fatal errors in production
           Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
         end
       end
 
       # Enable query analyzers
-      def begin!
-        analyzers = all_analyzers.select do |analyzer|
+      def begin!(analyzers = all_analyzers)
+        analyzers = analyzers.select do |analyzer|
           if analyzer.enabled?
             analyzer.begin!
 
             true
           end
-        rescue StandardError, QueryAnalyzers::Base::QueryAnalyzerError => e
+        rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
           Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
 
           false
@@ -88,7 +92,7 @@ module Gitlab
       def end!
         enabled_analyzers.select do |analyzer|
           analyzer.end!
-        rescue StandardError, QueryAnalyzers::Base::QueryAnalyzerError => e
+        rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
           Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
         end
 

lib/gitlab/database/query_analyzers/restrict_allowed_schemas.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module QueryAnalyzers
+      class RestrictAllowedSchemas < Base
+        UnsupportedSchemaError = Class.new(QueryAnalyzerError)
+        DDLNotAllowedError = Class.new(UnsupportedSchemaError)
+        DMLNotAllowedError = Class.new(UnsupportedSchemaError)
+        DMLAccessDeniedError = Class.new(UnsupportedSchemaError)
+
+        IGNORED_SCHEMAS = %i[gitlab_shared].freeze
+
+        class << self
+          def enabled?
+            true
+          end
+
+          def allowed_gitlab_schemas
+            self.context[:allowed_gitlab_schemas]
+          end
+
+          def allowed_gitlab_schemas=(value)
+            self.context[:allowed_gitlab_schemas] = value
+          end
+
+          def analyze(parsed)
+            # If list of schemas is empty, we allow only DDL changes
+            if self.dml_mode?
+              self.restrict_to_dml_only(parsed)
+            else
+              self.restrict_to_ddl_only(parsed)
+            end
+          end
+
+          def require_ddl_mode!(message = "")
+            return unless self.context
+
+            self.raise_dml_not_allowed_error(message) if self.dml_mode?
+          end
+
+          def require_dml_mode!(message = "")
+            return unless self.context
+
+            self.raise_ddl_not_allowed_error(message) if self.ddl_mode?
+          end
+
+          private
+
+          def restrict_to_ddl_only(parsed)
+            tables = self.dml_tables(parsed)
+            schemas = self.dml_schemas(tables)
+
+            if schemas.any?
+              self.raise_dml_not_allowed_error("Modifying of '#{tables}' (#{schemas.to_a}) with '#{parsed.sql}'")
+            end
+          end
+
+          def restrict_to_dml_only(parsed)
+            if parsed.pg.ddl_tables.any?
+              self.raise_ddl_not_allowed_error("Modifying of '#{parsed.pg.ddl_tables}' with '#{parsed.sql}'")
+            end
+
+            if parsed.pg.ddl_functions.any?
+              self.raise_ddl_not_allowed_error("Modifying of '#{parsed.pg.ddl_functions}' with '#{parsed.sql}'")
+            end
+
+            tables = self.dml_tables(parsed)
+            schemas = self.dml_schemas(tables)
+
+            if (schemas - self.allowed_gitlab_schemas).any?
+              raise DMLAccessDeniedError, "Select/DML queries (SELECT/UPDATE/DELETE) do access '#{tables}' (#{schemas.to_a}) " \
+                "which is outside of list of allowed schemas: '#{self.allowed_gitlab_schemas}'."
+            end
+          end
+
+          def dml_mode?
+            self.allowed_gitlab_schemas&.any?
+          end
+
+          def ddl_mode?
+            !self.dml_mode?
+          end
+
+          def dml_tables(parsed)
+            parsed.pg.select_tables + parsed.pg.dml_tables
+          end
+
+          def dml_schemas(tables)
+            extra_schemas = ::Gitlab::Database::GitlabSchema.table_schemas(tables)
+            extra_schemas.subtract(IGNORED_SCHEMAS)
+            extra_schemas
+          end
+
+          def raise_dml_not_allowed_error(message)
+            raise DMLNotAllowedError, "Select/DML queries (SELECT/UPDATE/DELETE) are disallowed in the DDL (structure) mode. #{message}"
+          end
+
+          def raise_ddl_not_allowed_error(message)
+            raise DDLNotAllowedError, "DDL queries (structure) are disallowed in the Select/DML (SELECT/UPDATE/DELETE) mode. #{message}"
+          end
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -92,6 +92,18 @@ RSpec.describe Gitlab::Database::LoadBalancing do
         end
       end
 
+      context 'when an invalid connection is used' do
+        it 'returns :unknown' do
+          expect(described_class.db_role_for_connection(:invalid)).to eq(:unknown)
+        end
+      end
+
+      context 'when a null connection is used' do
+        it 'returns :unknown' do
+          expect(described_class.db_role_for_connection(nil)).to eq(:unknown)
+        end
+      end
+
       context 'when a read connection is used' do
         it 'returns :replica' do
           load_balancer.read do |connection|

spec/lib/gitlab/database/migration_helpers/restrict_gitlab_schema_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::MigrationHelpers::RestrictGitlabSchema, query_analyzers: false, stub_feature_flags: false do
+  let(:schema_class) { Class.new(Gitlab::Database::Migration[1.0]).include(described_class) }
+
+  describe '#restrict_gitlab_migration' do
+    it 'invalid schema raises exception' do
+      expect { schema_class.restrict_gitlab_migration gitlab_schema: :gitlab_non_exisiting }
+        .to raise_error /Unknown 'gitlab_schema:/
+    end
+
+    it 'does configure allowed_gitlab_schema' do
+      schema_class.restrict_gitlab_migration gitlab_schema: :gitlab_main
+
+      expect(schema_class.allowed_gitlab_schemas).to eq(%i[gitlab_main])
+    end
+  end
+
+  context 'when executing migrations' do
+    using RSpec::Parameterized::TableSyntax
+
+    where do
+      {
+        "does create table in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            def change
+              create_table :_test_table do |t|
+                t.references :project, foreign_key: true, null: false
+                t.timestamps_with_timezone null: false
+              end
+            end
+          end,
+          query_matcher: /CREATE TABLE "_test_table"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does add column to projects in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            def change
+              add_column :projects, :__test_column, :integer
+            end
+          end,
+          query_matcher: /ALTER TABLE "projects" ADD "__test_column" integer/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does add column to ci_builds in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            def change
+              add_column :ci_builds, :__test_column, :integer
+            end
+          end,
+          query_matcher: /ALTER TABLE "ci_builds" ADD "__test_column" integer/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does add index to projects in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            def change
+              # Due to running in transactin we cannot use `add_concurrent_index`
+              add_index :projects, :hidden
+            end
+          end,
+          query_matcher: /CREATE INDEX/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does add index to ci_builds in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            def change
+              # Due to running in transactin we cannot use `add_concurrent_index`
+              add_index :ci_builds, :tag, where: "type = 'Ci::Build'", name: 'index_ci_builds_on_tag_and_type_eq_ci_build'
+            end
+          end,
+          query_matcher: /CREATE INDEX/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does create trigger in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            include Gitlab::Database::SchemaHelpers
+
+            def up
+              create_trigger_function('_test_trigger_function', replace: true) do
+                <<~SQL
+                  RETURN NULL;
+                SQL
+              end
+            end
+
+            def down
+              drop_function('_test_trigger_function')
+            end
+          end,
+          query_matcher: /CREATE OR REPLACE FUNCTION/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does create schema in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            include Gitlab::Database::SchemaHelpers
+
+            def up
+              execute("create schema __test_schema")
+            end
+
+            def down
+            end
+          end,
+          query_matcher: /create schema __test_schema/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_main: {
+              # This is not properly detected today since there are no helpers
+              # available to consider this as a DDL type of change
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does attach loose foreign key trigger in gitlab_main and gitlab_ci" => {
+          migration: ->(klass) do
+            include Gitlab::Database::MigrationHelpers::LooseForeignKeyHelpers
+
+            enable_lock_retries!
+
+            def up
+              track_record_deletions(:audit_events)
+            end
+
+            def down
+              untrack_record_deletions(:audit_events)
+            end
+          end,
+          query_matcher: /CREATE TRIGGER/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :ddl_not_allowed,
+              ci: :ddl_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does insert into software_licenses" => {
+          migration: ->(klass) do
+            def up
+              software_license_class.create!(name: 'aaa')
+            end
+
+            def down
+              software_license_class.where(name: 'aaa').delete_all
+            end
+
+            def software_license_class
+              Class.new(ActiveRecord::Base) do
+                self.table_name = 'software_licenses'
+              end
+            end
+          end,
+          query_matcher: /INSERT INTO "software_licenses"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does raise exception when accessing tables outside of gitlab_main" => {
+          migration: ->(klass) do
+            def up
+              ci_instance_variables_class.create!(variable_type: 1, key: 'aaa')
+            end
+
+            def down
+              ci_instance_variables_class.delete_all
+            end
+
+            def ci_instance_variables_class
+              Class.new(ActiveRecord::Base) do
+                self.table_name = 'ci_instance_variables'
+              end
+            end
+          end,
+          query_matcher: /INSERT INTO "ci_instance_variables"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :dml_access_denied,
+              ci: :skipped
+            }
+          }
+        },
+        "does allow modifying gitlab_shared" => {
+          migration: ->(klass) do
+            def up
+              detached_partitions_class.create!(drop_after: Time.current, table_name: '_test_table')
+            end
+
+            def down
+            end
+
+            def detached_partitions_class
+              Class.new(ActiveRecord::Base) do
+                self.table_name = 'detached_partitions'
+              end
+            end
+          end,
+          query_matcher: /INSERT INTO "detached_partitions"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_main: {
+              # TBD: This allow to selectively modify shared tables in context of a specific DB only
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does update data in batches of gitlab_main, but skips gitlab_ci" => {
+          migration: ->(klass) do
+            def up
+              update_column_in_batches(:projects, :archived, true) do |table, query|
+                query.where(table[:archived].eq(false)) # rubocop:disable CodeReuse/ActiveRecord
+              end
+            end
+
+            def down
+              # no-op
+            end
+          end,
+          query_matcher: /FROM "projects"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does not allow executing mixed DDL and DML migrations" => {
+          migration: ->(klass) do
+            def up
+              execute('UPDATE projects SET hidden=false')
+              add_index(:projects, :hidden, name: 'test_index')
+            end
+
+            def down
+              # no-op
+            end
+          end,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :ddl_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does schedule background migrations on gitlab_main" => {
+          migration: ->(klass) do
+            def up
+              queue_background_migration_jobs_by_range_at_intervals(
+                define_batchable_model('vulnerability_occurrences'),
+                'RemoveDuplicateVulnerabilitiesFindings',
+                2.minutes.to_i,
+                batch_size: 5_000
+              )
+            end
+
+            def down
+              # no-op
+            end
+          end,
+          query_matcher: /FROM "vulnerability_occurrences"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does support prepare_async_index" => {
+          migration: ->(klass) do
+            def up
+              prepare_async_index :projects, :hidden,
+                name: :index_projects_on_hidden
+            end
+
+            def down
+              unprepare_async_index_by_name :projects, :index_projects_on_hidden
+            end
+          end,
+          query_matcher: /INSERT INTO "postgres_async_indexes"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :success,
+              ci: :success
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_main: {
+              main: :dml_not_allowed,
+              ci: :skipped
+            }
+          }
+        },
+        "does raise exception when accessing current settings" => {
+          migration: ->(klass) do
+            def up
+              ApplicationSetting.last
+            end
+
+            def down
+            end
+          end,
+          query_matcher: /FROM "application_settings"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :success,
+              ci: :skipped
+            }
+          }
+        },
+        "does raise exception when accessing feature flags" => {
+          migration: ->(klass) do
+            def up
+              Feature.enabled?(:redis_hll_tracking, type: :ops, default_enabled: :yaml)
+            end
+
+            def down
+            end
+          end,
+          query_matcher: /FROM "features"/,
+          expected: {
+            no_gitlab_schema: {
+              main: :dml_not_allowed,
+              ci: :dml_not_allowed
+            },
+            gitlab_schema_gitlab_shared: {
+              main: :dml_access_denied,
+              ci: :dml_access_denied
+            },
+            gitlab_schema_gitlab_main: {
+              main: :success,
+              ci: :skipped
+            }
+          }
+        }
+      }
+    end
+
+    with_them do
+      let(:migration_class) { Class.new(schema_class, &migration) }
+
+      Gitlab::Database.database_base_models.each do |db_config_name, model|
+        context "for db_config_name=#{db_config_name}" do
+          around do |example|
+            with_reestablished_active_record_base do
+              reconfigure_db_connection(model: ActiveRecord::Base, config_model: model)
+
+              example.run
+            end
+          end
+
+          before do
+            allow_next_instance_of(migration_class) do |migration|
+              allow(migration).to receive(:transaction_open?).and_return(false)
+            end
+          end
+
+          %i[no_gitlab_schema gitlab_schema_gitlab_main gitlab_schema_gitlab_shared].each do |restrict_gitlab_migration|
+            context "while restrict_gitlab_migration=#{restrict_gitlab_migration}" do
+              it "does run migrate :up and :down" do
+                expected_result = expected.fetch(restrict_gitlab_migration)[db_config_name.to_sym]
+                skip "not configured" unless expected_result
+
+                case restrict_gitlab_migration
+                when :no_gitlab_schema
+                  # no-op
+                when :gitlab_schema_gitlab_main
+                  migration_class.restrict_gitlab_migration gitlab_schema: :gitlab_main
+                when :gitlab_schema_gitlab_shared
+                  migration_class.restrict_gitlab_migration gitlab_schema: :gitlab_shared
+                end
+
+                # In some cases (for :down) we ignore error and expect no other errors
+                case expected_result
+                when :success
+                  expect { migration_class.migrate(:up) }.to make_queries_matching(query_matcher)
+                  expect { migration_class.migrate(:down) }.not_to make_queries_matching(query_matcher)
+
+                when :dml_not_allowed
+                  expect { migration_class.migrate(:up) }.to raise_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DMLNotAllowedError)
+                  expect { ignore_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DMLNotAllowedError) { migration_class.migrate(:down) } }.not_to raise_error
+
+                when :dml_access_denied
+                  expect { migration_class.migrate(:up) }.to raise_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DMLAccessDeniedError)
+                  expect { ignore_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DMLAccessDeniedError) { migration_class.migrate(:down) } }.not_to raise_error
+
+                when :ddl_not_allowed
+                  expect { migration_class.migrate(:up) }.to raise_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DDLNotAllowedError)
+                  expect { ignore_error(Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas::DDLNotAllowedError) { migration_class.migrate(:down) } }.not_to raise_error
+
+                when :skipped
+                  expect { migration_class.migrate(:up) }.to raise_error(Gitlab::Database::MigrationHelpers::RestrictGitlabSchema::MigrationSkippedError)
+                  expect { migration_class.migrate(:down) }.to raise_error(Gitlab::Database::MigrationHelpers::RestrictGitlabSchema::MigrationSkippedError)
+                end
+              end
+            end
+          end
+
+          def ignore_error(error)
+            yield
+          rescue error
+          end
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/query_analyzer_spec.rb

@@ -4,6 +4,7 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
   let(:analyzer) { double(:query_analyzer) }
+  let(:user_analyzer) { double(:query_analyzer) }
   let(:disabled_analyzer) { double(:disabled_query_analyzer) }
 
   before do
@@ -53,6 +54,10 @@ RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
 
         expect { |b| described_class.instance.within(&b) }.to yield_control
       end
+
+      it 'raises exception when trying to re-define analyzers' do
+        expect { |b| described_class.instance.within([user_analyzer], &b) }.to raise_error /Query analyzers are already defined, cannot re-define them/
+      end
     end
 
     context 'when initializer is enabled' do
@@ -75,6 +80,18 @@ RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
         expect { |b| described_class.instance.within(&b) }.to yield_control
       end
     end
+
+    context 'when user analyzers are used' do
+      it 'calls begin! and end!' do
+        expect(analyzer).not_to receive(:begin!)
+        allow(user_analyzer).to receive(:enabled?).and_return(true)
+        allow(user_analyzer).to receive(:suppressed?).and_return(false)
+        expect(user_analyzer).to receive(:begin!)
+        expect(user_analyzer).to receive(:end!)
+
+        expect { |b| described_class.instance.within([user_analyzer], &b) }.to yield_control
+      end
+    end
   end
 
   describe '#process_sql' do

spec/lib/gitlab/database/query_analyzers/restrict_allowed_schemas_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::QueryAnalyzers::RestrictAllowedSchemas, query_analyzers: false do
+  let(:analyzer) { described_class }
+
+  context 'properly analyzes queries' do
+    using RSpec::Parameterized::TableSyntax
+
+    where do
+      examples = {
+        "for SELECT on projects" => {
+          sql: "SELECT 1 FROM projects",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :dml_not_allowed,
+            gitlab_main: :success,
+            gitlab_ci: :dml_access_denied # cross-schema access
+          }
+        },
+        "for INSERT" => {
+          sql: "INSERT INTO projects VALUES (1)",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :dml_not_allowed,
+            gitlab_main: :success,
+            gitlab_ci: :dml_access_denied # cross-schema access
+          }
+        },
+        "for CREATE INDEX" => {
+          sql: "CREATE INDEX index_projects_on_hidden ON projects (hidden)",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :success,
+            gitlab_main: :ddl_not_allowed,
+            gitlab_ci: :ddl_not_allowed
+          }
+        },
+        "for CREATE SCHEMA" => {
+          sql: "CREATE SCHEMA __test_schema",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :success,
+            # TODO: This is currently not properly detected
+            gitlab_main: :success,
+            gitlab_ci: :success
+          }
+        },
+        "for CREATE FUNCTION" => {
+          sql: "CREATE FUNCTION add(integer, integer) RETURNS integer AS 'select $1 + $2;' LANGUAGE SQL",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :success,
+            gitlab_main: :ddl_not_allowed,
+            gitlab_ci: :ddl_not_allowed
+          }
+        },
+        "for CREATE TRIGGER" => {
+          sql: "CREATE TRIGGER check_projects BEFORE UPDATE ON projects FOR EACH ROW EXECUTE PROCEDURE check_projects_update()",
+          expected_allowed_gitlab_schemas: {
+            no_schema: :success,
+            gitlab_main: :ddl_not_allowed,
+            gitlab_ci: :ddl_not_allowed
+          }
+        }
+      }
+
+      # Expands all examples into individual tests
+      examples.flat_map do |name, configuration|
+        configuration[:expected_allowed_gitlab_schemas].map do |allowed_gitlab_schema, expectation|
+          [
+            "#{name} for allowed_gitlab_schema=#{allowed_gitlab_schema}",
+            {
+              sql: configuration[:sql],
+              allowed_gitlab_schema: allowed_gitlab_schema, # nil, gitlab_main
+              expectation: expectation # success, dml_access_denied, ...
+            }
+          ]
+        end
+      end.to_h
+    end
+
+    with_them do
+      subject do
+        process_sql(sql) do
+          analyzer.allowed_gitlab_schemas = [allowed_gitlab_schema] unless allowed_gitlab_schema == :no_schema
+        end
+      end
+
+      it do
+        case expectation
+        when :success
+          expect { subject }.not_to raise_error
+        when :ddl_not_allowed
+          expect { subject }.to raise_error(described_class::DDLNotAllowedError)
+        when :dml_not_allowed
+          expect { subject }.to raise_error(described_class::DMLNotAllowedError)
+        when :dml_access_denied
+          expect { subject }.to raise_error(described_class::DMLAccessDeniedError)
+        else
+          raise "invalid expectation: #{expectation}"
+        end
+      end
+    end
+  end
+
+  describe '.require_ddl_mode!' do
+    subject { described_class.require_ddl_mode! }
+
+    it "when not configured does not raise exception" do
+      expect { subject }.not_to raise_error
+    end
+
+    it "when no schemas are configured does not raise exception (DDL mode)" do
+      with_analyzer do
+        expect { subject }.not_to raise_error
+      end
+    end
+
+    it "with schemas configured does raise exception (DML mode)" do
+      with_analyzer do
+        analyzer.allowed_gitlab_schemas = %i[gitlab_main]
+
+        expect { subject }.to raise_error(described_class::DMLNotAllowedError)
+      end
+    end
+  end
+
+  describe '.require_dml_mode!' do
+    subject { described_class.require_dml_mode! }
+
+    it "when not configured does not raise exception" do
+      expect { subject }.not_to raise_error
+    end
+
+    it "when no schemas are configured does raise exception (DDL mode)" do
+      with_analyzer do
+        expect { subject }.to raise_error(described_class::DDLNotAllowedError)
+      end
+    end
+
+    it "with schemas configured does raise exception (DML mode)" do
+      with_analyzer do
+        analyzer.allowed_gitlab_schemas = %i[gitlab_main]
+
+        expect { subject }.not_to raise_error
+      end
+    end
+  end
+
+  def with_analyzer
+    Gitlab::Database::QueryAnalyzer.instance.within([analyzer]) do
+      yield
+    end
+  end
+
+  def process_sql(sql, model = ActiveRecord::Base)
+    with_analyzer do
+      yield if block_given?
+
+      # Skip load balancer and retrieve connection assigned to model
+      Gitlab::Database::QueryAnalyzer.instance.process_sql(sql, model.retrieve_connection)
+    end
+  end
+end

spec/lib/gitlab/database_spec.rb

@@ -205,12 +205,12 @@ RSpec.describe Gitlab::Database do
     end
 
     context 'when the connection is LoadBalancing::ConnectionProxy' do
-      it 'returns nil' do
+      it 'returns primary_db_config' do
         lb_config = ::Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
         lb = ::Gitlab::Database::LoadBalancing::LoadBalancer.new(lb_config)
         proxy = ::Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
 
-        expect(described_class.db_config_for_connection(proxy)).to be_nil
+        expect(described_class.db_config_for_connection(proxy)).to eq(lb_config.primary_db_config)
       end
     end
 
@@ -229,7 +229,7 @@ RSpec.describe Gitlab::Database do
 
       # This is a ConnectionProxy
       expect(described_class.db_config_name(model.connection))
-        .to eq('unknown')
+        .to eq('main')
 
       # This is an actual connection
       expect(described_class.db_config_name(model.retrieve_connection))
@@ -245,6 +245,31 @@ RSpec.describe Gitlab::Database do
     end
   end
 
+  describe '.gitlab_schemas_for_connection' do
+    it 'does raise exception for invalid connection' do
+      expect { described_class.gitlab_schemas_for_connection(:invalid) }.to raise_error /key not found: "unknown"/
+    end
+
+    it 'does return a valid schema depending on a base model used', :request_store do
+      # This is currently required as otherwise the `Ci::Build.connection` == `Project.connection`
+      # ENV due to lib/gitlab/database/load_balancing/setup.rb:93
+      stub_env('GITLAB_USE_MODEL_LOAD_BALANCING', '1')
+      # FF due to lib/gitlab/database/load_balancing/configuration.rb:92
+      stub_feature_flags(force_no_sharing_primary_model: true)
+
+      expect(described_class.gitlab_schemas_for_connection(Project.connection)).to include(:gitlab_main, :gitlab_shared)
+      expect(described_class.gitlab_schemas_for_connection(Ci::Build.connection)).to include(:gitlab_ci, :gitlab_shared)
+    end
+
+    it 'does return gitlab_ci when a ActiveRecord::Base is using CI connection' do
+      with_reestablished_active_record_base do
+        reconfigure_db_connection(model: ActiveRecord::Base, config_model: Ci::Build)
+
+        expect(described_class.gitlab_schemas_for_connection(ActiveRecord::Base.connection)).to include(:gitlab_ci, :gitlab_shared)
+      end
+    end
+  end
+
   describe '#true_value' do
     it 'returns correct value' do
       expect(described_class.true_value).to eq "'t'"