Add latest changes from gitlab-org/gitlab@master

3 files changed, 94 insertions, 15 deletions

diff --git a/lib/gitlab/database/gitlab_schema.rb b/lib/gitlab/database/gitlab_schema.rb
index 0bd357b7730..31ceb898eee 100644
--- a/lib/gitlab/database/gitlab_schema.rb
+++ b/lib/gitlab/database/gitlab_schema.rb
@@ -87,24 +87,37 @@ module Gitlab
         # rubocop:enable Gitlab/DocUrl
       end
 
-      private_class_method def self.cross_access_allowed?(type, table_schemas)
+      def self.cross_joins_allowed?(table_schemas, all_tables)
+        return true unless table_schemas.many?
+
         table_schemas.any? do |schema|
-          extra_schemas = table_schemas - [schema]
-          extra_schemas -= Gitlab::Database.all_gitlab_schemas[schema]&.public_send(type) || [] # rubocop:disable GitlabSecurity/PublicSend
-          extra_schemas.empty?
+          schema_info = Gitlab::Database.all_gitlab_schemas[schema]
+          next false unless schema_info
+
+          schema_info.allow_cross_joins?(table_schemas, all_tables)
         end
       end
 
-      def self.cross_joins_allowed?(table_schemas)
-        table_schemas.empty? || self.cross_access_allowed?(:allow_cross_joins, table_schemas)
-      end
+      def self.cross_transactions_allowed?(table_schemas, all_tables)
+        return true unless table_schemas.many?
+
+        table_schemas.any? do |schema|
+          schema_info = Gitlab::Database.all_gitlab_schemas[schema]
+          next false unless schema_info
 
-      def self.cross_transactions_allowed?(table_schemas)
-        table_schemas.empty? || self.cross_access_allowed?(:allow_cross_transactions, table_schemas)
+          schema_info.allow_cross_transactions?(table_schemas, all_tables)
+        end
       end
 
-      def self.cross_foreign_key_allowed?(table_schemas)
-        self.cross_access_allowed?(:allow_cross_foreign_keys, table_schemas)
+      def self.cross_foreign_key_allowed?(table_schemas, all_tables)
+        return true if table_schemas.one?
+
+        table_schemas.any? do |schema|
+          schema_info = Gitlab::Database.all_gitlab_schemas[schema]
+          next false unless schema_info
+
+          schema_info.allow_cross_foreign_keys?(table_schemas, all_tables)
+        end
       end
 
       def self.dictionary_paths
diff --git a/lib/gitlab/database/gitlab_schema_info.rb b/lib/gitlab/database/gitlab_schema_info.rb
index 34b89cb9006..20d2b31a65c 100644
--- a/lib/gitlab/database/gitlab_schema_info.rb
+++ b/lib/gitlab/database/gitlab_schema_info.rb
@@ -2,6 +2,11 @@
 
 module Gitlab
   module Database
+    GitlabSchemaInfoAllowCross = Struct.new(
+      :specific_tables,
+      keyword_init: true
+    )
+
     GitlabSchemaInfo = Struct.new(
       :name,
       :description,
@@ -14,15 +19,76 @@ module Gitlab
       def initialize(*)
         super
         self.name = name.to_sym
-        self.allow_cross_joins = allow_cross_joins&.map(&:to_sym)&.freeze
-        self.allow_cross_transactions = allow_cross_transactions&.map(&:to_sym)&.freeze
-        self.allow_cross_foreign_keys = allow_cross_foreign_keys&.map(&:to_sym)&.freeze
+        self.allow_cross_joins = convert_array_to_hash(allow_cross_joins)
+        self.allow_cross_transactions = convert_array_to_hash(allow_cross_transactions)
+        self.allow_cross_foreign_keys = convert_array_to_hash(allow_cross_foreign_keys)
       end
 
       def self.load_file(yaml_file)
         content = YAML.load_file(yaml_file)
         new(**content.deep_symbolize_keys.merge(file_path: yaml_file))
       end
+
+      def allow_cross_joins?(table_schemas, all_tables)
+        allowed_schemas = allow_cross_joins || {}
+
+        allowed_for?(allowed_schemas, table_schemas, all_tables)
+      end
+
+      def allow_cross_transactions?(table_schemas, all_tables)
+        allowed_schemas = allow_cross_transactions || {}
+
+        allowed_for?(allowed_schemas, table_schemas, all_tables)
+      end
+
+      def allow_cross_foreign_keys?(table_schemas, all_tables)
+        allowed_schemas = allow_cross_foreign_keys || {}
+
+        allowed_for?(allowed_schemas, table_schemas, all_tables)
+      end
+
+      private
+
+      def allowed_for?(allowed_schemas, table_schemas, all_tables)
+        denied_schemas = table_schemas - [name]
+        denied_schemas -= allowed_schemas.keys
+        return false unless denied_schemas.empty?
+
+        all_tables.all? do |table|
+          table_schema = ::Gitlab::Database::GitlabSchema.table_schema!(table)
+          allowed_tables = allowed_schemas[table_schema]
+
+          allowed_tables.nil? || allowed_tables.specific_tables.include?(table)
+        end
+      end
+
+      # Convert from:
+      # - schema_a
+      # - schema_b:
+      #     specific_tables:
+      #     - table_b_of_schema_b
+      #     - table_c_of_schema_b
+      #
+      # To:
+      # { :schema_a => nil,
+      #   :schema_b => { specific_tables : [:table_b_of_schema_b, :table_c_of_schema_b] }
+      # }
+      #
+      def convert_array_to_hash(subject)
+        result = {}
+
+        subject&.each do |item|
+          if item.is_a?(Hash)
+            item.each do |key, value|
+              result[key.to_sym] = GitlabSchemaInfoAllowCross.new(value || {})
+            end
+          else
+            result[item.to_sym] = nil
+          end
+        end
+
+        result.freeze
+      end
     end
   end
 end
diff --git a/lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb b/lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb
index a9f2b963340..fb25cb70e57 100644
--- a/lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb
+++ b/lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb
@@ -113,7 +113,7 @@ module Gitlab
           schemas = ::Gitlab::Database::GitlabSchema.table_schemas!(all_tables)
           schemas += ApplicationRecord.gitlab_transactions_stack
 
-          unless ::Gitlab::Database::GitlabSchema.cross_transactions_allowed?(schemas)
+          unless ::Gitlab::Database::GitlabSchema.cross_transactions_allowed?(schemas, all_tables)
             messages = []
             messages << "Cross-database data modification of '#{schemas.to_a.join(", ")}' were detected within " \
                         "a transaction modifying the '#{all_tables.to_a.join(", ")}' tables. "