1use std::io::{Read, Write};
2
3use crate::limits::{MAX_PARAMS, MAX_SQL_BYTES};
4use crate::types::{ContainerConfig, Permission, QueryOptions, SessionConfig, SshConfig};
5use agent_first_data::{cli_parse_log_filters, cli_parse_output, OutputFormat};
6use clap::{Args, CommandFactory, Parser, Subcommand, ValueEnum};
7use serde_json::{json, Value};
8use std::collections::{btree_map::Entry, BTreeMap};
9
10const STARTUP_ENV_KEYS: &[&str] = &[
11 "AFPSQL_DSN_SECRET",
12 "AFPSQL_CONNINFO_SECRET",
13 "AFPSQL_HOST",
14 "AFPSQL_PORT",
15 "AFPSQL_USER",
16 "AFPSQL_DBNAME",
17 "AFPSQL_PASSWORD_SECRET",
18 "AFPSQL_SSH",
19 "AFPSQL_SSH_LOCAL_HOST",
20 "AFPSQL_SSH_LOCAL_PORT",
21 "AFPSQL_SSH_REMOTE_SOCKET",
22 "AFPSQL_SSH_SUDO_USER",
23 "AFPSQL_CONTAINER",
24 "AFPSQL_CONTAINER_DRIVER",
25 "AFPSQL_CONTAINER_RUNTIME",
26 "AFPSQL_CONTAINER_USER",
27 "AFPSQL_CONTAINER_NAMESPACE",
28 "AFPSQL_CONTAINER_CONTEXT",
29 "AFPSQL_CONTAINER_COMPOSE_FILE",
30 "AFPSQL_CONTAINER_COMPOSE_PROJECT",
31 "AFPSQL_CONTAINER_POD_CONTAINER",
32 "PGHOST",
33 "PGPORT",
34 "PGUSER",
35 "PGDATABASE",
36 "PGPASSWORD",
37 "PGSSLMODE",
38];
39
40pub enum Mode {
41 Cli(CliRequest),
42 Pipe(PipeInit),
43 PsqlAdmin(PsqlAdminRequest),
44 SkillAdmin(SkillAdminRequest),
45 PsqlUnsupported(PsqlUnsupportedRequest),
46}
47
48pub struct PipeInit {
49 pub output: OutputFormat,
50 pub session: SessionConfig,
51 pub log: Vec<String>,
52 pub startup_args: Value,
53 pub startup_env: Value,
54 pub startup_requested: bool,
55}
56
57#[derive(Debug, Clone)]
58pub struct PsqlAdminRequest {
59 pub action: PsqlAdminAction,
60 pub output: OutputFormat,
61}
62
63#[derive(Debug, Clone)]
64pub enum PsqlAdminAction {
65 Status { bin_dir: Option<String> },
66 Install { bin_dir: Option<String> },
67 Uninstall { bin_dir: Option<String> },
68}
69
70#[derive(Debug, Clone)]
71pub struct SkillAdminRequest {
72 pub action: SkillAdminAction,
73 pub output: OutputFormat,
74}
75
76#[derive(Debug, Clone)]
77pub enum SkillAdminAction {
78 Status(SkillAdminOptions),
79 Install(SkillAdminOptions),
80 Uninstall(SkillAdminOptions),
81}
82
83#[derive(Debug, Clone)]
84pub struct SkillAdminOptions {
85 pub agent: SkillAgentSelection,
86 pub scope: SkillScope,
87 pub skills_dir: Option<String>,
88 pub force: bool,
89}
90
91#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
92pub enum SkillAgentSelection {
93 All,
95 Codex,
97 #[value(name = "claude-code", alias = "claude")]
99 ClaudeCode,
100}
101
102#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
103pub enum SkillScope {
104 Personal,
106 Project,
108}
109
110pub struct CliRequest {
111 pub sql: String,
112 pub params: Vec<Value>,
113 pub options: QueryOptions,
114 pub session: SessionConfig,
115 pub output: OutputFormat,
116 pub output_file: Option<String>,
117 pub log_file: Option<String>,
118 pub log: Vec<String>,
119 pub startup_args: Value,
120 pub startup_env: Value,
121 pub startup_requested: bool,
122 pub dry_run: bool,
123 pub psql_mode: bool,
124}
125
126pub struct PsqlUnsupportedRequest {
127 pub reason: String,
128}
129
130#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
131enum RuntimeMode {
132 Cli,
133 Pipe,
134 #[value(name = "psql")]
135 Psql,
136}
137
138#[derive(Subcommand)]
139enum AfdCommand {
140 Psql(PsqlCommand),
142 Skill(SkillCommand),
144 Inspect(InspectCommand),
146}
147
148#[derive(Args)]
149struct InspectCommand {
150 #[command(subcommand)]
151 action: InspectAction,
152}
153
154#[derive(Subcommand)]
155enum InspectAction {
156 Databases,
158 Schemas,
160 Tables(InspectTablesArgs),
162 Views(InspectViewsArgs),
164 Table(InspectTableArgs),
166}
167
168#[derive(Args)]
169struct InspectTablesArgs {
170 #[arg(long = "schema", default_value = "public")]
172 schema: String,
173 #[arg(long = "like")]
175 like: Option<String>,
176}
177
178#[derive(Args)]
179struct InspectViewsArgs {
180 #[arg(long = "schema", default_value = "public")]
182 schema: String,
183 #[arg(long = "like")]
185 like: Option<String>,
186}
187
188#[derive(Args)]
189struct InspectTableArgs {
190 name: String,
192}
193
194#[derive(Args)]
195struct PsqlCommand {
196 #[command(subcommand)]
197 action: PsqlCliAction,
198}
199
200#[derive(Subcommand)]
201enum PsqlCliAction {
202 Status(PsqlPathArgs),
204 Install(PsqlPathArgs),
206 Uninstall(PsqlPathArgs),
208}
209
210#[derive(Args)]
211struct PsqlPathArgs {
212 #[arg(long = "bin-dir")]
214 bin_dir: Option<String>,
215}
216
217#[derive(Args)]
218struct SkillCommand {
219 #[command(subcommand)]
220 action: SkillCliAction,
221}
222
223#[derive(Subcommand)]
224enum SkillCliAction {
225 Status(SkillTargetArgs),
227 Install(SkillWriteArgs),
229 Uninstall(SkillWriteArgs),
231}
232
233#[derive(Args)]
234struct SkillTargetArgs {
235 #[arg(long = "agent", value_enum, default_value_t = SkillAgentSelection::All)]
237 agent: SkillAgentSelection,
238 #[arg(long = "scope", value_enum, default_value_t = SkillScope::Personal)]
240 scope: SkillScope,
241 #[arg(long = "skills-dir")]
243 skills_dir: Option<String>,
244}
245
246#[derive(Args)]
247struct SkillWriteArgs {
248 #[command(flatten)]
249 target: SkillTargetArgs,
250 #[arg(long)]
252 force: bool,
253}
254
255#[doc = r#"Agent-First PostgreSQL client.
256
257`afpsql` gives agents a reliable PostgreSQL contract: structured stdout
258events, first-class SSH/container transports, explicit write permissions,
259stable pipe sessions, and machine-readable failures.
260
261### Interface Policy
262
263- default mode is canonical agent-first CLI
264- `--mode psql` is argument translation only; runtime output stays JSONL
265- stdout carries protocol events; stderr is not a protocol channel
266- native CLI and pipe mode default to read-only transactions; writes require permission
267- SSH/container transports keep afpsql local instead of running human `psql` across boundaries
268
269### Query Sources and Parameters
270
271- use `--sql` for inline SQL or `--sql-file` for a file
272- use repeatable `--param N=value` for positional binds
273- placeholder count is validated from prepared-statement metadata, not by SQL text scanning
274
275### Connection Sources
276
277- `--dsn-secret` for a PostgreSQL URI
278- `--conninfo-secret` for libpq-style conninfo
279- or discrete `--host`, `--port`, `--user`, `--dbname`, `--password-secret`
280- add `--ssh user@server` when PostgreSQL is reachable only from the server boundary
281- add `--container TARGET` when PostgreSQL is reachable only from inside a container boundary
282- use named container scope flags instead of raw driver option passthrough
283- use `--container-driver docker|podman|nerdctl|compose|kubectl` for the exec syntax
284- combine `--ssh user@server --container TARGET` for containers on an SSH host
285- agent-first environment fallbacks: `AFPSQL_*`
286- PostgreSQL environment fallbacks: `PGHOST`, `PGPORT`, `PGUSER`, `PGDATABASE`, `PGPASSWORD`, `PGSSLMODE`
287
288### Result Shaping
289
290- default mode buffers a bounded inline result
291- use `--stream-rows` for large result sets, with `--batch-rows` and `--batch-bytes` to tune chunk size
292- `--output json|yaml|plain` changes rendering only, not the runtime schema
293
294### Examples
295
296```text
297afpsql --sql "select now() as now_rfc3339"
298afpsql --sql-file ./query.sql
299afpsql --sql "select * from users where id = $1" --param 1=123
300afpsql --dsn-secret-env DATABASE_URL --sql "select 1"
301afpsql --ssh user@server --host 127.0.0.1 --port 5432 --user app --dbname appdb --sql "select 1"
302afpsql --container pg-container --dsn-secret-env DATABASE_URL --sql "select 1"
303afpsql --ssh root@server --container app --host host.container.internal --port 5432 --user app --dbname appdb --sql "select 1"
304afpsql --mode psql -h 127.0.0.1 -p 5432 -U app -d appdb -c "select 1"
305afpsql --sql "select * from big_table" --stream-rows --batch-rows 1000
306afpsql --mode pipe
307afpsql psql status
308afpsql psql install
309afpsql skill status
310afpsql skill install
311```
312
313### Exit Codes
314
315- `0`: query completed successfully
316- `1`: SQL error or runtime error
317- `2`: invalid CLI arguments
318"#]
319#[derive(Parser)]
320#[command(name = "afpsql", version, verbatim_doc_comment)]
321pub struct AfdCli {
322 #[arg(long, allow_hyphen_values = true, help_heading = "Query")]
324 sql: Option<String>,
325 #[arg(long = "sql-file", allow_hyphen_values = true, help_heading = "Query")]
327 sql_file: Option<String>,
328 #[arg(long = "param", help_heading = "Query")]
330 param: Vec<String>,
331 #[arg(long = "stream-rows", help_heading = "Query")]
333 stream_rows: bool,
334 #[arg(long = "batch-rows", help_heading = "Query")]
336 batch_rows: Option<usize>,
337 #[arg(long = "batch-bytes", help_heading = "Query")]
339 batch_bytes: Option<usize>,
340 #[arg(long = "statement-timeout-ms", help_heading = "Query")]
342 statement_timeout_ms: Option<u64>,
343 #[arg(long = "lock-timeout-ms", help_heading = "Query")]
345 lock_timeout_ms: Option<u64>,
346 #[arg(long = "inline-max-rows", help_heading = "Query")]
348 inline_max_rows: Option<usize>,
349 #[arg(long = "inline-max-bytes", help_heading = "Query")]
351 inline_max_bytes: Option<usize>,
352 #[arg(long = "permission", value_parser = parse_permission_arg, help_heading = "Query")]
355 permission: Option<Permission>,
356 #[arg(long, help_heading = "Query")]
358 dry_run: bool,
359 #[arg(
362 long = "explain",
363 help_heading = "Query",
364 conflicts_with = "explain_analyze"
365 )]
366 explain: bool,
367 #[arg(long = "explain-analyze", help_heading = "Query")]
370 explain_analyze: bool,
371
372 #[arg(long = "dsn-secret", help_heading = "Connection")]
374 dsn_secret: Option<String>,
375 #[arg(long = "dsn-secret-env", help_heading = "Connection")]
377 dsn_secret_env: Option<String>,
378 #[arg(long = "conninfo-secret", help_heading = "Connection")]
380 conninfo_secret: Option<String>,
381 #[arg(long, help_heading = "Connection")]
383 host: Option<String>,
384 #[arg(long, help_heading = "Connection")]
386 port: Option<u16>,
387 #[arg(long, help_heading = "Connection")]
389 user: Option<String>,
390 #[arg(long, help_heading = "Connection")]
392 dbname: Option<String>,
393 #[arg(long = "password-secret", help_heading = "Connection")]
395 password_secret: Option<String>,
396 #[arg(long = "password-secret-env", help_heading = "Connection")]
398 password_secret_env: Option<String>,
399 #[arg(long = "ssh", help_heading = "SSH Transport")]
401 ssh: Option<String>,
402 #[arg(long = "ssh-option", help_heading = "SSH Transport")]
404 ssh_options: Vec<String>,
405 #[arg(long = "ssh-local-host", help_heading = "SSH Transport")]
407 ssh_local_host: Option<String>,
408 #[arg(long = "ssh-local-port", help_heading = "SSH Transport")]
410 ssh_local_port: Option<u16>,
411 #[arg(long = "ssh-remote-socket", help_heading = "SSH Transport")]
413 ssh_remote_socket: Option<String>,
414 #[arg(long = "ssh-sudo-user", help_heading = "SSH Transport")]
416 ssh_sudo_user: Option<String>,
417
418 #[arg(long = "container", help_heading = "Container Transport")]
420 container: Option<String>,
421 #[arg(long = "container-driver", help_heading = "Container Transport")]
423 container_driver: Option<String>,
424 #[arg(long = "container-runtime", help_heading = "Container Transport")]
426 container_runtime: Option<String>,
427 #[arg(long = "container-user", help_heading = "Container Transport")]
429 container_user: Option<String>,
430 #[arg(long = "container-namespace", help_heading = "Container Transport")]
432 container_namespace: Option<String>,
433 #[arg(long = "container-context", help_heading = "Container Transport")]
435 container_context: Option<String>,
436 #[arg(long = "container-compose-file", help_heading = "Container Transport")]
438 container_compose_files: Vec<String>,
439 #[arg(
441 long = "container-compose-project",
442 help_heading = "Container Transport"
443 )]
444 container_compose_project: Option<String>,
445 #[arg(long = "container-pod-container", help_heading = "Container Transport")]
447 container_pod_container: Option<String>,
448
449 #[arg(long, default_value = "json", global = true, help_heading = "Runtime")]
451 output: String,
452 #[arg(long = "log", value_delimiter = ',', help_heading = "Runtime")]
454 log: Vec<String>,
455 #[arg(long, value_enum, default_value_t = RuntimeMode::Cli, help_heading = "Runtime")]
457 mode: RuntimeMode,
458
459 #[command(subcommand)]
460 command: Option<AfdCommand>,
461}
462
463pub fn parse_args() -> Result<Mode, String> {
464 let raw: Vec<String> = std::env::args().collect();
465 if is_psql_mode_requested(&raw) {
466 return parse_psql_mode(&raw);
467 }
468 let startup_requested = startup_requested_from_raw(&raw);
469
470 if top_level_help_requested(&raw) {
472 let _ = writeln!(
473 std::io::stdout(),
474 "{}",
475 agent_first_data::cli_render_help(&AfdCli::command(), &[])
476 );
477 std::process::exit(0);
478 }
479 if top_level_help_markdown_requested(&raw) {
481 let _ = writeln!(
482 std::io::stdout(),
483 "{}",
484 agent_first_data::cli_render_help_markdown(&AfdCli::command(), &[])
485 );
486 std::process::exit(0);
487 }
488
489 let cli = match AfdCli::try_parse_from(&raw) {
490 Ok(c) => c,
491 Err(e) => {
492 use clap::error::ErrorKind;
493 if matches!(e.kind(), ErrorKind::DisplayVersion | ErrorKind::DisplayHelp) {
494 let _ = writeln!(std::io::stdout(), "{e}");
495 std::process::exit(0);
496 }
497 return Err(e.to_string());
498 }
499 };
500 let output = parse_output(&cli.output)?;
501 let log = parse_log_categories(&cli.log);
502 let dsn_secret = resolve_secret_value(
503 "--dsn-secret",
504 cli.dsn_secret,
505 cli.dsn_secret_env.as_deref(),
506 )?;
507 let password_secret = resolve_secret_value(
508 "--password-secret",
509 cli.password_secret,
510 cli.password_secret_env.as_deref(),
511 )?;
512 let session = SessionConfig {
513 dsn_secret,
514 conninfo_secret: cli.conninfo_secret,
515 host: cli.host,
516 port: cli.port,
517 user: cli.user,
518 dbname: cli.dbname,
519 password_secret,
520 ssh: SshConfig {
521 destination: cli.ssh.or_else(|| std::env::var("AFPSQL_SSH").ok()),
522 options: cli.ssh_options,
523 local_host: cli
524 .ssh_local_host
525 .or_else(|| std::env::var("AFPSQL_SSH_LOCAL_HOST").ok()),
526 local_port: cli.ssh_local_port.or_else(|| {
527 std::env::var("AFPSQL_SSH_LOCAL_PORT")
528 .ok()
529 .and_then(|v| v.parse().ok())
530 }),
531 remote_socket: cli
532 .ssh_remote_socket
533 .or_else(|| std::env::var("AFPSQL_SSH_REMOTE_SOCKET").ok()),
534 sudo_user: cli
535 .ssh_sudo_user
536 .or_else(|| std::env::var("AFPSQL_SSH_SUDO_USER").ok()),
537 },
538 container: ContainerConfig {
539 target: cli
540 .container
541 .or_else(|| std::env::var("AFPSQL_CONTAINER").ok()),
542 driver: cli
543 .container_driver
544 .or_else(|| std::env::var("AFPSQL_CONTAINER_DRIVER").ok()),
545 runtime: cli
546 .container_runtime
547 .or_else(|| std::env::var("AFPSQL_CONTAINER_RUNTIME").ok()),
548 user: cli
549 .container_user
550 .or_else(|| std::env::var("AFPSQL_CONTAINER_USER").ok()),
551 namespace: cli
552 .container_namespace
553 .or_else(|| std::env::var("AFPSQL_CONTAINER_NAMESPACE").ok()),
554 context: cli
555 .container_context
556 .or_else(|| std::env::var("AFPSQL_CONTAINER_CONTEXT").ok()),
557 compose_files: resolve_container_compose_files(cli.container_compose_files),
558 compose_project: cli
559 .container_compose_project
560 .or_else(|| std::env::var("AFPSQL_CONTAINER_COMPOSE_PROJECT").ok()),
561 pod_container: cli
562 .container_pod_container
563 .or_else(|| std::env::var("AFPSQL_CONTAINER_POD_CONTAINER").ok()),
564 },
565 };
566 let mode_name = match cli.mode {
567 RuntimeMode::Cli => "cli",
568 RuntimeMode::Pipe => "pipe",
569 RuntimeMode::Psql => "psql",
570 };
571 let startup_env = startup_env_snapshot();
572
573 if let Some(command) = cli.command {
574 return match command {
575 AfdCommand::Psql(psql) => Ok(Mode::PsqlAdmin(PsqlAdminRequest {
576 action: psql_admin_action(psql.action),
577 output,
578 })),
579 AfdCommand::Skill(skill) => Ok(Mode::SkillAdmin(SkillAdminRequest {
580 action: skill_admin_action(skill.action),
581 output,
582 })),
583 AfdCommand::Inspect(inspect) => {
584 let (sql, params) = build_inspect_sql(inspect.action);
585 let startup_args = startup_args(mode_name, Some(&sql), None, params.len());
586 Ok(Mode::Cli(CliRequest {
587 sql,
588 params,
589 options: QueryOptions::default(),
590 session,
591 output,
592 output_file: None,
593 log_file: None,
594 log,
595 startup_args,
596 startup_env,
597 startup_requested,
598 dry_run: false,
599 psql_mode: false,
600 }))
601 }
602 };
603 }
604
605 match cli.mode {
606 RuntimeMode::Pipe => {
607 return Ok(Mode::Pipe(PipeInit {
608 output,
609 session,
610 log: log.clone(),
611 startup_args: startup_args(mode_name, None, None, 0),
612 startup_env,
613 startup_requested,
614 }));
615 }
616 RuntimeMode::Cli | RuntimeMode::Psql => {}
617 }
618
619 let startup_sql_file = cli.sql_file.clone();
620 let user_sql = load_sql(cli.sql, cli.sql_file)?;
621 let params = parse_params(&cli.param)?;
622 let sql = if cli.explain {
623 wrap_explain_sql(&user_sql, false)
624 } else if cli.explain_analyze {
625 wrap_explain_sql(&user_sql, true)
626 } else {
627 user_sql
628 };
629 let startup_args = startup_args(
630 mode_name,
631 Some(&sql),
632 startup_sql_file.as_deref(),
633 params.len(),
634 );
635
636 let options = QueryOptions {
637 stream_rows: cli.stream_rows,
638 batch_rows: cli.batch_rows,
639 batch_bytes: cli.batch_bytes,
640 statement_timeout_ms: cli.statement_timeout_ms,
641 lock_timeout_ms: cli.lock_timeout_ms,
642 permission: cli.permission,
643 inline_max_rows: cli.inline_max_rows,
644 inline_max_bytes: cli.inline_max_bytes,
645 };
646
647 Ok(Mode::Cli(CliRequest {
648 sql,
649 params,
650 options,
651 session,
652 output,
653 output_file: None,
654 log_file: None,
655 log,
656 startup_args,
657 startup_env,
658 startup_requested,
659 dry_run: cli.dry_run,
660 psql_mode: false,
661 }))
662}
663
664fn parse_psql_mode(raw: &[String]) -> Result<Mode, String> {
665 let startup_requested = startup_requested_from_raw(raw);
666 let mut state = PsqlModeState::default();
667
668 let mut i = 1usize;
669 while i < raw.len() {
670 let arg = raw[i].as_str();
671 if arg == "--" {
672 i += 1;
673 while i < raw.len() {
674 state.positionals.push(raw[i].clone());
675 i += 1;
676 }
677 break;
678 }
679 if arg.starts_with("--") {
680 parse_psql_long_arg(raw, &mut i, &mut state)?;
681 continue;
682 }
683 if arg.starts_with('-') && arg.len() > 1 {
684 parse_psql_short_arg(raw, &mut i, &mut state)?;
685 continue;
686 }
687 state.positionals.push(raw[i].clone());
688 i += 1;
689 }
690
691 if let Some(reason) = state.interactive_reason {
692 return Ok(Mode::PsqlUnsupported(PsqlUnsupportedRequest { reason }));
693 }
694
695 apply_psql_positionals(&mut state)?;
696 if state.list_databases {
697 state.sql = Some(psql_list_databases_sql());
698 state.sql_file = None;
699 }
700 if state.sql.is_none() && state.sql_file.is_none() {
701 return Ok(Mode::PsqlUnsupported(PsqlUnsupportedRequest {
702 reason: "no -c/--command, -f/--file, or -l/--list was provided".to_string(),
703 }));
704 }
705
706 let dsn_secret = resolve_secret_value(
707 "--dsn-secret",
708 state.dsn_secret,
709 state.dsn_secret_env.as_deref(),
710 )?;
711 let password_secret = resolve_secret_value(
712 "--password-secret",
713 state.password_secret,
714 state.password_secret_env.as_deref(),
715 )?;
716 let session = SessionConfig {
717 dsn_secret,
718 conninfo_secret: state.conninfo_secret,
719 host: state.host,
720 port: state.port,
721 user: state.user,
722 dbname: state.dbname,
723 password_secret,
724 ssh: SshConfig::default(),
725 container: ContainerConfig {
726 target: state
727 .container
728 .or_else(|| std::env::var("AFPSQL_CONTAINER").ok()),
729 driver: state
730 .container_driver
731 .or_else(|| std::env::var("AFPSQL_CONTAINER_DRIVER").ok()),
732 runtime: state
733 .container_runtime
734 .or_else(|| std::env::var("AFPSQL_CONTAINER_RUNTIME").ok()),
735 user: state
736 .container_user
737 .or_else(|| std::env::var("AFPSQL_CONTAINER_USER").ok()),
738 namespace: state
739 .container_namespace
740 .or_else(|| std::env::var("AFPSQL_CONTAINER_NAMESPACE").ok()),
741 context: state
742 .container_context
743 .or_else(|| std::env::var("AFPSQL_CONTAINER_CONTEXT").ok()),
744 compose_files: resolve_container_compose_files(state.container_compose_files),
745 compose_project: state
746 .container_compose_project
747 .or_else(|| std::env::var("AFPSQL_CONTAINER_COMPOSE_PROJECT").ok()),
748 pod_container: state
749 .container_pod_container
750 .or_else(|| std::env::var("AFPSQL_CONTAINER_POD_CONTAINER").ok()),
751 },
752 };
753
754 let startup_sql_file = state.sql_file.clone();
755 let sql = load_sql(state.sql, state.sql_file)?;
756 let params = parse_params(&state.params_kv)?;
757 let startup_args = psql_startup_args(PsqlStartupArgs {
758 mode: "psql",
759 sql: Some(&sql),
760 sql_file: startup_sql_file,
761 param_count: params.len(),
762 });
763 Ok(Mode::Cli(CliRequest {
764 sql,
765 params,
766 options: QueryOptions {
767 permission: Some(if session.uses_container_transport() {
768 Permission::ContainerWrite
769 } else {
770 Permission::Write
771 }),
772 ..Default::default()
773 },
774 session,
775 output: state.output,
776 output_file: state.output_file,
777 log_file: state.log_file,
778 log: parse_log_categories(&state.log_entries),
779 startup_args,
780 startup_env: startup_env_snapshot(),
781 startup_requested,
782 dry_run: false,
783 psql_mode: true,
784 }))
785}
786
787struct PsqlModeState {
788 sql: Option<String>,
789 sql_file: Option<String>,
790 host: Option<String>,
791 port: Option<u16>,
792 user: Option<String>,
793 dbname: Option<String>,
794 dsn_secret: Option<String>,
795 dsn_secret_env: Option<String>,
796 conninfo_secret: Option<String>,
797 password_secret: Option<String>,
798 password_secret_env: Option<String>,
799 container: Option<String>,
800 container_driver: Option<String>,
801 container_runtime: Option<String>,
802 container_user: Option<String>,
803 container_namespace: Option<String>,
804 container_context: Option<String>,
805 container_compose_files: Vec<String>,
806 container_compose_project: Option<String>,
807 container_pod_container: Option<String>,
808 params_kv: Vec<String>,
809 output: OutputFormat,
810 log_entries: Vec<String>,
811 output_file: Option<String>,
812 log_file: Option<String>,
813 list_databases: bool,
814 positionals: Vec<String>,
815 interactive_reason: Option<String>,
816}
817
818impl Default for PsqlModeState {
819 fn default() -> Self {
820 Self {
821 sql: None,
822 sql_file: None,
823 host: None,
824 port: None,
825 user: None,
826 dbname: None,
827 dsn_secret: None,
828 dsn_secret_env: None,
829 conninfo_secret: None,
830 password_secret: None,
831 password_secret_env: None,
832 container: None,
833 container_driver: None,
834 container_runtime: None,
835 container_user: None,
836 container_namespace: None,
837 container_context: None,
838 container_compose_files: vec![],
839 container_compose_project: None,
840 container_pod_container: None,
841 params_kv: vec![],
842 output: OutputFormat::Json,
843 log_entries: vec![],
844 output_file: None,
845 log_file: None,
846 list_databases: false,
847 positionals: vec![],
848 interactive_reason: None,
849 }
850 }
851}
852
853impl PsqlModeState {
854 fn set_sql(&mut self, sql: String, flag: &str) -> Result<(), String> {
855 if self.sql.is_some() || self.sql_file.is_some() {
856 return Err(format!(
857 "psql mode currently supports only one -c/--command or -f/--file source; repeated source at {flag}"
858 ));
859 }
860 self.sql = Some(sql);
861 Ok(())
862 }
863
864 fn set_sql_file(&mut self, path: String, flag: &str) -> Result<(), String> {
865 if self.sql.is_some() || self.sql_file.is_some() {
866 return Err(format!(
867 "psql mode currently supports only one -c/--command or -f/--file source; repeated source at {flag}"
868 ));
869 }
870 self.sql_file = Some(path);
871 Ok(())
872 }
873}
874
875fn parse_psql_long_arg(
876 raw: &[String],
877 i: &mut usize,
878 state: &mut PsqlModeState,
879) -> Result<(), String> {
880 let arg = raw[*i].as_str();
881 if arg == "--mode" {
882 let value = take_arg_value(raw, i, "--mode")?;
883 if value != "psql" {
884 return Err(format!(
885 "unsupported psql-mode argument: --mode {value}; only --mode psql is allowed with psql translation"
886 ));
887 }
888 return Ok(());
889 }
890 if let Some(value) = arg.strip_prefix("--mode=") {
891 if value != "psql" {
892 return Err(format!(
893 "unsupported psql-mode argument: {arg}; only --mode=psql is allowed with psql translation"
894 ));
895 }
896 *i += 1;
897 return Ok(());
898 }
899
900 if arg == "--help" || arg.starts_with("--help=") {
901 emit_psql_mode_help();
902 std::process::exit(0);
903 }
904 if arg == "--version" {
905 emit_psql_mode_version();
906 std::process::exit(0);
907 }
908
909 match long_name(arg) {
910 "--command" => {
911 let value = take_long_arg_value(raw, i, "--command")?;
912 state.set_sql(value, "--command")
913 }
914 "--file" => {
915 let value = take_long_arg_value(raw, i, "--file")?;
916 state.set_sql_file(value, "--file")
917 }
918 "--host" => {
919 state.host = Some(take_long_arg_value(raw, i, "--host")?);
920 Ok(())
921 }
922 "--port" => {
923 state.port = Some(parse_port(
924 &take_long_arg_value(raw, i, "--port")?,
925 "--port",
926 )?);
927 Ok(())
928 }
929 "--username" | "--user" => {
930 state.user = Some(take_long_arg_value(raw, i, long_name(arg))?);
931 Ok(())
932 }
933 "--dbname" => {
934 apply_dbname_value(state, take_long_arg_value(raw, i, "--dbname")?);
935 Ok(())
936 }
937 "--set" | "--variable" => {
938 let value = take_long_arg_value(raw, i, long_name(arg))?;
939 add_psql_variable(state, value)
940 }
941 "--list" => {
942 state.list_databases = true;
943 *i += 1;
944 Ok(())
945 }
946 "--no-password"
947 | "--no-psqlrc"
948 | "--no-readline"
949 | "--quiet"
950 | "--echo-all"
951 | "--echo-errors"
952 | "--echo-queries"
953 | "--echo-hidden"
954 | "--no-align"
955 | "--csv"
956 | "--html"
957 | "--tuples-only"
958 | "--expanded"
959 | "--field-separator-zero"
960 | "--record-separator-zero"
961 | "--single-transaction" => {
962 *i += 1;
963 Ok(())
964 }
965 "--field-separator" | "--record-separator" | "--pset" | "--table-attr" => {
966 let _ = take_long_arg_value(raw, i, long_name(arg))?;
967 Ok(())
968 }
969 "--password" => {
970 state.interactive_reason =
971 Some("--password/-W requests an interactive password prompt".to_string());
972 *i += 1;
973 Ok(())
974 }
975 "--single-step" => {
976 state.interactive_reason =
977 Some("--single-step/-s requires interactive command confirmation".to_string());
978 *i += 1;
979 Ok(())
980 }
981 "--single-line" => {
982 state.interactive_reason =
983 Some("--single-line/-S is a human-interactive input mode".to_string());
984 *i += 1;
985 Ok(())
986 }
987 "--dsn-secret" => {
988 state.dsn_secret = Some(take_long_arg_value(raw, i, "--dsn-secret")?);
989 Ok(())
990 }
991 "--dsn-secret-env" => {
992 state.dsn_secret_env = Some(take_long_arg_value(raw, i, "--dsn-secret-env")?);
993 Ok(())
994 }
995 "--conninfo-secret" => {
996 state.conninfo_secret = Some(take_long_arg_value(raw, i, "--conninfo-secret")?);
997 Ok(())
998 }
999 "--password-secret" => {
1000 state.password_secret = Some(take_long_arg_value(raw, i, "--password-secret")?);
1001 Ok(())
1002 }
1003 "--password-secret-env" => {
1004 state.password_secret_env = Some(take_long_arg_value(raw, i, "--password-secret-env")?);
1005 Ok(())
1006 }
1007 "--container" => {
1008 state.container = Some(take_long_arg_value(raw, i, "--container")?);
1009 Ok(())
1010 }
1011 "--container-driver" => {
1012 state.container_driver = Some(take_long_arg_value(raw, i, "--container-driver")?);
1013 Ok(())
1014 }
1015 "--container-runtime" => {
1016 state.container_runtime = Some(take_long_arg_value(raw, i, "--container-runtime")?);
1017 Ok(())
1018 }
1019 "--container-user" => {
1020 state.container_user = Some(take_long_arg_value(raw, i, "--container-user")?);
1021 Ok(())
1022 }
1023 "--container-namespace" => {
1024 state.container_namespace = Some(take_long_arg_value(raw, i, "--container-namespace")?);
1025 Ok(())
1026 }
1027 "--container-context" => {
1028 state.container_context = Some(take_long_arg_value(raw, i, "--container-context")?);
1029 Ok(())
1030 }
1031 "--container-compose-file" => {
1032 state.container_compose_files.push(take_long_arg_value(
1033 raw,
1034 i,
1035 "--container-compose-file",
1036 )?);
1037 Ok(())
1038 }
1039 "--container-compose-project" => {
1040 state.container_compose_project =
1041 Some(take_long_arg_value(raw, i, "--container-compose-project")?);
1042 Ok(())
1043 }
1044 "--container-pod-container" => {
1045 state.container_pod_container =
1046 Some(take_long_arg_value(raw, i, "--container-pod-container")?);
1047 Ok(())
1048 }
1049 "--output" => {
1050 let value = take_long_arg_value(raw, i, "--output")?;
1051 if let Ok(format) = parse_output(&value) {
1052 state.output = format;
1053 } else {
1054 state.output_file = Some(value);
1055 }
1056 Ok(())
1057 }
1058 "--output-format" => {
1059 let value = take_long_arg_value(raw, i, long_name(arg))?;
1060 state.output = parse_output(&value)?;
1061 Ok(())
1062 }
1063 "--log" => {
1064 let values = take_long_arg_value(raw, i, "--log")?;
1065 add_log_entries(state, &values);
1066 Ok(())
1067 }
1068 "--log-file" => {
1069 state.log_file = Some(take_long_arg_value(raw, i, "--log-file")?);
1070 Ok(())
1071 }
1072 _ => Err(format!("unsupported psql-mode argument: {arg}")),
1073 }
1074}
1075
1076fn parse_psql_short_arg(
1077 raw: &[String],
1078 i: &mut usize,
1079 state: &mut PsqlModeState,
1080) -> Result<(), String> {
1081 let arg = raw[*i].as_str();
1082 let mut offset = 1usize;
1083 while offset < arg.len() {
1084 let flag = arg.as_bytes()[offset] as char;
1085 offset += 1;
1086 match flag {
1087 '?' => {
1088 emit_psql_mode_help();
1089 std::process::exit(0);
1090 }
1091 'V' => {
1092 emit_psql_mode_version();
1093 std::process::exit(0);
1094 }
1095 'c' => {
1096 let value = take_short_arg_value(raw, i, arg, offset, "-c")?;
1097 return state.set_sql(value, "-c");
1098 }
1099 'f' => {
1100 let value = take_short_arg_value(raw, i, arg, offset, "-f")?;
1101 return state.set_sql_file(value, "-f");
1102 }
1103 'h' => {
1104 state.host = Some(take_short_arg_value(raw, i, arg, offset, "-h")?);
1105 return Ok(());
1106 }
1107 'p' => {
1108 let value = take_short_arg_value(raw, i, arg, offset, "-p")?;
1109 state.port = Some(parse_port(&value, "-p")?);
1110 return Ok(());
1111 }
1112 'U' => {
1113 state.user = Some(take_short_arg_value(raw, i, arg, offset, "-U")?);
1114 return Ok(());
1115 }
1116 'd' => {
1117 apply_dbname_value(state, take_short_arg_value(raw, i, arg, offset, "-d")?);
1118 return Ok(());
1119 }
1120 'v' => {
1121 let value = take_short_arg_value(raw, i, arg, offset, "-v")?;
1122 return add_psql_variable(state, value);
1123 }
1124 'F' | 'P' | 'R' | 'T' => {
1125 let _ = take_short_arg_value(raw, i, arg, offset, &format!("-{flag}"))?;
1126 return Ok(());
1127 }
1128 'L' => {
1129 state.log_file = Some(take_short_arg_value(raw, i, arg, offset, "-L")?);
1130 return Ok(());
1131 }
1132 'o' => {
1133 state.output_file = Some(take_short_arg_value(raw, i, arg, offset, "-o")?);
1134 return Ok(());
1135 }
1136 'l' => state.list_databases = true,
1137 'W' => {
1138 state.interactive_reason =
1139 Some("--password/-W requests an interactive password prompt".to_string());
1140 }
1141 's' => {
1142 state.interactive_reason =
1143 Some("--single-step/-s requires interactive command confirmation".to_string());
1144 }
1145 'S' => {
1146 state.interactive_reason =
1147 Some("--single-line/-S is a human-interactive input mode".to_string());
1148 }
1149 'a' | 'A' | 'b' | 'e' | 'E' | 'H' | 'n' | 'q' | 't' | 'w' | 'x' | 'X' | 'z' | '0'
1150 | '1' => {}
1151 _ => return Err(format!("unsupported psql-mode argument: -{flag}")),
1152 }
1153 }
1154 *i += 1;
1155 Ok(())
1156}
1157
1158fn long_name(arg: &str) -> &str {
1159 arg.split_once('=').map(|(name, _)| name).unwrap_or(arg)
1160}
1161
1162fn take_arg_value(raw: &[String], i: &mut usize, flag: &str) -> Result<String, String> {
1163 *i += 1;
1164 let value = raw
1165 .get(*i)
1166 .ok_or_else(|| format!("{flag} requires value"))?
1167 .clone();
1168 *i += 1;
1169 Ok(value)
1170}
1171
1172fn take_long_arg_value(raw: &[String], i: &mut usize, flag: &str) -> Result<String, String> {
1173 let arg = raw[*i].as_str();
1174 if let Some((_, value)) = arg.split_once('=') {
1175 *i += 1;
1176 return Ok(value.to_string());
1177 }
1178 take_arg_value(raw, i, flag)
1179}
1180
1181fn take_short_arg_value(
1182 raw: &[String],
1183 i: &mut usize,
1184 arg: &str,
1185 offset: usize,
1186 flag: &str,
1187) -> Result<String, String> {
1188 if offset < arg.len() {
1189 let value = arg[offset..].to_string();
1190 *i += 1;
1191 return Ok(value);
1192 }
1193 take_arg_value(raw, i, flag)
1194}
1195
1196fn parse_port(value: &str, flag: &str) -> Result<u16, String> {
1197 value.parse().map_err(|_| format!("invalid {flag} port"))
1198}
1199
1200fn add_log_entries(state: &mut PsqlModeState, values: &str) {
1201 for part in values.split(',') {
1202 let trimmed = part.trim();
1203 if !trimmed.is_empty() {
1204 state.log_entries.push(trimmed.to_string());
1205 }
1206 }
1207}
1208
1209fn add_psql_variable(state: &mut PsqlModeState, value: String) -> Result<(), String> {
1210 let name = value
1211 .split_once('=')
1212 .map(|(name, _)| name)
1213 .unwrap_or(value.as_str());
1214 if name.parse::<usize>().is_ok() {
1215 if value.contains('=') {
1216 state.params_kv.push(value);
1217 return Ok(());
1218 }
1219 return Err(format!("invalid param '{value}', expected N=value"));
1220 }
1221 if is_psql_behavior_variable(name) {
1222 return Ok(());
1223 }
1224 Err(format!(
1225 "invalid or unsupported psql variable '{name}'; afpsql supports numeric -v N=value bind parameters, not client-side :name interpolation"
1226 ))
1227}
1228
1229fn is_psql_behavior_variable(name: &str) -> bool {
1230 matches!(
1231 name.to_ascii_uppercase().as_str(),
1232 "ON_ERROR_STOP"
1233 | "ON_ERROR_ROLLBACK"
1234 | "QUIET"
1235 | "ECHO"
1236 | "ECHO_HIDDEN"
1237 | "FETCH_COUNT"
1238 | "VERBOSITY"
1239 | "SHOW_CONTEXT"
1240 | "HISTCONTROL"
1241 | "HISTFILE"
1242 | "HISTSIZE"
1243 | "IGNOREEOF"
1244 | "PAGER"
1245 | "COLUMNS"
1246 )
1247}
1248
1249fn apply_psql_positionals(state: &mut PsqlModeState) -> Result<(), String> {
1250 let positionals = std::mem::take(&mut state.positionals);
1251 for value in positionals {
1252 if is_postgres_uri(&value) {
1253 state.dsn_secret = Some(value);
1254 continue;
1255 }
1256 if looks_like_conninfo(&value) {
1257 state.conninfo_secret = Some(value);
1258 continue;
1259 }
1260 if state.dbname.is_none() {
1261 state.dbname = Some(value);
1262 continue;
1263 }
1264 if state.user.is_none() {
1265 state.user = Some(value);
1266 continue;
1267 }
1268 return Err(format!("too many positional psql arguments: {value}"));
1269 }
1270 Ok(())
1271}
1272
1273fn apply_dbname_value(state: &mut PsqlModeState, value: String) {
1274 if is_postgres_uri(&value) {
1275 state.dsn_secret = Some(value);
1276 } else if looks_like_conninfo(&value) {
1277 state.conninfo_secret = Some(value);
1278 } else {
1279 state.dbname = Some(value);
1280 }
1281}
1282
1283fn is_postgres_uri(value: &str) -> bool {
1284 value.starts_with("postgresql://") || value.starts_with("postgres://")
1285}
1286
1287fn looks_like_conninfo(value: &str) -> bool {
1288 value.contains('=')
1289}
1290
1291fn psql_list_databases_sql() -> String {
1292 "select datname as name from pg_catalog.pg_database where datallowconn order by datname"
1293 .to_string()
1294}
1295
1296fn emit_psql_mode_version() {
1297 let _ = writeln!(
1298 std::io::stdout(),
1299 "psql (afpsql wrapper) {}",
1300 env!("CARGO_PKG_VERSION")
1301 );
1302}
1303
1304fn emit_psql_mode_help() {
1305 let _ = writeln!(
1306 std::io::stdout(),
1307 "psql (afpsql wrapper) {}\n\
1308Usage:\n psql [OPTION]... [DBNAME [USERNAME]]\n\n\
1309Supported non-interactive forms:\n -c, --command=SQL\n -f, --file=FILE\n -l, --list\n -h/-p/-U/-d and --host/--port/--username/--dbname\n -v N=value, --set N=value for positional bind parameters\n\n\
1310Output routing:\n -o, --output=FILE writes structured output to FILE\n -L, --log-file=FILE tees structured output to FILE\n --output-format=json|yaml|plain changes afpsql rendering\n\n\
1311Human-interactive psql modes and psql meta-commands are not supported by this wrapper.",
1312 env!("CARGO_PKG_VERSION")
1313 );
1314}
1315
1316fn top_level_help_requested(raw: &[String]) -> bool {
1317 raw.len() == 2 && matches!(raw.get(1).map(String::as_str), Some("--help" | "-h"))
1318}
1319
1320fn top_level_help_markdown_requested(raw: &[String]) -> bool {
1321 let mut i = 1usize;
1322 while i < raw.len() {
1323 let arg = raw[i].as_str();
1324 if arg == "--" {
1325 break;
1326 }
1327 if arg == "--help-markdown" {
1328 return true;
1329 }
1330 if arg == "--mode" {
1331 i += 2;
1332 continue;
1333 }
1334 if top_level_arg_consumes_value(arg) {
1335 i += if arg.contains('=') { 1 } else { 2 };
1336 continue;
1337 }
1338 if arg.starts_with('-') {
1339 i += 1;
1340 continue;
1341 }
1342 break;
1343 }
1344 false
1345}
1346
1347fn psql_admin_action(action: PsqlCliAction) -> PsqlAdminAction {
1348 match action {
1349 PsqlCliAction::Status(args) => PsqlAdminAction::Status {
1350 bin_dir: args.bin_dir,
1351 },
1352 PsqlCliAction::Install(args) => PsqlAdminAction::Install {
1353 bin_dir: args.bin_dir,
1354 },
1355 PsqlCliAction::Uninstall(args) => PsqlAdminAction::Uninstall {
1356 bin_dir: args.bin_dir,
1357 },
1358 }
1359}
1360
1361fn skill_admin_action(action: SkillCliAction) -> SkillAdminAction {
1362 match action {
1363 SkillCliAction::Status(args) => SkillAdminAction::Status(skill_options(args, false)),
1364 SkillCliAction::Install(args) => {
1365 SkillAdminAction::Install(skill_options(args.target, args.force))
1366 }
1367 SkillCliAction::Uninstall(args) => {
1368 SkillAdminAction::Uninstall(skill_options(args.target, args.force))
1369 }
1370 }
1371}
1372
1373fn skill_options(args: SkillTargetArgs, force: bool) -> SkillAdminOptions {
1374 SkillAdminOptions {
1375 agent: args.agent,
1376 scope: args.scope,
1377 skills_dir: args.skills_dir,
1378 force,
1379 }
1380}
1381
1382fn is_psql_mode_requested(raw: &[String]) -> bool {
1383 let mut i = 1usize;
1384 while i < raw.len() {
1385 let arg = raw[i].as_str();
1386 if arg == "--" {
1387 break;
1388 }
1389 if arg == "--mode" {
1390 if let Some(v) = raw.get(i + 1) {
1391 return v == "psql";
1392 }
1393 return false;
1394 }
1395 if arg == "--mode=psql" {
1396 return true;
1397 }
1398 if top_level_arg_consumes_value(arg) {
1399 i += if arg.contains('=') { 1 } else { 2 };
1400 continue;
1401 }
1402 if arg.starts_with('-') {
1403 i += 1;
1404 continue;
1405 }
1406 break;
1407 }
1408 false
1409}
1410
1411fn top_level_arg_consumes_value(arg: &str) -> bool {
1412 let name = arg.split_once('=').map(|(name, _)| name).unwrap_or(arg);
1413 matches!(
1414 name,
1415 "--sql"
1416 | "--sql-file"
1417 | "--param"
1418 | "--batch-rows"
1419 | "--batch-bytes"
1420 | "--statement-timeout-ms"
1421 | "--lock-timeout-ms"
1422 | "--inline-max-rows"
1423 | "--inline-max-bytes"
1424 | "--permission"
1425 | "--dsn-secret"
1426 | "--dsn-secret-env"
1427 | "--conninfo-secret"
1428 | "--host"
1429 | "--port"
1430 | "--user"
1431 | "--dbname"
1432 | "--password-secret"
1433 | "--password-secret-env"
1434 | "--ssh"
1435 | "--ssh-option"
1436 | "--ssh-local-host"
1437 | "--ssh-local-port"
1438 | "--ssh-remote-socket"
1439 | "--ssh-sudo-user"
1440 | "--container"
1441 | "--container-driver"
1442 | "--container-runtime"
1443 | "--container-user"
1444 | "--container-namespace"
1445 | "--container-context"
1446 | "--container-compose-file"
1447 | "--container-compose-project"
1448 | "--container-pod-container"
1449 | "--output"
1450 | "--log"
1451 )
1452}
1453
1454fn resolve_container_compose_files(cli_files: Vec<String>) -> Vec<String> {
1455 if !cli_files.is_empty() {
1456 return cli_files;
1457 }
1458 std::env::var("AFPSQL_CONTAINER_COMPOSE_FILE")
1459 .ok()
1460 .map(|value| {
1461 value
1462 .split(':')
1463 .filter(|part| !part.is_empty())
1464 .map(std::string::ToString::to_string)
1465 .collect()
1466 })
1467 .unwrap_or_default()
1468}
1469
1470fn load_sql(sql: Option<String>, sql_file: Option<String>) -> Result<String, String> {
1471 match (sql, sql_file) {
1472 (Some(s), None) => validate_sql_size(s),
1473 (None, Some(path)) if path == "-" => {
1474 let stdin = std::io::stdin();
1475 read_limited_sql(stdin.lock(), "read --sql-file -")
1476 }
1477 (None, Some(path)) => {
1478 let metadata =
1479 std::fs::metadata(&path).map_err(|e| format!("read --sql-file failed: {e}"))?;
1480 if metadata.is_file() && metadata.len() > MAX_SQL_BYTES as u64 {
1481 return Err(sql_size_error());
1482 }
1483 let file =
1484 std::fs::File::open(&path).map_err(|e| format!("read --sql-file failed: {e}"))?;
1485 read_limited_sql(file, "read --sql-file")
1486 }
1487 (Some(_), Some(_)) => Err("--sql and --sql-file are mutually exclusive".to_string()),
1488 (None, None) => Err("one of --sql or --sql-file is required".to_string()),
1489 }
1490}
1491
1492fn read_limited_sql<R: Read>(reader: R, context: &str) -> Result<String, String> {
1493 let mut buf = Vec::new();
1494 let mut limited = reader.take(MAX_SQL_BYTES as u64 + 1);
1495 limited
1496 .read_to_end(&mut buf)
1497 .map_err(|e| format!("{context} failed: {e}"))?;
1498 if buf.len() > MAX_SQL_BYTES {
1499 return Err(sql_size_error());
1500 }
1501 String::from_utf8(buf).map_err(|e| format!("{context} failed: {e}"))
1502}
1503
1504fn validate_sql_size(sql: String) -> Result<String, String> {
1505 if sql.len() > MAX_SQL_BYTES {
1506 return Err(sql_size_error());
1507 }
1508 Ok(sql)
1509}
1510
1511fn sql_size_error() -> String {
1512 format!("sql exceeds maximum size; maximum SQL size is {MAX_SQL_BYTES} bytes")
1513}
1514
1515fn parse_output(v: &str) -> Result<OutputFormat, String> {
1516 cli_parse_output(v)
1517}
1518
1519fn parse_permission_arg(v: &str) -> Result<Permission, String> {
1520 v.parse()
1521}
1522
1523fn parse_log_categories(entries: &[String]) -> Vec<String> {
1524 cli_parse_log_filters(entries)
1525}
1526
1527fn startup_requested_from_raw(raw: &[String]) -> bool {
1528 let mut i = 1usize;
1529 while i < raw.len() {
1530 if raw[i] == "--log" {
1531 if let Some(values) = raw.get(i + 1) {
1532 for part in values.split(',') {
1533 let v = part.trim().to_ascii_lowercase();
1534 if matches!(v.as_str(), "startup" | "all" | "*") {
1535 return true;
1536 }
1537 }
1538 }
1539 i += 2;
1540 continue;
1541 }
1542 if let Some(values) = raw[i].strip_prefix("--log=") {
1543 for part in values.split(',') {
1544 let v = part.trim().to_ascii_lowercase();
1545 if matches!(v.as_str(), "startup" | "all" | "*") {
1546 return true;
1547 }
1548 }
1549 }
1550 i += 1;
1551 }
1552 false
1553}
1554
1555fn startup_env_snapshot() -> Value {
1556 Value::Array(
1557 STARTUP_ENV_KEYS
1558 .iter()
1559 .map(|key| {
1560 json!({
1561 "key": key,
1562 "present": std::env::var_os(key).is_some(),
1563 })
1564 })
1565 .collect(),
1566 )
1567}
1568
1569fn startup_args(
1570 mode: &str,
1571 sql: Option<&str>,
1572 sql_file: Option<&str>,
1573 param_count: usize,
1574) -> Value {
1575 json!({
1576 "mode": mode,
1577 "sql": startup_sql_summary(sql, sql_file),
1578 "param_count": param_count,
1579 })
1580}
1581
1582fn startup_sql_summary(sql: Option<&str>, sql_file: Option<&str>) -> Value {
1583 let Some(sql) = sql else {
1584 return json!({
1585 "present": false,
1586 "source": "none",
1587 "bytes": 0,
1588 "chars": 0,
1589 "operation": null,
1590 });
1591 };
1592 json!({
1593 "present": true,
1594 "source": if sql_file.is_some() { "file" } else { "inline" },
1595 "bytes": sql.len(),
1596 "chars": sql.chars().count(),
1597 "operation": sql_operation(sql),
1598 })
1599}
1600
1601fn sql_operation(sql: &str) -> Option<String> {
1602 let sql = trim_leading_sql_comments(sql);
1603 let token: String = sql
1604 .chars()
1605 .skip_while(|c| c.is_whitespace())
1606 .take_while(|c| c.is_ascii_alphabetic() || *c == '_')
1607 .collect();
1608 if token.is_empty() {
1609 None
1610 } else {
1611 Some(token.to_ascii_lowercase())
1612 }
1613}
1614
1615fn trim_leading_sql_comments(mut sql: &str) -> &str {
1616 loop {
1617 sql = sql.trim_start();
1618 if let Some(rest) = sql.strip_prefix("--") {
1619 sql = rest.split_once('\n').map(|(_, rest)| rest).unwrap_or("");
1620 continue;
1621 }
1622 if let Some(rest) = sql.strip_prefix("/*") {
1623 let Some((_, after)) = rest.split_once("*/") else {
1624 return "";
1625 };
1626 sql = after;
1627 continue;
1628 }
1629 return sql;
1630 }
1631}
1632
1633struct PsqlStartupArgs<'a> {
1634 mode: &'a str,
1635 sql: Option<&'a str>,
1636 sql_file: Option<String>,
1637 param_count: usize,
1638}
1639
1640fn psql_startup_args(args: PsqlStartupArgs<'_>) -> Value {
1641 startup_args(
1642 args.mode,
1643 args.sql,
1644 args.sql_file.as_deref(),
1645 args.param_count,
1646 )
1647}
1648
1649fn resolve_secret_value(
1650 flag_name: &str,
1651 direct: Option<String>,
1652 env_name: Option<&str>,
1653) -> Result<Option<String>, String> {
1654 match (direct, env_name) {
1655 (Some(_), Some(_)) => Err(format!(
1656 "{flag_name} and {flag_name}-env are mutually exclusive"
1657 )),
1658 (Some(value), None) => Ok(Some(value)),
1659 (None, Some(name)) => {
1660 if name.is_empty() {
1661 return Err(format!(
1662 "{flag_name}-env requires a non-empty variable name"
1663 ));
1664 }
1665 std::env::var(name).map(Some).map_err(|_| {
1666 format!("{flag_name}-env references unset environment variable: {name}")
1667 })
1668 }
1669 (None, None) => Ok(None),
1670 }
1671}
1672
1673pub fn parse_params(entries: &[String]) -> Result<Vec<Value>, String> {
1674 if entries.len() > MAX_PARAMS {
1675 return Err(format!("too many params; maximum params is {MAX_PARAMS}"));
1676 }
1677
1678 let mut by_index: BTreeMap<usize, Value> = BTreeMap::new();
1679 for entry in entries {
1680 let (idx, raw) = split_index_value(entry)?;
1681 if idx == 0 {
1682 return Err("param index must start at 1".to_string());
1683 }
1684 if idx > MAX_PARAMS {
1685 return Err(format!(
1686 "parameter index {idx} exceeds maximum params {MAX_PARAMS}"
1687 ));
1688 }
1689 match by_index.entry(idx) {
1690 Entry::Vacant(slot) => {
1691 slot.insert(parse_param_value(raw));
1692 }
1693 Entry::Occupied(_) => return Err(format!("duplicate parameter index {idx}")),
1694 }
1695 }
1696 if by_index.is_empty() {
1697 return Ok(vec![]);
1698 }
1699 let max = by_index.keys().max().copied().unwrap_or(0);
1700 for i in 1..=max {
1701 if !by_index.contains_key(&i) {
1702 return Err(format!("missing parameter index {i}"));
1703 }
1704 }
1705 Ok(by_index.into_values().collect())
1706}
1707
1708fn split_index_value(entry: &str) -> Result<(usize, &str), String> {
1709 let mut parts = entry.splitn(2, '=');
1710 let left = parts.next().unwrap_or_default();
1711 let right = parts
1712 .next()
1713 .ok_or_else(|| format!("invalid param '{entry}', expected N=value"))?;
1714 let idx = left
1715 .parse::<usize>()
1716 .map_err(|_| format!("invalid param index in '{entry}'"))?;
1717 Ok((idx, right))
1718}
1719
1720fn parse_param_value(v: &str) -> Value {
1721 if v == "null" {
1722 return Value::Null;
1723 }
1724 if v == "true" {
1725 return Value::Bool(true);
1726 }
1727 if v == "false" {
1728 return Value::Bool(false);
1729 }
1730 Value::String(v.to_string())
1734}
1735
1736fn wrap_explain_sql(user_sql: &str, analyze: bool) -> String {
1737 let body = user_sql.trim_end_matches([';', ' ', '\n', '\t', '\r']);
1738 if analyze {
1739 format!("explain (analyze true, format json, buffers true) {body}")
1740 } else {
1741 format!("explain (format json) {body}")
1742 }
1743}
1744
1745fn build_inspect_sql(action: InspectAction) -> (String, Vec<Value>) {
1746 match action {
1747 InspectAction::Databases => (
1748 "select datname as database, \
1749 pg_catalog.pg_get_userbyid(datdba) as owner, \
1750 pg_catalog.pg_encoding_to_char(encoding) as encoding \
1751 from pg_catalog.pg_database \
1752 where not datistemplate \
1753 order by datname"
1754 .to_string(),
1755 vec![],
1756 ),
1757 InspectAction::Schemas => (
1758 "select schema_name, schema_owner \
1759 from information_schema.schemata \
1760 where schema_name not in ('pg_catalog', 'information_schema') \
1761 and schema_name not like 'pg_toast%' \
1762 and schema_name not like 'pg_temp_%' \
1763 order by schema_name"
1764 .to_string(),
1765 vec![],
1766 ),
1767 InspectAction::Tables(args) => {
1768 let mut sql = String::from(
1769 "select table_schema as schema, table_name as name, table_type as kind \
1770 from information_schema.tables \
1771 where table_schema = $1",
1772 );
1773 let mut params = vec![Value::String(args.schema)];
1774 if let Some(pattern) = args.like {
1775 sql.push_str(" and table_name like $2");
1776 params.push(Value::String(pattern));
1777 }
1778 sql.push_str(" order by table_name");
1779 (sql, params)
1780 }
1781 InspectAction::Views(args) => {
1782 let mut sql = String::from(
1783 "select table_schema as schema, table_name as name \
1784 from information_schema.views \
1785 where table_schema = $1",
1786 );
1787 let mut params = vec![Value::String(args.schema)];
1788 if let Some(pattern) = args.like {
1789 sql.push_str(" and table_name like $2");
1790 params.push(Value::String(pattern));
1791 }
1792 sql.push_str(" order by table_name");
1793 (sql, params)
1794 }
1795 InspectAction::Table(args) => {
1796 let (schema, name) = match args.name.split_once('.') {
1797 Some((s, n)) => (s.to_string(), n.to_string()),
1798 None => ("public".to_string(), args.name),
1799 };
1800 (
1801 "select column_name as name, data_type as type, \
1802 is_nullable = 'YES' as nullable, \
1803 column_default as default, \
1804 ordinal_position as position \
1805 from information_schema.columns \
1806 where table_schema = $1 and table_name = $2 \
1807 order by ordinal_position"
1808 .to_string(),
1809 vec![Value::String(schema), Value::String(name)],
1810 )
1811 }
1812 }
1813}
1814
1815#[cfg(test)]
1816#[path = "../tests/support/unit_cli.rs"]
1817mod tests;