vectra-client 0.2.2 → 0.3.1

data/lib/vectra/pool.rb

@@ -0,0 +1,256 @@
+# frozen_string_literal: true
+
+require "concurrent"
+
+module Vectra
+  # Connection pool with warmup support
+  #
+  # Provides connection pooling for database providers with configurable
+  # pool size, timeout, and connection warmup.
+  #
+  # @example Basic usage
+  #   pool = Vectra::Pool.new(size: 5, timeout: 5) { create_connection }
+  #   pool.warmup(3) # Pre-create 3 connections
+  #
+  #   pool.with_connection do |conn|
+  #     conn.execute("SELECT 1")
+  #   end
+  #
+  class Pool
+    class TimeoutError < Vectra::Error; end
+    class PoolExhaustedError < Vectra::Error; end
+
+    attr_reader :size, :timeout
+
+    # Initialize connection pool
+    #
+    # @param size [Integer] maximum pool size
+    # @param timeout [Integer] checkout timeout in seconds
+    # @yield connection factory block
+    def initialize(size:, timeout: 5, &block)
+      raise ArgumentError, "Connection factory block required" unless block_given?
+
+      @size = size
+      @timeout = timeout
+      @factory = block
+      @pool = Concurrent::Array.new
+      @checked_out = Concurrent::AtomicFixnum.new(0)
+      @mutex = Mutex.new
+      @condition = ConditionVariable.new
+      @shutdown = false
+    end
+
+    # Warmup the pool by pre-creating connections
+    #
+    # @param count [Integer] number of connections to create (default: pool size)
+    # @return [Integer] number of connections created
+    def warmup(count = nil)
+      count ||= size
+      count = [count, size].min
+      created = 0
+
+      count.times do
+        break if @pool.size >= size
+
+        conn = create_connection
+        if conn
+          @pool << conn
+          created += 1
+        end
+      end
+
+      created
+    end
+
+    # Execute block with a connection from the pool
+    #
+    # @yield [connection] the checked out connection
+    # @return [Object] result of the block
+    def with_connection
+      conn = checkout
+      begin
+        yield conn
+      ensure
+        checkin(conn)
+      end
+    end
+
+    # Checkout a connection from the pool
+    #
+    # @return [Object] a connection
+    # @raise [TimeoutError] if checkout times out
+    # @raise [PoolExhaustedError] if pool is exhausted
+    def checkout
+      raise PoolExhaustedError, "Pool has been shutdown" if @shutdown
+
+      deadline = Time.now + timeout
+
+      @mutex.synchronize do
+        loop do
+          # Try to get an existing connection
+          conn = @pool.pop
+          if conn
+            @checked_out.increment
+            return conn if healthy?(conn)
+
+            # Connection is unhealthy, discard and try again
+            close_connection(conn)
+            next
+          end
+
+          # Try to create a new connection if under limit
+          if @checked_out.value + @pool.size < size
+            conn = create_connection
+            if conn
+              @checked_out.increment
+              return conn
+            end
+          end
+
+          # Wait for a connection to be returned
+          remaining = deadline - Time.now
+          raise TimeoutError, "Connection checkout timed out after #{timeout}s" if remaining <= 0
+
+          @condition.wait(@mutex, remaining)
+        end
+      end
+    end
+
+    # Return a connection to the pool
+    #
+    # @param connection [Object] connection to return
+    def checkin(connection)
+      return if @shutdown
+
+      @mutex.synchronize do
+        @checked_out.decrement
+        if healthy?(connection) && @pool.size < size
+          @pool << connection
+        else
+          close_connection(connection)
+        end
+        @condition.signal
+      end
+    end
+
+    # Shutdown the pool, closing all connections
+    #
+    # @return [void]
+    def shutdown
+      @shutdown = true
+      @mutex.synchronize do
+        while (conn = @pool.pop)
+          close_connection(conn)
+        end
+      end
+    end
+
+    # Get pool statistics
+    #
+    # @return [Hash] pool stats
+    def stats
+      {
+        size: size,
+        available: @pool.size,
+        checked_out: @checked_out.value,
+        total_created: @pool.size + @checked_out.value,
+        shutdown: @shutdown
+      }
+    end
+
+    # Check if pool is healthy (public method)
+    #
+    # @return [Boolean]
+    def pool_healthy?
+      !@shutdown && (@pool.size + @checked_out.value).positive?
+    end
+
+    private
+
+    # Internal health check for individual connections
+    def healthy?(conn)
+      return false if conn.nil?
+      return true unless conn.respond_to?(:status)
+
+      # For PG connections, check status. Otherwise assume healthy.
+      if defined?(PG::CONNECTION_OK)
+        conn.status == PG::CONNECTION_OK
+      else
+        # If PG not loaded, assume connection is healthy if it exists
+        true
+      end
+    rescue StandardError
+      false
+    end
+
+    def create_connection
+      @factory.call
+    rescue StandardError => e
+      Vectra.configuration.logger&.error("Pool: Failed to create connection: #{e.message}")
+      nil
+    end
+
+    def close_connection(conn)
+      conn.close if conn.respond_to?(:close)
+    rescue StandardError => e
+      Vectra.configuration.logger&.warn("Pool: Error closing connection: #{e.message}")
+    end
+  end
+
+  # Pooled connection module for pgvector
+  module PooledConnection
+    # Get a pooled connection
+    #
+    # @return [Pool] connection pool
+    def connection_pool
+      @connection_pool ||= create_pool
+    end
+
+    # Warmup the connection pool
+    #
+    # @param count [Integer] number of connections to pre-create
+    # @return [Integer] connections created
+    def warmup_pool(count = nil)
+      connection_pool.warmup(count)
+    end
+
+    # Execute with pooled connection
+    #
+    # @yield [connection] database connection
+    def with_pooled_connection(&)
+      connection_pool.with_connection(&)
+    end
+
+    # Shutdown the connection pool
+    def shutdown_pool
+      @connection_pool&.shutdown
+      @connection_pool = nil
+    end
+
+    # Get pool statistics
+    #
+    # @return [Hash]
+    def pool_stats
+      return { status: "not_initialized" } unless @connection_pool
+
+      connection_pool.stats
+    end
+
+    private
+
+    def create_pool
+      pool_size = config.pool_size || 5
+      pool_timeout = config.pool_timeout || 5
+
+      Pool.new(size: pool_size, timeout: pool_timeout) do
+        create_raw_connection
+      end
+    end
+
+    def create_raw_connection
+      require "pg"
+      conn_params = parse_connection_params
+      PG.connect(conn_params)
+    end
+  end
+end

data/lib/vectra/rate_limiter.rb

@@ -0,0 +1,304 @@
+# frozen_string_literal: true
+
+module Vectra
+  # Proactive rate limiter using token bucket algorithm
+  #
+  # Throttles requests BEFORE sending to prevent rate limit errors from providers.
+  # Uses token bucket algorithm for smooth rate limiting with burst support.
+  #
+  # @example Basic usage
+  #   limiter = Vectra::RateLimiter.new(
+  #     requests_per_second: 10,
+  #     burst_size: 20
+  #   )
+  #
+  #   limiter.acquire do
+  #     client.query(...)
+  #   end
+  #
+  # @example With client wrapper
+  #   client = Vectra::RateLimitedClient.new(
+  #     Vectra::Client.new(...),
+  #     requests_per_second: 10
+  #   )
+  #   client.query(...)  # Automatically rate limited
+  #
+  class RateLimiter
+    class RateLimitExceededError < Vectra::Error
+      attr_reader :wait_time
+
+      def initialize(wait_time:)
+        @wait_time = wait_time
+        super("Rate limit exceeded. Retry after #{wait_time.round(2)} seconds")
+      end
+    end
+
+    attr_reader :requests_per_second, :burst_size
+
+    # Initialize rate limiter
+    #
+    # @param requests_per_second [Float] Sustained request rate
+    # @param burst_size [Integer] Maximum burst capacity (default: 2x RPS)
+    def initialize(requests_per_second:, burst_size: nil)
+      @requests_per_second = requests_per_second.to_f
+      @burst_size = burst_size || (requests_per_second * 2).to_i
+      @tokens = @burst_size.to_f
+      @last_refill = Time.now
+      @mutex = Mutex.new
+    end
+
+    # Acquire a token and execute block
+    #
+    # @param wait [Boolean] Wait for token if not available (default: true)
+    # @param timeout [Float] Maximum wait time in seconds (default: 30)
+    # @yield Block to execute after acquiring token
+    # @return [Object] Result of block
+    # @raise [RateLimitExceededError] If wait is false and no token available
+    def acquire(wait: true, timeout: 30, &)
+      acquired = try_acquire(wait: wait, timeout: timeout)
+
+      unless acquired
+        wait_time = time_until_token
+        raise RateLimitExceededError.new(wait_time: wait_time)
+      end
+
+      yield
+    end
+
+    # Try to acquire a token without blocking
+    #
+    # @return [Boolean] true if token acquired
+    def try_acquire(wait: false, timeout: 30)
+      deadline = Time.now + timeout
+
+      loop do
+        @mutex.synchronize do
+          refill_tokens
+          if @tokens >= 1
+            @tokens -= 1
+            return true
+          end
+        end
+
+        return false unless wait
+        return false if Time.now >= deadline
+
+        # Wait a bit before retrying
+        sleep([time_until_token, 0.1].min)
+      end
+    end
+
+    # Get current token count
+    #
+    # @return [Float]
+    def available_tokens
+      @mutex.synchronize do
+        refill_tokens
+        @tokens
+      end
+    end
+
+    # Get time until next token is available
+    #
+    # @return [Float] Seconds until next token
+    def time_until_token
+      @mutex.synchronize do
+        refill_tokens
+        return 0 if @tokens >= 1
+
+        tokens_needed = 1 - @tokens
+        tokens_needed / @requests_per_second
+      end
+    end
+
+    # Get rate limiter statistics
+    #
+    # @return [Hash]
+    def stats
+      @mutex.synchronize do
+        refill_tokens
+        {
+          requests_per_second: @requests_per_second,
+          burst_size: @burst_size,
+          available_tokens: @tokens.round(2),
+          time_until_token: @tokens >= 1 ? 0 : ((1 - @tokens) / @requests_per_second).round(3)
+        }
+      end
+    end
+
+    # Reset the rate limiter to full capacity
+    #
+    # @return [void]
+    def reset!
+      @mutex.synchronize do
+        @tokens = @burst_size.to_f
+        @last_refill = Time.now
+      end
+    end
+
+    private
+
+    def refill_tokens
+      now = Time.now
+      elapsed = now - @last_refill
+      @last_refill = now
+
+      # Add tokens based on elapsed time
+      new_tokens = elapsed * @requests_per_second
+      @tokens = [@tokens + new_tokens, @burst_size.to_f].min
+    end
+  end
+
+  # Client wrapper with automatic rate limiting
+  #
+  # @example
+  #   client = Vectra::RateLimitedClient.new(
+  #     Vectra::Client.new(provider: :pinecone, ...),
+  #     requests_per_second: 10,
+  #     burst_size: 20
+  #   )
+  #
+  #   client.query(...)  # Automatically waits if rate limited
+  #
+  class RateLimitedClient
+    RATE_LIMITED_METHODS = [:upsert, :query, :fetch, :update, :delete].freeze
+
+    attr_reader :client, :limiter
+
+    # Initialize rate limited client
+    #
+    # @param client [Client] The underlying Vectra client
+    # @param requests_per_second [Float] Request rate limit
+    # @param burst_size [Integer] Burst capacity
+    # @param wait [Boolean] Wait for rate limit (default: true)
+    def initialize(client, requests_per_second:, burst_size: nil, wait: true)
+      @client = client
+      @limiter = RateLimiter.new(
+        requests_per_second: requests_per_second,
+        burst_size: burst_size
+      )
+      @wait = wait
+    end
+
+    # Rate-limited upsert
+    def upsert(...)
+      with_rate_limit { client.upsert(...) }
+    end
+
+    # Rate-limited query
+    def query(...)
+      with_rate_limit { client.query(...) }
+    end
+
+    # Rate-limited fetch
+    def fetch(...)
+      with_rate_limit { client.fetch(...) }
+    end
+
+    # Rate-limited update
+    def update(...)
+      with_rate_limit { client.update(...) }
+    end
+
+    # Rate-limited delete
+    def delete(...)
+      with_rate_limit { client.delete(...) }
+    end
+
+    # Pass through other methods without rate limiting
+    def method_missing(method, *, **, &)
+      if client.respond_to?(method)
+        client.public_send(method, *, **, &)
+      else
+        super
+      end
+    end
+
+    def respond_to_missing?(method, include_private = false)
+      client.respond_to?(method, include_private) || super
+    end
+
+    # Get rate limiter stats
+    #
+    # @return [Hash]
+    def rate_limit_stats
+      limiter.stats
+    end
+
+    private
+
+    def with_rate_limit(&)
+      limiter.acquire(wait: @wait, &)
+    end
+  end
+
+  # Per-provider rate limiter registry
+  #
+  # @example
+  #   # Configure rate limits per provider
+  #   Vectra::RateLimiterRegistry.configure(:pinecone, requests_per_second: 100)
+  #   Vectra::RateLimiterRegistry.configure(:qdrant, requests_per_second: 50)
+  #
+  #   # Get limiter for provider
+  #   limiter = Vectra::RateLimiterRegistry[:pinecone]
+  #   limiter.acquire { ... }
+  #
+  module RateLimiterRegistry
+    class << self
+      # Configure rate limiter for provider
+      #
+      # @param provider [Symbol] Provider name
+      # @param requests_per_second [Float] Request rate
+      # @param burst_size [Integer] Burst capacity
+      # @return [RateLimiter]
+      def configure(provider, requests_per_second:, burst_size: nil)
+        limiters[provider.to_sym] = RateLimiter.new(
+          requests_per_second: requests_per_second,
+          burst_size: burst_size
+        )
+      end
+
+      # Get limiter for provider
+      #
+      # @param provider [Symbol] Provider name
+      # @return [RateLimiter, nil]
+      def [](provider)
+        limiters[provider.to_sym]
+      end
+
+      # Get all configured limiters
+      #
+      # @return [Hash<Symbol, RateLimiter>]
+      def all
+        limiters.dup
+      end
+
+      # Get stats for all limiters
+      #
+      # @return [Hash<Symbol, Hash>]
+      def stats
+        limiters.transform_values(&:stats)
+      end
+
+      # Reset all limiters
+      #
+      # @return [void]
+      def reset_all!
+        limiters.each_value(&:reset!)
+      end
+
+      # Clear all limiters
+      #
+      # @return [void]
+      def clear!
+        @limiters = {}
+      end
+
+      private
+
+      def limiters
+        @limiters ||= {}
+      end
+    end
+  end
+end

data/lib/vectra/streaming.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+module Vectra
+  # Streaming query results for large datasets
+  #
+  # Provides lazy enumeration over query results with automatic pagination,
+  # reducing memory usage for large result sets.
+  #
+  # @example Stream through results
+  #   stream = Vectra::Streaming.new(client)
+  #   stream.query_each(index: 'my-index', vector: query_vec, total: 1000) do |match|
+  #     process(match)
+  #   end
+  #
+  # @example As lazy enumerator
+  #   results = stream.query_stream(index: 'my-index', vector: query_vec, total: 1000)
+  #   results.take(50).each { |m| puts m.id }
+  #
+  class Streaming
+    DEFAULT_PAGE_SIZE = 100
+
+    attr_reader :client, :page_size
+
+    # Initialize streaming query handler
+    #
+    # @param client [Client] the Vectra client
+    # @param page_size [Integer] results per page (default: 100)
+    def initialize(client, page_size: DEFAULT_PAGE_SIZE)
+      @client = client
+      @page_size = [page_size, 1].max
+    end
+
+    # Stream query results with a block
+    #
+    # @param index [String] the index name
+    # @param vector [Array<Float>] query vector
+    # @param total [Integer] total results to fetch
+    # @param namespace [String, nil] optional namespace
+    # @param filter [Hash, nil] metadata filter
+    # @yield [Match] each match result
+    # @return [Integer] total results yielded
+    def query_each(index:, vector:, total:, namespace: nil, filter: nil, &block)
+      return 0 unless block_given?
+
+      count = 0
+      query_stream(index: index, vector: vector, total: total, namespace: namespace, filter: filter).each do |match|
+        block.call(match)
+        count += 1
+      end
+      count
+    end
+
+    # Create a lazy enumerator for streaming results
+    #
+    # @param index [String] the index name
+    # @param vector [Array<Float>] query vector
+    # @param total [Integer] total results to fetch
+    # @param namespace [String, nil] optional namespace
+    # @param filter [Hash, nil] metadata filter
+    # @return [Enumerator::Lazy] lazy enumerator of results
+    def query_stream(index:, vector:, total:, namespace: nil, filter: nil)
+      Enumerator.new do |yielder|
+        fetched = 0
+        seen_ids = Set.new
+
+        while fetched < total
+          batch_size = [page_size, total - fetched].min
+
+          result = client.query(
+            index: index,
+            vector: vector,
+            top_k: batch_size,
+            namespace: namespace,
+            filter: filter,
+            include_metadata: true
+          )
+
+          break if result.empty?
+
+          result.each do |match|
+            # Skip duplicates (some providers may return overlapping results)
+            next if seen_ids.include?(match.id)
+
+            seen_ids.add(match.id)
+            yielder << match
+            fetched += 1
+
+            break if fetched >= total
+          end
+
+          # If we got fewer results than requested, we've exhausted the index
+          break if result.size < batch_size
+        end
+      end.lazy
+    end
+
+    # Scan all vectors in an index (provider-dependent)
+    #
+    # @param index [String] the index name
+    # @param namespace [String, nil] optional namespace
+    # @param batch_size [Integer] IDs per batch
+    # @yield [Vector] each vector
+    # @return [Integer] total vectors scanned
+    # @note Not all providers support efficient scanning
+    def scan_all(index:, namespace: nil, batch_size: 1000)
+      return 0 unless block_given?
+
+      count = 0
+      offset = 0
+
+      loop do
+        # This is a simplified scan - actual implementation depends on provider
+        stats = client.stats(index: index, namespace: namespace)
+        total = stats[:total_vector_count] || 0
+
+        break if offset >= total
+
+        # Fetch IDs in batches (this is provider-specific)
+        # For now, we return what we can
+        break if offset.positive? # Only one iteration for basic implementation
+
+        offset += batch_size
+        count = total
+      end
+
+      count
+    end
+  end
+
+  # Streaming result wrapper with additional metadata
+  class StreamingResult
+    include Enumerable
+
+    attr_reader :enumerator, :metadata
+
+    def initialize(enumerator, metadata = {})
+      @enumerator = enumerator
+      @metadata = metadata
+    end
+
+    def each(&)
+      enumerator.each(&)
+    end
+
+    def take(n)
+      enumerator.take(n)
+    end
+
+    def to_a
+      enumerator.to_a
+    end
+  end
+end

data/lib/vectra/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Vectra
-  VERSION = "0.2.2"
+  VERSION = "0.3.1"
 end