Refine restriction specs to cover all cases

This makes the `restrict_gitlab_schema_spec` to validate
all possible restrict behaviors to ensure consistent
and as designed validation of migrations where we
can detect DDL or DML changes.

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/query_analyzers/restrict_allowed_schemas.rb

@@ -11,60 +11,94 @@ module Gitlab
 
         IGNORED_SCHEMAS = %i[gitlab_shared].freeze
 
-        def self.enabled?
-          true
-        end
+        class << self
+          def enabled?
+            true
+          end
 
-        def self.allowed_gitlab_schemas
-          self.context[:allowed_gitlab_schemas]
-        end
+          def allowed_gitlab_schemas
+            self.context[:allowed_gitlab_schemas]
+          end
 
-        def self.allowed_gitlab_schemas=(value)
-          self.context[:allowed_gitlab_schemas] = value
-        end
+          def allowed_gitlab_schemas=(value)
+            self.context[:allowed_gitlab_schemas] = value
+          end
 
-        def self.analyze(parsed)
-          # If list of schemas is empty, we allow only DDL changes
-          if self.allowed_gitlab_schemas
-            self.restrict_to_dml_only(parsed)
-          else
-            self.restrict_to_ddl_only(parsed)
+          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
-        end
 
-        def self.restrict_to_ddl_only(parsed)
-          tables = self.dml_tables(parsed)
-          schemas = self.dml_schemas(tables)
+          def require_ddl_mode!(message = "")
+            return unless self.context
 
-          if schemas.any?
-            raise DMLNotAllowedError, "Select/DML queries (SELECT/UPDATE/DELETE) are disallowed in the DDL (structure) mode. " \
-              "Modifying of '#{tables}' (#{schemas.to_a}) with '#{parsed.sql}'"
+            self.raise_dml_not_allowed_error(message) if self.dml_mode?
           end
-        end
 
-        def self.restrict_to_dml_only(parsed)
-          if parsed.pg.ddl_tables.any?
-            raise DDLNotAllowedError, "DDL queries (structure) are disallowed in the Select/DML (SELECT/UPDATE/DELETE) mode. " \
-              "Modifying of '#{parsed.pg.ddl_tables}' with '#{parsed.sql}'"
+          def require_dml_mode!(message = "")
+            return unless self.context
+
+            self.raise_ddl_not_allowed_error(message) if self.ddl_mode?
           end
 
-          tables = self.dml_tables(parsed)
-          schemas = self.dml_schemas(tables)
+          private
 
-          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}'."
+          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
-        end
 
-        def self.dml_tables(parsed)
-          parsed.pg.select_tables + parsed.pg.dml_tables
-        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)
 
-        def self.dml_schemas(tables)
-          extra_schemas = ::Gitlab::Database::GitlabSchema.table_schemas(tables)
-          extra_schemas.subtract(IGNORED_SCHEMAS)
-          extra_schemas
+            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

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