Use ActiveRecord configuration to read from replica for reporting

app/jobs/reports/base_report.rb

@@ -31,8 +31,12 @@ def report_timeout
       IdentityConfig.store.report_timeout
     end
 
-    def transaction_with_timeout
-      Db::EstablishConnection::ReadReplicaConnection.new.call do
+    def transaction_with_timeout(rails_env = Rails.env)
+      # rspec-rails's use_transactional_tests does not seem to act as expected when switching
+      # connections mid-test, so we just skip for now :[
+      return yield if rails_env.test?
+
+      ApplicationRecord.connected_to(role: :reading, shard: :read_replica) do
         ActiveRecord::Base.transaction do
           quoted_timeout = ActiveRecord::Base.connection.quote(report_timeout)
           ActiveRecord::Base.connection.execute("SET LOCAL statement_timeout = #{quoted_timeout}")

app/models/application_record.rb

@@ -1,3 +1,12 @@
 class ApplicationRecord < ActiveRecord::Base
   self.abstract_class = true
+
+  connects_to shards: {
+    default: { writing: :primary, reading: :primary },
+    read_replica: {
+      # writing to the read_replica won't work, but AR needs to have something here
+      writing: :read_replica,
+      reading: :read_replica,
+    },
+  }
 end

config/application.rb

@@ -23,8 +23,23 @@ class Application < Rails::Application
     )
     IdentityConfig.build_store(configuration)
 
+    console do
+      if ENV['ALLOW_CONSOLE_DB_WRITE_ACCESS'] != 'true' &&
+         IdentityConfig.store.database_readonly_username.present? &&
+         IdentityConfig.store.database_readonly_password.present?
+        warn <<-EOS.squish
+          WARNING: Loading database a configuration with the readonly database user.
+          If you wish to make changes to records in the database set
+          ALLOW_CONSOLE_DB_WRITE_ACCESS to "true" in the environment
+        EOS
+
+        ActiveRecord::Base.establish_connection :read_replica
+      end
+    end
+
     config.load_defaults '6.1'
     config.active_record.belongs_to_required_by_default = false
+    config.active_record.legacy_connection_handling = false
     config.assets.unknown_asset_fallback = true
 
     if IdentityConfig.store.ruby_workers_enabled

config/database.yml

@@ -1,5 +1,3 @@
-<% require 'production_database_configuration' %>
-
 postgresql: &postgresql
   adapter: postgresql
   encoding: utf8
@@ -24,22 +22,49 @@ defaults: &defaults
     statement_timeout: <%= IdentityConfig.store.database_statement_timeout %>
 
 development:
-  <<: *defaults
+  primary:
+    <<: *defaults
+  read_replica:
+    <<: *defaults
+    replica: true
 
 test:
-  <<: *defaults
-  pool: 10
-  checkout_timeout: 10
-  database: <%= ENV['POSTGRES_DB'] || "upaya_test#{ENV['TEST_ENV_NUMBER']}" %>
-  user: <%= ENV['POSTGRES_USER'] %>
-  password: <%= ENV['POSTGRES_PASSWORD'] %>
+  primary: &test
+    <<: *defaults
+    pool: 10
+    checkout_timeout: 10
+    database: <%= ENV['POSTGRES_DB'] || "upaya_test#{ENV['TEST_ENV_NUMBER']}" %>
+    user: <%= ENV['POSTGRES_USER'] %>
+    password: <%= ENV['POSTGRES_PASSWORD'] %>
+  read_replica:
+    <<: *test
+    replica: true
+
+<%
+  pool = if Identity::Hostdata.instance_role == 'worker'
+    IdentityConfig.store.good_job_max_threads + IdentityConfig.store.database_pool_idp
+  else
+    IdentityConfig.store.database_pool_idp
+  end
+%>
 
 production:
-  <<: *defaults
-  database: <%= IdentityConfig.store.database_name %>
-  username: <%= ProductionDatabaseConfiguration.username %>
-  host: <%= ProductionDatabaseConfiguration.host %>
-  password: <%= ProductionDatabaseConfiguration.password %>
-  pool: <%= ProductionDatabaseConfiguration.pool %>
-  sslmode: 'verify-full'
-  sslrootcert: '/usr/local/share/aws/rds-combined-ca-bundle.pem'
+  primary:
+    <<: *defaults
+    database: <%= IdentityConfig.store.database_name %>
+    username: <%= IdentityConfig.store.database_username %>
+    host: <%= IdentityConfig.store.database_host %>
+    password: <%= IdentityConfig.store.database_password %>
+    pool: <%= pool %>
+    sslmode: 'verify-full'
+    sslrootcert: '/usr/local/share/aws/rds-combined-ca-bundle.pem'
+  read_replica:
+    <<: *defaults
+    database: <%= IdentityConfig.store.database_name %>
+    username: <%= IdentityConfig.store.database_readonly_username %>
+    host: <%= IdentityConfig.store.database_read_replica_host %>
+    password: <%= IdentityConfig.store.database_readonly_password %>
+    pool: <%= pool %>
+    sslmode: 'verify-full'
+    sslrootcert: '/usr/local/share/aws/rds-combined-ca-bundle.pem'
+    replica: true

spec/jobs/reports/base_report_spec.rb

@@ -0,0 +1,22 @@
+require 'rails_helper'
+
+RSpec.describe Reports::BaseReport do
+  subject(:report) { described_class.new }
+
+  describe '#transaction_with_timeout' do
+    let(:rails_env) { ActiveSupport::StringInquirer.new('production') }
+    let(:report_timeout) { 999 }
+
+    before do
+      allow(IdentityConfig.store).to receive(:report_timeout).and_return(report_timeout)
+    end
+
+    it 'sets the statement_timeout inside a transaction' do
+      result = report.send(:transaction_with_timeout, rails_env) do
+        ActiveRecord::Base.connection.execute('SHOW statement_timeout')
+      end
+
+      expect(result.first['statement_timeout']).to eq("#{report_timeout}ms")
+    end
+  end
+end