tobias 0.3.0 → 0.4.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4b50326ed51ec5388f5f63528cbfe0cda2f27db4e505b0cdf1349f1cc9313311
-  data.tar.gz: e76ddcd08800005f89ddb772953df045ad623fcfc13e114090bdfcdd49c0dabb
+  metadata.gz: 6e8024f2a5efa0e4ffe41d2973817f1906fccfb32193782b86e55f822a82f5e7
+  data.tar.gz: 5b6b27888ed7341c76fea8503bc8e577662949903089970930fced41367583e0
 SHA512:
-  metadata.gz: 456a7a52310fcac42a5755ad02aeb464bb429742888a195569fb5d5df7f33416d583cef442eca92b035af9807e7ec5d17c626b0eda20b9d43998fc4c99ccb6f3
-  data.tar.gz: e8c5d112db7cd2e87f37455d90b25a6dfe53eba27e60a492c4fbc4798a941bdc9b69532b813eb39f4c69b68c34812e09112ec19a80485f90e526869e2d0c1794
+  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: 10
     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,7 +55,7 @@ module Tobias
         raise "Script not found at: #{script}"
       end
 
-      container = Container.new(code)
+      container = Container.new(code, database)
       results = {}
 
       parsed = TTY::Markdown.parse(<<~MARKDOWN)

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

data/lib/tobias/evaluations/base.rb

@@ -1,32 +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
+      attr_reader :database, :container, :options
 
-      def initialize(database, container)
+      def initialize(database, container, options)
         @database = database
         @container = container
+        @options = options
       end
 
-      def run(options, &block)
-        results = {}
+      def run(&block)
+        results = Concurrent::Array.new
 
+        container.run_setup
         container.queries.each do |name, query|
-          results.merge!(run_each(name, query, options))
+          result = run_each(name, query)
+          results << result if result
         end
 
         to_markdown(results)
+      ensure
+        container.run_teardown
       end
 
-      def run_each(query, options)
+      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

@@ -3,57 +3,55 @@
 module Tobias
   module Evaluations
     class WorkMem < Base
-
       def work_mems
-        Tobias::WorkMem.valid_for(database)
+        @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
-        "Optional work_mem settings"
+        "Optimal work_mem settings"
       end
 
-      def run_each(name, query, options)
+      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("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
+            container.run_query(query)
 
-            query_result = database.instance_eval(&query)
-            options[:iterations].to_i.times do
-              database.run(query_result.sql)
-            end
-
-            stats = database[:pg_stat_database].where(datname: Sequel.function(:current_database)).first
+            stats = database[:pg_stat_database].
+              where(datname: Sequel.function(:current_database)).
+              first
 
             if stats[:temp_files] == 0 && stats[:temp_bytes] == 0
-              return { name => value }
+              return Result.new(name: name, value: value)
             end
           end
         end
 
-        # TODO: Add a warning message or something.
-        return { name => nil }
+        # 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)
-        current_work_mem = Tobias::WorkMem.from_sql(database.fetch("SHOW work_mem").first[:work_mem])
-
         <<~MARKDOWN
           ## #{description}
 
-          | Query | Required work_mem |
-          |-------|-------------------|
-          #{results.map { |name, work_mem| "| #{name} | #{work_mem.to_sql} |" }.join("\n")}
+          #{render_table(headers: ["Query", "Required work_mem"], body: results.map { |r| [r.name, r.value.to_sql] })}
 
-          Your application will need to run with at least #{results.values.max.to_sql} of work_mem.
+          I see that your current `work_mem` setting is `#{current_work_mem.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.values.max.to_sql}';
+          ALTER SYSTEM SET work_mem = '#{results.max.value.to_sql}';
           SELECT pg_reload_conf();
           ```
         MARKDOWN

data/lib/tobias/evaluations.rb

@@ -4,7 +4,7 @@ module Tobias
   module Evaluations
     def self.run(database, container, options)
       results = []
-      results << WorkMem.new(database, container).run(options)
+      results << WorkMem.new(database, container, options).run
       results
     end
   end

data/lib/tobias/version.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module Tobias
-  VERSION = "0.3.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.rb

@@ -11,15 +11,22 @@ 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__)
 

metadata

@@ -1,13 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: tobias
 version: !ruby/object:Gem::Version
-  version: 0.3.0
+  version: 0.4.0
 platform: ruby
 authors:
 - Jon Daniel
+autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-08-04 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:
@@ -251,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
@@ -265,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: []