Merge branch 'cross-database-modification-txn-in-wrong-db' into 'master'

Detect modification in wrong DB transaction

See merge request gitlab-org/gitlab!76717

app/models/application_record.rb

@@ -4,6 +4,7 @@ class ApplicationRecord < ActiveRecord::Base
   include DatabaseReflection
   include Transactions
   include LegacyBulkInsert
+  include CrossDatabaseModification
 
   self.abstract_class = true
 

app/models/concerns/cross_database_modification.rb

+# frozen_string_literal: true
+
+module CrossDatabaseModification
+  extend ActiveSupport::Concern
+
+  class TransactionStackTrackRecord
+    def initialize(subject, gitlab_schema)
+      @subject = subject
+      @gitlab_schema = gitlab_schema
+      @subject.gitlab_transactions_stack.push(gitlab_schema)
+    end
+
+    def done!
+      unless @done
+        @done = true
+        @subject.gitlab_transactions_stack.pop
+      end
+
+      true
+    end
+
+    def trigger_transactional_callbacks?
+      false
+    end
+
+    def before_committed!
+    end
+
+    def rolledback!(force_restore_state: false, should_run_callbacks: true)
+      done!
+    end
+
+    def committed!(should_run_callbacks: true)
+      done!
+    end
+  end
+
+  included do
+    private_class_method :gitlab_schema
+  end
+
+  class_methods do
+    def gitlab_transactions_stack
+      Thread.current[:gitlab_transactions_stack] ||= []
+    end
+
+    def transaction(**options, &block)
+      if track_gitlab_schema_in_current_transaction?
+        super(**options) do
+          # Hook into current transaction to ensure that once
+          # the `COMMIT` is executed the `gitlab_transactions_stack`
+          # will be allowing to execute `after_commit_queue`
+          record = TransactionStackTrackRecord.new(self, gitlab_schema)
+
+          begin
+            connection.current_transaction.add_record(record)
+
+            yield
+          ensure
+            record.done!
+          end
+        end
+      else
+        super(**options, &block)
+      end
+    end
+
+    def track_gitlab_schema_in_current_transaction?
+      return false unless Feature::FlipperFeature.table_exists?
+
+      Feature.enabled?(:track_gitlab_schema_in_current_transaction, default_enabled: :yaml)
+    rescue ActiveRecord::NoDatabaseError, PG::ConnectionBad
+      false
+    end
+
+    def gitlab_schema
+      case self.name
+      when 'ActiveRecord::Base', 'ApplicationRecord'
+        :gitlab_main
+      when 'Ci::ApplicationRecord'
+        :gitlab_ci
+      else
+        Gitlab::Database::GitlabSchema.table_schema(table_name) if table_name
+      end
+    end
+  end
+end

config/feature_flags/development/track_gitlab_schema_in_current_transaction.yml

+---
+name: track_gitlab_schema_in_current_transaction
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/76717
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/349944
+milestone: '14.8'
+type: development
+group: group::sharding
+default_enabled: false

ee/app/services/security/ingestion/ingest_report_slice_service.rb

@@ -32,9 +32,11 @@ module Security
       end
 
       def execute
+        ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.allow_cross_database_modification_within_transaction(url: 'https://gitlab.com/gitlab-org/gitlab/-/issues/350092') do
           ApplicationRecord.transaction do
             TASKS.each { |task| execute_task(task) }
           end
+        end
 
         @finding_maps.map(&:vulnerability_id)
       end

ee/spec/models/ci/pipeline_spec.rb

@@ -361,11 +361,9 @@ RSpec.describe Ci::Pipeline do
     end
 
     context 'when pipeline project has downstream subscriptions' do
-      let(:pipeline) { create(:ci_empty_pipeline, project: create(:project, :public)) }
-
-      before do
-        pipeline.project.downstream_projects << create(:project)
-      end
+      let(:downstream_project) { create(:project) }
+      let(:project) { create(:project, :public, downstream_projects: [downstream_project]) }
+      let(:pipeline) { create(:ci_empty_pipeline, project: project) }
 
       context 'when pipeline runs on a tag' do
         before do

ee/spec/models/project_spec.rb

@@ -3231,9 +3231,8 @@ RSpec.describe Project do
 
   describe '#upstream_projects' do
     it 'returns the upstream projects' do
-      primary_project = create(:project, :public)
       upstream_project = create(:project, :public)
-      primary_project.upstream_projects << upstream_project
+      primary_project = create(:project, :public, upstream_projects: [upstream_project])
 
       with_cross_joins_prevented do
         expect(primary_project.upstream_projects).to eq([upstream_project])
@@ -3243,9 +3242,8 @@ RSpec.describe Project do
 
   describe '#upstream_projects_count' do
     it 'returns the upstream projects count' do
-      primary_project = create(:project, :public)
       upstream_projects = create_list(:project, 2, :public)
-      primary_project.upstream_projects = upstream_projects
+      primary_project = create(:project, :public, upstream_projects: upstream_projects)
 
       with_cross_joins_prevented do
         expect(primary_project.upstream_projects_count).to eq(2)
@@ -3255,9 +3253,8 @@ RSpec.describe Project do
 
   describe '#downstream_projects' do
     it 'returns the downstream projects' do
-      primary_project = create(:project, :public)
       downstream_project = create(:project, :public)
-      primary_project.downstream_projects << downstream_project
+      primary_project = create(:project, :public, downstream_projects: [downstream_project])
 
       with_cross_joins_prevented do
         expect(primary_project.downstream_projects).to eq([downstream_project])
@@ -3267,9 +3264,8 @@ RSpec.describe Project do
 
   describe '#downstream_projects_count' do
     it 'returns the downstream projects count' do
-      primary_project = create(:project, :public)
       downstream_projects = create_list(:project, 2, :public)
-      primary_project.downstream_projects = downstream_projects
+      primary_project = create(:project, :public, downstream_projects: downstream_projects)
 
       with_cross_joins_prevented do
         expect(primary_project.downstream_projects_count).to eq(2)

lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb

@@ -87,6 +87,8 @@ module Gitlab
           all_tables = context[:modified_tables_by_db].values.map(&:to_a).flatten
           schemas = ::Gitlab::Database::GitlabSchema.table_schemas(all_tables)
 
+          schemas += ApplicationRecord.gitlab_transactions_stack
+
           if schemas.many?
             message = "Cross-database data modification of '#{schemas.to_a.join(", ")}' were detected within " \
               "a transaction modifying the '#{all_tables.to_a.join(", ")}' tables." \

lib/gitlab/import_export/project/object_builder.rb

@@ -14,10 +14,12 @@ module Gitlab
       # It also adds some logic around Group Labels/Milestones for edge cases.
       class ObjectBuilder < Base::ObjectBuilder
         def self.build(*args)
+          ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.allow_cross_database_modification_within_transaction(url: 'https://gitlab.com/gitlab-org/gitlab/-/issues/350091') do
             ::Project.transaction do
               super
             end
           end
+        end
 
         def initialize(klass, attributes)
           super

spec/graphql/resolvers/package_pipelines_resolver_spec.rb

@@ -15,8 +15,9 @@ RSpec.describe Resolvers::PackagePipelinesResolver do
     subject { resolve(described_class, obj: package, args: args, ctx: { current_user: user }) }
 
     before do
-      package.pipelines = pipelines
-      package.save!
+      pipelines.each do |pipeline|
+        create(:package_build_info, package: package, pipeline: pipeline)
+      end
     end
 
     it { is_expected.to contain_exactly(*pipelines) }

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

@@ -14,23 +14,25 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
     Gitlab::Database::QueryAnalyzer.instance.within { example.run }
   end
 
-  shared_examples 'successful examples' do
+  shared_examples 'successful examples' do |model:|
+    let(:model) { model }
+
     context 'outside transaction' do
       it { expect { run_queries }.not_to raise_error }
     end
 
-    context 'within transaction' do
+    context "within #{model} transaction" do
       it do
-        Project.transaction do
+        model.transaction do
           expect { run_queries }.not_to raise_error
         end
       end
     end
 
-    context 'within nested transaction' do
+    context "within nested #{model} transaction" do
       it do
-        Project.transaction(requires_new: true) do
-          Project.transaction(requires_new: true) do
+        model.transaction(requires_new: true) do
+          model.transaction(requires_new: true) do
             expect { run_queries }.not_to raise_error
           end
         end
@@ -38,13 +40,26 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
     end
   end
 
+  shared_examples 'cross-database modification errors' do |model:|
+    let(:model) { model }
+
+    context "within #{model} transaction" do
+      it 'raises error' do
+        model.transaction do
+          expect { run_queries }.to raise_error /Cross-database data modification/
+        end
+      end
+    end
+  end
+
   context 'when CI and other tables are read in a transaction' do
     def run_queries
       pipeline.reload
       project.reload
     end
 
-    include_examples 'successful examples'
+    include_examples 'successful examples', model: Project
+    include_examples 'successful examples', model: Ci::Pipeline
   end
 
   context 'when only CI data is modified' do
@@ -53,7 +68,9 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
       project.reload
     end
 
-    include_examples 'successful examples'
+    include_examples 'successful examples', model: Ci::Pipeline
+
+    include_examples 'cross-database modification errors', model: Project
   end
 
   context 'when other data is modified' do
@@ -62,7 +79,9 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
       project.touch
     end
 
-    include_examples 'successful examples'
+    include_examples 'successful examples', model: Project
+
+    include_examples 'cross-database modification errors', model: Ci::Pipeline
   end
 
   context 'when both CI and other data is modified' do
@@ -144,7 +163,9 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
         project.save!
       end
 
-      include_examples 'successful examples'
+      include_examples 'successful examples', model: Ci::Pipeline
+
+      include_examples 'cross-database modification errors', model: Project
     end
 
     describe '.allow_cross_database_modification_within_transaction' do

spec/models/concerns/cross_database_modification_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe CrossDatabaseModification do
+  describe '.transaction' do
+    context 'feature flag disabled' do
+      before do
+        stub_feature_flags(track_gitlab_schema_in_current_transaction: false)
+      end
+
+      it 'does not add to gitlab_transactions_stack' do
+        ApplicationRecord.transaction do
+          expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+
+          Project.first
+        end
+
+        expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+      end
+    end
+
+    context 'feature flag is not yet setup' do
+      before do
+        allow(Feature::FlipperFeature).to receive(:table_exists?).and_raise(ActiveRecord::NoDatabaseError)
+      end
+
+      it 'does not add to gitlab_transactions_stack' do
+        ApplicationRecord.transaction do
+          expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+
+          Project.first
+        end
+
+        expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+      end
+    end
+
+    it 'adds the current gitlab schema to gitlab_transactions_stack', :aggregate_failures do
+      ApplicationRecord.transaction do
+        expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_main)
+
+        Project.first
+      end
+
+      expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+
+      Ci::ApplicationRecord.transaction do
+        expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_ci)
+
+        Project.first
+      end
+
+      expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+
+      Project.transaction do
+        expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_main)
+
+        Project.first
+      end
+
+      expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+
+      Ci::Pipeline.transaction do
+        expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_ci)
+
+        Project.first
+      end
+
+      expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+
+      ApplicationRecord.transaction do
+        expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_main)
+
+        Ci::Pipeline.transaction do
+          expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_main, :gitlab_ci)
+
+          Project.first
+        end
+      end
+
+      expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+    end
+
+    it 'yields' do
+      expect { |block| ApplicationRecord.transaction(&block) }.to yield_control
+    end
+  end
+end

spec/requests/api/graphql/packages/package_spec.rb

@@ -145,8 +145,9 @@ RSpec.describe 'package details' do
       let(:pipeline_gids) { pipelines.sort_by(&:id).map(&:to_gid).map(&:to_s).reverse }
 
       before do
-        composer_package.pipelines = pipelines
-        composer_package.save!
+        pipelines.each do |pipeline|
+          create(:package_build_info, package: composer_package, pipeline: pipeline)
+        end
       end
 
       def run_query(args)