From 95a749c3db9eccfef9f66f636e861e5be0d06059 Mon Sep 17 00:00:00 2001
From: Philipp Thun <philipp.thun@sap.com>
Date: Mon, 6 May 2024 13:49:31 +0200
Subject: [PATCH] Fix the concurrent statement timeout

- Fix the 20240314131908_add_user_guid_to_jobs_table migration (don't
  use 'with_concurrent_timeout' in a block - e.g. alter_table {}).
- Also fix migration_concurrent_statement_timeout_spec (be less specific
  about the original timeout).
- Also fix the concurrent_timeout_in_milliseconds calculation.
---
 ...240314131908_add_user_guid_to_jobs_table.rb | 18 ++++++------------
 lib/cloud_controller/db.rb                     | 15 ++++++++++-----
 ...ration_concurrent_statement_timeout_spec.rb |  2 +-
 3 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/db/migrations/20240314131908_add_user_guid_to_jobs_table.rb b/db/migrations/20240314131908_add_user_guid_to_jobs_table.rb
index 3bac5f73c4..fb869ef17a 100644
--- a/db/migrations/20240314131908_add_user_guid_to_jobs_table.rb
+++ b/db/migrations/20240314131908_add_user_guid_to_jobs_table.rb
@@ -4,12 +4,9 @@
 
   up do
     if database_type == :postgres
-      db = self
-      alter_table :jobs do
-        add_column :user_guid, String, size: 255, if_not_exists: true
-        VCAP::Migration.with_concurrent_timeout(db) do
-          add_index :user_guid, name: :jobs_user_guid_index, if_not_exists: true, concurrently: true
-        end
+      add_column :jobs, :user_guid, String, size: 255, if_not_exists: true
+      VCAP::Migration.with_concurrent_timeout(self) do
+        add_index :jobs, :user_guid, name: :jobs_user_guid_index, if_not_exists: true, concurrently: true
       end
 
     elsif database_type == :mysql
@@ -24,13 +21,10 @@
 
   down do
     if database_type == :postgres
-      db = self
-      alter_table :jobs do
-        VCAP::Migration.with_concurrent_timeout(db) do
-          drop_index :user_guid, name: :jobs_user_guid_index, if_exists: true, concurrently: true
-        end
-        drop_column :user_guid, if_exists: true
+      VCAP::Migration.with_concurrent_timeout(self) do
+        drop_index :jobs, :user_guid, name: :jobs_user_guid_index, if_exists: true, concurrently: true
       end
+      drop_column :jobs, :user_guid, if_exists: true
     end
 
     if database_type == :mysql
diff --git a/lib/cloud_controller/db.rb b/lib/cloud_controller/db.rb
index f255df7fd0..93e405a64d 100644
--- a/lib/cloud_controller/db.rb
+++ b/lib/cloud_controller/db.rb
@@ -224,15 +224,20 @@ def self.uuid_function(migration)
     # Concurrent migrations can take a long time to run, so this helper can be used to override 'max_migration_statement_runtime_in_seconds' for a specific migration.
     # REF: https://www.postgresql.org/docs/current/sql-createindex.html#SQL-CREATEINDEX-CONCURRENTLY
     def self.with_concurrent_timeout(db, &block)
-      concurrent_timeout_seconds = VCAP::CloudController::Config.config&.get(:migration_psql_concurrent_statement_timeout_in_seconds) || PSQL_DEFAULT_STATEMENT_TIMEOUT
-
-      if concurrent_timeout_seconds && db.database_type == :postgres
+      concurrent_timeout_in_seconds = VCAP::CloudController::Config.config&.get(:migration_psql_concurrent_statement_timeout_in_seconds)
+      concurrent_timeout_in_milliseconds = if concurrent_timeout_in_seconds.nil? || concurrent_timeout_in_seconds <= 0
+                                             PSQL_DEFAULT_STATEMENT_TIMEOUT
+                                           else
+                                             concurrent_timeout_in_seconds * 1000
+                                           end
+
+      if db.database_type == :postgres
         original_timeout = db.fetch("select setting from pg_settings where name = 'statement_timeout'").first[:setting]
-        db.run("SET statement_timeout TO #{concurrent_timeout_seconds * 1000}")
+        db.run("SET statement_timeout TO #{concurrent_timeout_in_milliseconds}")
       end
       block.call
     ensure
-      db.run("SET statement_timeout TO #{original_timeout}") if original_timeout && db.database_type == :postgres
+      db.run("SET statement_timeout TO #{original_timeout}") if original_timeout
     end
 
     def self.logger
diff --git a/spec/migrations/migration_concurrent_statement_timeout_spec.rb b/spec/migrations/migration_concurrent_statement_timeout_spec.rb
index 9df06c2ee2..1a9ed112ed 100644
--- a/spec/migrations/migration_concurrent_statement_timeout_spec.rb
+++ b/spec/migrations/migration_concurrent_statement_timeout_spec.rb
@@ -33,7 +33,7 @@
     skip if db.database_type != :postgres
     expect { Sequel::Migrator.run(db, tmp_migrations_dir, allow_missing_migration_files: true) }.not_to raise_error
     expect(db).to have_received(:run).exactly(2).times
+    expect(db).to have_received(:run).with(/SET statement_timeout TO \d+/).twice
     expect(db).to have_received(:run).with('SET statement_timeout TO 1899000').once
-    expect(db).to have_received(:run).with('SET statement_timeout TO 30000').once
   end
 end