From ee844c79bc442e6236c3fc84a2ecee2be2a93c99 Mon Sep 17 00:00:00 2001
From: Manoj M J <mmj@gitlab.com>
Date: Tue, 5 Mar 2024 05:17:40 +0000
Subject: [PATCH] Modify helper to fetch more details of BBM
---
keeps/helpers/postgres_ai.rb | 14 +++++++++++++-
spec/keeps/helpers/postgres_ai_spec.rb | 24 +++++++++++++++++++++++-
2 files changed, 36 insertions(+), 2 deletions(-)
diff --git a/keeps/helpers/postgres_ai.rb b/keeps/helpers/postgres_ai.rb
index e4833e1403822..266a416ad75ca 100644
--- a/keeps/helpers/postgres_ai.rb
+++ b/keeps/helpers/postgres_ai.rb
@@ -11,7 +11,8 @@ 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, gitlab_schema
+ SELECT id, created_at, updated_at, finished_at, started_at, status, job_class_name,
+ gitlab_schema, total_tuple_count
FROM batched_background_migrations
WHERE job_class_name = $1::text
SQL
@@ -19,6 +20,17 @@ def fetch_background_migration_status(job_class_name)
pg_client.exec_params(query, [job_class_name])
end
+ def fetch_migrated_tuple_count(batched_background_migration_id)
+ query = <<~SQL
+ SELECT SUM("batched_background_migration_jobs"."batch_size")
+ FROM "batched_background_migration_jobs"
+ WHERE "batched_background_migration_jobs"."batched_background_migration_id" = #{batched_background_migration_id}
+ AND ("batched_background_migration_jobs"."status" IN (3))
+ SQL
+
+ pg_client.exec_params(query)
+ end
+
private
def connection_string
diff --git a/spec/keeps/helpers/postgres_ai_spec.rb b/spec/keeps/helpers/postgres_ai_spec.rb
index f7f1191522453..b6e460fd4c28a 100644
--- a/spec/keeps/helpers/postgres_ai_spec.rb
+++ b/spec/keeps/helpers/postgres_ai_spec.rb
@@ -39,7 +39,8 @@
let(:job_class_name) { 'ExampleJob' }
let(:query) do
<<~SQL
- SELECT id, created_at, updated_at, finished_at, started_at, status, job_class_name, gitlab_schema
+ SELECT id, created_at, updated_at, finished_at, started_at, status, job_class_name,
+ gitlab_schema, total_tuple_count
FROM batched_background_migrations
WHERE job_class_name = $1::text
SQL
@@ -54,4 +55,25 @@
expect(result).to eq(query_response)
end
end
+
+ describe '#fetch_migrated_tuple_count' do
+ let(:batched_background_migration_id) { 100 }
+ let(:query) do
+ <<~SQL
+ SELECT SUM("batched_background_migration_jobs"."batch_size")
+ FROM "batched_background_migration_jobs"
+ WHERE "batched_background_migration_jobs"."batched_background_migration_id" = 100
+ AND ("batched_background_migration_jobs"."status" IN (3))
+ SQL
+ end
+
+ let(:query_response) { double }
+
+ subject(:result) { described_class.new.fetch_migrated_tuple_count(batched_background_migration_id) }
+
+ it 'fetches data from Postgres AI' do
+ expect(pg_client).to receive(:exec_params).with(query).and_return(query_response)
+ expect(result).to eq(query_response)
+ end
+ end
end
--
GitLab