Setting the file. One moment.
Psql Connect · Aurora Dsql · aws/agent-toolkit-for-aws · Skills Docs
Back to the top of the page 10
Setup DevOps Agent
33
AWS Deployment
scripts/psql-connect.sh
scripts/ psql-connect.sh
Shell · 332 lines · 12 KB
15
set
-euo
pipefail
16
17 # psql-connect.sh - Connect to Aurora DSQL using psql with IAM auth
18 #
19 # Usage: ./psql-connect.sh [CLUSTER_ID|--cluster CLUSTER_ID] [--region REGION] [--user USER] [--admin] [--ai-model MODEL_ID] [--command "SQL" | --script PATH]
20 #
21 # Examples:
22 # ./psql-connect.sh --cluster abc123def456 --ai-model claude-opus-4-6
23 # ./psql-connect.sh abc123def456 --ai-model claude-opus-4-6 --region us-west-2
24 # ./psql-connect.sh --cluster abc123def456 --admin
25 # ./psql-connect.sh --cluster abc123def456 --command "SELECT * FROM entities LIMIT 5"
26 # ./psql-connect.sh --cluster abc123def456 --script ./migration.sql # multi-statement file
27
28 CLUSTER_ID = "${ CLUSTER :- }"
29 REGION = "${ REGION :- ${ AWS_REGION :- us-east-1 }}"
30 # Note: avoid using bare `USER` here — bash sets it automatically to the login
31 # user, and overwriting it would clobber that for child processes.
32 DB_USER_NAME = "${ DB_USER :- admin }"
33 ADMIN = false
34 COMMAND = ""
35 SCRIPT_FILE = ""
36 AI_MODEL = ""
37 SKIP_CERT_VERIFY = false
38
39 # require_value FLAG NEXT — validate that a value-taking flag has a non-empty,
40 # non-flag argument following it. Aborts with a clean error otherwise.
41 require_value () {
42 local flag = " $1 "
43 local next = " ${2 :- } "
44 if [[ -z " $next " ]]; then
45 echo "Error: $flag requires a value." >&2
46 exit 1
47 fi
48 if [[ " $next " == - * ]]; then
49 echo "Error: $flag requires a value, got ' $next ' (looks like another flag)." >&2
50 exit 1
51 fi
52 }
53
54 # Track which source supplied CLUSTER_ID so positional+--cluster mismatch is
55 # caught instead of silently letting the last writer win.
56 CLUSTER_FROM_FLAG = ""
57 CLUSTER_FROM_POSITIONAL = ""
58
59 # set_cluster SOURCE VALUE — record the cluster ID from a specific source and
60 # reject conflicting values from a different source.
61 set_cluster () {
62 local src = " $1 "
63 local val = " $2 "
64 case " $src " in
65 flag )
66 if [[ -n " $CLUSTER_FROM_POSITIONAL " && " $CLUSTER_FROM_POSITIONAL " != " $val " ]]; then
67 echo "Error: cluster id supplied by both --cluster (' $val ') and positional (' $CLUSTER_FROM_POSITIONAL '); they disagree." >&2
68 exit 1
69 fi
70 CLUSTER_FROM_FLAG = " $val "
71 ;;
72 positional )
73 if [[ -n " $CLUSTER_FROM_FLAG " && " $CLUSTER_FROM_FLAG " != " $val " ]]; then
74 echo "Error: cluster id supplied by both positional (' $val ') and --cluster (' $CLUSTER_FROM_FLAG '); they disagree." >&2
75 exit 1
76 fi
77 CLUSTER_FROM_POSITIONAL = " $val "
78 ;;
79 esac
80 CLUSTER_ID = " $val "
81 }
82
83 # Parse arguments
84 while [[ $# -gt 0 ]]; do
85 case $1 in
86 --region )
87 require_value " $1 " " ${2 :- } "
88 REGION = " $2 "
89 shift 2
90 ;;
91 --user )
92 require_value " $1 " " ${2 :- } "
93 DB_USER_NAME = " $2 "
94 shift 2
95 ;;
96 --admin )
97 ADMIN = true
98 shift
99 ;;
100 --command | -c )
101 require_value " $1 " " ${2 :- } "
102 COMMAND = " $2 "
103 shift 2
104 ;;
105 --script | -f )
106 require_value " $1 " " ${2 :- } "
107 SCRIPT_FILE = " $2 "
108 shift 2
109 ;;
110 --cluster )
111 require_value " $1 " " ${2 :- } "
112 set_cluster flag " $2 "
113 shift 2
114 ;;
115 --ai-model )
116 require_value " $1 " " ${2 :- } "
117 AI_MODEL = " $2 "
118 shift 2
119 ;;
120 --skip-cert-verify )
121 SKIP_CERT_VERIFY = true
122 shift
123 ;;
124 -- )
125 # End-of-options sentinel — remaining args are positional.
126 shift
127 while [[ $# -gt 0 ]]; do
128 set_cluster positional " $1 "
129 shift
130 done
131 break
132 ;;
133 -h | --help )
134 echo "Usage: $0 [CLUSTER_ID|--cluster CLUSTER_ID] [--region REGION] [--user USER] [--admin] [--command SQL | --script PATH]"
135 echo ""
136 echo "Connect to Aurora DSQL using psql with IAM authentication."
137 echo ""
138 echo "Arguments:"
139 echo " CLUSTER_ID Cluster identifier (positional, or via --cluster, or \$ CLUSTER env var)"
140 echo ""
141 echo "Options:"
142 echo " --cluster ID Cluster identifier (alternative to positional argument)"
143 echo " --region REGION AWS region (default: \$ REGION or \$ AWS_REGION or us-east-1)"
144 echo " --user USER Database user (default: \$ DB_USER or 'admin')"
145 echo " --admin Generate IAM admin auth token (uses generate-db-connect-admin-auth-token)"
146 echo " --command SQL, -c Execute one SQL statement and exit (single-statement; chained semicolons rejected)"
147 echo " --script PATH, -f Run a multi-statement SQL file via 'psql -f' (no semicolon guard)"
148 echo " --ai-model ID AI model identifier appended to application_name (e.g. claude-opus-4-6)"
149 echo " --skip-cert-verify Downgrade TLS to sslmode=require (encrypt only; vulnerable to MITM)."
150 echo " Do NOT use in production."
151 echo " -h, --help Show this help message"
152 echo ""
153 echo "Environment Variables:"
154 echo " CLUSTER Default cluster identifier"
155 echo " REGION Default AWS region"
156 echo " DB_USER Default database user"
157 exit 0
158 ;;
159 - * )
160 echo "Unknown option: $1 " >&2
161 exit 1
162 ;;
163 *)
164 set_cluster positional " $1 "
165 shift
166 ;;
167 esac
168 done
169
170 # Validate cluster ID — trim surrounding whitespace and enforce DSQL's
171 # alphanumeric format. Catches `--cluster ""`, `--cluster " "`, and accidental
172 # slashes/dots in the ID before they reach the AWS CLI or psql.
173 CLUSTER_ID = "${ CLUSTER_ID # "${ CLUSTER_ID %% [ ! [ : space : ]] * }"}"
174 CLUSTER_ID = "${ CLUSTER_ID % "${ CLUSTER_ID ##* [ ! [ : space : ]]}"}"
175 if [[ -z " $CLUSTER_ID " ]]; then
176 echo "Error: CLUSTER_ID is required. Set \$ CLUSTER env var or pass as argument." >&2
177 echo "" >&2
178 echo "Usage: $0 [CLUSTER_ID|--cluster CLUSTER_ID] [options]" >&2
179 echo " or: export CLUSTER=abc123 && $0 [options]" >&2
180 exit 1
181 fi
182 if [[ ! " $CLUSTER_ID " =~ ^[a-z0-9]+$ ]]; then
183 echo "Error: CLUSTER_ID ' $CLUSTER_ID ' is invalid (DSQL cluster IDs are lowercase alphanumeric)." >&2
184 exit 1
185 fi
186
187 # Build endpoint
188 ENDPOINT = "${ CLUSTER_ID }.dsql.${ REGION }.on.aws"
189
190 # Generate auth token. Capture stderr alongside stdout so an aws CLI failure
191 # (expired creds, missing dsql:DbConnect, wrong region) surfaces a useful
192 # message — under `set -e` the bare command-substitution would otherwise abort
193 # the script before the empty-token guard below could fire.
194 echo "Generating IAM auth token for $ENDPOINT ..." >&2
195
196 if [[ " $ADMIN " == "true" ]]; then
197 TOKEN_CMD = ( aws dsql generate-db-connect-admin-auth-token --hostname " $ENDPOINT " --region " $REGION " )
198 else
199 TOKEN_CMD = ( aws dsql generate-db-connect-auth-token --hostname " $ENDPOINT " --region " $REGION " )
200 fi
201
202 if ! TOKEN = $( "${ TOKEN_CMD [ @ ]}" 2>&1 ); then
203 echo "Error: Failed to generate auth token (aws CLI exited non-zero)." >&2
204 echo " Command: ${ TOKEN_CMD [ * ]}" >&2
205 echo " Output: $TOKEN " >&2
206 exit 1
207 fi
208
209 # Check if token generation was successful
210 if [[ -z " $TOKEN " ]]; then
211 echo "Error: Failed to generate auth token (empty result). Check your AWS credentials." >&2
212 exit 1
213 fi
214
215 echo "Connecting to $ENDPOINT as $DB_USER_NAME ..." >&2
216 echo "" >&2
217
218 # DSQL requires TLS and rejects non-TLS connections. Default to verify-full
219 # which validates the server certificate against DSQL's CA, preventing MITM
220 # attacks. Point sslrootcert at the OS trust store so users don't need a
221 # per-user ~/.postgresql/root.crt. Use --skip-cert-verify to downgrade to
222 # require (encrypt only).
223 # See https://docs.aws.amazon.com/aurora-dsql/latest/userguide/accessing-psql.html
224 if [[ " $SKIP_CERT_VERIFY " == "true" ]]; then
225 echo "WARNING: Certificate verification disabled. Connection is vulnerable to MITM attacks. Do NOT use in production." >&2
226 export PGSSLMODE = require
227 else
228 export PGSSLMODE = verify-full
229 # libpq defaults to ~/.postgresql/root.crt — fall back to the OS trust store
230 # when the user has not provisioned a personal CA bundle. Honor any caller-
231 # supplied PGSSLROOTCERT (e.g., a corporate bundle) by not overwriting it.
232 : "${ PGSSLROOTCERT := system }"
233 export PGSSLROOTCERT
234 fi
235
236 # Set application_name with AI model identifier if provided
237 PGAPPNAME = "dsql-skill"
238 if [[ -n " $AI_MODEL " ]]; then
239 # Validate: allow only alphanumeric, hyphens, underscores, and dots
240 if [[ ! " $AI_MODEL " =~ ^[a-zA-Z0-9._-]+$ ]]; then
241 echo "Error: --ai-model must contain only alphanumeric characters, hyphens, underscores, and dots." >&2
242 exit 1
243 fi
244 PGAPPNAME = "dsql-skill/${ AI_MODEL }"
245 fi
246 export PGAPPNAME
247
248 # Sanitize --command input: reject multi-statement chaining and comment injection.
249 # psql -c runs a single command; allow at most ONE trailing semicolon.
250 # This is a defense-in-depth measure — callers should also validate inputs.
251 # Limitations: does not handle escaped quotes (\' or ''), dollar-quoted strings
252 # ($$...$$), or all edge cases. For complex queries, use --script PATH instead
253 # to pipe a multi-statement file via stdin without the semicolon guard.
254 if [[ -n " $COMMAND " && -n " $SCRIPT_FILE " ]]; then
255 echo "Error: --command and --script are mutually exclusive." >&2
256 exit 1
257 fi
258
259 if [[ -n " $COMMAND " ]]; then
260 # Reject whitespace-only --command early so the user gets a clear error
261 # rather than psql's downstream syntax message.
262 if [[ -z "${ COMMAND // [[ : space : ]] / }" ]]; then
263 echo "Error: --command is whitespace-only." >&2
264 exit 1
265 fi
266 # Reject newlines — sed processes the strip-quotes pipeline line by line, so a
267 # newline-spanning literal would defeat the multi-statement detector. Use
268 # --script for SQL that needs to span multiple lines.
269 if [[ " $COMMAND " == * $' \n ' * ]]; then
270 echo "Error: --command does not support newlines. Use --script PATH for multi-line SQL." >&2
271 exit 1
272 fi
273 # Reject dollar-quoting which can interfere with single-quote stripping
274 if echo " $COMMAND " | grep -qE '\$\$|\$[a-zA-Z_][a-zA-Z0-9_]*\$' ; then
275 echo "Error: Dollar-quoting is not supported in --command. Use --script PATH for SQL with dollar-quoted strings." >&2
276 exit 1
277 fi
278
279 # Reject multi-statement chaining (semicolons outside string/identifier
280 # literals, ignoring an optional trailing whitespace/semicolon at the end).
281 # Strip in this order: (1) collapse SQL-standard doubled-quote escapes ('')
282 # so the next pass treats them as empty literals; (2) strip single-quoted
283 # string literals; (3) strip double-quoted identifiers; (4) trim a single
284 # trailing semicolon. Any semicolon that survives is genuine statement
285 # chaining.
286 stripped = $( echo " $COMMAND " \
287 | sed "s/''//g" \
288 | sed "s/'[^']*'//g" \
289 | sed 's/"[^"]*"//g' \
290 | sed -E 's/[[:space:]]*;[[:space:]]*$//' )
291 if echo " $stripped " | grep -q ';' ; then
292 echo "Error: Multiple SQL statements are not allowed in --command. Use --script PATH for multi-statement input." >&2
293 exit 1
294 fi
295 # Reject SQL comment sequences that could hide injected code
296 if echo " $stripped " | grep -qE -- '--|/\*' ; then
297 echo "Error: SQL comments (-- or /*) are not allowed in --command. Use --script PATH if you need comments." >&2
298 exit 1
299 fi
300
301 # Execute command and exit
302 exec env PGPASSWORD=" $TOKEN " psql \
303 -h " $ENDPOINT " \
304 -U " $DB_USER_NAME " \
305 -d postgres \
306 -c " $COMMAND "
307 elif [[ -n " $SCRIPT_FILE " ]]; then
308 # Multi-statement file mode — no semicolon guard. Caller is responsible for
309 # the contents of the file; build the SQL with safe_query.build() upstream
310 # whenever values come from untrusted input.
311 if [[ ! -f " $SCRIPT_FILE " ]]; then
312 echo "Error: --script path ' $SCRIPT_FILE ' is not a regular file." >&2
313 exit 1
314 fi
315 if [[ ! -r " $SCRIPT_FILE " ]]; then
316 echo "Error: --script file ' $SCRIPT_FILE ' is not readable." >&2
317 exit 1
318 fi
319 exec env PGPASSWORD=" $TOKEN " psql \
320 -P pager=off \
321 -v ON_ERROR_STOP= 1 \
322 -h " $ENDPOINT " \
323 -U " $DB_USER_NAME " \
324 -d postgres \
325 -f " $SCRIPT_FILE "
326 else
327 # Interactive session
328 exec env PGPASSWORD=" $TOKEN " psql \
329 -h " $ENDPOINT " \
330 -U " $DB_USER_NAME " \
331 -d postgres
332 fi