diff --git a/keeps/helpers/postgres_ai.rb b/keeps/helpers/postgres_ai.rb index e4833e1403822760fdf89c41c33a4421fea07f54..266a416ad75ca270a24a41a9b195939aaeace72d 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 f7f11915224531371439f49c3dcac66353ddca89..b6e460fd4c28aa5f52dbd0ac3d25e2c613e0d938 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