diff --git a/keeps/helpers/postgres_ai.rb b/keeps/helpers/postgres_ai.rb
index c4cee13325c245372003b8d7aeb1f8b791dc397f..e4833e1403822760fdf89c41c33a4421fea07f54 100644
--- a/keeps/helpers/postgres_ai.rb
+++ b/keeps/helpers/postgres_ai.rb
@@ -11,7 +11,7 @@ def initialize
 
       def fetch_background_migration_status(job_class_name)
         query = <<~SQL
-        SELECT id, created_at, updated_at, finished_at, started_at, status, job_class_name
+        SELECT id, created_at, updated_at, finished_at, started_at, status, job_class_name, gitlab_schema
         FROM batched_background_migrations
         WHERE job_class_name = $1::text
         SQL
diff --git a/keeps/overdue_finalize_background_migration.rb b/keeps/overdue_finalize_background_migration.rb
index 9c1da8a8cb25a72a4c73684c21db1782f77d6509..f56789b6da87e7a3ffddb767b7faf2f209e4bc01 100644
--- a/keeps/overdue_finalize_background_migration.rb
+++ b/keeps/overdue_finalize_background_migration.rb
@@ -81,7 +81,7 @@ def each_change
         migration_file = generator.invoke_all.first
         change.changed_files = [migration_file]
 
-        add_ensure_call_to_migration(migration_file, queue_method_node, job_name)
+        add_ensure_call_to_migration(migration_file, queue_method_node, job_name, migration_record)
         ::Gitlab::Housekeeper::Shell.execute('rubocop', '-a', migration_file)
 
         digest = Digest::SHA256.hexdigest(generator.migration_number)
@@ -126,7 +126,7 @@ def last_migration_for_job(job_name)
       nil
     end
 
-    def add_ensure_call_to_migration(file, queue_method_node, job_name)
+    def add_ensure_call_to_migration(file, queue_method_node, job_name, migration_record)
       source = RuboCop::ProcessedSource.new(File.read(file), 3.1)
       ast = source.ast
       source_buffer = source.buffer
@@ -140,7 +140,7 @@ def add_ensure_call_to_migration(file, queue_method_node, job_name)
       column_name = queue_method_node.children[4]
       job_arguments = queue_method_node.children[5..].select { |s| s.type != :hash } # All remaining non-keyword args
 
-      gitlab_schema = ::Gitlab::Database::GitlabSchema.table_schema(table_name.value.to_s)
+      gitlab_schema = migration_record.gitlab_schema
 
       added_content = <<~RUBY.strip
       disable_ddl_transaction!
diff --git a/spec/keeps/helpers/postgres_ai_spec.rb b/spec/keeps/helpers/postgres_ai_spec.rb
index e3003be34609cf0e4b99e0e7c976ffcd18bd3ff4..f7f11915224531371439f49c3dcac66353ddca89 100644
--- a/spec/keeps/helpers/postgres_ai_spec.rb
+++ b/spec/keeps/helpers/postgres_ai_spec.rb
@@ -39,7 +39,7 @@
     let(:job_class_name) { 'ExampleJob' }
     let(:query) do
       <<~SQL
-      SELECT id, created_at, updated_at, finished_at, started_at, status, job_class_name
+      SELECT id, created_at, updated_at, finished_at, started_at, status, job_class_name, gitlab_schema
       FROM batched_background_migrations
       WHERE job_class_name = $1::text
       SQL