Stream query results (LG-5351)

Gemfile

@@ -40,6 +40,7 @@ gem 'local_time'
 gem 'lograge', '>= 0.11.2'
 gem 'lru_redux'
 gem 'maxminddb'
+gem 'multiset'
 gem 'net-sftp'
 gem 'newrelic_rpm', '~> 7.0'
 gem 'pg'

Gemfile.lock

@@ -368,6 +368,7 @@ GEM
     minitest (5.14.4)
     msgpack (1.4.2)
     multipart-post (2.1.1)
+    multiset (0.5.3)
     nenv (0.3.0)
     net-sftp (2.1.2)
       net-ssh (>= 2.6.5)
@@ -727,6 +728,7 @@ DEPENDENCIES
   lograge (>= 0.11.2)
   lru_redux
   maxminddb
+  multiset
   net-sftp
   newrelic_rpm (~> 7.0)
   nokogiri (~> 1.12.5)

app/jobs/reports/agency_invoice_iaa_supplement_report.rb

@@ -11,13 +11,8 @@ class AgencyInvoiceIaaSupplementReport < BaseReport
 
     def perform(_date)
       raw_results = iaas.flat_map do |iaa|
-        [:sum, :unique, :new_unique].flat_map do |aggregate|
-          transaction_with_timeout do
-            Db::MonthlySpAuthCount::UniqueMonthlyAuthCountsByIaa.call(
-              iaa: iaa,
-              aggregate: aggregate,
-            )
-          end.to_a
+        transaction_with_timeout do
+          Db::MonthlySpAuthCount::UniqueMonthlyAuthCountsByIaa.call(**iaa)
         end
       end
 
@@ -49,30 +44,30 @@ def iaas
 
     # Turns ial1/ial2 rows into ial1/ial2 columns
     def combine_by_iaa_month(raw_results)
-      raw_results.group_by { |r| [r['key'], r['year_month']] }.
+      raw_results.group_by { |r| [r[:key], r[:year_month]] }.
         transform_values do |grouped|
-          key = grouped.first['key']
-          iaa_start_date = grouped.first['iaa_start_date']
-          iaa_end_date = grouped.first['iaa_end_date']
-          year_month = grouped.first['year_month']
+          key = grouped.first[:key]
+          iaa_start_date = grouped.first[:iaa_start_date]
+          iaa_end_date = grouped.first[:iaa_end_date]
+          year_month = grouped.first[:year_month]
 
           {
             iaa: key,
             iaa_start_date: iaa_start_date,
             iaa_end_date: iaa_end_date,
             year_month: year_month,
-            ial1_total_auth_count: extract(grouped, 'total_auth_count', ial: 1),
-            ial2_total_auth_count: extract(grouped, 'total_auth_count', ial: 2),
-            ial1_unique_users: extract(grouped, 'unique_users', ial: 1),
-            ial2_unique_users: extract(grouped, 'unique_users', ial: 2),
-            ial1_new_unique_users: extract(grouped, 'new_unique_users', ial: 1),
-            ial2_new_unique_users: extract(grouped, 'new_unique_users', ial: 2),
+            ial1_total_auth_count: extract(grouped, :total_auth_count, ial: 1),
+            ial2_total_auth_count: extract(grouped, :total_auth_count, ial: 2),
+            ial1_unique_users: extract(grouped, :unique_users, ial: 1),
+            ial2_unique_users: extract(grouped, :unique_users, ial: 2),
+            ial1_new_unique_users: extract(grouped, :new_unique_users, ial: 1),
+            ial2_new_unique_users: extract(grouped, :new_unique_users, ial: 2),
           }
         end.values
     end
 
     def extract(arr, key, ial:)
-      arr.find { |elem| elem['ial'] == ial && elem[key] }&.dig(key) || 0
+      arr.find { |elem| elem[:ial] == ial && elem[key] }&.dig(key) || 0
     end
   end
 end

app/jobs/reports/agency_invoice_issuer_supplement_report.rb

@@ -11,13 +11,10 @@ class AgencyInvoiceIssuerSupplementReport < BaseReport
 
     def perform(_date)
       raw_results = service_providers.flat_map do |service_provider|
-        [:sum, :unique].flat_map do |aggregate|
-          transaction_with_timeout do
-            Db::MonthlySpAuthCount::TotalMonthlyAuthCountsWithinIaaWindow.call(
-              service_provider: service_provider,
-              aggregate: aggregate,
-            )
-          end.to_a
+        transaction_with_timeout do
+          Db::MonthlySpAuthCount::TotalMonthlyAuthCountsWithinIaaWindow.call(
+            service_provider: service_provider,
+          )
         end
       end
 
@@ -35,30 +32,30 @@ def service_providers
 
     # Turns ial1/ial2 rows into ial1/ial2 columns
     def combine_by_issuer_month(raw_results)
-      raw_results.group_by { |r| [r['issuer'], r['year_month']] }.
+      raw_results.group_by { |r| [r[:issuer], r[:year_month]] }.
         transform_values do |grouped|
-          issuer = grouped.first['issuer']
-          iaa = grouped.first['iaa']
-          iaa_start_date = grouped.first['iaa_start_date']
-          iaa_end_date = grouped.first['iaa_end_date']
-          year_month = grouped.first['year_month']
+          issuer = grouped.first[:issuer]
+          iaa = grouped.first[:iaa]
+          iaa_start_date = grouped.first[:iaa_start_date]
+          iaa_end_date = grouped.first[:iaa_end_date]
+          year_month = grouped.first[:year_month]
 
           {
             issuer: issuer,
             iaa: iaa,
             iaa_start_date: iaa_start_date,
             iaa_end_date: iaa_end_date,
             year_month: year_month,
-            ial1_total_auth_count: extract(grouped, 'total_auth_count', ial: 1),
-            ial2_total_auth_count: extract(grouped, 'total_auth_count', ial: 2),
-            ial1_unique_users: extract(grouped, 'unique_users', ial: 1),
-            ial2_unique_users: extract(grouped, 'unique_users', ial: 2),
+            ial1_total_auth_count: extract(grouped, :total_auth_count, ial: 1),
+            ial2_total_auth_count: extract(grouped, :total_auth_count, ial: 2),
+            ial1_unique_users: extract(grouped, :unique_users, ial: 1),
+            ial2_unique_users: extract(grouped, :unique_users, ial: 2),
           }
         end.values
     end
 
     def extract(arr, key, ial:)
-      arr.find { |elem| elem['ial'] == ial && elem[key] }&.dig(key) || 0
+      arr.find { |elem| elem[:ial] == ial && elem[key] }&.dig(key) || 0
     end
   end
 end

app/jobs/reports/query_helpers.rb

@@ -7,5 +7,20 @@ def quote(value)
         ActiveRecord::Base.connection.quote(value)
       end
     end
+
+    # Wrapper around PG::Result#stream_each, runs a query and yields each row to the block
+    # as it is returned from the DB
+    # @param [String] SQL query
+    # @yieldparam [Hash] row
+    def stream_query(query)
+      ActiveRecord::Base.logger.debug(query) # using send_query skips ActiveRecord logging
+      connection = ActiveRecord::Base.connection.raw_connection
+      connection.send_query(query)
+      connection.set_single_row_mode
+      connection.get_result.stream_each do |row|
+        yield row
+      end
+      connection.get_result
+    end
   end
 end

app/services/db/monthly_sp_auth_count/total_monthly_auth_counts_within_iaa_window.rb

@@ -10,9 +10,8 @@ module TotalMonthlyAuthCountsWithinIaaWindow
       module_function
 
       # @param [ServiceProvider] service_provider
-      # @param [Symbol] aggregate (one of :sum, :unique)
       # @return [PG::Result,Array]
-      def call(service_provider:, aggregate:)
+      def call(service_provider:)
         return [] if !service_provider.iaa_start_date || !service_provider.iaa_end_date
 
         iaa_range = (service_provider.iaa_start_date..service_provider.iaa_end_date)
@@ -28,51 +27,57 @@ def call(service_provider:, aggregate:)
         # - full months from monthly_sp_auth_counts
         # - partial months by aggregating sp_return_logs
         # The results are rows with [user_id, ial, auth_count, year_month]
-        subquery = [
+        queries = [
           *full_month_subquery(issuer: issuer, full_months: full_months),
           *partial_month_subqueries(issuer: issuer, partial_months: partial_months),
-        ].join(' UNION ALL ')
+        ]
 
-        select_clause = case aggregate
-        when :sum
-          <<~SQL
-            SUM(billing_month_logs.auth_count)::bigint AS total_auth_count
-          SQL
-        when :unique
-          <<~SQL
-            COUNT(DISTINCT billing_month_logs.user_id) AS unique_users
-          SQL
-        else
-          raise "unknown aggregate=#{aggregate}"
+        ial_to_year_month_to_users = Hash.new do |ial_h, ial_k|
+          ial_h[ial_k] = Hash.new { |ym_h, ym_k| ym_h[ym_k] = Multiset.new }
         end
 
-        params = {
-          iaa_start_date: quote(iaa_range.begin),
-          iaa_end_date: quote(iaa_range.end),
-          iaa: quote(service_provider.iaa),
-          issuer: quote(issuer),
-          subquery: subquery,
-          select_clause: select_clause,
-        }
-
-        sql = format(<<~SQL, params)
-          WITH subquery AS (%{subquery})
-          SELECT
-            billing_month_logs.year_month
-          , billing_month_logs.ial
-          , %{iaa_start_date} AS iaa_start_date
-          , %{iaa_end_date} AS iaa_end_date
-          , %{issuer} AS issuer
-          , %{iaa} AS iaa
-          , %{select_clause}
-          FROM
-            subquery billing_month_logs
-          GROUP BY
-            billing_month_logs.year_month
-          , billing_month_logs.ial
-        SQL
+        queries.each do |query|
+          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
+
+        rows = []
+
+        ial_to_year_month_to_users.each do |ial, year_month_to_users|
+          prev_seen_users = Set.new
+
+          year_months = year_month_to_users.keys.sort
+
+          year_months.each do |year_month|
+            year_month_users = year_month_to_users[year_month]
+
+            auth_count = year_month_users.count
+            unique_users = year_month_users.uniq.to_set
+
+            new_unique_users = unique_users - prev_seen_users
+            prev_seen_users |= unique_users
+
+            rows << {
+              issuer: service_provider.issuer,
+              iaa: service_provider.iaa,
+              ial: ial,
+              year_month: year_month,
+              iaa_start_date: iaa_range.begin.to_s,
+              iaa_end_date: iaa_range.end.to_s,
+              total_auth_count: auth_count,
+              unique_users: unique_users.count,
+              new_unique_users: new_unique_users.count,
+            }
+          end
+        end
 
-        ActiveRecord::Base.connection.execute(sql)
+        rows
       end
 
       # @return [String]