blob: dc3062a075df65370fd9cc7a6b8ec6e82beee648 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
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
|
#!/usr/bin/env bash
set -e
ROOT_DIR="${LLM_ROOT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)}"
# @meta require-tools usql
# @env LLM_AGENT_VAR_DSN! The database connection url. e.g. pgsql://user:pass@host:port
# @cmd Execute a SELECT query
# @option --query! SELECT SQL query to execute
read_query() {
if ! grep -qi '^select' <<<"$argc_query"; then
echo "error: only SELECT query is allowed" >&2
exit 1
fi
_run_sql "$argc_query"
}
# @cmd Execute an SQL query
# @option --query! SQL query to execute
write_query() {
"$ROOT_DIR/utils/guard_operation.sh" "Execute SQL?"
_run_sql "$argc_query"
}
# @cmd List all tables
list_tables() {
_run_sql "\dt+"
}
# @cmd Get the schema information for a specific table
# @option --table-name! Name of the table to describe
describe_table() {
_run_sql "\d $argc_table_name"
}
_run_sql() {
usql "$LLM_AGENT_VAR_DSN" -c "$1" >> "$LLM_OUTPUT"
}
# See more details at https://github.com/sigoden/argc
eval "$(argc --argc-eval "$0" "$@")"
|
