diff --git a/app/jobs/reports/agency_invoice_iaa_supplement_report.rb b/app/jobs/reports/agency_invoice_iaa_supplement_report.rb index db7371964cd..bd222154b83 100644 --- a/app/jobs/reports/agency_invoice_iaa_supplement_report.rb +++ b/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) diff --git a/app/jobs/reports/base_report.rb b/app/jobs/reports/base_report.rb index 20ccec76a29..c984d3d0219 100644 --- a/app/jobs/reports/base_report.rb +++ b/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:) diff --git a/app/jobs/reports/combined_invoice_supplement_report.rb b/app/jobs/reports/combined_invoice_supplement_report.rb index 3cedf95f830..45d2dbe0c5d 100644 --- a/app/jobs/reports/combined_invoice_supplement_report.rb +++ b/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 diff --git a/app/services/db/monthly_sp_auth_count/total_monthly_auth_counts_within_iaa_window.rb b/app/services/db/monthly_sp_auth_count/total_monthly_auth_counts_within_iaa_window.rb index be97f93686c..1efccde96f3 100644 --- a/app/services/db/monthly_sp_auth_count/total_monthly_auth_counts_within_iaa_window.rb +++ b/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 diff --git a/app/services/db/monthly_sp_auth_count/unique_monthly_auth_counts_by_iaa.rb b/app/services/db/monthly_sp_auth_count/unique_monthly_auth_counts_by_iaa.rb index 968474d40e1..a8748d0a167 100644 --- a/app/services/db/monthly_sp_auth_count/unique_monthly_auth_counts_by_iaa.rb +++ b/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