Add configurable CCDB connection parameters

lib/cloud_controller/config_schemas/api_schema.rb

@@ -121,7 +121,8 @@ class ApiSchema < VCAP::Config
             optional(:connection_expiration_random_delay) => Integer,
             optional(:ssl_verify_hostname) => bool,
             optional(:ca_cert_path) => String,
-            optional(:enable_paginate_window) => bool
+            optional(:enable_paginate_window) => bool,
+            optional(:connection_parameters) => Hash
           },
 
           optional(:redis) => {

lib/cloud_controller/config_schemas/worker_schema.rb

@@ -59,7 +59,8 @@ class WorkerSchema < VCAP::Config
             log_db_queries: bool,
             ssl_verify_hostname: bool,
             connection_validation_timeout: Integer,
-            optional(:ca_cert_path) => String
+            optional(:ca_cert_path) => String,
+            optional(:connection_parameters) => Hash
           },
 
           staging: {

lib/cloud_controller/db_connection/postgres_options_factory.rb

@@ -1,6 +1,9 @@
 module VCAP::CloudController
   module DbConnection
     class PostgresOptionsFactory
+      SQL_CONNECTION_PARAMETERS = %i[statement_timeout idle_in_transaction_session_timeout].freeze
+      LIBPQ_CONNECTION_PARAMETERS = %i[keepalives keepalives_idle keepalives_interval keepalives_count].freeze
+
       def self.build(opts)
         options = {}
 
@@ -9,9 +12,17 @@ def self.build(opts)
           options[:sslmode] = opts[:ssl_verify_hostname] ? 'verify-full' : 'verify-ca'
         end
 
-        options[:after_connect] = proc do |connection|
-          connection.exec("SET time zone 'UTC'")
+        connection_parameters = opts[:connection_parameters] || {}
+
+        sql_params = connection_parameters.slice(*SQL_CONNECTION_PARAMETERS)
+        connect_sqls = ["SET time zone 'UTC'"]
+        sql_params.each do |key, value|
+          connect_sqls << "SET #{key} TO '#{value}'"
         end
+        options[:connect_sqls] = connect_sqls
+
+        libpq_params = connection_parameters.slice(*LIBPQ_CONNECTION_PARAMETERS)
+        options.merge!(libpq_params.transform_keys(&:to_sym))
 
         options
       end

spec/unit/lib/cloud_controller/db_connection/options_factory_spec.rb

@@ -46,10 +46,8 @@
         let(:adapter) { 'postgres' }
 
         it 'returns postgres-specific options' do
-          connection = double('connection', exec: '')
           postgres_options = VCAP::CloudController::DbConnection::OptionsFactory.build(required_options)
-          postgres_options[:after_connect].call(connection)
-          expect(connection).to have_received(:exec).with("SET time zone 'UTC'")
+          expect(postgres_options[:connect_sqls]).to include("SET time zone 'UTC'")
         end
       end
     end

spec/unit/lib/cloud_controller/db_connection/postgres_options_factory_spec.rb

@@ -17,10 +17,8 @@
       )
     end
 
-    it 'sets the timezone via a Proc' do
-      connection = double('connection', exec: '')
-      postgres_options[:after_connect].call(connection)
-      expect(connection).to have_received(:exec).with("SET time zone 'UTC'")
+    it 'sets the timezone via connect_sqls' do
+      expect(postgres_options[:connect_sqls]).to include("SET time zone 'UTC'")
     end
 
     describe 'when the CA cert path is not set' do
@@ -57,5 +55,102 @@
         end
       end
     end
+
+    describe 'connection_parameters' do
+      context 'when connection_parameters is not set' do
+        let(:postgres_options) do
+          VCAP::CloudController::DbConnection::PostgresOptionsFactory.build(
+            database: { adapter: 'postgres' }
+          )
+        end
+
+        it 'only sets the timezone in connect_sqls' do
+          expect(postgres_options[:connect_sqls]).to eq(["SET time zone 'UTC'"])
+        end
+
+        it 'does not include any keepalive options in the returned hash' do
+          expect(postgres_options[:keepalives]).to be_nil
+          expect(postgres_options[:keepalives_idle]).to be_nil
+          expect(postgres_options[:keepalives_interval]).to be_nil
+          expect(postgres_options[:keepalives_count]).to be_nil
+        end
+      end
+
+      context 'when connection_parameters contains SQL params' do
+        let(:postgres_options) do
+          VCAP::CloudController::DbConnection::PostgresOptionsFactory.build(
+            database: { adapter: 'postgres' },
+            connection_parameters: {
+              statement_timeout: '3600000',
+              idle_in_transaction_session_timeout: '600000'
+            }
+          )
+        end
+
+        it 'sets the SQL params via connect_sqls' do
+          expect(postgres_options[:connect_sqls]).to include("SET time zone 'UTC'")
+          expect(postgres_options[:connect_sqls]).to include("SET statement_timeout TO '3600000'")
+          expect(postgres_options[:connect_sqls]).to include("SET idle_in_transaction_session_timeout TO '600000'")
+        end
+
+        it 'does not put SQL params into the returned options hash' do
+          expect(postgres_options[:statement_timeout]).to be_nil
+          expect(postgres_options[:idle_in_transaction_session_timeout]).to be_nil
+        end
+      end
+
+      context 'when connection_parameters contains libpq keepalive params' do
+        let(:postgres_options) do
+          VCAP::CloudController::DbConnection::PostgresOptionsFactory.build(
+            database: { adapter: 'postgres' },
+            connection_parameters: {
+              keepalives: 1,
+              keepalives_idle: 30,
+              keepalives_interval: 10,
+              keepalives_count: 3
+            }
+          )
+        end
+
+        it 'merges keepalive params into the returned options hash' do
+          expect(postgres_options[:keepalives]).to eq(1)
+          expect(postgres_options[:keepalives_idle]).to eq(30)
+          expect(postgres_options[:keepalives_interval]).to eq(10)
+          expect(postgres_options[:keepalives_count]).to eq(3)
+        end
+
+        it 'does not SET keepalive params via connect_sqls' do
+          expect(postgres_options[:connect_sqls]).not_to include(match(/SET keepalives/))
+        end
+      end
+
+      context 'when connection_parameters contains both SQL and libpq params' do
+        let(:postgres_options) do
+          VCAP::CloudController::DbConnection::PostgresOptionsFactory.build(
+            database: { adapter: 'postgres' },
+            connection_parameters: {
+              statement_timeout: '3600000',
+              keepalives: 1,
+              keepalives_idle: 30,
+              keepalives_interval: 10,
+              keepalives_count: 3
+            }
+          )
+        end
+
+        it 'sets SQL params via connect_sqls and merges libpq params into options hash' do
+          expect(postgres_options[:connect_sqls]).to include("SET statement_timeout TO '3600000'")
+          expect(postgres_options[:keepalives]).to eq(1)
+          expect(postgres_options[:keepalives_idle]).to eq(30)
+          expect(postgres_options[:keepalives_interval]).to eq(10)
+          expect(postgres_options[:keepalives_count]).to eq(3)
+        end
+
+        it 'does not mix up the two kinds: SQL params not in options hash, libpq params not SET via SQL' do
+          expect(postgres_options[:statement_timeout]).to be_nil
+          expect(postgres_options[:connect_sqls]).not_to include(match(/SET keepalives/))
+        end
+      end
+    end
   end
 end