Introduce `ActiveRecord::Base.gitlab_schema` indicator

app/models/application_record.rb

 # frozen_string_literal: true
 
 class ApplicationRecord < ActiveRecord::Base
+  self.gitlab_schema = :gitlab_main
   self.abstract_class = true
 
   alias_method :reset, :reload

app/models/ci/application_record.rb

@@ -2,6 +2,7 @@
 
 module Ci
   class ApplicationRecord < ::ApplicationRecord
+    self.gitlab_schema = :gitlab_ci
     self.abstract_class = true
 
     def self.table_name_prefix

config/initializers/00_active_record_gitlab_schema.rb

+# frozen_string_literal: true
+
+# This parameter describes a virtual context to indicate
+# table affinity to other tables.
+#
+# Table affinity limits cross-joins, cross-modifications,
+# foreign keys and validates relationship between tables
+#
+# By default it is undefined
+ActiveRecord::Base.class_attribute :gitlab_schema, default: nil

config/initializers/0_acts_as_taggable.rb

@@ -13,3 +13,7 @@ raise "Counter cache is not disabled" if
 ActsAsTaggableOn::Tagging.include IgnorableColumns
 ActsAsTaggableOn::Tagging.ignore_column :id_convert_to_bigint, remove_with: '14.2', remove_after: '2021-08-22'
 ActsAsTaggableOn::Tagging.ignore_column :taggable_id_convert_to_bigint, remove_with: '14.2', remove_after: '2021-08-22'
+
+# The tags and taggings are supposed to be part of `gitlab_ci`
+ActsAsTaggableOn::Tag.gitlab_schema = :gitlab_ci
+ActsAsTaggableOn::Tagging.gitlab_schema = :gitlab_ci

spec/support/database/ci_tables.rb

-# frozen_string_literal: true
-
-# This module stores the CI-related database tables which are
-# going to be moved to a separate database.
-module Database
-  module CiTables
-    def self.include?(name)
-      ci_tables.include?(name)
-    end
-
-    def self.ci_tables
-      @@ci_tables ||= Set.new.tap do |tables| # rubocop:disable Style/ClassVars
-        tables.merge(Ci::ApplicationRecord.descendants.map(&:table_name).compact)
-
-        # It was decided that taggings/tags are best placed with CI
-        # https://gitlab.com/gitlab-org/gitlab/-/issues/333413
-        tables.add('taggings')
-        tables.add('tags')
-      end
-    end
-  end
-end

spec/support/database/gitlab_schema.rb

+# frozen_string_literal: true
+
+# This module gathes information about table to schema mapping
+# to understand table affinity
+module Database
+  module GitlabSchema
+    def self.table_schemas(tables)
+      tables.map { |table| table_schema(table) }.to_set
+    end
+
+    def self.table_schema(name)
+      tables_to_schema[name] || :undefined
+    end
+
+    def self.tables_to_schema
+      @tables_to_schema ||= all_classes_with_schema.to_h do |klass|
+        [klass.table_name, klass.gitlab_schema]
+      end
+    end
+
+    def self.all_classes_with_schema
+      ActiveRecord::Base.descendants.reject(&:abstract_class?).select(&:gitlab_schema?) # rubocop:disable Database/MultipleDatabases
+    end
+  end
+end

spec/support/database/prevent_cross_database_modification.rb

@@ -75,17 +75,17 @@ module Database
       return if cross_database_context[:transaction_depth_by_db].values.all?(&:zero?)
 
       tables = PgQuery.parse(sql).dml_tables
-
       return if tables.empty?
 
       cross_database_context[:modified_tables_by_db][database].merge(tables)
 
       all_tables = cross_database_context[:modified_tables_by_db].values.map(&:to_a).flatten
+      schemas = Database::GitlabSchema.table_schemas(all_tables)
 
-      unless PreventCrossJoins.only_ci_or_only_main?(all_tables)
+      if schemas.many?
         raise Database::PreventCrossDatabaseModification::CrossDatabaseModificationAcrossUnsupportedTablesError,
-          "Cross-database data modification queries (CI and Main) were detected within " \
-          "a transaction '#{all_tables.join(", ")}' discovered"
+          "Cross-database data modification of '#{schemas.to_a.join(", ")}' were detected within " \
+          "a transaction modifying the '#{all_tables.to_a.join(", ")}'"
       end
     end
   end

spec/support/database/prevent_cross_joins.rb

@@ -27,20 +27,15 @@ module Database
       # PgQuery might fail in some cases due to limited nesting:
       # https://github.com/pganalyze/pg_query/issues/209
       tables = PgQuery.parse(sql).tables
+      schemas = Database::GitlabSchema.table_schemas(tables)
 
-      unless only_ci_or_only_main?(tables)
+      if schemas.many?
         raise CrossJoinAcrossUnsupportedTablesError,
-          "Unsupported cross-join across '#{tables.join(", ")}' discovered " \
+          "Unsupported cross-join across '#{tables.join(", ")}' modifying '#{schemas.to_a.join(", ")}' discovered " \
           "when executing query '#{sql}'"
       end
     end
 
-    # Returns true if a set includes only CI tables, or includes only non-CI tables
-    def self.only_ci_or_only_main?(tables)
-      tables.all? { |table| CiTables.include?(table) } ||
-        tables.none? { |table| CiTables.include?(table) }
-    end
-
     module SpecHelpers
       def with_cross_joins_prevented
         subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |event|

spec/support_specs/database/prevent_cross_database_modification_spec.rb

@@ -66,7 +66,7 @@ RSpec.describe 'Database::PreventCrossDatabaseModification' do
             pipeline.touch
           end
         end
-      end.to raise_error /Cross-database data modification queries/
+      end.to raise_error /Cross-database data modification/
     end
   end
 
@@ -84,7 +84,7 @@ RSpec.describe 'Database::PreventCrossDatabaseModification' do
       context 'when data modification happens in a transaction' do
         it 'raises error' do
           Project.transaction do
-            expect { run_queries }.to raise_error /Cross-database data modification queries/
+            expect { run_queries }.to raise_error /Cross-database data modification/
           end
         end
 
@@ -93,7 +93,7 @@ RSpec.describe 'Database::PreventCrossDatabaseModification' do
             Project.transaction(requires_new: true) do
               project.touch
               Project.transaction(requires_new: true) do
-                expect { pipeline.touch }.to raise_error /Cross-database data modification queries/
+                expect { pipeline.touch }.to raise_error /Cross-database data modification/
               end
             end
           end
@@ -127,7 +127,7 @@ RSpec.describe 'Database::PreventCrossDatabaseModification' do
           ApplicationRecord.transaction do
             create(:ci_pipeline)
           end
-        end.to raise_error /Cross-database data modification queries/
+        end.to raise_error /Cross-database data modification/
       end
 
       it 'skips raising error on factory creation' do