Class: RailsAiContext::Tools::Query

Inherits:
BaseTool
  • Object
show all
Defined in:
lib/rails_ai_context/tools/query.rb

Defined Under Namespace

Classes: ResultProxy

Constant Summary collapse

BLOCKED_KEYWORDS =

── Layer 1: SQL validation ─────────────────────────────────────

/\b(INSERT|UPDATE|DELETE|DROP|ALTER|TRUNCATE|CREATE|GRANT|REVOKE|SET|COPY|MERGE|REPLACE)\b/i
BLOCKED_CLAUSES = 
/\bFOR\s+(UPDATE|SHARE|NO\s+KEY\s+UPDATE)\b/i
BLOCKED_SHOWS = 
/\bSHOW\s+(GRANTS|PROCESSLIST|BINLOG|SLAVE|MASTER|REPLICAS)\b/i
SELECT_INTO = 
/\bSELECT\b[^;]*\bINTO\b/i
MULTI_STATEMENT = 
/;\s*\S/
ALLOWED_PREFIX = 
/\A\s*(SELECT|WITH|SHOW|EXPLAIN|DESCRIBE|DESC)\b/i
TAUTOLOGY_PATTERNS =

SQL injection tautology patterns: OR 1=1, OR true, OR ”=”, UNION SELECT, etc.

[
  /\bOR\s+1\s*=\s*1\b/i,
  /\bOR\s+true\b/i,
  /\bOR\s+'[^']*'\s*=\s*'[^']*'/i,
  /\bOR\s+"[^"]*"\s*=\s*"[^"]*"/i,
  /\bOR\s+\d+\s*=\s*\d+/i,
  /\bUNION\s+(ALL\s+)?SELECT\b/i
].freeze
HARD_ROW_CAP = 
1000

Constants inherited from BaseTool

BaseTool::SESSION_CONTEXT, BaseTool::SHARED_CACHE

Class Method Summary collapse

Methods inherited from BaseTool

abstract!, abstract?, cache_key, cached_context, config, extract_method_source_from_file, extract_method_source_from_string, find_closest_match, fuzzy_find_key, inherited, not_found_response, paginate, rails_app, registered_tools, reset_all_caches!, reset_cache!, session_queries, session_record, session_reset!, set_call_params, text_response

Class Method Details

.call(sql: nil, limit: nil, format: "table", explain: false, server_context: nil, **_extra) ⇒ Object 



63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
 
# File 'lib/rails_ai_context/tools/query.rb', line 63

def self.call(sql: nil, limit: nil, format: "table", explain: false, server_context: nil, **_extra)
  set_call_params(sql: sql&.truncate(60))
  # ── Environment guard ───────────────────────────────────────
  unless config.allow_query_in_production || !Rails.env.production?
    return text_response(
      "rails_query is disabled in production for data privacy. " \
      "Set config.allow_query_in_production = true to override."
    )
  end

  # ── ActiveRecord guard (api-only apps) ──────────────────────
  # Must come BEFORE any code that rescues ActiveRecord::* — Ruby
  # resolves rescue class constants at raise time, and `rescue
  # ActiveRecord::ConnectionNotEstablished` crashes with NameError
  # on apps where ActiveRecord is not loaded (e.g.
  # `rails new --api --skip-active-record`).
  unless defined?(ActiveRecord::Base)
    return text_response(
      "Database queries are unavailable: ActiveRecord is not loaded in this app. " \
      "This happens on API-only apps created with `rails new --api --skip-active-record`. " \
      "rails_query requires a database connection to function."
    )
  end

  # ── Layer 1: SQL validation ─────────────────────────────────
  valid, error = validate_sql(sql)
  return text_response(error) unless valid

  # ── EXPLAIN mode ────────────────────────────────────────────
  if explain
    return execute_explain(sql.strip, config.query_timeout)
  end

  # Resolve row limit
  row_limit = limit ? [ limit.to_i, HARD_ROW_CAP ].min : config.query_row_limit
  row_limit = [ row_limit, 1 ].max
  timeout_seconds = config.query_timeout

  # ── Layers 2-3: Execute with DB-level safety + row limit ────
  result = execute_safely(sql.strip, row_limit, timeout_seconds)

  # ── Layer 4: Redact sensitive columns ───────────────────────
  redacted = redact_results(result)

  # ── Format output ───────────────────────────────────────────
  output = case format
  when "csv"
    format_csv(redacted)
  else
    format_table(redacted)
  end

  text_response(output)
rescue ActiveRecord::ConnectionNotEstablished, ActiveRecord::NoDatabaseError => e
  text_response("Database unavailable: #{clean_error_message(e.message)}\n\n**Troubleshooting:**\n- Check `config/database.yml` for correct host/port/credentials\n- Try `RAILS_ENV=test` if the development DB is remote\n- Run `bin/rails db:create` if the database doesn't exist yet")
rescue ActiveRecord::StatementInvalid => e
  if e.message.match?(/timeout|statement_timeout|MAX_EXECUTION_TIME/i)
    text_response("Query exceeded #{config.query_timeout} second timeout. Simplify the query or add indexes.")
  elsif e.message.match?(/could not find|does not exist|Unknown database/i)
    text_response("Database not found: #{clean_error_message(e.message)}\n\n**Troubleshooting:**\n- Run `bin/rails db:create` to create the database\n- Check `config/database.yml` for the correct database name\n- Try `RAILS_ENV=test` if the development DB is remote")
  else
    text_response("SQL error: #{clean_error_message(e.message)}")
  end
rescue => e
  text_response("Query failed: #{clean_error_message(e.message)}")
end

.strip_sql_comments(sql) ⇒ Object

── SQL comment stripping ───────────────────────────────────────



131
132
133
134
135
136
137
 
# File 'lib/rails_ai_context/tools/query.rb', line 131

def self.strip_sql_comments(sql)
  sql
    .gsub(/\/\*.*?\*\//m, " ")   # Block comments: /* ... */
    .gsub(/--[^\n]*/, " ")        # Line comments: -- ...
    .gsub(/^\s*#[^\n]*/m, " ")   # MySQL-style comments: # at line start only
    .squeeze(" ").strip
end

.validate_sql(sql) ⇒ Object

── SQL validation (Layer 1) ────────────────────────────────────



140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
 
# File 'lib/rails_ai_context/tools/query.rb', line 140

def self.validate_sql(sql)
  return [ false, "SQL query is required." ] if sql.nil? || sql.strip.empty?

  cleaned = strip_sql_comments(sql)

  # Check multi-statement and clause patterns first — they provide more
  # specific error messages than the generic keyword blocker.
  return [ false, "Blocked: multiple statements (no semicolons)" ] if cleaned.match?(MULTI_STATEMENT)
  return [ false, "Blocked: FOR UPDATE/SHARE clause" ] if cleaned.match?(BLOCKED_CLAUSES)
  return [ false, "Blocked: sensitive SHOW command" ] if cleaned.match?(BLOCKED_SHOWS)
  return [ false, "Blocked: SELECT INTO creates a table" ] if cleaned.match?(SELECT_INTO)

  # Check for SQL injection tautology patterns (OR 1=1, UNION SELECT, etc.)
  tautology = TAUTOLOGY_PATTERNS.find { |p| cleaned.match?(p) }
  return [ false, "Blocked: SQL injection pattern detected (#{cleaned[tautology]})" ] if tautology

  # Check blocked keywords before the allowed-prefix fallback so that
  # INSERT/UPDATE/DELETE/DROP etc. get a specific "Blocked" error
  # rather than the generic "Only SELECT... allowed" message.
  if (m = cleaned.match(BLOCKED_KEYWORDS))
    return [ false, "Blocked: contains #{m[0]}" ]
  end

  return [ false, "Only SELECT, WITH, SHOW, EXPLAIN, DESCRIBE allowed" ] unless cleaned.match?(ALLOWED_PREFIX)

  [ true, nil ]
end