Merge branch 'audit-event-export-improvement-adam-hegyi' into 'master'

Improve audit events preloading logic

See merge request gitlab-org/gitlab!44958

app/models/audit_event.rb

@@ -4,6 +4,7 @@ class AuditEvent < ApplicationRecord
   include CreatedAtFilterable
   include IgnorableColumns
   include BulkInsertSafe
+  include EachBatch
 
   PARALLEL_PERSISTENCE_COLUMNS = [
     :author_name,
@@ -54,7 +55,9 @@ class AuditEvent < ApplicationRecord
   end
 
   def initialize_details
-    self.details = {} if details.nil?
+    return unless self.has_attribute?(:details)
+
+    self.details = {} if details&.nil?
   end
 
   def author_name

ee/app/services/audit_events/export_csv_service.rb

@@ -20,7 +20,7 @@ module AuditEvents
 
     def data
       events = AuditLogFinder.new(finder_params).execute
-      Gitlab::Audit::Events::Preloader.preload!(events)
+      Gitlab::Audit::Events::Preloader.new(events)
     end
 
     def finder_params

ee/changelogs/unreleased/266976-audit-event-export-improvement-for-csv-export.yml

+---
+title: Improve audit events preloading logic for CSV export
+merge_request: 44958
+author:
+type: performance

ee/lib/audit/details.rb

@@ -109,7 +109,7 @@ module Audit
     end
 
     def detail_value
-      @details.each_value.first
+      String(@details.each_value.first)
     end
 
     def oauth_label

ee/lib/gitlab/audit/events/preloader.rb

@@ -12,6 +12,25 @@ module Gitlab
             end
           end
         end
+
+        def initialize(audit_events)
+          @audit_events = audit_events
+        end
+
+        def find_each(&block)
+          @audit_events.each_batch(column: :created_at) do |relation|
+            relation.each do |audit_event|
+              audit_event.lazy_author
+              audit_event.lazy_entity
+            end
+
+            relation.each do |audit_event|
+              yield(audit_event)
+            end
+
+            BatchLoader::Executor.clear_current
+          end
+        end
       end
     end
   end

ee/spec/lib/gitlab/audit/events/preloader_spec.rb

@@ -3,11 +3,10 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Audit::Events::Preloader do
-  describe '.preload!' do
-    let_it_be(:audit_events) { create_list(:audit_event, 2) }
-
-    let(:audit_events_relation) { AuditEvent.where(id: audit_events.map(&:id)) }
+  let_it_be(:audit_events) { create_list(:audit_event, 2) }
+  let(:audit_events_relation) { AuditEvent.where(id: audit_events.map(&:id)) }
 
+  describe '.preload!' do
     subject { described_class.preload!(audit_events_relation) }
 
     it 'returns an ActiveRecord::Relation' do
@@ -15,12 +14,6 @@ RSpec.describe Gitlab::Audit::Events::Preloader do
     end
 
     it 'preloads associated records' do
-      log = ActiveRecord::QueryRecorder.new do
-        subject.map do |event|
-          [event.author_name, event.lazy_entity.name]
-        end
-      end
-
       # Expected queries when requesting for AuditEvent with associated records
       #
       # 1. On the audit_events table
@@ -30,7 +23,39 @@ RSpec.describe Gitlab::Audit::Events::Preloader do
       # 3. On the users table for entity name
       #    SELECT "users".* FROM "users" WHERE "users"."id" IN (2, 4)
       #
-      expect(log.count).to eq(3)
+      expect do
+        subject.map do |event|
+          [event.author_name, event.lazy_entity.name]
+        end
+      end.not_to exceed_query_limit(3)
+    end
+  end
+
+  describe '#find_each' do
+    let(:preloader) { described_class.new(audit_events_relation) }
+
+    it 'yields a list audit events' do
+      expect { |b| preloader.find_each(&b) }.to yield_successive_args(*audit_events)
+    end
+
+    it 'loads audit events in batches with preloaded associated records' do
+      # Expected queries when requesting for AuditEvent with associated records
+      #
+      # 1. Get the start of created_at value in 1k batch
+      #    SELECT "audit_events"."created_at" FROM "audit_events" WHERE "audit_events"."id" IN (1, 2) ORDER BY "audit_events"."created_at" ASC LIMIT 1
+      # 2. Get the end of created_at value in 1k batch
+      #    SELECT "audit_events"."created_at" FROM "audit_events" WHERE "audit_events"."id" IN (1, 2) AND "audit_events"."created_at" >= '2020-10-15 04:51:06.392709' ORDER BY "audit_events"."created_at" ASC LIMIT 1 OFFSET 1000
+      # 3. Get the audit_events in 1k batch
+      #    SELECT "audit_events".* FROM "audit_events" WHERE "audit_events"."id" IN (1, 2) AND "audit_events"."created_at" >= '2020-10-15 04:51:06.392709'
+      # 4. On the users table for author_name
+      #    SELECT "users"."id", "users"."name", "users"."username" FROM "users" WHERE "users"."id" IN (1, 3) ORDER BY "users"."id" ASC LIMIT 1000
+      # 5. On the users table for entity name
+      #    SELECT "users".* FROM "users" WHERE "users"."id" IN (2, 4) ORDER BY "users"."id" ASC LIMIT 1000
+      expect do
+        preloader.find_each do |event|
+          [event.author_name, event.lazy_entity.name]
+        end
+      end.not_to exceed_query_limit(5)
     end
   end
 end