sqlitesweep 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 1873157a16c9f24d344824160bbcc81bfd02bff047a3d7f8aa293abf69a5784e
+  data.tar.gz: 105e3556be154e3c16d329ce14d78a7cb5cd890a9bbc7771447ce014ee2a91e0
+SHA512:
+  metadata.gz: 2128e4c7ab40f54c45873d5c91472e5807a2011818fa4f9946819015e3931ecaacc53d8458420ab24843b05e8b4bb9b261f1ea1b6bd58265a767a1e3c9f66059
+  data.tar.gz: 62bb09f2f215fcfb6865557768c2779adddc53265470fbe8bccb7636676be1956fb664300aea63b0daccde7e558ecc37aed68c918dbe277cd38ea7c916f1f9a1

data/README.md

@@ -0,0 +1,111 @@
+# SQLiteSweep
+
+Query millions of SQLite databases across remote hosts via SSH, aggregating results in real-time. Designed for multi-tenant apps where each tenant has their own SQLite database.
+
+## Installation
+
+```bash
+gem install sqlitesweep
+```
+
+Or add to your Gemfile:
+
+```ruby
+gem "sqlitesweep"
+```
+
+## Usage
+
+```bash
+sqlitesweep \
+  -q 'SELECT count(*) FROM products WHERE price > 0' \
+  -a sum \
+  -s 'rails runner "Account.active.find_each { |a| puts a.db_uri }"' \
+  -c 16
+```
+
+The `-s` (source) command should output one database URI per line. URIs can be:
+
+- Local paths: `/data/tenants/acme.sqlite3`
+- File URIs: `file:///data/tenants/acme.sqlite3`
+- SSH URIs: `ssh://deploy@web1.example.com/data/tenants/acme.sqlite3`
+
+Results go to stdout, progress goes to stderr — so it's pipe-friendly:
+
+```bash
+sqlitesweep -q "SELECT count(*) FROM users" -a sum -s "cat db_uris.txt" > result.txt
+```
+
+### Actions
+
+| Action | Description |
+|--------|-------------|
+| `sum` | Sum the first column across all databases (default) |
+| `average` / `avg` | Average the first column across all databases |
+| `list` | Write all rows to a JSONL file, print the file path |
+
+### Options
+
+| Flag | Long | Default | Description |
+|------|------|---------|-------------|
+| `-q` | `--query` | *(required)* | SQL query to execute on each database |
+| `-a` | `--action` | `sum` | `sum`, `average`/`avg`, or `list` |
+| `-s` | `--source` | *(required)* | Shell command that outputs database URIs |
+| `-c` | `--concurrency` | `8` | Max parallel query workers |
+| | `--max-ssh` | `50` | Max SSH master connections |
+| | `--no-live` | `false` | Disable live progress display |
+| | `--batch-size` | `4` | Databases to query per SSH call |
+| | `--ssh-timeout` | `10` | SSH connect timeout (seconds) |
+| | `--query-timeout` | `30` | Per-query timeout (seconds) |
+
+## How it works
+
+- **Local databases** are queried directly via the `sqlite3` gem
+- **Remote databases** are queried by shelling out to `ssh` + `sqlite3` on the remote host
+- SSH connections use `ControlMaster` multiplexing — one master connection per host, shared across queries
+- Multiple databases on the same host are batched into a single SSH call (configurable via `--batch-size`)
+- A thread pool (via `concurrent-ruby`) runs queries in parallel with back-pressure
+
+### Requirements
+
+- Ruby >= 4.0
+- `sqlite3` available on remote hosts (for SSH queries)
+- SSH agent or key-based auth configured (BatchMode is enforced — no password prompts)
+
+## Development
+
+```bash
+bundle install
+bundle exec rake test
+```
+
+Run the integration test:
+
+```bash
+bundle exec ruby test/integration/test_local_sweep.rb
+```
+
+Watch the live progress display with a slow-drip demo:
+
+```bash
+ruby test/integration/harness/live_demo.rb
+ruby test/integration/harness/live_demo.rb --count 50 --delay 0.5
+ruby test/integration/harness/live_demo.rb --action list
+```
+
+Run the benchmark:
+
+```bash
+ruby test/integration/harness/sweep_bench.rb --db-count 1000
+```
+
+### Docker-based SSH testing
+
+```bash
+docker compose -f test/integration/harness/docker-compose.yml up -d
+ruby test/integration/harness/sweep_bench.rb --docker --db-count 500
+```
+
+## License
+
+MIT

data/exe/sqlitesweep

@@ -0,0 +1,5 @@
+#!/usr/bin/env ruby
+
+require "sqlitesweep"
+
+SQLiteSweep::CLI.run(ARGV)

data/lib/sqlitesweep/aggregator.rb

@@ -0,0 +1,86 @@
+module SQLiteSweep
+  # Thread-safe accumulator for query results. Supports three modes:
+  #
+  #   :sum     - Running total of the first column value from each database.
+  #   :average - Running total divided by the number of databases queried.
+  #   :list    - Delegates to a ResultFile that streams rows to a JSONL file.
+  #
+  # For :sum and :average, the first value of the first row from each query
+  # result is extracted and added to a running total. This means your query
+  # should return a single numeric value (e.g. SELECT count(*) FROM ...).
+  #
+  # All methods are mutex-protected and safe to call from multiple worker threads.
+  #
+  # @example
+  #   agg = Aggregator.new(:sum)
+  #   agg.add(Result.new(rows: [{"count" => 10}], source: "db1"))
+  #   agg.add(Result.new(rows: [{"count" => 20}], source: "db2"))
+  #   agg.value  # => "30"
+  #   agg.count  # => 2
+  #
+  class Aggregator
+    # @return [Integer] Number of databases successfully queried.
+    attr_reader :count
+
+    # @return [Integer] Number of databases that produced errors.
+    attr_reader :error_count
+
+    # @param action [Symbol] One of :sum, :average, or :list.
+    # @param result_file [ResultFile, nil] Required for :list action. Receives
+    #   streamed results so they don't accumulate in memory.
+    def initialize(action, result_file: nil)
+      @action = action
+      @result_file = result_file
+      @mutex = Mutex.new
+      @total = 0.0
+      @count = 0
+      @error_count = 0
+    end
+
+    # Records a successful query result.
+    #
+    # @param result [Result] The query result to aggregate.
+    def add(result)
+      @mutex.synchronize do
+        case @action
+        when :sum, :average
+          result.rows.each do |row|
+            value = row.values.first
+            @total += value.to_f
+          end
+        when :list
+          @result_file&.write(result)
+        end
+        @count += 1
+      end
+    end
+
+    # Records a failed query (increments error count).
+    def record_error
+      @mutex.synchronize { @error_count += 1 }
+    end
+
+    # Returns the final aggregated value as a string.
+    #
+    # @return [String] For :sum, the total. For :average, the mean.
+    #   For :list, the path to the JSONL result file.
+    def value
+      case @action
+      when :sum
+        format_number(@total)
+      when :average
+        @count > 0 ? format_number(@total / @count) : "0"
+      when :list
+        @result_file&.path
+      end
+    end
+
+    private
+
+    # Formats a number as a clean string: integers without decimal point,
+    # floats with their natural precision.
+    def format_number(n)
+      n == n.to_i ? n.to_i.to_s : n.to_s
+    end
+  end
+end

data/lib/sqlitesweep/cli.rb

@@ -0,0 +1,89 @@
+require "optparse"
+
+module SQLiteSweep
+  # Command-line interface. Parses ARGV into a Config and launches the Runner.
+  #
+  # Designed for pipe-friendly usage: results go to stdout, progress/errors
+  # go to stderr. This means you can do:
+  #
+  #   sqlitesweep -q "SELECT count(*) FROM users" -a sum -s "cat uris.txt" > result.txt
+  #
+  class CLI
+    # Parses command-line arguments and runs the sweep.
+    #
+    # @param argv [Array<String>] Command-line arguments (typically ARGV).
+    # @return [void]
+    def self.run(argv)
+      options = {}
+      parser = OptionParser.new do |opts|
+        opts.banner = "Usage: sqlitesweep [options]"
+
+        opts.on("-q", "--query QUERY", "SQL query to execute on each database (required)") do |v|
+          options[:query] = v
+        end
+
+        opts.on("-a", "--action ACTION", "sum, average/avg, or list (default: sum)") do |v|
+          options[:action] = v
+        end
+
+        opts.on("-s", "--source COMMAND", "Shell command that outputs database URIs (required)") do |v|
+          options[:source] = v
+        end
+
+        opts.on("-c", "--concurrency N", Integer, "Max parallel query workers (default: 8)") do |v|
+          options[:concurrency] = v
+        end
+
+        opts.on("--max-ssh N", Integer, "Max SSH master connections (default: 50)") do |v|
+          options[:max_ssh] = v
+        end
+
+        opts.on("--no-live", "Disable live progress display") do
+          options[:live] = false
+        end
+
+        opts.on("--batch-size N", Integer, "Databases to query per SSH call (default: 4)") do |v|
+          options[:batch_size] = v
+        end
+
+        opts.on("--ssh-timeout N", Integer, "SSH connect timeout in seconds (default: 10)") do |v|
+          options[:ssh_timeout] = v
+        end
+
+        opts.on("--query-timeout N", Integer, "Per-query timeout in seconds (default: 30)") do |v|
+          options[:query_timeout] = v
+        end
+
+        opts.on("-h", "--help", "Show this help") do
+          $stderr.puts opts
+          exit 0
+        end
+
+        opts.on("--version", "Show version") do
+          puts "sqlitesweep #{VERSION}"
+          exit 0
+        end
+      end
+
+      parser.parse!(argv)
+
+      unless options[:query]
+        $stderr.puts "Error: --query is required"
+        $stderr.puts parser
+        exit 1
+      end
+
+      unless options[:source]
+        $stderr.puts "Error: --source is required"
+        $stderr.puts parser
+        exit 1
+      end
+
+      config = Config.new(**options)
+      Runner.new(config).run
+    rescue ConfigError => e
+      $stderr.puts "Error: #{e.message}"
+      exit 1
+    end
+  end
+end

data/lib/sqlitesweep/config.rb

@@ -0,0 +1,58 @@
+module SQLiteSweep
+  # Immutable configuration for a sweep run. Built from CLI flags and passed
+  # to every component.
+  #
+  # Uses Data.define for a frozen value object — once created, a Config
+  # cannot be modified.
+  #
+  # @example
+  #   config = Config.new(
+  #     query: "SELECT count(*) FROM users",
+  #     source: "cat uris.txt",
+  #     action: :sum,
+  #     concurrency: 16
+  #   )
+  #   config.action     # => :sum
+  #   config.concurrency # => 16
+  #
+  Config = Data.define(
+    :query,          # SQL query to run on every database
+    :action,         # :sum, :average, or :list
+    :source,         # Shell command that outputs database URIs (one per line)
+    :concurrency,    # Max worker threads in the pool
+    :max_ssh,        # Max simultaneous SSH ControlMaster connections
+    :live,           # Whether to show live ANSI progress on stderr
+    :batch_size,     # Number of databases to query per SSH call
+    :ssh_timeout,    # SSH connect timeout in seconds
+    :query_timeout   # Per-query (or per-batch) timeout in seconds
+  ) do
+    def initialize(
+      query:,
+      source:,
+      action: :sum,
+      concurrency: 8,
+      max_ssh: 50,
+      live: $stderr.tty?,
+      batch_size: 4,
+      ssh_timeout: 10,
+      query_timeout: 30
+    )
+      action = action.to_sym
+      action = :average if action == :avg
+      unless %i[sum average list].include?(action)
+        raise ConfigError, "Unknown action: #{action}. Must be sum, average/avg, or list"
+      end
+      super(
+        query: query,
+        action: action,
+        source: source,
+        concurrency: concurrency,
+        max_ssh: max_ssh,
+        live: live,
+        batch_size: batch_size,
+        ssh_timeout: ssh_timeout,
+        query_timeout: query_timeout
+      )
+    end
+  end
+end

data/lib/sqlitesweep/database_uri.rb

@@ -0,0 +1,89 @@
+require "uri"
+
+module SQLiteSweep
+  # Parses and represents a database location. Supports three URI formats:
+  #
+  #   - Local path:  /data/tenants/acme.sqlite3
+  #   - File URI:    file:///data/tenants/acme.sqlite3
+  #   - SSH URI:     ssh://deploy@web1.example.com/data/tenants/acme.sqlite3
+  #
+  # For SSH URIs, the user portion is optional. The host and path are extracted
+  # and used to build SSH commands.
+  #
+  # @example
+  #   uri = DatabaseURI.new("ssh://deploy@web1/data/db.sqlite3")
+  #   uri.remote?          # => true
+  #   uri.ssh_destination  # => "deploy@web1"
+  #   uri.path             # => "/data/db.sqlite3"
+  #
+  #   uri = DatabaseURI.new("/tmp/local.sqlite3")
+  #   uri.local?           # => true
+  #   uri.path             # => "/tmp/local.sqlite3"
+  #
+  class DatabaseURI
+    attr_reader :user, :host, :path
+
+    def initialize(uri_string)
+      uri_string = uri_string.strip
+      if uri_string.start_with?("ssh://")
+        parsed = URI.parse(uri_string)
+        @user = parsed.user
+        @host = parsed.host
+        @path = parsed.path
+        raise ConfigError, "SSH URI missing host: #{uri_string}" if @host.nil? || @host.empty?
+        raise ConfigError, "SSH URI missing path: #{uri_string}" if @path.nil? || @path.empty?
+      elsif uri_string.start_with?("file://")
+        parsed = URI.parse(uri_string)
+        @user = nil
+        @host = nil
+        @path = parsed.path
+      else
+        @user = nil
+        @host = nil
+        @path = uri_string
+      end
+    end
+
+    # Returns true if this database is on a remote host (accessed via SSH).
+    def remote?
+      !@host.nil?
+    end
+
+    # Returns true if this database is on the local filesystem.
+    def local?
+      @host.nil?
+    end
+
+    # Returns the SSH destination string (e.g. "deploy@web1" or "web1").
+    # Used as the target for ssh commands. Returns nil for local URIs.
+    def ssh_destination
+      return nil unless remote?
+      @user ? "#{@user}@#{@host}" : @host
+    end
+
+    # Returns a key that uniquely identifies the remote host (user@host).
+    # Used by HostBatcher to group databases on the same host into batches.
+    # Returns nil for local URIs.
+    def host_key
+      ssh_destination
+    end
+
+    def to_s
+      if remote?
+        "ssh://#{ssh_destination}#{@path}"
+      else
+        @path
+      end
+    end
+
+    def ==(other)
+      other.is_a?(DatabaseURI) && @user == other.user && @host == other.host && @path == other.path
+    end
+
+    def hash
+      [@user, @host, @path].hash
+    end
+
+    alias_method :eql?, :==
+  end
+end

data/lib/sqlitesweep/display.rb

@@ -0,0 +1,77 @@
+module SQLiteSweep
+  # Live progress display on stderr. Shows a single-line status that updates
+  # in-place using ANSI escape codes (\r to return to line start, \e[2K to
+  # clear the line).
+  #
+  # Output format:
+  #   Queried: 14523 | Errors: 3 | Rate: 847/s | Elapsed: 17.1s | Result: 2847291
+  #
+  # When live mode is disabled (--no-live or non-TTY stderr), all rendering
+  # is silently skipped. This keeps piped output clean.
+  #
+  # The display auto-refreshes every 250ms via a background timer thread,
+  # and can also be manually refreshed after each query completes.
+  #
+  class Display
+    # @param aggregator [Aggregator] Source of count/error/value data.
+    # @param live [Boolean] Whether to render live output.
+    def initialize(aggregator, live: true)
+      @aggregator = aggregator
+      @live = live
+      @start_time = nil
+      @mutex = Mutex.new
+      @timer = nil
+    end
+
+    # Starts the display timer. Must be called before any refresh calls.
+    #
+    # @param start_time [Float] Monotonic clock timestamp (from Process.clock_gettime).
+    def start(start_time)
+      @start_time = start_time
+      return unless @live
+
+      @timer = Thread.new do
+        loop do
+          sleep 0.25
+          refresh
+        end
+      rescue
+        # timer thread exits silently on kill
+      end
+    end
+
+    # Manually triggers a display refresh. Called by workers after each
+    # query completes to keep the display responsive.
+    def refresh
+      return unless @live
+      @mutex.synchronize { render }
+    end
+
+    # Stops the timer, renders one final update, and prints a newline
+    # so subsequent output starts on a fresh line.
+    def finish
+      @timer&.kill
+      @timer&.join(1)
+      if @live
+        @mutex.synchronize do
+          render
+          $stderr.write("\n")
+        end
+      end
+    end
+
+    private
+
+    def render
+      elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - @start_time
+      count = @aggregator.count
+      errors = @aggregator.error_count
+      rate = elapsed > 0 ? (count / elapsed).round(0).to_i : 0
+
+      value_display = @aggregator.value
+
+      line = "  Queried: #{count} | Errors: #{errors} | Rate: #{rate}/s | Elapsed: #{elapsed.round(1)}s | Result: #{value_display}"
+      $stderr.write("\r\e[2K#{line}")
+    end
+  end
+end

data/lib/sqlitesweep/errors.rb

@@ -0,0 +1,19 @@
+module SQLiteSweep
+  # Base error class for all SQLiteSweep errors.
+  class Error < StandardError; end
+
+  # Raised when a SQL query fails against a database (local or remote).
+  class QueryError < Error; end
+
+  # Raised when the source command (the -s flag) fails or exits non-zero.
+  class SourceError < Error; end
+
+  # Raised when an SSH connection or master setup fails.
+  class SSHError < Error; end
+
+  # Raised for invalid configuration (bad action name, missing fields, etc.).
+  class ConfigError < Error; end
+
+  # Raised when an SSH command or query exceeds its timeout.
+  class TimeoutError < Error; end
+end

data/lib/sqlitesweep/host_batcher.rb

@@ -0,0 +1,127 @@
+module SQLiteSweep
+  # Groups remote database URIs by host and flushes them as batches.
+  #
+  # Instead of making one SSH round-trip per database, HostBatcher accumulates
+  # URIs for the same host and submits them as a single batch to the worker pool.
+  # Each batch becomes one SSH command that queries multiple databases sequentially
+  # on the remote host.
+  #
+  # Batches are flushed when either:
+  #   - The batch reaches --batch-size (default 4)
+  #   - A timeout fires (200ms) to avoid stalling on slow source streams
+  #
+  # Local URIs should NOT be sent here — they bypass batching entirely and go
+  # straight to the worker pool (see Runner).
+  #
+  # @example
+  #   batcher = HostBatcher.new(config, pool, aggregator, display, ssh_manager)
+  #   batcher.add(uri)      # accumulates by host
+  #   batcher.flush_all     # drains any remaining partial batches
+  #
+  class HostBatcher
+    # How long to wait before flushing an incomplete batch (seconds).
+    # Prevents URIs from sitting in the buffer when the source stream is slow.
+    FLUSH_TIMEOUT = 0.2
+
+    def initialize(config, pool, aggregator, display, ssh_manager)
+      @config = config
+      @pool = pool
+      @aggregator = aggregator
+      @display = display
+      @ssh_manager = ssh_manager
+      @batch_size = config.batch_size
+      @mutex = Mutex.new
+      @batches = {} # host_key => [uri, ...]
+      @timers = {}  # host_key => Thread
+      @cancelled = false
+    end
+
+    # Adds a remote URI to the appropriate host batch.
+    # May trigger an immediate flush if the batch is full.
+    #
+    # @param uri [DatabaseURI] A remote database URI.
+    def add(uri)
+      return if @cancelled
+
+      host_key = uri.host_key
+      flush_batch = nil
+
+      @mutex.synchronize do
+        @batches[host_key] ||= []
+        @batches[host_key] << uri
+
+        if @batches[host_key].length >= @batch_size
+          flush_batch = @batches.delete(host_key)
+          cancel_timer(host_key)
+        else
+          start_timer(host_key) unless @timers[host_key]
+        end
+      end
+
+      submit_batch(flush_batch) if flush_batch
+    end
+
+    # Flushes all pending batches regardless of size. Called at the end of a
+    # sweep to drain any partial batches that haven't reached batch_size.
+    def flush_all
+      batches_to_flush = nil
+      @mutex.synchronize do
+        @timers.each_value(&:kill)
+        @timers.clear
+        batches_to_flush = @batches.dup
+        @batches.clear
+      end
+
+      batches_to_flush.each_value { |batch| submit_batch(batch) }
+    end
+
+    # Cancels all pending batches and timers. Used during shutdown.
+    def cancel
+      @mutex.synchronize do
+        @cancelled = true
+        @timers.each_value(&:kill)
+        @timers.clear
+        @batches.clear
+      end
+    end
+
+    private
+
+    # Starts a timeout timer for a host. If the batch hasn't been flushed
+    # by the time the timer fires, it flushes whatever is accumulated.
+    def start_timer(host_key)
+      @timers[host_key] = Thread.new do
+        sleep FLUSH_TIMEOUT
+        batch = @mutex.synchronize do
+          @timers.delete(host_key)
+          @batches.delete(host_key)
+        end
+        submit_batch(batch) if batch
+      end
+    end
+
+    def cancel_timer(host_key)
+      timer = @timers.delete(host_key)
+      timer&.kill
+    end
+
+    # Submits a batch of URIs (all for the same host) to the worker pool.
+    # The worker creates a Query::Remote and executes all queries in a
+    # single SSH command.
+    def submit_batch(uris)
+      return if uris.nil? || uris.empty?
+
+      @pool.submit do
+        remote_query = Query::Remote.new(@config, @ssh_manager)
+        begin
+          results = remote_query.execute_batch(uris)
+          results.each { |result| @aggregator.add(result) }
+        rescue QueryError => e
+          uris.length.times { @aggregator.record_error }
+          $stderr.puts "\n#{e.message}" unless @config.live
+        end
+        @display.refresh
+      end
+    end
+  end
+end

data/lib/sqlitesweep/query/base.rb

@@ -0,0 +1,27 @@
+module SQLiteSweep
+  module Query
+    # Abstract base class for query executors.
+    #
+    # Subclasses must implement #execute(uri) which runs the configured SQL
+    # query against the given database and returns a Result.
+    #
+    # Two implementations:
+    #   - Query::Local  — uses the sqlite3 gem for direct file access
+    #   - Query::Remote — shells out via SSH to run sqlite3 on a remote host
+    #
+    class Base
+      def initialize(config)
+        @config = config
+      end
+
+      # Executes the configured query against the given database.
+      #
+      # @param uri [DatabaseURI] The database to query.
+      # @return [Result] The query result.
+      # @raise [QueryError] If the query fails.
+      def execute(uri)
+        raise NotImplementedError
+      end
+    end
+  end
+end