Class: OxAiWorkers::Tool::Database

Inherits:
Object
  • Object
show all
Includes:
DependencyHelper, LoadI18n, OxAiWorkers::ToolDefinition
Defined in:
lib/oxaiworkers/tool/database.rb

Overview

Connects to a database, executes SQL queries, and outputs DB schema for Agents to use

Gem requirements:

gem "sequel", "~> 5.68.0"

Usage:

database = OxAiWorkers::Tool::Database.new(connection_string: "postgres://user:password@localhost:5432/db_name")

Instance Attribute Summary collapse

Attributes included from LoadI18n

#locale

Attributes included from OxAiWorkers::ToolDefinition

#white_list

Instance Method Summary collapse

Methods included from LoadI18n

#store_locale, #with_locale

Methods included from DependencyHelper

#depends_on

Methods included from OxAiWorkers::ToolDefinition

#define_function, #full_function_name, #function_schemas, #init_white_list_with, #tool_name

Constructor Details

#initialize(connection_string:, tables: [], exclude_tables: [], only: nil) ⇒ Database

Establish a database connection

Parameters:

  • connection_string (String) —

    Database connection info, e.g. ‘postgres://user:password@localhost:5432/db_name’

  • tables (Array<Symbol>) (defaults to: []) —

    The tables to use. Will use all if empty.

  • except_tables (Array<Symbol>) —

    The tables to exclude. Will exclude none if empty.

Raises:

  • (StandardError) 


25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
 
# File 'lib/oxaiworkers/tool/database.rb', line 25

def initialize(connection_string:, tables: [], exclude_tables: [], only: nil)
  depends_on 'sequel'

  raise StandardError, 'connection_string parameter cannot be blank' if connection_string.empty?

  store_locale

  init_white_list_with only

  define_function :list_tables,
                  description: I18n.t('oxaiworkers.tool.database_tool.list_tables.description')

  define_function :describe_tables,
                  description: I18n.t('oxaiworkers.tool.database_tool.describe_tables.description') do
    property :tables, type: 'string',
                      description: I18n.t('oxaiworkers.tool.database_tool.describe_tables.tables'),
                      required: true
  end

  define_function :dump_schema,
                  description: I18n.t('oxaiworkers.tool.database_tool.dump_schema.description')

  define_function :execute,
                  description: I18n.t('oxaiworkers.tool.database_tool.execute.description') do
    property :input, type: 'string',
                     description: I18n.t('oxaiworkers.tool.database_tool.execute.input'),
                     required: true
  end

  @db = Sequel.connect(connection_string)
  @requested_tables = tables
  @excluded_tables = exclude_tables
end

Instance Attribute Details

#db ⇒ Object (readonly)

Returns the value of attribute db.



17
18
19
 
# File 'lib/oxaiworkers/tool/database.rb', line 17

def db
  @db
end

#excluded_tables ⇒ Object (readonly)

Returns the value of attribute excluded_tables.



17
18
19
 
# File 'lib/oxaiworkers/tool/database.rb', line 17

def excluded_tables
  @excluded_tables
end

#requested_tables ⇒ Object (readonly)

Returns the value of attribute requested_tables.



17
18
19
 
# File 'lib/oxaiworkers/tool/database.rb', line 17

def requested_tables
  @requested_tables
end

Instance Method Details

#describe_table(table, schema) ⇒ Object 



88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
 
# File 'lib/oxaiworkers/tool/database.rb', line 88

def describe_table(table, schema)
  primary_key_columns = []
  primary_key_column_count = db.schema(table).count { |column| column[1][:primary_key] == true }

  schema << "CREATE TABLE #{table}(\n"
  db.schema(table).each do |column|
    schema << "#{column[0]} #{column[1][:type]}"
    if column[1][:primary_key] == true
      schema << ' PRIMARY KEY' if primary_key_column_count == 1
    else
      primary_key_columns << column[0]
    end
    schema << ",\n" unless column == db.schema(table).last && primary_key_column_count == 1
  end
  schema << "PRIMARY KEY (#{primary_key_columns.join(',')})" if primary_key_column_count > 1
  db.foreign_key_list(table).each do |fk|
    schema << ",\n" if fk == db.foreign_key_list(table).first
    schema << "FOREIGN KEY (#{fk[:columns][0]}) REFERENCES #{fk[:table]}(#{fk[:key][0]})"
    schema << ",\n" unless fk == db.foreign_key_list(table).last
  end
  schema << ");\n"
end

#describe_tables(tables:) ⇒ String

Database Tool: Returns the schema for a list of tables

Parameters:

  • tables (String) —

    The tables to describe.

Returns:

  • (String) —

    Database schema for the tables



68
69
70
71
72
73
74
 
# File 'lib/oxaiworkers/tool/database.rb', line 68

def describe_tables(tables:)
  schema = ''
  tables.split(',').each do |table|
    describe_table(table, schema)
  end
  schema
end

#dump_schema ⇒ String

Database Tool: Returns the database schema

Returns:

  • (String) —

    Database schema



79
80
81
82
83
84
85
86
 
# File 'lib/oxaiworkers/tool/database.rb', line 79

def dump_schema
  OxAiWorkers.logger.debug('Dumping schema tables and keys', for: self.class)
  schema = ''
  db.tables.each do |table|
    describe_table(table, schema)
  end
  schema
end

#execute(input:) ⇒ Array

Database Tool: Executes a SQL query and returns the results

Parameters:

  • input (String) —

    SQL query to be executed

Returns:

  • (Array) —

    Results from the SQL query



115
116
117
118
119
120
121
122
 
# File 'lib/oxaiworkers/tool/database.rb', line 115

def execute(input:)
  OxAiWorkers.logger.info("Executing \"#{input}\"", for: self.class)

  db[input].to_a
rescue Sequel::DatabaseError => e
  OxAiWorkers.logger.info(e.message, for: self.class)
  e.message
end

#list_tables ⇒ Object

Database Tool: Returns a list of tables in the database



60
61
62
 
# File 'lib/oxaiworkers/tool/database.rb', line 60

def list_tables
  db.tables
end