Skip to main content

agent_first_psql/
cli.rs

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