tobias 0.2.0 → 0.4.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 342518e2c3730a589319c0d16967f7e9fdc80affde011c3820ef2c1034fcbc3e
-  data.tar.gz: adcdc3152a07d4e1203ad6ceff26901e999f5a8110e883675aa20c9e894ae932
+  metadata.gz: 6e8024f2a5efa0e4ffe41d2973817f1906fccfb32193782b86e55f822a82f5e7
+  data.tar.gz: 5b6b27888ed7341c76fea8503bc8e577662949903089970930fced41367583e0
 SHA512:
-  metadata.gz: af1341883626cda785ba564368baad44e0cec9e55751bb7c00d204b1e12cf32ebc3817002f2c7fd81e7b57a6ddd2a78be74b11fa0e1f3f837741552f8064f338
-  data.tar.gz: 58d79c4c82d5bf1cd21e42c6770dcc454c52bfab5a22a5fa1d15ed6dfe8b3abfd340d7a31b07c6205b7955787214395814176cd997307e2dad083256c29c5139
+  metadata.gz: 2e289f868dad239629020799f64c05be1f5658f2facd905d917ce162bcccbefc957e979c19bdb36ec37e1c21eb89ae26d7b2f2a458f11b8d1d335eb5b9a9a37c
+  data.tar.gz: fece84c5f4c2d65961c091ee4bb7e2879b43646e60a8c470079d20529541757e0e61d9f4a2163e0615660c38b62c758504bdde9009a35e426f0578d598089468

data/bin/tobias

@@ -3,5 +3,8 @@
 
 require_relative "../lib/tobias"
 
-RubyVM::YJIT.enable
+if defined?(RubyVM::YJIT)
+  RubyVM::YJIT.enable
+end
+
 Tobias::CLI.start(ARGV, exit_on_failure: true)

data/lib/tobias/cli.rb

@@ -8,14 +8,46 @@ module Tobias
       true
     end
 
+    desc "recommend", "recommend a work_mem setting for a database"
+    option :database_url, type: :string, required: true
+    option :debug, type: :boolean, default: false
+    def recommend
+      database = Sequel.connect(options[:database_url])
+      database.extension :pgvector
+      database.loggers << Logger.new(STDERR) if options[:debug]
+
+      parsed = TTY::Markdown.parse(<<~MARKDOWN)
+        # @tobias is thinking...
+      MARKDOWN
+      puts parsed
+
+      work_mem = Tobias::WorkMem.valid_for(database).sort_by(&:amount).reverse.first
+
+      parsed = TTY::Markdown.parse(<<~MARKDOWN)
+        # @tobias has sent you a new message
+
+        I've reviewed your database by analyzing your shared buffers and connection limits and
+        recommend setting `work_mem` to `#{work_mem.to_sql}`. To apply my recommendation, run the following SQL:
+
+        ```sql
+        ALTER SYSTEM SET work_mem = '#{work_mem.to_sql}';
+        SELECT pg_reload_conf();
+        ```
+
+        Regards,
+        ~ Tobias
+      MARKDOWN
+      puts parsed
+    end
+
     desc "profile SCRIPT", "profile"
     option :database_url, type: :string, required: true
-    option :iterations, type: :numeric, default: 100
     option :debug, type: :boolean, default: false
     def profile(script)
-      database = Sequel.connect(options[:database_url])
+      database = Sequel.connect(options[:database_url], max_connections: Etc.nprocessors + 2)
       database.loggers << Logger.new(STDERR) if options[:debug]
       database.extension :pg_json
+      database.extension :pgvector
 
       if File.exist?(script)
         code = File.read(script)
@@ -23,8 +55,7 @@ module Tobias
         raise "Script not found at: #{script}"
       end
 
-      container = Container.new(code)
-      work_mems = WorkMem.valid_for(database)
+      container = Container.new(code, database)
       results = {}
 
       parsed = TTY::Markdown.parse(<<~MARKDOWN)
@@ -33,11 +64,7 @@ module Tobias
       puts parsed
 
       thinking_time = Benchmark.realtime do
-        container.queries.each do |name, block|
-          work_mem = Evaluation.new(database, work_mems).run(options, &block)
-
-          results[name] = work_mem
-        end
+        results = Evaluations.run(database, container, options)
       end
 
       parsed = TTY::Markdown.parse(<<~MARKDOWN)
@@ -45,18 +72,7 @@ module Tobias
 
         I thought about your queries for precisely #{thinking_time.round(2)} seconds and here is what I recommend:
 
-        | Query | Required work_mem |
-        |-------|-------------------|
-        #{results.map { |name, work_mem| "| #{name} | #{work_mem.to_sql} |" }.join("\n")}
-
-        Your application will need to run with at least #{results.values.max.to_sql} of work_mem.
-
-        To apply my recommendations, run the following SQL:
-
-        ```sql
-        ALTER SYSTEM SET work_mem = '#{results.values.max.to_sql}';
-        SELECT pg_reload_conf();
-        ```
+        #{results.join("\n")}
 
         Regards,
         ~ Tobias

data/lib/tobias/container.rb

@@ -1,20 +1,95 @@
 # frozen_string_literal: true
 
+require "open3"
+
 module Tobias
   class Container
-    def initialize(code)
+    def initialize(code, database)
       @code = code
-      @queries = {}
+      @database = database
+      @queries = Concurrent::Hash.new
+      @sql = Concurrent::Hash.new
+      @options = Concurrent::Hash.new
+      @setup = Proc.new { }
+      @teardown = Proc.new { }
+      @load_data = Proc.new { }
+      @helpers = Module.new
 
       eval(code, binding, __FILE__, __LINE__)
     end
 
+    module DefaultHelpers
+      def db
+        @database
+      end
+
+      def run_parallel(list = Etc.nprocessors.times, &block)
+        db.disconnect
+
+        Parallel.each(list, in_processes: Etc.nprocessors) do |item|
+          instance_exec(item, &block)
+        end
+      end
+    end
+
+    def run_setup
+      @database.run("CREATE EXTENSION IF NOT EXISTS pg_stat_statements")
+      run_action(@setup)
+    end
+
+    def run_query(query)
+      @database.run(run_action(query).sql)
+    end
+
+    def run_teardown
+      run_action(@teardown)
+    end
+
+    def options
+      Struct.new(*@options.keys).new(*@options.values)
+    end
+
+    def run_action(action)
+      helpers = @helpers
+
+      class_eval do
+        include DefaultHelpers
+        include helpers
+      end
+
+      instance_eval(&action)
+    end
+
     def queries
       @queries
     end
 
-    def query(name, &block)
-      @queries[name] = block
+    def option(name, default = nil, &block)
+      @options[name] = block || default
+    end
+
+    def helpers(&block)
+      @helpers.class_eval(&block) if block_given?
+    end
+
+    def setup(&block)
+      @setup = block
+    end
+
+    def teardown(&block)
+      @teardown = block
+    end
+
+    def load_data(&block)
+      @load_data = block
+    end
+
+    def query(name, sql = nil, &block)
+      if sql.is_a?(String)
+        @queries[name] = Proc.new { sql }
+      else
+        @queries[name] = block || Proc.new { raise "No SQL provided for query '#{name}'" }
+      end
     end
   end
-end
+end

data/lib/tobias/evaluations/base.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+class MarkdownTableBorder < TTY::Table::Border
+  def_border do
+    left         "|"
+    center       "|"
+    right        "|"
+    bottom       " "
+    bottom_mid   " "
+    bottom_left  " "
+    bottom_right " "
+  end
+end
+
+module Tobias
+  module Evaluations
+    Result = Struct.new(:name, :value, keyword_init: true) do
+      def <=>(other)
+        value <=> other.value
+      end
+    end
+
+    class Base
+      attr_reader :database, :container, :options
+
+      def initialize(database, container, options)
+        @database = database
+        @container = container
+        @options = options
+      end
+
+      def run(&block)
+        results = Concurrent::Array.new
+
+        container.run_setup
+        container.queries.each do |name, query|
+          result = run_each(name, query)
+          results << result if result
+        end
+
+        to_markdown(results)
+      ensure
+        container.run_teardown
+      end
+
+      def run_each(query)
+        raise NotImplementedError
+      end
+
+      def to_markdown(results)
+        raise NotImplementedError
+      end
+
+      def render_table(headers:, body:)
+        table = TTY::Table.new(header: headers)
+        body.each do |row|
+          table << row
+        end
+
+        table.render_with(MarkdownTableBorder)
+      end
+    end
+  end
+end

data/lib/tobias/evaluations/work_mem.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module Tobias
+  module Evaluations
+    class WorkMem < Base
+      def work_mems
+        @work_mems ||= Tobias::WorkMem.valid_for(database)
+      end
+
+      def current_work_mem
+        @current_work_mem ||= Tobias::WorkMem.from_sql(database.fetch("SHOW work_mem").first[:work_mem])
+      end
+
+      def description
+        "Optimal work_mem settings"
+      end
+
+      def run_each(name, query)
+        database.run("CREATE EXTENSION IF NOT EXISTS pg_stat_statements")
+
+        work_mems.each do |value|
+          database.transaction do
+            database.run("SET LOCAL work_mem = '#{value.to_sql}'")
+            database.select(Sequel.function(:pg_stat_reset)).first
+            container.run_query(query)
+
+            stats = database[:pg_stat_database].
+              where(datname: Sequel.function(:current_database)).
+              first
+
+            if stats[:temp_files] == 0 && stats[:temp_bytes] == 0
+              return Result.new(name: name, value: value)
+            end
+          end
+        end
+
+        # Fallback to the highest work_mem setting if no results are found.
+        Result.new(name: name, value: work_mems.last)
+      end
+
+      def to_markdown(results)
+        <<~MARKDOWN
+          ## #{description}
+
+          #{render_table(headers: ["Query", "Required work_mem"], body: results.map { |r| [r.name, r.value.to_sql] })}
+
+          I see that your current `work_mem` setting is `#{current_work_mem.to_sql}`.
+
+          Your application will need to run with at least `#{results.max.value.to_sql}` of `work_mem`.
+
+          To apply my recommendations, run the following SQL:
+
+          ```sql
+          ALTER SYSTEM SET work_mem = '#{results.max.value.to_sql}';
+          SELECT pg_reload_conf();
+          ```
+        MARKDOWN
+      end
+    end
+  end
+end

data/lib/tobias/evaluations.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Tobias
+  module Evaluations
+    def self.run(database, container, options)
+      results = []
+      results << WorkMem.new(database, container, options).run
+      results
+    end
+  end
+end

data/lib/tobias/version.rb

@@ -1,8 +1,7 @@
 # frozen_string_literal: true
 
 module Tobias
-  VERSION = "0.2.0"
-
+  VERSION = "0.4.0"
   SUMMARY = "Tobias is a tool to help you find the optimal work_mem for your queries."
   DESCRIPTION = "Tobias is a tool to help you find the optimal work_mem for your queries."
 end

data/lib/tobias/work_mem.rb

@@ -4,6 +4,21 @@ module Tobias
   class WorkMem
     attr_reader :amount
 
+    def self.from_sql(sql)
+      case sql
+      when /^\d+B$/
+        new(sql.to_i)
+      when /^\d+kB$/
+        new(sql.to_i * 1024)
+      when /^\d+MB$/
+        new(sql.to_i * 1024 * 1024)
+      when /^\d+GB$/
+        new(sql.to_i * 1024 * 1024 * 1024)
+      else
+        raise "Invalid work_mem setting: #{sql}"
+      end
+    end
+
     def initialize(amount)
       @amount = amount
     end
@@ -52,6 +67,7 @@ module Tobias
       [
         new(64.kilobytes),
         new(128.kilobytes),
+        new(256.kilobytes),
         new(512.kilobytes),
         new(1.megabyte),
         new(4.megabytes),
@@ -101,4 +117,4 @@ module Tobias
       self.all.select { |work_mem| work_mem.amount < bytes_per_connection.to_i }
     end
   end
-end
+end

data/lib/tobias.rb

@@ -11,21 +11,31 @@ if RUBY_ENGINE == "ruby"
   end
 end
 
+# See: https://github.com/ged/ruby-pg/issues/538#issuecomment-1591629049
+ENV["PGGSSENCMODE"] = "disable"
+
 require "bundler/setup"
 Bundler.require(:default)
 
 require "thor"
 require "active_support/all"
 require "sequel"
+require "pgvector"
 require "enumerable-stats"
 require "benchmark"
+require "parquet"
+require "parallel"
 require "tty-markdown"
+require "tty-table"
 
 $LOAD_PATH.unshift File.dirname(__FILE__)
 
+require "tobias/evaluations"
+require "tobias/evaluations/base"
+require "tobias/evaluations/work_mem"
+
 module Tobias
   autoload :CLI, "tobias/cli"
   autoload :Container, "tobias/container"
-  autoload :Evaluation, "tobias/evaluation"
   autoload :WorkMem, "tobias/work_mem"
-end
+end

metadata

@@ -1,13 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: tobias
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.4.0
 platform: ruby
 authors:
 - Jon Daniel
+autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-08-03 00:00:00.000000000 Z
+date: 2025-08-23 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activesupport
@@ -189,6 +190,66 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: 5.76.0
+- !ruby/object:Gem::Dependency
+  name: parquet
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.7'
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.7.3
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.7'
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.7.3
+- !ruby/object:Gem::Dependency
+  name: pgvector
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.3'
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.3.0
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.3'
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.3.0
+- !ruby/object:Gem::Dependency
+  name: parallel
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.20'
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 1.20.0
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.20'
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 1.20.0
 - !ruby/object:Gem::Dependency
   name: thor
   requirement: !ruby/object:Gem::Requirement
@@ -229,6 +290,26 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: 0.7.0
+- !ruby/object:Gem::Dependency
+  name: tty-table
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.12'
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.12.0
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.12'
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 0.12.0
 description: Tobias is a tool to help you find the optimal work_mem for your queries.
 email: binarycleric@gmail.com
 executables:
@@ -240,7 +321,9 @@ files:
 - lib/tobias.rb
 - lib/tobias/cli.rb
 - lib/tobias/container.rb
-- lib/tobias/evaluation.rb
+- lib/tobias/evaluations.rb
+- lib/tobias/evaluations/base.rb
+- lib/tobias/evaluations/work_mem.rb
 - lib/tobias/version.rb
 - lib/tobias/work_mem.rb
 homepage: https://github.com/binarycleric/tobias
@@ -249,6 +332,7 @@ licenses:
 metadata:
   source_code_uri: https://github.com/binarycleric/tobias
   rubygems_mfa_required: 'true'
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -263,7 +347,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.6.2
+rubygems_version: 3.5.22
+signing_key:
 specification_version: 4
 summary: Tobias is a tool to help you find the optimal work_mem for your queries.
 test_files: []

data/lib/tobias/evaluation.rb

@@ -1,36 +0,0 @@
-# frozen_string_literal: true
-
-module Tobias
-  class Evaluation
-    attr_reader :database
-
-    def initialize(database, work_mems)
-      @database = database
-      @work_mems = work_mems
-    end
-
-    def run(options, &block)
-      @work_mems.each do |value|
-        database.transaction do
-          database.run("CREATE EXTENSION IF NOT EXISTS pg_stat_statements")
-          database.run("SET LOCAL work_mem = '#{value.to_sql}'")
-          database.select(Sequel.function(:pg_stat_reset)).first
-          database.instance_eval(&block)
-
-          query = database.instance_eval(&block)
-          options[:iterations].to_i.times do
-            database.run(query.sql)
-          end
-
-          stats = database[:pg_stat_database].where(datname: Sequel.function(:current_database)).first
-
-          if stats[:temp_files] == 0 && stats[:temp_bytes] == 0
-            return value
-          end
-        end
-      end
-
-      raise "No work_mem found."
-    end
-  end
-end