Update long-running report queries to reset connection

app/jobs/reports/agency_invoice_iaa_supplement_report.rb

@@ -11,14 +11,12 @@ class AgencyInvoiceIaaSupplementReport < BaseReport
 
     def perform(_date)
       raw_results = IaaReportingHelper.iaas.flat_map do |iaa|
-        transaction_with_timeout do
-          Db::MonthlySpAuthCount::UniqueMonthlyAuthCountsByIaa.call(
-            key: iaa.key,
-            issuers: iaa.issuers,
-            start_date: iaa.start_date,
-            end_date: iaa.end_date,
-          )
-        end
+        Db::MonthlySpAuthCount::UniqueMonthlyAuthCountsByIaa.call(
+          key: iaa.key,
+          issuers: iaa.issuers,
+          start_date: iaa.start_date,
+          end_date: iaa.end_date,
+        )
       end
 
       results = combine_by_iaa_month(raw_results)

app/jobs/reports/base_report.rb

@@ -7,6 +7,20 @@ class BaseReport < ApplicationJob
     # We use good_job's concurrency features to cancel "extra" or duplicative runs of the same job
     discard_on GoodJob::ActiveJobExtensions::Concurrency::ConcurrencyExceededError
 
+    def self.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?
+
+      ActiveRecord::Base.connected_to(role: :reading, shard: :read_replica) do
+        ActiveRecord::Base.transaction do
+          quoted_timeout = ActiveRecord::Base.connection.quote(IdentityConfig.store.report_timeout)
+          ActiveRecord::Base.connection.execute("SET LOCAL statement_timeout = #{quoted_timeout}")
+          yield
+        end
+      end
+    end
+
     private
 
     def public_bucket_name
@@ -27,22 +41,8 @@ def end_of_today
       Time.zone.now.end_of_day
     end
 
-    def report_timeout
-      IdentityConfig.store.report_timeout
-    end
-
-    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?
-
-      ActiveRecord::Base.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}")
-          yield
-        end
-      end
+    def transaction_with_timeout(...)
+      self.class.transaction_with_timeout(...)
     end
 
     def save_report(report_name, body, extension:)

app/jobs/reports/combined_invoice_supplement_report.rb

@@ -15,26 +15,22 @@ def perform(_date)
       iaas = IaaReportingHelper.iaas
 
       by_iaa_results = iaas.flat_map do |iaa|
-        transaction_with_timeout do
-          Db::MonthlySpAuthCount::UniqueMonthlyAuthCountsByIaa.call(
-            key: iaa.key,
-            issuers: iaa.issuers,
-            start_date: iaa.start_date,
-            end_date: iaa.end_date,
-          )
-        end
+        Db::MonthlySpAuthCount::UniqueMonthlyAuthCountsByIaa.call(
+          key: iaa.key,
+          issuers: iaa.issuers,
+          start_date: iaa.start_date,
+          end_date: iaa.end_date,
+        )
       end
 
       by_issuer_results = iaas.flat_map do |iaa|
         iaa.issuers.flat_map do |issuer|
-          transaction_with_timeout do
-            Db::MonthlySpAuthCount::TotalMonthlyAuthCountsWithinIaaWindow.call(
-              issuer: issuer,
-              iaa_start_date: iaa.start_date,
-              iaa_end_date: iaa.end_date,
-              iaa: iaa.key,
-            )
-          end
+          Db::MonthlySpAuthCount::TotalMonthlyAuthCountsWithinIaaWindow.call(
+            issuer: issuer,
+            iaa_start_date: iaa.start_date,
+            iaa_end_date: iaa.end_date,
+            iaa: iaa.key,
+          )
         end
       end
 

app/services/db/monthly_sp_auth_count/total_monthly_auth_counts_within_iaa_window.rb

@@ -33,16 +33,21 @@ def call(issuer:, iaa:, iaa_start_date:, iaa_end_date:)
 
           with_retries(
             max_tries: 3,
-            rescue: PG::TRSerializationFailure,
-            handler: proc { ial_to_year_month_to_users = temp_copy },
+            rescue: [PG::TRSerializationFailure, PG::UnableToSend],
+            handler: proc do
+              ial_to_year_month_to_users = temp_copy
+              ActiveRecord::Base.connection.reconnect!
+            end,
           ) do
-            stream_query(query) do |row|
-              user_id = row['user_id']
-              year_month = row['year_month']
-              auth_count = row['auth_count']
-              ial = row['ial']
-
-              ial_to_year_month_to_users[ial][year_month].add(user_id, auth_count)
+            Reports::BaseReport.transaction_with_timeout do
+              stream_query(query) do |row|
+                user_id = row['user_id']
+                year_month = row['year_month']
+                auth_count = row['auth_count']
+                ial = row['ial']
+
+                ial_to_year_month_to_users[ial][year_month].add(user_id, auth_count)
+              end
             end
           end
         end
@@ -111,6 +116,10 @@ def build_queries(issuer:, months:)
           SQL
         end
       end
+
+      def default_call
+        yield
+      end
     end
   end
 end

app/services/db/monthly_sp_auth_count/unique_monthly_auth_counts_by_iaa.rb

@@ -29,16 +29,21 @@ def call(key:, issuers:, start_date:, end_date:)
 
           with_retries(
             max_tries: 3,
-            rescue: PG::TRSerializationFailure,
-            handler: proc { ial_to_year_month_to_users = temp_copy },
+            rescue: [PG::TRSerializationFailure, PG::UnableToSend],
+            handler: proc do
+              ial_to_year_month_to_users = temp_copy
+              ActiveRecord::Base.connection.reconnect!
+            end,
           ) do
-            stream_query(query) do |row|
-              user_id = row['user_id']
-              year_month = row['year_month']
-              auth_count = row['auth_count']
-              ial = row['ial']
-
-              ial_to_year_month_to_users[ial][year_month].add(user_id, auth_count)
+            Reports::BaseReport.transaction_with_timeout do
+              stream_query(query) do |row|
+                user_id = row['user_id']
+                year_month = row['year_month']
+                auth_count = row['auth_count']
+                ial = row['ial']
+
+                ial_to_year_month_to_users[ial][year_month].add(user_id, auth_count)
+              end
             end
           end
         end