search-engine-for-typesense 1.0.0

data/lib/search_engine/sources/active_record_source.rb

@@ -0,0 +1,185 @@
+# frozen_string_literal: true
+
+module SearchEngine
+  module Sources
+    # ActiveRecord-backed source adapter that yields arrays of records in batches.
+    #
+    # Uses ORM batch APIs (`in_batches`/`find_in_batches`) honoring batch_size and an
+    # optional scope proc. Ensures stable memory by disabling query cache and using
+    # readonly relations. Does not accumulate results across batches.
+    #
+    # @example
+    #   src = SearchEngine::Sources::ActiveRecordSource.new(model: ::Product, scope: -> { where(active: true) }, batch_size: 2000)
+    #   src.each_batch(partition: nil, cursor: nil) { |rows| ... }
+    #
+    # @note Emits "search_engine.source.batch_fetched" and "search_engine.source.error".
+    class ActiveRecordSource
+      include Base
+
+      # @param model [Class] ActiveRecord model class
+      # @param scope [Proc, nil] optional proc evaluated in the context of model.all (returns a Relation)
+      # @param batch_size [Integer, nil] override batch size (defaults from config)
+      # @param use_transaction [Boolean, nil] wrap in read-only transaction (best-effort)
+      # @param readonly [Boolean, nil] mark relations readonly (default true)
+      def initialize(model:, scope: nil, batch_size: nil, use_transaction: nil, readonly: nil)
+        @model = model
+        @scope_proc = scope
+        cfg = SearchEngine.config.sources.active_record
+        @batch_size = (batch_size || cfg.batch_size).to_i
+        @use_transaction = use_transaction.nil? ? truthy?(cfg.use_transaction) : truthy?(use_transaction)
+        @readonly = !readonly.nil? ? truthy?(readonly) : truthy?(cfg.readonly)
+      end
+
+      # Iterate over batches of records.
+      #
+      # @param partition [Object, nil] optional opaque partition (e.g., id range)
+      # @param cursor [Object, nil] optional opaque cursor (e.g., last id)
+      # @yieldparam rows [Array<Object>] array of model instances
+      # @return [Enumerator] when no block is given
+      def each_batch(partition: nil, cursor: nil, &block)
+        return enum_for(:each_batch, partition: partition, cursor: cursor) unless block_given?
+
+        relation = base_relation
+        relation = apply_partition(relation, partition) if partition
+        relation = apply_cursor(relation, cursor) if cursor
+
+        idx = 0
+        started = monotonic_ms
+        run_with_connection do
+          run_in_readonly_txn_if_needed do
+            without_query_cache do
+              dispatch_batches(relation, started, idx, partition, cursor, &block)
+            end
+          end
+        end
+      rescue StandardError => error
+        instrument_error(source: 'active_record', error: error, partition: partition, cursor: cursor,
+                         adapter_options: { batch_size: @batch_size }
+        )
+        raise
+      end
+
+      private
+
+      def truthy?(val)
+        val == true
+      end
+
+      def dispatch_batches(relation, started, idx, partition, cursor)
+        if relation.respond_to?(:in_batches)
+          relation.in_batches(of: @batch_size) do |batch_scope|
+            batch_scope = mark_readonly(batch_scope)
+            rows = batch_scope.to_a
+            duration = monotonic_ms - started
+            instrument_batch_fetched(source: 'active_record', batch_index: idx, rows_count: rows.size,
+                                     duration_ms: duration, partition: partition, cursor: cursor,
+                                     adapter_options: { batch_size: @batch_size }
+            )
+            yield(rows)
+            idx += 1
+            started = monotonic_ms
+          end
+        elsif relation.respond_to?(:find_in_batches)
+          relation.find_in_batches(batch_size: @batch_size) do |rows|
+            rows = rows.map { |r| r }
+            duration = monotonic_ms - started
+            instrument_batch_fetched(source: 'active_record', batch_index: idx, rows_count: rows.size,
+                                     duration_ms: duration, partition: partition, cursor: cursor,
+                                     adapter_options: { batch_size: @batch_size }
+            )
+            yield(rows)
+            idx += 1
+            started = monotonic_ms
+          end
+        else
+          # Last resort: materialize and slice
+          records = relation.to_a
+          records.each_slice(@batch_size) do |rows|
+            duration = monotonic_ms - started
+            instrument_batch_fetched(source: 'active_record', batch_index: idx, rows_count: rows.size,
+                                     duration_ms: duration, partition: partition, cursor: cursor,
+                                     adapter_options: { batch_size: @batch_size }
+            )
+            yield(rows)
+            idx += 1
+            started = monotonic_ms
+          end
+        end
+      end
+
+      def base_relation
+        rel = @model.all
+        rel = rel.instance_exec(&@scope_proc) if @scope_proc.respond_to?(:call)
+        mark_readonly(rel)
+      end
+
+      def mark_readonly(rel)
+        if @readonly && rel.respond_to?(:readonly)
+          rel.readonly(true)
+        else
+          rel
+        end
+      end
+
+      def apply_partition(rel, partition)
+        case partition
+        when Range
+          pk = rel.klass.primary_key
+          rel.where(pk => partition)
+        when Hash
+          rel.where(partition)
+        else
+          rel
+        end
+      end
+
+      def apply_cursor(rel, cursor)
+        return rel unless cursor
+
+        pk = rel.klass.primary_key
+        if cursor.is_a?(Hash)
+          rel.where(cursor)
+        else
+          rel.where(arel_table(rel)[pk].gt(cursor)).order(pk => :asc)
+        end
+      end
+
+      def arel_table(rel)
+        rel.klass.arel_table
+      end
+
+      def run_with_connection(&block)
+        if defined?(ActiveRecord::Base)
+          ActiveRecord::Base.connection_pool.with_connection(&block)
+        else
+          yield
+        end
+      end
+
+      def run_in_readonly_txn_if_needed
+        if @use_transaction && defined?(ActiveRecord::Base)
+          ActiveRecord::Base.connection.transaction(requires_new: true) do
+            if ActiveRecord::Base.connection.respond_to?(:execute)
+              begin
+                ActiveRecord::Base.connection.execute('SET TRANSACTION READ ONLY')
+              rescue StandardError
+                # best-effort; ignore if not supported
+              end
+            end
+            yield
+          end
+        else
+          yield
+        end
+      end
+
+      def without_query_cache(&block)
+        if defined?(ActiveRecord::Base) && ActiveRecord::Base.respond_to?(:connection)
+          ActiveRecord::Base.connection.uncached(&block)
+        else
+          yield
+        end
+      end
+    end
+  end
+end

data/lib/search_engine/sources/base.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module SearchEngine
+  module Sources
+    # Internal helpers shared by source adapters.
+    #
+    # Provides small utilities for instrumentation and for returning
+    # an Enumerator when a block is not supplied.
+    module Base
+      private
+
+      def monotonic_ms
+        SearchEngine::Instrumentation.monotonic_ms
+      end
+
+      def instrument_batch_fetched(source:, batch_index:, rows_count:, duration_ms:, partition: nil, cursor: nil,
+                                   adapter_options: {})
+        return unless defined?(ActiveSupport::Notifications)
+
+        payload = {
+          source: source,
+          batch_index: batch_index,
+          rows_count: rows_count,
+          duration_ms: duration_ms
+        }
+        payload[:partition] = partition unless partition.nil?
+        payload[:cursor] = cursor unless cursor.nil?
+        if defined?(SearchEngine::Observability)
+          payload[:adapter_options] = SearchEngine::Observability.redact(adapter_options)
+        end
+        SearchEngine::Instrumentation.instrument('search_engine.source.batch_fetched', payload) {}
+      end
+
+      def instrument_error(source:, error:, partition: nil, cursor: nil, adapter_options: {})
+        return unless defined?(ActiveSupport::Notifications)
+
+        payload = {
+          source: source,
+          error_class: error.class.name,
+          message: error.message.to_s[0, 200]
+        }
+        payload[:partition] = partition unless partition.nil?
+        payload[:cursor] = cursor unless cursor.nil?
+        if defined?(SearchEngine::Observability)
+          payload[:adapter_options] = SearchEngine::Observability.redact(adapter_options)
+        end
+        if error.respond_to?(:to_h)
+          h = error.to_h
+          payload[:error_hint] = h[:hint] if h[:hint]
+          payload[:error_doc] = h[:doc] if h[:doc]
+        end
+        SearchEngine::Instrumentation.instrument('search_engine.source.error', payload) {}
+      end
+
+      def enum_for_each_batch(partition:, cursor:)
+        return to_enum(:each_batch, partition: partition, cursor: cursor) unless block_given?
+
+        yield
+      end
+    end
+  end
+end

data/lib/search_engine/sources/lambda_source.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module SearchEngine
+  module Sources
+    # Adapter that delegates batch enumeration to a provided callable.
+    #
+    # The callable is expected to implement `call(cursor:, partition:)` and return either
+    # an Enumerator or yield arrays of rows. Shapes are application-defined.
+    #
+    # @example
+    #   src = SearchEngine::Sources::LambdaSource.new(->(cursor:, partition:) { [[row1, row2]] })
+    #   src.each_batch { |rows| ... }
+    #
+    # @see `https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/indexer`
+    class LambdaSource
+      include Base
+
+      # @param callable [#call] object responding to call(cursor:, partition:)
+      # @raise [ArgumentError] when callable does not respond to :call
+      def initialize(callable)
+        raise ArgumentError, 'callable must respond to :call(cursor:, partition:)' unless callable.respond_to?(:call)
+
+        @callable = callable
+      end
+
+      # Enumerate batches produced by the callable.
+      # @param partition [Object, nil]
+      # @param cursor [Object, nil]
+      # @yieldparam rows [Array]
+      # @return [Enumerator]
+      def each_batch(partition: nil, cursor: nil)
+        return enum_for(:each_batch, partition: partition, cursor: cursor) unless block_given?
+
+        started = monotonic_ms
+        begin
+          enum = @callable.call(cursor: cursor, partition: partition)
+          Array(enum).each do |rows|
+            duration = monotonic_ms - started
+            instrument_batch_fetched(source: 'lambda', batch_index: nil, rows_count: Array(rows).size,
+                                     duration_ms: duration, partition: partition, cursor: cursor,
+                                     adapter_options: { callable: @callable.class.name }
+            )
+            yield(rows)
+            started = monotonic_ms
+          end
+        rescue StandardError => error
+          instrument_error(source: 'lambda', error: error, partition: partition, cursor: cursor,
+                           adapter_options: { callable: @callable.class.name }
+          )
+          raise
+        end
+      end
+    end
+  end
+end

data/lib/search_engine/sources/sql_source.rb

@@ -0,0 +1,196 @@
+# frozen_string_literal: true
+
+module SearchEngine
+  module Sources
+    # Stream rows from SQL using ActiveRecord connection, with PG cursor support when available.
+    #
+    # Supports large result sets via server-side cursors on PostgreSQL connections and falls
+    # back to paginated SELECT for other adapters. Yields arrays of rows in batches.
+    #
+    # @example
+    #   src = SearchEngine::Sources::SqlSource.new(sql: "SELECT id, name FROM users", fetch_size: 2000)
+    #   src.each_batch { |rows| ... }
+    #
+    # @note Emits "search_engine.source.batch_fetched" and "search_engine.source.error".
+    # @see `https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/indexer`
+    class SqlSource
+      include Base
+
+      # @param sql [String] SQL statement to execute
+      # @param binds [Array, Hash, nil] optional parameters for the SQL
+      # @param fetch_size [Integer, nil] batch size hint
+      # @param row_shape [Symbol, nil] :hash, :array or :auto (PG only)
+      # @param statement_timeout_ms [Integer, nil] optional statement timeout (PG only)
+      def initialize(sql:, binds: nil, fetch_size: nil, row_shape: nil, statement_timeout_ms: nil)
+        @sql = sql.to_s
+        @binds = binds
+        cfg = SearchEngine.config.sources.sql
+        @fetch_size = (fetch_size || cfg.fetch_size).to_i
+        @row_shape = row_shape || :auto
+        @statement_timeout_ms = statement_timeout_ms
+      end
+
+      # Iterate over batches produced by the SQL query.
+      # @param partition [Object, nil]
+      # @param cursor [Object, nil]
+      # @yieldparam rows [Array<Hash>, Array<Array>]
+      # @return [Enumerator]
+      def each_batch(partition: nil, cursor: nil, &block)
+        return enum_for(:each_batch, partition: partition, cursor: cursor) unless block_given?
+
+        run_with_connection do
+          if postgres_connection?(ActiveRecord::Base.connection.raw_connection)
+            stream_postgres(ActiveRecord::Base.connection.raw_connection, partition: partition, cursor: cursor, &block)
+          else
+            stream_generic(nil, partition: partition, cursor: cursor, &block)
+          end
+        end
+      rescue StandardError => error
+        instrument_error(source: 'sql', error: error, partition: partition, cursor: cursor,
+                         adapter_options: { fetch_size: @fetch_size, row_shape: @row_shape }
+        )
+        raise
+      end
+
+      private
+
+      def run_with_connection
+        unless defined?(ActiveRecord::Base)
+          raise SearchEngine::Errors::InvalidParams, 'SqlSource requires ActiveRecord connection'
+        end
+
+        ActiveRecord::Base.connection_pool.with_connection do |ar_conn|
+          raw = raw_connection(ar_conn)
+          yield raw
+        end
+      end
+
+      def raw_connection(ar_conn)
+        if ar_conn.respond_to?(:raw_connection)
+          ar_conn.raw_connection
+        else
+          ar_conn
+        end
+      end
+
+      def postgres_connection?(conn)
+        (defined?(PG) && conn.is_a?(PG::Connection)) || conn.class.name.include?('PG')
+      end
+
+      def stream_postgres(conn, partition:, cursor:)
+        cursor_name = "se_cursor_#{object_id}"
+        sql, params = build_sql_and_params(partition: partition, cursor: cursor)
+        started = monotonic_ms
+        begin
+          set_statement_timeout(conn, @statement_timeout_ms) if @statement_timeout_ms
+          # Use unnamed prepared statement + DECLARE CURSOR for streaming
+          conn.exec('BEGIN READ ONLY')
+          conn.exec_params("DECLARE #{cursor_name} NO SCROLL CURSOR FOR #{sql}", params)
+          idx = 0
+          loop do
+            res = conn.exec("FETCH FORWARD #{@fetch_size} FROM #{cursor_name}")
+            break if res.ntuples.zero?
+
+            rows = rows_from_pg_result(res)
+            duration = monotonic_ms - started
+            instrument_batch_fetched(source: 'sql', batch_index: idx, rows_count: rows.size, duration_ms: duration,
+                                     partition: partition, cursor: cursor,
+                                     adapter_options: { fetch_size: @fetch_size, row_shape: @row_shape }
+            )
+            yield rows
+            idx += 1
+            started = monotonic_ms
+          end
+        ensure
+          begin
+            conn.exec("CLOSE #{cursor_name}")
+          rescue StandardError
+            # ignore
+          end
+          begin
+            conn.exec('COMMIT')
+          rescue StandardError
+            # ignore
+          end
+          reset_statement_timeout(conn) if @statement_timeout_ms
+        end
+      end
+
+      def rows_from_pg_result(res)
+        if @row_shape == :hash || @row_shape == :auto
+          res.to_a
+        else
+          res.values
+        end
+      end
+
+      def stream_generic(_conn, partition:, cursor:)
+        # Fallback: try ActiveRecord select_all with pagination via placeholders
+        ar_conn = ActiveRecord::Base.connection
+        sql, params = build_sql_and_params(partition: partition, cursor: cursor)
+        idx = 0
+        started = monotonic_ms
+        loop do
+          chunked_sql = sql_with_limit(sql, @fetch_size, idx)
+          rows = ar_conn.exec_query(chunked_sql, 'SqlSource', params_for_ar(params)).to_a
+          break if rows.empty?
+
+          duration = monotonic_ms - started
+          instrument_batch_fetched(source: 'sql', batch_index: idx, rows_count: rows.size, duration_ms: duration,
+                                   partition: partition, cursor: cursor,
+                                   adapter_options: { fetch_size: @fetch_size, row_shape: :hash }
+          )
+          yield rows
+          idx += 1
+          started = monotonic_ms
+        end
+      end
+
+      def sql_with_limit(base_sql, fetch_size, page_idx)
+        offset = page_idx * fetch_size
+        "SELECT * FROM (#{base_sql}) se_sub LIMIT #{Integer(fetch_size)} OFFSET #{Integer(offset)}"
+      end
+
+      def params_for_ar(params)
+        return [] if params.nil?
+
+        return params if params.is_a?(Array)
+
+        []
+      end
+
+      def build_sql_and_params(**)
+        # For safety, do not mutate the original SQL. Bind support is adapter-specific.
+        # We support a simple Hash expansion for named placeholders in PG `exec_params` style: $1, $2 ...
+        sql = @sql.dup
+        binds = []
+        case @binds
+        when Array
+          binds = @binds.dup
+        when Hash
+          raise SearchEngine::Errors::InvalidParams,
+                'SqlSource with Hash binds is not supported; use positional binds (Array)'
+        when nil
+          # noop
+        else
+          raise SearchEngine::Errors::InvalidParams, 'SqlSource binds must be an Array or nil'
+        end
+
+        # Partition/cursor semantics are adapter/domain-specific; callers should incorporate placeholders.
+        [sql, binds]
+      end
+
+      def set_statement_timeout(conn, ms)
+        conn.exec_params('SET LOCAL statement_timeout = $1', [Integer(ms)])
+      rescue StandardError
+        # ignore
+      end
+
+      def reset_statement_timeout(conn)
+        conn.exec('RESET statement_timeout')
+      rescue StandardError
+        # ignore
+      end
+    end
+  end
+end

data/lib/search_engine/sources.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module SearchEngine
+  # Factory and DSL for building data source adapters that yield batches.
+  #
+  # Usage via symbol and options:
+  #   SearchEngine::Sources.build(:active_record, model: ::Product, scope: -> { where(active: true) }, batch_size: 2000)
+  #   SearchEngine::Sources.build(:sql, sql: "SELECT * FROM products WHERE active = TRUE", fetch_size: 2000)
+  #
+  # Usage via block (lambda source):
+  #   SearchEngine::Sources.build(:lambda) do |cursor:, partition:|
+  #     Enumerator.new { |y| external_api.each_page(cursor) { |rows| y << rows } }
+  #   end
+  #
+  # All adapters implement `each_batch(partition:, cursor:)` and return an Enumerator
+  # when no block is provided.
+  module Sources
+    # Build a source adapter from a symbol and options or from a block.
+    #
+    # @param type [Symbol] :active_record, :sql, or :lambda
+    # @param options [Hash] adapter-specific options
+    # @yield for :lambda sources, a block taking (cursor:, partition:) and returning an Enumerator
+    # @return [Object] adapter responding to `each_batch(partition:, cursor:)`
+    def self.build(type, **options, &block)
+      case type.to_sym
+      when :active_record
+        model = options[:model]
+        unless model.is_a?(Class)
+          raise SearchEngine::Errors::InvalidParams,
+                'active_record source requires :model (ActiveRecord class). See ' \
+                'https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/indexer.'
+        end
+
+        scope = options[:scope]
+        batch_size = options[:batch_size]
+        readonly = options[:readonly]
+        use_txn = options[:use_transaction]
+        ActiveRecordSource.new(model: model, scope: scope, batch_size: batch_size, use_transaction: use_txn,
+                               readonly: readonly
+        )
+      when :sql
+        sql = options[:sql]
+        unless sql.is_a?(String) && !sql.strip.empty?
+          raise SearchEngine::Errors::InvalidParams,
+                'sql source requires :sql (String). See ' \
+                'https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/indexer.'
+        end
+
+        binds = options[:binds]
+        fetch_size = options[:fetch_size]
+        row_shape = options[:row_shape]
+        stmt_timeout = options[:statement_timeout_ms]
+        SqlSource.new(sql: sql, binds: binds, fetch_size: fetch_size, row_shape: row_shape,
+                      statement_timeout_ms: stmt_timeout
+        )
+      when :lambda
+        callable = block || options[:callable]
+        unless callable
+          raise SearchEngine::Errors::InvalidParams,
+                'lambda source requires a block or :callable. See ' \
+                'https://nikita-shkoda.mintlify.app/projects/search-engine-for-typesense/indexer.'
+        end
+
+        LambdaSource.new(callable)
+      else
+        raise SearchEngine::Errors::InvalidParams,
+              "unknown source type: #{type.inspect}. Supported: :active_record, :sql, :lambda"
+      end
+    end
+  end
+end