activeadmin_batched_export 0.2.0 → 0.3.1

data/lib/activeadmin/batched_export/controller_methods.rb

@@ -1,12 +1,16 @@
 # frozen_string_literal: true
 
-require "builder"
-require "csv"
+require "activeadmin/batched_export/chunk_renderer"
+require "activeadmin/batched_export/errors"
+require "activeadmin/batched_export/export_cursor"
+require "activeadmin/batched_export/keyset_page"
+require "activeadmin/batched_export/snapshot_page"
 
 module ActiveAdmin
   module BatchedExport
     module ControllerMethods
       extend ActiveSupport::Concern
+      include ChunkRenderer
 
       def batched_export
         authorize! ActiveAdmin::Authorization::READ, active_admin_config.resource_class
@@ -18,29 +22,65 @@ module ActiveAdmin
           return render(json: batched_export_meta(export_format))
         end
 
-        batch_page = params[:batch_page].to_i
-        if batch_page.positive?
-          if request.format.symbol != export_format
-            head :not_acceptable
-            return
-          end
-          ensure_batch_download_format_allowed!(export_format)
-          page = page_relation(batch_page)
-          if page.out_of_range?
-            head :not_found
-            return
-          end
-          begin
-            body = batched_export_batch_body(export_format, batch_page, page: page)
-          rescue ExportMacroCatalog::UnknownMacroError
-            return render(
-              plain: I18n.t("active_admin.batched_export_page.unknown_macro"),
-              status: :unprocessable_content
-            )
-          end
-          return render(plain: body, content_type: batch_content_type(export_format))
+        return render_export_batch(export_format) if export_batch_request?(export_format)
+        return head(:not_acceptable) if %i[csv json xml].include?(request.format.symbol)
+
+        render_export_workspace(export_format)
+      end
+
+      private
+
+      def normalized_export_format
+        format_name = params[:export_format].to_s.downcase
+        format_name = "csv" if format_name.blank?
+        unless %w[csv xml json].include?(format_name)
+          render(plain: "Invalid export format", status: :bad_request)
+          return nil
         end
+        format_name.to_sym
+      end
 
+      def export_batch_request?(export_format)
+        %i[csv json xml].include?(request.format.symbol) &&
+          request.format.symbol == export_format &&
+          params[:export_meta].blank?
+      end
+
+      def render_export_batch(export_format)
+        ensure_batch_download_format_allowed!(export_format)
+        refuse_over_max_export_rows! if starting_export_walk?
+        field, direction = batched_export_sort_pair
+        cursor = decode_export_cursor(field: field, direction: direction)
+        page = snapshot_export_page(field: field, direction: direction, cursor: cursor)
+        response.set_header(ExportCursor::NEXT_HEADER, page.next_cursor) if page.next_cursor
+        response.set_header(SnapshotPage::HEADER, page.snapshot_token) if page.snapshot_token
+        body = batched_export_batch_body(export_format, page.records, first_page: page.first_page)
+        render(plain: body, content_type: batch_content_type(export_format))
+      rescue ExportCursor::Invalid
+        render(plain: "Invalid export cursor", status: :bad_request)
+      rescue InvalidExportSnapshotError
+        render(
+          plain: I18n.t("active_admin.batched_export_page.unavailable_session"),
+          status: :bad_request
+        )
+      rescue ExportTooLargeError
+        render(
+          plain: I18n.t("active_admin.batched_export_page.over_max_rows"),
+          status: :bad_request
+        )
+      rescue UnresolvableExportColumnsError
+        render(
+          plain: I18n.t("active_admin.batched_export_page.select_at_least_one_column"),
+          status: :bad_request
+        )
+      rescue ExportMacroCatalog::UnknownMacroError
+        render(
+          plain: I18n.t("active_admin.batched_export_page.unknown_macro"),
+          status: :unprocessable_content
+        )
+      end
+
+      def render_export_workspace(export_format)
         ensure_batch_download_format_allowed!(export_format)
         @batched_export_format = export_format
         @batched_export_meta_url = batched_export_url_for(
@@ -57,20 +97,10 @@ module ActiveAdmin
         render "active_admin/batched_export/workspace", layout: "active_admin"
       end
 
-      private
-
-      def normalized_export_format
-        format_name = params[:export_format].to_s.downcase
-        format_name = "csv" if format_name.blank?
-        unless %w[csv xml json].include?(format_name)
-          render(plain: "Invalid export format", status: :bad_request)
-          return nil
-        end
-        format_name.to_sym
-      end
-
       def batched_export_url_for(request_format:, extra_params: {})
-        query = request.query_parameters.except(:format, :commit, :page, :batch_page, :export_meta)
+        query = request.query_parameters.except(
+          :format, :commit, :page, :batch_page, :export_meta, :export_cursor, :export_snapshot
+        )
         hash = query.respond_to?(:to_unsafe_h) ? query.to_unsafe_h : query.to_h
         hash = hash.merge(extra_params.stringify_keys)
         url_for(action: :batched_export, format: request_format, params: hash, only_path: true)
@@ -92,16 +122,34 @@ module ActiveAdmin
         first_page = paginate(base, 1, effective_batch_size)
         total_count = first_page.total_count
         total_batches = total_count.zero? ? 0 : first_page.total_pages
+        cap = BatchedExport.config.max_export_rows
         {
           export_format: export_format,
           total_count: total_count,
           total_batches: total_batches,
           batch_size: effective_batch_size,
           filename: export_filename(export_format),
-          large_export: total_count >= BatchedExport.config.large_export_row_threshold
+          large_export: total_count >= BatchedExport.config.large_export_row_threshold,
+          over_max: cap.present? && total_count > cap,
+          max_export_rows: cap
         }
       end
 
+      def export_filtered_count
+        paginate(find_collection(except: [:pagination]), 1, effective_batch_size).total_count
+      end
+
+      def starting_export_walk?
+        params[:export_snapshot].blank? && params[:export_cursor].blank?
+      end
+
+      def refuse_over_max_export_rows!
+        cap = BatchedExport.config.max_export_rows
+        return if cap.blank?
+
+        raise ExportTooLargeError if export_filtered_count > cap
+      end
+
       def export_filename(format_symbol)
         filename_proc = active_admin_config.batched_export_filename_proc
         if filename_proc
@@ -112,15 +160,6 @@ module ActiveAdmin
         "#{base}-#{Time.zone.now.to_date}.#{format_symbol}"
       end
 
-      def batched_export_batch_body(export_format, batch_page, page: nil)
-        case export_format
-        when :csv then batched_csv_chunk(batch_page, page: page)
-        when :json then batched_json_chunk(batch_page, page: page)
-        when :xml then batched_xml_chunk(batch_page, page: page)
-        else ""
-        end
-      end
-
       def batch_content_type(export_format)
         case export_format
         when :csv then "text/csv; charset=utf-8"
@@ -130,84 +169,46 @@ module ActiveAdmin
         end
       end
 
-      def batched_csv_chunk(batch_page, page: nil)
-        builder = active_admin_config.csv_builder
-        options = builder.options.dup
-        csv_options = options.except(:encoding_options, :humanize_name, :byte_order_mark)
-        columns = batched_export_filter_columns(builder.exec_columns(view_context))
-        buffer = +""
-        byte_order_mark = options[:byte_order_mark]
-        buffer << byte_order_mark if batch_page == 1 && byte_order_mark
-        if batch_page == 1 && options.fetch(:column_names, true)
-          header_line = columns.map do |column|
-            ActiveAdmin::Sanitizer.sanitize(builder.send(:encode, column.name, options))
-          end
-          buffer << CSV.generate_line(header_line, **csv_options)
-        end
-        paginated_export_rows(batch_page, page: page) do |resource|
-          row = builder.build_row(resource, columns, options)
-          row = apply_export_macros(row, columns, resource)
-          buffer << CSV.generate_line(row, **csv_options)
+      def batched_export_sort_pair
+        model = active_admin_config.resource_class
+        order_param = params[:order].presence || active_admin_config.sort_order
+        clause = ActiveAdmin::OrderClause.new(active_admin_config, order_param)
+        field = clause.valid? ? clause.field.to_s.split(".").last : nil
+        if field && model.column_names.include?(field)
+          [field, clause.order.to_s]
+        else
+          [model.primary_key.to_s, "desc"]
         end
-        buffer
       end
 
-      def batched_json_chunk(batch_page, page: nil)
-        builder = active_admin_config.csv_builder
-        options = builder.options
-        columns = batched_export_filter_columns(builder.exec_columns(view_context))
-        names = columns.map(&:name)
-        rows = []
-        paginated_export_rows(batch_page, page: page) do |resource|
-          row = apply_export_macros(builder.build_row(resource, columns, options), columns, resource)
-          rows << names.zip(row).to_h
-        end
-        rows.to_json
-      end
+      def decode_export_cursor(field:, direction:)
+        raw = params[:export_cursor]
+        return nil if raw.blank?
 
-      def batched_xml_chunk(batch_page, page: nil)
-        builder = active_admin_config.csv_builder
-        options = builder.options
-        columns = batched_export_filter_columns(builder.exec_columns(view_context))
-        xml = Builder::XmlMarkup.new(indent: 0)
-        paginated_export_rows(batch_page, page: page) do |resource|
-          row = apply_export_macros(builder.build_row(resource, columns, options), columns, resource)
-          xml.batch do
-            xml.record do
-              columns.each_with_index do |column, index|
-                xml.field("name" => column.name) { xml.text!(row[index].to_s) }
-              end
-            end
-          end
+        cursor = ExportCursor.decode(raw)
+        unless cursor.field == field && cursor.direction == direction
+          raise ExportCursor::Invalid, "mismatch"
         end
-        xml.target!
-      end
-
-      def apply_export_macros(row, columns, resource)
-        ExportMacroResolver.apply(
-          row: row,
-          columns: columns,
-          resource: resource,
-          resource_settings: active_admin_config.batched_export_settings,
-          registry: merged_macro_registry
-        )
-      end
 
-      def merged_macro_registry
-        BatchedExport.config.registered_macros.merge(ExportMacroCatalog.global_registry)
+        cursor
       end
 
-      def paginated_export_rows(batch_page, page: nil)
-        (page || page_relation(batch_page)).each do |resource|
-          yield apply_decorator(resource)
-        end
+      def snapshot_export_page(field:, direction:, cursor:)
+        SnapshotPage.fetch(
+          export_collection,
+          model: active_admin_config.resource_class,
+          field: field,
+          direction: direction,
+          cursor: cursor,
+          snapshot_param: params[:export_snapshot],
+          limit: effective_batch_size
+        )
       end
 
-      def page_relation(page)
+      def export_collection
         collection = find_collection(except: [:pagination])
         includes_list = active_admin_config.batched_export_includes
-        collection = collection.includes(includes_list) if includes_list.present?
-        paginate(collection, page, effective_batch_size)
+        includes_list.present? ? collection.includes(includes_list) : collection
       end
 
       def effective_batch_size
@@ -269,7 +270,9 @@ module ActiveAdmin
         return columns if indices.empty?
 
         resolved = indices.filter_map { |index| columns[index] }
-        resolved.presence || columns
+        raise UnresolvableExportColumnsError if resolved.empty?
+
+        resolved
       end
     end
   end

data/lib/activeadmin/batched_export/engine.rb

@@ -12,13 +12,23 @@ module ActiveAdmin
         require "activeadmin/batched_export/export_macro_resolver"
         require "activeadmin/batched_export/configuration"
         require "activeadmin/batched_export/resource_extension"
+        require "activeadmin/batched_export/errors"
+        require "activeadmin/batched_export/export_cursor"
+        require "activeadmin/batched_export/keyset_page"
+        require "activeadmin/batched_export/snapshot_row"
+        require "activeadmin/batched_export/snapshot_page"
+        require "activeadmin/batched_export/row_sanitizer"
+        require "activeadmin/batched_export/chunk_renderer"
         require "activeadmin/batched_export/controller_methods"
         require "activeadmin/batched_export/install"
       end
 
       initializer "activeadmin_batched_export.assets" do |app|
         assets_path = root.join("app/assets")
-        app.config.importmap.cache_sweepers << assets_path.join("controllers") if app.config.respond_to?(:importmap)
+        next unless app.config.respond_to?(:importmap)
+
+        app.config.importmap.cache_sweepers << assets_path.join("controllers")
+        app.config.importmap.cache_sweepers << assets_path.join("javascripts")
       end
 
       initializer "activeadmin_batched_export.importmap", after: :load_config_initializers do
@@ -27,6 +37,8 @@ module ActiveAdmin
         pin_controller = proc do |importmap|
           importmap.pin "controllers/activeadmin_batched_export/batched_export_controller",
             to: "activeadmin_batched_export/batched_export_controller.js"
+          importmap.pin "activeadmin_batched_export/chunk_assembly",
+            to: "activeadmin_batched_export/chunk_assembly.mjs"
         end
 
         Rails.application.importmap.draw(&pin_controller)

data/lib/activeadmin/batched_export/errors.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module ActiveAdmin
+  module BatchedExport
+    class UnresolvableExportColumnsError < StandardError; end
+    class ExportTooLargeError < StandardError; end
+    class InvalidExportSnapshotError < StandardError; end
+  end
+end

data/lib/activeadmin/batched_export/export_cursor.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require "base64"
+require "json"
+
+module ActiveAdmin
+  module BatchedExport
+    class ExportCursor
+      Invalid = Class.new(StandardError)
+      NEXT_HEADER = "X-Batched-Export-Next"
+
+      Payload = Struct.new(:field, :direction, :primary_key, :sort_value, :position)
+
+      def self.encode(field:, direction:, primary_key:, sort_value:, position: nil)
+        payload = {
+          "f" => field,
+          "d" => direction,
+          "k" => dump_value(primary_key),
+          "s" => dump_value(sort_value)
+        }
+        payload["p"] = position unless position.nil?
+        Base64.urlsafe_encode64(JSON.generate(payload), padding: false)
+      end
+
+      def self.decode(raw)
+        raise Invalid, "blank" if raw.nil? || raw.to_s.empty?
+
+        payload_from_hash(JSON.parse(Base64.urlsafe_decode64(raw.to_s)))
+      rescue ArgumentError, JSON::ParserError, KeyError, TypeError
+        raise Invalid
+      end
+
+      def self.dump_value(value)
+        case value
+        when Time, DateTime then value.iso8601(6)
+        when Date then value.iso8601
+        else value
+        end
+      end
+
+      def self.payload_from_hash(hash)
+        direction = hash.fetch("d")
+        raise Invalid, "direction" unless %w[asc desc].include?(direction)
+
+        Payload.new(
+          hash.fetch("f").to_s,
+          direction,
+          hash.fetch("k"),
+          hash.fetch("s"),
+          hash["p"]
+        )
+      end
+      private_class_method :payload_from_hash
+    end
+  end
+end

data/lib/activeadmin/batched_export/keyset_page.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module ActiveAdmin
+  module BatchedExport
+    class KeysetPage
+      def self.records(relation, model:, field:, direction:, cursor:, limit:)
+        scoped = ordered(relation, model: model, field: field, direction: direction)
+        scoped = apply_after(scoped, model: model, field: field, direction: direction, cursor: cursor) if cursor
+        scoped.limit(limit).to_a
+      end
+
+      def self.ordered(relation, model:, field:, direction:)
+        table = model.arel_table
+        primary_key = model.primary_key.to_s
+        primary_key_order = arel_direction(table[primary_key], direction)
+        return relation.except(:order).order(primary_key_order) if field == primary_key
+
+        relation.except(:order).order(
+          table[field].eq(nil),
+          arel_direction(table[field], direction),
+          primary_key_order
+        )
+      end
+
+      def self.next_cursor(records, field:, direction:, primary_key:, limit:, position: nil)
+        return nil if records.empty? || records.length < limit
+
+        last_record = records.last
+        ExportCursor.encode(
+          field: field,
+          direction: direction,
+          primary_key: last_record.public_send(primary_key),
+          sort_value: last_record.public_send(field),
+          position: position
+        )
+      end
+
+      def self.apply_after(relation, model:, field:, direction:, cursor:)
+        table = model.arel_table
+        primary_key = model.primary_key
+        primary_key_value = cast(model, primary_key, cursor.primary_key)
+        if field == primary_key.to_s
+          return after_primary_key(relation, table[primary_key], primary_key_value, direction)
+        end
+
+        after_sort_and_primary_key(
+          relation,
+          table[field],
+          cast(model, field, cursor.sort_value),
+          table[primary_key],
+          primary_key_value,
+          direction
+        )
+      end
+
+      def self.arel_direction(column, direction)
+        (direction == "desc") ? column.desc : column.asc
+      end
+
+      def self.after_primary_key(relation, column, value, direction)
+        comparator = (direction == "desc") ? :lt : :gt
+        relation.where(column.public_send(comparator, value))
+      end
+
+      def self.after_sort_and_primary_key(relation, sort_column, sort_value, pk_column, pk_value, direction)
+        pk_past = pk_column.public_send((direction == "desc") ? :lt : :gt, pk_value)
+        return relation.where(sort_column.eq(nil).and(pk_past)) if sort_value.nil?
+
+        sort_past = sort_column.public_send((direction == "desc") ? :lt : :gt, sort_value)
+        relation.where(
+          sort_past.and(sort_column.not_eq(nil))
+            .or(sort_column.eq(sort_value).and(pk_past))
+            .or(sort_column.eq(nil))
+        )
+      end
+
+      def self.cast(model, name, raw)
+        model.type_for_attribute(name).cast(raw)
+      end
+      private_class_method :apply_after, :arel_direction, :after_primary_key,
+        :after_sort_and_primary_key, :cast
+    end
+  end
+end

data/lib/activeadmin/batched_export/row_sanitizer.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module ActiveAdmin
+  module BatchedExport
+    module RowSanitizer
+      def self.apply(row)
+        row.map { |cell| ActiveAdmin::Sanitizer.sanitize(cell) }
+      end
+    end
+  end
+end

data/lib/activeadmin/batched_export/snapshot_page.rb

@@ -0,0 +1,140 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require "activeadmin/batched_export/errors"
+require "activeadmin/batched_export/export_cursor"
+require "activeadmin/batched_export/keyset_page"
+require "activeadmin/batched_export/snapshot_row"
+
+module ActiveAdmin
+  module BatchedExport
+    class SnapshotPage
+      HEADER = "X-Batched-Export-Snapshot"
+      INSERT_SLICE = 500
+      Page = Struct.new(:records, :snapshot_token, :next_cursor, :first_page)
+
+      def self.fetch(collection, model:, field:, direction:, cursor:, snapshot_param:, limit:)
+        if snapshot_param.present?
+          resume(collection, model: model, field: field, direction: direction,
+            cursor: cursor, snapshot_param: snapshot_param, limit: limit)
+        elsif SnapshotRow.available? && cursor.nil?
+          freeze_first(collection, model: model, field: field, direction: direction, limit: limit)
+        else
+          live(collection, model: model, field: field, direction: direction, cursor: cursor, limit: limit)
+        end
+      end
+
+      def self.freeze_first(collection, model:, field:, direction:, limit:)
+        SnapshotRow.expire_stale!(ttl: BatchedExport.config.snapshot_ttl)
+        # Full id list in process memory on freeze. Stream into insert_all when that ceiling matters.
+        ids = KeysetPage.ordered(collection, model: model, field: field, direction: direction)
+          .pluck(model.primary_key)
+        return Page.new(records: [], snapshot_token: nil, next_cursor: nil, first_page: true) if ids.empty?
+
+        token = insert_ids(ids, model.name)
+        walk(collection, model: model, field: field, direction: direction,
+          token: token, after_position: 0, limit: limit, first_page: true)
+      end
+
+      def self.resume(collection, model:, field:, direction:, cursor:, snapshot_param:, limit:)
+        raise InvalidExportSnapshotError unless SnapshotRow.available?
+        raise InvalidExportSnapshotError if cursor.nil? || cursor.position.nil?
+
+        token = authentic_token(snapshot_param, model.name)
+        walk(collection, model: model, field: field, direction: direction,
+          token: token, after_position: cursor.position, limit: limit, first_page: false)
+      end
+
+      def self.live(collection, model:, field:, direction:, cursor:, limit:)
+        records = KeysetPage.records(
+          collection, model: model, field: field, direction: direction, cursor: cursor, limit: limit
+        )
+        Page.new(
+          records: records,
+          snapshot_token: nil,
+          next_cursor: KeysetPage.next_cursor(
+            records, field: field, direction: direction, primary_key: model.primary_key, limit: limit
+          ),
+          first_page: cursor.nil?
+        )
+      end
+
+      def self.insert_ids(ids, resource_type)
+        token = SecureRandom.urlsafe_base64(32)
+        created_at = Time.current
+        ids.each_slice(INSERT_SLICE).with_index do |slice, slice_index|
+          SnapshotRow.insert_all(slice.each_with_index.map { |record_id, offset|
+            {
+              token: token, resource_type: resource_type, record_id: record_id,
+              position: (slice_index * INSERT_SLICE) + offset + 1, created_at: created_at
+            }
+          })
+        end
+        token
+      end
+
+      def self.authentic_token(raw, resource_type)
+        anchor = SnapshotRow.where(token: raw).order(:position).first
+        raise InvalidExportSnapshotError if anchor.nil?
+        raise InvalidExportSnapshotError if SnapshotRow.expired?(anchor.created_at, ttl: BatchedExport.config.snapshot_ttl)
+        raise InvalidExportSnapshotError unless anchor.resource_type == resource_type
+
+        anchor.token
+      end
+
+      def self.walk(collection, model:, field:, direction:, token:, after_position:, limit:, first_page:)
+        records, consumed = fill_live(collection, model: model, token: token,
+          after_position: after_position, limit: limit)
+        if SnapshotRow.remaining_after?(token, consumed) && records.length == limit
+          return Page.new(
+            records: records, snapshot_token: token, first_page: first_page,
+            next_cursor: KeysetPage.next_cursor(
+              records, field: field, direction: direction,
+              primary_key: model.primary_key, limit: limit, position: consumed
+            )
+          )
+        end
+
+        SnapshotRow.delete_token!(token)
+        Page.new(records: records, snapshot_token: nil, next_cursor: nil, first_page: first_page)
+      end
+
+      def self.fill_live(collection, model:, token:, after_position:, limit:)
+        live = []
+        consumed = after_position
+        position = after_position
+        loop do
+          window = SnapshotRow.window_after(token, position, limit)
+          break if window.empty?
+
+          found = live_by_id(collection, model.primary_key, window)
+          consumed = append_live(window, found, live, limit)
+          break if live.length == limit
+
+          position = window.last.position
+        end
+        [live, consumed]
+      end
+
+      def self.live_by_id(collection, primary_key, rows)
+        collection.where(primary_key => rows.map(&:record_id))
+          .index_by { |record| record.public_send(primary_key).to_i }
+      end
+
+      def self.append_live(window, found, live, limit)
+        consumed = nil
+        window.each do |row|
+          consumed = row.position
+          record = found[row.record_id.to_i]
+          next unless record
+
+          live << record
+          break if live.length == limit
+        end
+        consumed
+      end
+      private_class_method :freeze_first, :resume, :live, :insert_ids, :authentic_token,
+        :walk, :fill_live, :live_by_id, :append_live
+    end
+  end
+end

data/lib/activeadmin/batched_export/snapshot_row.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require "active_record"
+
+module ActiveAdmin
+  module BatchedExport
+    class SnapshotRow < ::ActiveRecord::Base
+      self.table_name = "active_admin_batched_export_snapshot_rows"
+      self.record_timestamps = false
+
+      def self.available?
+        table_exists?
+      end
+
+      def self.expire_stale!(ttl:)
+        # Global sweep on freeze start. Upgrade to a per-token job when volume matters.
+        where("created_at < ?", cutoff(ttl)).delete_all
+      end
+
+      def self.expired?(created_at, ttl:)
+        created_at < cutoff(ttl)
+      end
+
+      def self.cutoff(ttl)
+        ttl.respond_to?(:ago) ? ttl.ago : Time.current - ttl
+      end
+
+      def self.window_after(token, position, limit)
+        where(token: token).where("position > ?", position).order(:position).limit(limit).to_a
+      end
+
+      def self.remaining_after?(token, position)
+        where(token: token).where("position > ?", position).exists?
+      end
+
+      def self.delete_token!(token)
+        where(token: token).delete_all
+      end
+    end
+  end
+end