15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
|
# File 'lib/ask/rails/tools/query_database.rb', line 15
def execute(sql:, limit: 50)
sql = sql.strip
if WRITE_STATEMENTS.match?(sql)
return Ask::Result.failure(
"Only SELECT queries are allowed. Write statements (#{sql.match(WRITE_STATEMENTS)[1]}) are rejected in all environments."
)
end
if ::Rails.env.production? && !sql.match?(/\A\s*SELECT\b/i)
return Ask::Result.failure(
"Only SELECT queries are allowed in the production environment."
)
end
pool = ActiveRecord::Base.connection_pool
pool.with_connection do |conn|
limited_sql = sql.match?(/\bLIMIT\b/i) ? sql : "#{sql.chomp(';')} LIMIT #{limit.to_i}"
result = conn.exec_query(limited_sql)
columns = result.columns
rows = result.rows.first(limit.to_i).map { |row| build_row(row, columns) }
{
columns: columns,
rows: rows,
count: rows.size,
truncated: result.rows.size > limit.to_i
}
end
rescue ActiveRecord::StatementInvalid => e
Ask::Result.failure("SQL error: #{e.message}")
rescue ActiveRecord::ConnectionNotEstablished => e
Ask::Result.failure("Database not connected: #{e.message}. Verify the database is running and Rails is connected.")
end
|