#434

backend/app/controllers/counts_exports_controller.rb

@@ -1,6 +1,29 @@
 class CountsExportsController < ApplicationController
   wrap_parameters false
 
+  def prepare
+    locale = resolve_locale!
+    payload = request.request_parameters.merge(request.query_parameters)
+    exporter = CountsExporter.new(payload, language_code: locale.language_code)
+    session[:counts_export_owner] ||= SecureRandom.hex(32)
+    token = CountsExportDownload.prepare(exporter, session[:counts_export_owner])
+    response.headers['Cache-Control'] = 'private, no-store'
+    render json: { token: }
+  rescue ArgumentError, TypeError, KeyError, JSON::ParserError
+    head :bad_request
+  end
+
+  def download
+    file = CountsExportDownload.find(params[:token], session[:counts_export_owner])
+    return head :not_found unless file
+
+    response.headers['Cache-Control'] = 'private, no-store'
+    send_file file.fetch('path'),
+              filename: "counts-#{ file.fetch('resource') }.#{ file.fetch('format') }",
+              type: file['format'] == 'json' ? 'application/json' : 'text/csv',
+              disposition: 'attachment'
+  end
+
   def index
     locale = resolve_locale!
     payload = request.request_parameters.except('authenticity_token')

backend/app/controllers/counts_tables_controller.rb

@@ -1,7 +1,20 @@
 class CountsTablesController < ApplicationController
+  def definition
+    resolve_locale!
+    definition = CountsResourceDefinition.new(params[:resource])
+    render json: { fields: definition.metadata,
+                   axes: { latest: definition.axes(false), history: definition.axes(true) },
+                   grains: CountsAxes::GRAINS }
+  rescue ArgumentError
+    head :bad_request
+  end
+
   def index
     resolve_locale!
     render json: CountsTableQuery.new(request.query_parameters).call
+  rescue CountsTableQuery::TooManyColumns
+    render json: { error: 'too_many_columns', max_columns: CountsTableQuery::MAX_COLUMNS },
+           status: :unprocessable_entity
   rescue ArgumentError, TypeError, KeyError, JSON::ParserError
     head :bad_request
   end

backend/app/services/counts_export_download.rb

@@ -0,0 +1,51 @@
+require 'tempfile'
+require 'fileutils'
+
+# Finish generation before reporting readiness; downloads never hold the full file in RAM.
+class CountsExportDownload
+  EXPIRY = 10.minutes
+
+  def self.directory = Rails.root.join('tmp', 'counts_exports')
+  def self.verifier = Rails.application.message_verifier('counts_export_download')
+
+  def self.prepare(exporter, owner)
+    FileUtils.mkdir_p(directory, mode: 0o700)
+    Dir.glob(directory.join('counts-*.export')).each do |path|
+      File.unlink(path) if File.mtime(path) < Time.current - EXPIRY
+    rescue Errno::ENOENT
+      # Another request may have removed an expired download.
+    end
+    file = Tempfile.create(['counts-', '.part'], directory)
+    complete = false
+    path = file.path
+    begin
+      file.binmode
+      exporter.body.each { file.write(_1) }
+      file.close
+      finished_path = path.sub(/\.part\z/, '.export')
+      File.rename(path, finished_path)
+      path = finished_path
+      token = verifier.generate({ 'file' => File.basename(path),
+                          'resource' => exporter.resource, 'format' => exporter.format,
+                          'owner' => owner }, expires_in: EXPIRY, purpose: 'counts_export')
+      complete = true
+      token
+    ensure
+      file.close unless file.closed?
+      File.unlink(path) if !complete && File.exist?(path)
+    end
+  end
+
+  def self.find(token, owner)
+    metadata = verifier.verified(token, purpose: 'counts_export')
+    return unless metadata.is_a?(Hash) && metadata['owner'] == owner
+    return unless metadata['file']&.match?(/\Acounts-[a-zA-Z0-9_-]+\.export\z/)
+    return unless CountsResourceDefinition::DEFINITIONS.key?(metadata['resource'])
+    return unless ['json', 'csv'].include?(metadata['format'])
+
+    path = directory.join(metadata['file'])
+    return unless path.file?
+
+    metadata.merge('path' => path.to_s)
+  end
+end

backend/app/services/counts_table_query.rb

@@ -1,5 +1,8 @@
 # Common entrypoint for six version resources; legacy endpoints keep their contracts.
 class CountsTableQuery
+  MAX_COLUMNS = 100
+  class TooManyColumns < ArgumentError; end
+
   KEYS = ['resource', 'definition', 'mode', 'at', 'execution_at', 'from', 'through',
           'timezone', 'rows', 'cols', 'conditions', 'drill', 'tags', 'tag_refs', 'match',
           'tag', 'kind', 'resource_id', 'user_id', 'page', 'limit', 'sort', 'locale'].freeze
@@ -45,12 +48,15 @@ class CountsTableQuery
     { page:, limit:, execution_at: population.at.iso8601(6), at: @params['at'].presence,
       definition: population.history ? 'history' : 'latest',
       timezone: population.timezone.identifier, rows: axes.rows, cols: axes.columns,
-      fields: @definition.metadata, count: population.scope.except(:select).count }
+      fields: @definition.metadata,
+      count: population.scope.except(:select).optimizer_hints('MAX_EXECUTION_TIME(10000)').count }
   end
 
   def summary
     expressions = axes.expressions
-    selected = expressions.each_with_index.map { |expression, index| "#{ expression } AS a#{ index }" }
+    selected = expressions.each_with_index.map do |expression, index|
+      "#{ expression } AS a#{ index }"
+    end
     source = population.scope
     source = source.joins(axes.joins) unless axes.joins.empty?
     sql = if expressions.empty?
@@ -61,6 +67,7 @@ class CountsTableQuery
                 .to_sql
           end
     prefix = "WITH cells AS (#{ sql })"
+    column_values = limited_columns(prefix)
     row_keys = axes.rows.each_index.map { "a#{ _1 }" }
     # Page whole row tuples so a pivot row is never split across display pages.
     row_selection = row_keys.empty? ? '1 AS singleton' : row_keys.join(', ')
@@ -84,29 +91,28 @@ class CountsTableQuery
     order_rows = row_sort.map(&:last) + row_keys.map { "#{ _1 } IS NULL, #{ _1 }" }
     order_rows = ['singleton ASC'] if order_rows.empty?
     group_sql = "SELECT #{ ([row_selection] + extras).join(', ') } FROM cells #{ row_group }"
-    total = @connection.select_value(
-        "#{ prefix } SELECT COUNT(*) FROM (#{ group_sql }) row_groups").to_i
-    join = row_keys.empty? ? 'TRUE' : row_keys.map { "cells.#{ _1 } <=> page_rows.#{ _1 }" }.join(' AND ')
+    total = @connection.select_value(<<~SQL).to_i
+      #{ prefix }
+      SELECT /*+ MAX_EXECUTION_TIME(10000) */ COUNT(*) FROM (#{ group_sql }) row_groups
+    SQL
+    join = row_keys.map { "cells.#{ _1 } <=> page_rows.#{ _1 }" }.join(' AND ')
+    join = 'TRUE' if row_keys.empty?
     order_cols = @sort.filter_map do |key, direction|
-      "cells.#{ key } #{ direction.upcase }" if key.match?(/\Aa[0-9]+\z/) && !row_keys.include?(key)
+      next unless key.match?(/\Aa[0-9]+\z/) && !row_keys.include?(key)
+
+      "cells.#{ key } #{ direction.upcase }"
     end
+    order = ['page_rows.position'] + order_cols +
+        axes.all.each_index.map { "cells.a#{ _1 }" }
     records = @connection.select_all(<<~SQL)
       #{ prefix }, row_groups AS (#{ group_sql }),
       numbered AS (SELECT row_groups.*,
         ROW_NUMBER() OVER (ORDER BY #{ order_rows.join(', ') }) AS position FROM row_groups),
       page_rows AS (SELECT * FROM numbered
         WHERE position > #{ (page - 1) * limit } AND position <= #{ page * limit })
-      SELECT cells.* FROM cells JOIN page_rows ON #{ join }
-      ORDER BY #{ (['page_rows.position'] + order_cols + axes.all.each_index.map { "cells.a#{ _1 }" }).join(', ') }
+      SELECT /*+ MAX_EXECUTION_TIME(10000) */ cells.* FROM cells JOIN page_rows ON #{ join }
+      ORDER BY #{ order.join(', ') }
     SQL
-    column_keys = axes.columns.each_index.map { "a#{ axes.rows.length + _1 }" }
-    column_values = if column_keys.empty?
-                      [[]]
-                    else
-                      @connection.select_rows(
-                          "#{ prefix } SELECT DISTINCT #{ column_keys.join(', ') } FROM cells " +
-                          "ORDER BY #{ column_keys.map { "#{ _1 } IS NULL, #{ _1 }" }.join(', ') }")
-                    end
     { group_count: total, column_values:, cells: records.map do |row|
         { values: axes.all.each_index.map { row["a#{ _1 }"] }, count: row['count'].to_i }
       end, axis_labels: axis_labels(records, column_values) }
@@ -120,9 +126,11 @@ class CountsTableQuery
         .limit(limit).offset((page - 1) * limit)
     rows = query.to_a
     if ['posts', 'tags'].include?(@definition.resource)
+      point = population.history ? population.at :
+          population.timestamp(@params['at']) || population.at
       values = CountsVersionRepr.many(rows, resource: @definition.resource,
           language_code: locale.language_code,
-          at: population.history ? population.at : population.timestamp(@params['at']) || population.at)
+          at: point)
       values.each_with_index { |value, index| value[:event_type] = rows[index].event_type }
     else
       users = User.where(id: rows.map(&:created_by_user_id)).pluck(:id, :name).to_h
@@ -155,6 +163,21 @@ class CountsTableQuery
 
   def locale = @locale ||= Locale.find_by!(code: @params.fetch('locale', 'ja'))
 
+  def limited_columns(prefix)
+    keys = axes.columns.each_index.map { "a#{ axes.rows.length + _1 }" }
+    return [[]] if keys.empty?
+
+    order = keys.map { "#{ _1 } IS NULL, #{ _1 }" }.join(', ')
+    values = @connection.select_rows(<<~SQL)
+      #{ prefix }
+      SELECT /*+ MAX_EXECUTION_TIME(10000) */ DISTINCT #{ keys.join(', ') } FROM cells
+      ORDER BY #{ order } LIMIT #{ MAX_COLUMNS + 1 }
+    SQL
+    raise TooManyColumns if values.length > MAX_COLUMNS
+
+    values
+  end
+
   def measure_values(key)
     value = JSON.parse([key.delete_prefix('measure_')].pack('H*'))
     unless (value.is_a?(Array) && value.length == axes.columns.length &&
@@ -171,7 +194,9 @@ class CountsTableQuery
                 'original_created_at' => 'original_created_from' }
     field = aliases.fetch(key, key)
     expression = @definition.expression(field)
-    @definition.fields[field] == :json ? "CAST(#{ expression } AS CHAR) COLLATE utf8mb4_bin" : expression
+    return expression unless @definition.fields[field] == :json
+
+    "CAST(#{ expression } AS CHAR) COLLATE utf8mb4_bin"
   end
 
   def sort_sql
@@ -208,7 +233,8 @@ class CountsTableQuery
     external = values.grep(/\Aexternal:/).map { _1.split(':').last.to_i }
     point = population.timestamp(@params['at']) || population.at
     labels = {}
-    [[TagVersion, internal, 'internal'], [NicoTagVersion, external, 'external']].each do |model, ids, kind|
+    sources = [[TagVersion, internal, 'internal'], [NicoTagVersion, external, 'external']]
+    sources.each do |model, ids, kind|
       next if ids.empty?
 
       table = model.table_name