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::secret_config::{SecretConfigRef, resolve_config_secret};
5use crate::types::{ContainerConfig, Permission, QueryOptions, SessionConfig, SshConfig};
6use agent_first_data::{LogFilters, OutputFormat, cli_parse_log_filters, cli_parse_output};
7use clap::{Args, CommandFactory, FromArgMatches, Parser, Subcommand, ValueEnum};
8use serde_json::{Value, json};
9use std::collections::{BTreeMap, btree_map::Entry};
10
11const STARTUP_ENV_KEYS: &[&str] = &[
12    "AFPSQL_DSN_SECRET",
13    "AFPSQL_CONNINFO_SECRET",
14    "AFPSQL_HOST",
15    "AFPSQL_PORT",
16    "AFPSQL_USER",
17    "AFPSQL_DBNAME",
18    "AFPSQL_PASSWORD_SECRET",
19    "AFPSQL_SSH",
20    "AFPSQL_SSH_LOCAL_HOST",
21    "AFPSQL_SSH_LOCAL_PORT",
22    "AFPSQL_SSH_REMOTE_SOCKET",
23    "AFPSQL_SSH_SUDO_USER",
24    "AFPSQL_CONTAINER",
25    "AFPSQL_CONTAINER_DRIVER",
26    "AFPSQL_CONTAINER_RUNTIME",
27    "AFPSQL_CONTAINER_USER",
28    "AFPSQL_CONTAINER_NAMESPACE",
29    "AFPSQL_CONTAINER_CONTEXT",
30    "AFPSQL_CONTAINER_COMPOSE_FILE",
31    "AFPSQL_CONTAINER_COMPOSE_PROJECT",
32    "AFPSQL_CONTAINER_POD_CONTAINER",
33    "PGHOST",
34    "PGPORT",
35    "PGUSER",
36    "PGDATABASE",
37    "PGPASSWORD",
38    "PGSSLMODE",
39];
40
41pub enum Mode {
42    Cli(CliRequest),
43    Pipe(PipeInit),
44    PsqlAdmin(PsqlAdminRequest),
45    SkillAdmin(SkillAdminRequest),
46    PsqlUnsupported(PsqlUnsupportedRequest),
47}
48
49pub struct PipeInit {
50    pub output: OutputFormat,
51    pub session: SessionConfig,
52    pub log: LogFilters,
53    pub startup_args: Value,
54    pub startup_env: Value,
55    pub startup_requested: bool,
56}
57
58#[derive(Debug, Clone)]
59pub struct PsqlAdminRequest {
60    pub action: PsqlAdminAction,
61    pub output: OutputFormat,
62}
63
64#[derive(Debug, Clone)]
65pub enum PsqlAdminAction {
66    Status { bin_dir: Option<String> },
67    Install { bin_dir: Option<String> },
68    Uninstall { bin_dir: Option<String> },
69}
70
71#[derive(Debug, Clone)]
72pub struct SkillAdminRequest {
73    pub action: SkillAdminAction,
74    pub output: OutputFormat,
75}
76
77#[derive(Debug, Clone)]
78pub enum SkillAdminAction {
79    Status(SkillAdminOptions),
80    Install(SkillAdminOptions),
81    Uninstall(SkillAdminOptions),
82}
83
84#[derive(Debug, Clone)]
85pub struct SkillAdminOptions {
86    pub agent: SkillAgentSelection,
87    pub scope: SkillScope,
88    pub skills_dir: Option<String>,
89    pub force: bool,
90}
91
92#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
93pub enum SkillAgentSelection {
94    /// Manage every agent that supports the requested scope.
95    All,
96    /// Manage the Codex local skill under $CODEX_HOME/skills.
97    Codex,
98    /// Manage the Claude Code skill under ~/.claude/skills or .claude/skills.
99    #[value(name = "claude-code", alias = "claude")]
100    ClaudeCode,
101    /// Manage the opencode skill under ~/.config/opencode/skills or .opencode/skills.
102    Opencode,
103    /// Manage the Hermes skill under $HERMES_HOME/skills or ~/.hermes/skills.
104    Hermes,
105}
106
107#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
108pub enum SkillScope {
109    /// Install under the user-level skills directory.
110    Personal,
111    /// Install under the current workspace's skills directory.
112    Workspace,
113}
114
115pub struct CliRequest {
116    pub sql: String,
117    pub params: Vec<Value>,
118    pub options: QueryOptions,
119    pub session: SessionConfig,
120    pub output: OutputFormat,
121    pub log: LogFilters,
122    pub startup_args: Value,
123    pub startup_env: Value,
124    pub startup_requested: bool,
125    pub dry_run: bool,
126    pub psql_mode: bool,
127}
128
129pub struct PsqlUnsupportedRequest {
130    pub reason: String,
131}
132
133#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
134enum RuntimeMode {
135    Cli,
136    Pipe,
137    #[value(name = "psql")]
138    Psql,
139}
140
141#[derive(Subcommand)]
142enum AfdCommand {
143    /// Manage the local psql wrapper for afpsql --mode psql.
144    Psql(PsqlCommand),
145    /// Manage Agent-First PSQL skills for Codex, Claude Code, opencode, and Hermes.
146    Skill(SkillCommand),
147    /// Schema discovery: inspect databases, schemas, tables, indexes, or snapshots.
148    Inspect(InspectCommand),
149}
150
151#[derive(Args)]
152struct InspectCommand {
153    #[command(subcommand)]
154    action: InspectAction,
155}
156
157#[derive(Subcommand)]
158enum InspectAction {
159    /// List databases on the connected server with size, encoding, and connection facts.
160    Databases(InspectDatabasesArgs),
161    /// Summarize the connected database: schema/table/view/sequence counts and size.
162    Database,
163    /// List user-visible schemas.
164    Schemas,
165    /// Export full schema metadata for one schema.
166    Schema(InspectSchemaArgs),
167    /// Export a stable full-schema snapshot for machine consumption.
168    Snapshot(InspectSchemaArgs),
169    /// List tables in a schema with owner, estimated rows, and size.
170    Tables(InspectTablesArgs),
171    /// List views (regular and materialized) in a schema with owner.
172    Views(InspectViewsArgs),
173    /// List indexes with definitions, size, validity, and optional usage stats.
174    Indexes(InspectIndexesArgs),
175    /// Describe a table's columns: types, nullability, defaults, primary key, comments.
176    Table(InspectTableArgs),
177}
178
179#[derive(Args)]
180struct InspectDatabasesArgs {
181    /// Include template databases (template0/template1) in the listing.
182    #[arg(long = "all")]
183    all: bool,
184}
185
186#[derive(Args)]
187struct InspectTablesArgs {
188    /// Schema to filter on. Defaults to `public`.
189    #[arg(long = "schema", default_value = "public")]
190    schema: String,
191    /// Optional `LIKE` pattern matched against the table name (use `%` as wildcard).
192    #[arg(long = "like")]
193    like: Option<String>,
194}
195
196#[derive(Args)]
197struct InspectSchemaArgs {
198    /// Schema to inspect. Defaults to `public`.
199    #[arg(long = "schema", default_value = "public")]
200    schema: String,
201    /// Optional `LIKE` pattern matched against relation names (use `%` as wildcard).
202    #[arg(long = "like")]
203    like: Option<String>,
204}
205
206#[derive(Args)]
207struct InspectViewsArgs {
208    /// Schema to filter on. Defaults to `public`.
209    #[arg(long = "schema", default_value = "public")]
210    schema: String,
211    /// Optional `LIKE` pattern matched against the view name (use `%` as wildcard).
212    #[arg(long = "like")]
213    like: Option<String>,
214}
215
216#[derive(Args)]
217struct InspectIndexesArgs {
218    /// Schema to filter on. Defaults to `public`.
219    #[arg(long = "schema", default_value = "public")]
220    schema: String,
221    /// Optional table name to filter on. Accepts `schema.table` to override --schema.
222    #[arg(long = "table")]
223    table: Option<String>,
224    /// Include PostgreSQL's built-in pg_stat_user_indexes usage counters.
225    #[arg(long = "stats")]
226    stats: bool,
227}
228
229#[derive(Args)]
230struct InspectTableArgs {
231    /// Table name. Accepts `schema.table`; defaults to `public.NAME` when unqualified.
232    name: String,
233    /// Include relation, constraints, indexes, triggers, and sequence/default metadata.
234    #[arg(long = "full")]
235    full: bool,
236}
237
238#[derive(Args)]
239struct PsqlCommand {
240    #[command(subcommand)]
241    action: PsqlCliAction,
242}
243
244#[derive(Subcommand)]
245enum PsqlCliAction {
246    /// Show whether the afpsql-managed psql wrapper is installed and active.
247    Status(PsqlPathArgs),
248    /// Install an afpsql-managed psql wrapper.
249    Install(PsqlPathArgs),
250    /// Remove an afpsql-managed psql wrapper.
251    Uninstall(PsqlPathArgs),
252}
253
254#[derive(Args)]
255struct PsqlPathArgs {
256    /// Directory that contains the psql wrapper. Defaults to the afpsql executable directory.
257    #[arg(long = "bin-dir")]
258    bin_dir: Option<String>,
259}
260
261#[derive(Args)]
262struct SkillCommand {
263    #[command(subcommand)]
264    action: SkillCliAction,
265}
266
267#[derive(Subcommand)]
268enum SkillCliAction {
269    /// Show whether the Agent-First PSQL skill is installed, valid, and up to date.
270    Status(SkillTargetArgs),
271    /// Install the Agent-First PSQL skill.
272    Install(SkillWriteArgs),
273    /// Remove an afpsql-managed Agent-First PSQL skill.
274    Uninstall(SkillWriteArgs),
275}
276
277#[derive(Args)]
278struct SkillTargetArgs {
279    /// Agent to manage. Defaults to all personal skill targets.
280    #[arg(long = "agent", value_enum, default_value_t = SkillAgentSelection::All)]
281    agent: SkillAgentSelection,
282    /// Skill scope.
283    #[arg(long = "scope", value_enum, default_value_t = SkillScope::Personal)]
284    scope: SkillScope,
285    /// Directory that contains skill folders. Requires an explicit single --agent.
286    #[arg(long = "skills-dir")]
287    skills_dir: Option<String>,
288}
289
290#[derive(Args)]
291struct SkillWriteArgs {
292    #[command(flatten)]
293    target: SkillTargetArgs,
294    /// Overwrite or remove an unmanaged Agent-First PSQL skill at the target path.
295    #[arg(long)]
296    force: bool,
297}
298
299#[doc = r#"`afpsql` gives agents a reliable PostgreSQL contract: structured stdout
300events, first-class SSH/container transports, explicit write permissions,
301stable pipe sessions, and machine-readable failures.
302
303### Interface Policy
304
305- default mode is canonical agent-first CLI
306- `--mode psql` is argument translation only; runtime output stays JSONL
307- stdout carries protocol events; stderr is not a protocol channel
308- native CLI and pipe mode default to read-only transactions; writes require permission
309- SSH/container transports keep afpsql local instead of running human `psql` across boundaries
310
311### Modes
312
313- default (native CLI): one SQL action per process — a single agent step
314- `--mode pipe`: a long-lived JSONL session with `id` correlation and named sessions for multi-step work
315- `--mode psql`: run existing `psql` scripts unchanged — flags are translated, runtime output stays JSONL
316
317### Query Sources and Parameters
318
319- use `--sql` for inline SQL or `--sql-file` for a file
320- use repeatable `--param N=value` for positional binds
321- placeholder count is validated from prepared-statement metadata, not by SQL text scanning
322
323### Connection Sources
324
325- `--dsn-secret` for a PostgreSQL URI
326- `--conninfo-secret` for libpq-style conninfo
327- or discrete `--host`, `--port`, `--user`, `--dbname`, `--password-secret`
328- every `*-secret` flag has a `*-secret-env` partner that reads the value from a named environment variable
329- every secret slot also has a `*-secret-config FILE DOT_PATH` source for JSON, TOML, YAML, or dotenv
330- add `--ssh user@server` when PostgreSQL is reachable only from the server boundary
331- add `--container TARGET` when PostgreSQL is reachable only from inside a container boundary
332- use named container scope flags instead of raw driver option passthrough
333- use `--container-driver docker|podman|nerdctl|compose|kubectl` for the exec syntax
334- combine `--ssh user@server --container TARGET` for containers on an SSH host
335- agent-first environment fallbacks: `AFPSQL_*`
336- PostgreSQL environment fallbacks: `PGHOST`, `PGPORT`, `PGUSER`, `PGDATABASE`, `PGPASSWORD`, `PGSSLMODE`
337
338### Result Shaping
339
340- default mode buffers a bounded inline result
341- use `--stream-rows` for large result sets, with `--batch-rows` and `--batch-bytes` to tune chunk size
342- `--output json|yaml|plain` changes rendering only, not the runtime schema
343
344### Examples
345
346```text
347afpsql --sql "select now() as now_rfc3339"
348afpsql --sql-file ./query.sql
349afpsql --sql 'select * from users where id = $1' --param 1=123
350afpsql --dsn-secret-env DATABASE_URL --sql "select 1"
351afpsql --dsn-secret-config config.yaml database.url --sql "select 1"
352afpsql --ssh user@server --host 127.0.0.1 --port 5432 --user app --dbname appdb --sql "select 1"
353afpsql --container pg-container --dsn-secret-env DATABASE_URL --sql "select 1"
354afpsql --ssh root@server --container app --host host.container.internal --port 5432 --user app --dbname appdb --sql "select 1"
355afpsql --mode psql -h 127.0.0.1 -p 5432 -U app -d appdb -c "select 1"
356afpsql --sql "select * from big_table" --stream-rows --batch-rows 1000
357afpsql --mode pipe
358afpsql psql status
359afpsql psql install
360afpsql skill status
361afpsql skill install
362```
363
364### Exit Codes
365
366- `0`: query completed successfully
367- `1`: SQL error or runtime error
368- `2`: invalid CLI arguments
369"#]
370#[derive(Parser)]
371#[command(
372    name = env!("DISPLAY_NAME"),
373    bin_name = "afpsql",
374    version,
375    verbatim_doc_comment,
376    about = env!("CARGO_PKG_DESCRIPTION"),
377)]
378pub struct AfdCli {
379    /// Inline SQL string to execute.
380    #[arg(long, allow_hyphen_values = true, help_heading = "Query")]
381    sql: Option<String>,
382    /// Read SQL from a file.
383    #[arg(long = "sql-file", allow_hyphen_values = true, help_heading = "Query")]
384    sql_file: Option<String>,
385    /// Positional bind parameter in `N=value` form. Repeat for additional parameters.
386    #[arg(long = "param", help_heading = "Query")]
387    param: Vec<String>,
388    /// Stream large result sets as `result_rows` batches instead of a single inline result.
389    #[arg(long = "stream-rows", help_heading = "Query")]
390    stream_rows: bool,
391    /// Maximum rows per streamed batch.
392    #[arg(long = "batch-rows", help_heading = "Query")]
393    batch_rows: Option<usize>,
394    /// Soft byte target per streamed batch.
395    #[arg(long = "batch-bytes", help_heading = "Query")]
396    batch_bytes: Option<usize>,
397    /// Per-query statement timeout in milliseconds.
398    #[arg(long = "statement-timeout-ms", help_heading = "Query")]
399    statement_timeout_ms: Option<u64>,
400    /// Per-query lock timeout in milliseconds.
401    #[arg(long = "lock-timeout-ms", help_heading = "Query")]
402    lock_timeout_ms: Option<u64>,
403    /// Maximum inline rows before returning `result_too_large`.
404    #[arg(long = "inline-max-rows", help_heading = "Query")]
405    inline_max_rows: Option<usize>,
406    /// Maximum inline payload bytes before returning `result_too_large`.
407    #[arg(long = "inline-max-bytes", help_heading = "Query")]
408    inline_max_bytes: Option<usize>,
409    /// Query permission. Defaults to read, ssh-read with --ssh, or
410    /// container-read with --container.
411    #[arg(long = "permission", value_enum, help_heading = "Query")]
412    permission: Option<Permission>,
413    /// Preview the query without executing it
414    #[arg(long, help_heading = "Query")]
415    dry_run: bool,
416    /// Wrap the query in EXPLAIN (FORMAT JSON) and return the plan tree instead
417    /// of executing the user's SQL.
418    #[arg(
419        long = "explain",
420        help_heading = "Query",
421        conflicts_with = "explain_analyze"
422    )]
423    explain: bool,
424    /// Wrap the query in EXPLAIN (ANALYZE, FORMAT JSON, BUFFERS). The underlying
425    /// SQL actually runs; writes require the matching write permission.
426    #[arg(long = "explain-analyze", help_heading = "Query")]
427    explain_analyze: bool,
428
429    /// PostgreSQL DSN URI. Redacted in structured output.
430    #[arg(
431        long = "dsn-secret",
432        global = true,
433        help_heading = "Connection",
434        conflicts_with_all = ["dsn_secret_env", "dsn_secret_config"]
435    )]
436    dsn_secret: Option<String>,
437    /// Read PostgreSQL DSN URI from an environment variable.
438    #[arg(
439        long = "dsn-secret-env",
440        global = true,
441        help_heading = "Connection",
442        conflicts_with = "dsn_secret_config"
443    )]
444    dsn_secret_env: Option<String>,
445    /// Read PostgreSQL DSN URI from FILE at DOT_PATH.
446    #[arg(
447        long = "dsn-secret-config",
448        global = true,
449        help_heading = "Connection",
450        value_names = ["FILE", "DOT_PATH"],
451        num_args = 2
452    )]
453    dsn_secret_config: Option<Vec<String>>,
454    /// libpq-style conninfo string. Redacted in structured output.
455    #[arg(
456        long = "conninfo-secret",
457        global = true,
458        help_heading = "Connection",
459        conflicts_with_all = ["conninfo_secret_env", "conninfo_secret_config"]
460    )]
461    conninfo_secret: Option<String>,
462    /// Read libpq-style conninfo string from an environment variable.
463    #[arg(
464        long = "conninfo-secret-env",
465        global = true,
466        help_heading = "Connection",
467        conflicts_with = "conninfo_secret_config"
468    )]
469    conninfo_secret_env: Option<String>,
470    /// Read libpq-style conninfo from FILE at DOT_PATH.
471    #[arg(
472        long = "conninfo-secret-config",
473        global = true,
474        help_heading = "Connection",
475        value_names = ["FILE", "DOT_PATH"],
476        num_args = 2
477    )]
478    conninfo_secret_config: Option<Vec<String>>,
479    /// PostgreSQL host.
480    #[arg(long, global = true, help_heading = "Connection")]
481    host: Option<String>,
482    /// PostgreSQL port.
483    #[arg(long, global = true, help_heading = "Connection")]
484    port: Option<u16>,
485    /// PostgreSQL user name.
486    #[arg(long, global = true, help_heading = "Connection")]
487    user: Option<String>,
488    /// PostgreSQL database name.
489    #[arg(long, global = true, help_heading = "Connection")]
490    dbname: Option<String>,
491    /// PostgreSQL password. Redacted in structured output.
492    #[arg(
493        long = "password-secret",
494        global = true,
495        help_heading = "Connection",
496        conflicts_with_all = ["password_secret_env", "password_secret_config"]
497    )]
498    password_secret: Option<String>,
499    /// Read PostgreSQL password from an environment variable.
500    #[arg(
501        long = "password-secret-env",
502        global = true,
503        help_heading = "Connection",
504        conflicts_with = "password_secret_config"
505    )]
506    password_secret_env: Option<String>,
507    /// Read PostgreSQL password from FILE at DOT_PATH.
508    #[arg(
509        long = "password-secret-config",
510        global = true,
511        help_heading = "Connection",
512        value_names = ["FILE", "DOT_PATH"],
513        num_args = 2
514    )]
515    password_secret_config: Option<Vec<String>>,
516    /// Open an SSH transport to USER@HOST before connecting to PostgreSQL.
517    #[arg(long = "ssh", global = true, help_heading = "SSH Transport")]
518    ssh: Option<String>,
519    /// SSH hop to reach before the final --ssh destination. Repeat for multiple hops.
520    #[arg(long = "ssh-via", global = true, help_heading = "SSH Transport")]
521    ssh_via: Vec<String>,
522    /// Additional OpenSSH -o option. Repeat for multiple options.
523    #[arg(long = "ssh-option", global = true, help_heading = "SSH Transport")]
524    ssh_options: Vec<String>,
525    /// Local bind host for the SSH tunnel.
526    #[arg(long = "ssh-local-host", global = true, help_heading = "SSH Transport")]
527    ssh_local_host: Option<String>,
528    /// Local bind port for the SSH tunnel. Defaults to an ephemeral port.
529    #[arg(long = "ssh-local-port", global = true, help_heading = "SSH Transport")]
530    ssh_local_port: Option<u16>,
531    /// Explicit remote PostgreSQL Unix socket path for SSH forwarding.
532    #[arg(
533        long = "ssh-remote-socket",
534        global = true,
535        help_heading = "SSH Transport"
536    )]
537    ssh_remote_socket: Option<String>,
538    /// Remote OS user for sudo -n Unix-socket bridge mode; requires an explicit socket.
539    #[arg(long = "ssh-sudo-user", global = true, help_heading = "SSH Transport")]
540    ssh_sudo_user: Option<String>,
541
542    /// Run a container exec stdio bridge in TARGET before connecting to PostgreSQL.
543    #[arg(
544        long = "container",
545        global = true,
546        help_heading = "Container Transport"
547    )]
548    container: Option<String>,
549    /// Container exec driver: docker, podman, nerdctl, compose, or kubectl.
550    #[arg(
551        long = "container-driver",
552        global = true,
553        help_heading = "Container Transport"
554    )]
555    container_driver: Option<String>,
556    /// Runtime command for the selected container driver. Defaults to the driver command.
557    #[arg(
558        long = "container-runtime",
559        global = true,
560        help_heading = "Container Transport"
561    )]
562    container_runtime: Option<String>,
563    /// OS user passed to drivers that support exec user selection.
564    #[arg(
565        long = "container-user",
566        global = true,
567        help_heading = "Container Transport"
568    )]
569    container_user: Option<String>,
570    /// Kubernetes namespace for kubectl exec.
571    #[arg(
572        long = "container-namespace",
573        global = true,
574        help_heading = "Container Transport"
575    )]
576    container_namespace: Option<String>,
577    /// Docker or Kubernetes context for the selected driver.
578    #[arg(
579        long = "container-context",
580        global = true,
581        help_heading = "Container Transport"
582    )]
583    container_context: Option<String>,
584    /// Compose file passed before compose exec. Repeat for multiple files.
585    #[arg(
586        long = "container-compose-file",
587        global = true,
588        help_heading = "Container Transport"
589    )]
590    container_compose_files: Vec<String>,
591    /// Compose project name passed before compose exec.
592    #[arg(
593        long = "container-compose-project",
594        global = true,
595        help_heading = "Container Transport"
596    )]
597    container_compose_project: Option<String>,
598    /// Kubernetes container name for multi-container pods.
599    #[arg(
600        long = "container-pod-container",
601        global = true,
602        help_heading = "Container Transport"
603    )]
604    container_pod_container: Option<String>,
605
606    /// Output format: json (default), yaml, or plain.
607    #[arg(
608        long,
609        short = 'o',
610        default_value = "json",
611        global = true,
612        help_heading = "Runtime"
613    )]
614    output: String,
615    /// Redirect stdout bytes to this file.
616    #[arg(
617        long = "stdout-file",
618        value_name = "PATH",
619        global = true,
620        help_heading = "Runtime"
621    )]
622    stdout_file: Option<String>,
623    /// Redirect stderr bytes to this file.
624    #[arg(
625        long = "stderr-file",
626        value_name = "PATH",
627        global = true,
628        help_heading = "Runtime"
629    )]
630    stderr_file: Option<String>,
631    /// Diagnostic log categories (comma-separated). Categories: startup,
632    /// connect, query, transport, mode; or an exact event name like
633    /// `query.error`; or `all` for everything.
634    #[arg(
635        long = "log",
636        value_delimiter = ',',
637        global = true,
638        help_heading = "Runtime"
639    )]
640    log: Vec<String>,
641    /// Runtime mode: canonical cli, pipe, or `psql` translation mode.
642    #[arg(long, value_enum, default_value_t = RuntimeMode::Cli, help_heading = "Runtime")]
643    mode: RuntimeMode,
644
645    #[command(subcommand)]
646    command: Option<AfdCommand>,
647}
648
649pub fn parse_args(bin_name: &str) -> Result<Mode, String> {
650    let raw: Vec<String> = std::env::args().collect();
651    if is_psql_mode_requested(&raw) {
652        return parse_psql_mode(&raw);
653    }
654    let startup_requested = startup_requested_from_raw(&raw);
655
656    match agent_first_data::cli_handle_version_or_continue(
657        &raw,
658        bin_name,
659        env!("CARGO_PKG_VERSION"),
660    ) {
661        Ok(Some(version)) => {
662            let _ = write!(std::io::stdout(), "{version}");
663            std::process::exit(0);
664        }
665        Ok(None) => {}
666        Err(err) => {
667            let stdout = std::io::stdout();
668            let mut emitter = agent_first_data::CliEmitter::new(stdout.lock(), OutputFormat::Json);
669            let _ = emitter.emit_error("cli_error", &err.to_string());
670            std::process::exit(2);
671        }
672    }
673
674    match agent_first_data::cli_handle_help_or_continue(
675        &raw,
676        &command_for_bin(bin_name),
677        &agent_first_data::HelpConfig::human_cli_default(),
678    ) {
679        Ok(Some(help)) => {
680            let _ = write!(std::io::stdout(), "{help}");
681            std::process::exit(0);
682        }
683        Ok(None) => {}
684        Err(err) => {
685            let stdout = std::io::stdout();
686            let mut emitter = agent_first_data::CliEmitter::new(stdout.lock(), OutputFormat::Json);
687            let _ = emitter.emit_error("cli_error", &err.to_string());
688            std::process::exit(2);
689        }
690    }
691
692    let cli = match command_for_bin(bin_name)
693        .try_get_matches_from(&raw)
694        .and_then(|matches| AfdCli::from_arg_matches(&matches))
695    {
696        Ok(c) => c,
697        Err(e) => {
698            use clap::error::ErrorKind;
699            if matches!(e.kind(), ErrorKind::DisplayVersion | ErrorKind::DisplayHelp) {
700                let _ = writeln!(std::io::stdout(), "{e}");
701                std::process::exit(0);
702            }
703            return Err(e.to_string());
704        }
705    };
706    let _stream_redirect_args = (&cli.stdout_file, &cli.stderr_file);
707    let output = parse_output(&cli.output)?;
708    let log = parse_log_categories(&cli.log);
709    let dsn_config = SecretConfigRef::from_values("--dsn-secret-config", cli.dsn_secret_config)?;
710    let conninfo_config =
711        SecretConfigRef::from_values("--conninfo-secret-config", cli.conninfo_secret_config)?;
712    let password_config =
713        SecretConfigRef::from_values("--password-secret-config", cli.password_secret_config)?;
714    let connection_sources = connection_source_metadata([
715        (
716            "dsn",
717            cli.dsn_secret.is_some(),
718            cli.dsn_secret_env.as_deref(),
719            dsn_config.as_ref(),
720        ),
721        (
722            "conninfo",
723            cli.conninfo_secret.is_some(),
724            cli.conninfo_secret_env.as_deref(),
725            conninfo_config.as_ref(),
726        ),
727        (
728            "password",
729            cli.password_secret.is_some(),
730            cli.password_secret_env.as_deref(),
731            password_config.as_ref(),
732        ),
733    ]);
734    let dsn_secret = resolve_secret_value(
735        "--dsn-secret",
736        cli.dsn_secret,
737        cli.dsn_secret_env.as_deref(),
738        dsn_config.as_ref(),
739    )?;
740    let password_secret = resolve_secret_value(
741        "--password-secret",
742        cli.password_secret,
743        cli.password_secret_env.as_deref(),
744        password_config.as_ref(),
745    )?;
746    let conninfo_secret = resolve_secret_value(
747        "--conninfo-secret",
748        cli.conninfo_secret,
749        cli.conninfo_secret_env.as_deref(),
750        conninfo_config.as_ref(),
751    )?;
752    let session = SessionConfig {
753        dsn_secret,
754        conninfo_secret,
755        host: cli.host,
756        port: cli.port,
757        user: cli.user,
758        dbname: cli.dbname,
759        password_secret,
760        ssh: SshConfig {
761            destination: cli.ssh.or_else(|| std::env::var("AFPSQL_SSH").ok()),
762            via: if cli.ssh_via.is_empty() {
763                parse_csv_env("AFPSQL_SSH_VIA")
764            } else {
765                cli.ssh_via
766            },
767            options: cli.ssh_options,
768            local_host: cli
769                .ssh_local_host
770                .or_else(|| std::env::var("AFPSQL_SSH_LOCAL_HOST").ok()),
771            local_port: cli.ssh_local_port.or_else(|| {
772                std::env::var("AFPSQL_SSH_LOCAL_PORT")
773                    .ok()
774                    .and_then(|v| v.parse().ok())
775            }),
776            remote_socket: cli
777                .ssh_remote_socket
778                .or_else(|| std::env::var("AFPSQL_SSH_REMOTE_SOCKET").ok()),
779            sudo_user: cli
780                .ssh_sudo_user
781                .or_else(|| std::env::var("AFPSQL_SSH_SUDO_USER").ok()),
782        },
783        container: ContainerConfig {
784            target: cli
785                .container
786                .or_else(|| std::env::var("AFPSQL_CONTAINER").ok()),
787            driver: cli
788                .container_driver
789                .or_else(|| std::env::var("AFPSQL_CONTAINER_DRIVER").ok()),
790            runtime: cli
791                .container_runtime
792                .or_else(|| std::env::var("AFPSQL_CONTAINER_RUNTIME").ok()),
793            user: cli
794                .container_user
795                .or_else(|| std::env::var("AFPSQL_CONTAINER_USER").ok()),
796            namespace: cli
797                .container_namespace
798                .or_else(|| std::env::var("AFPSQL_CONTAINER_NAMESPACE").ok()),
799            context: cli
800                .container_context
801                .or_else(|| std::env::var("AFPSQL_CONTAINER_CONTEXT").ok()),
802            compose_files: resolve_container_compose_files(cli.container_compose_files),
803            compose_project: cli
804                .container_compose_project
805                .or_else(|| std::env::var("AFPSQL_CONTAINER_COMPOSE_PROJECT").ok()),
806            pod_container: cli
807                .container_pod_container
808                .or_else(|| std::env::var("AFPSQL_CONTAINER_POD_CONTAINER").ok()),
809        },
810    };
811    let mode_name = match cli.mode {
812        RuntimeMode::Cli => "cli",
813        RuntimeMode::Pipe => "pipe",
814        RuntimeMode::Psql => "psql",
815    };
816    let startup_env = startup_env_snapshot();
817
818    if let Some(command) = cli.command {
819        return match command {
820            AfdCommand::Psql(psql) => Ok(Mode::PsqlAdmin(PsqlAdminRequest {
821                action: psql_admin_action(psql.action),
822                output,
823            })),
824            AfdCommand::Skill(skill) => Ok(Mode::SkillAdmin(SkillAdminRequest {
825                action: skill_admin_action(skill.action),
826                output,
827            })),
828            AfdCommand::Inspect(inspect) => {
829                let (sql, params) = build_inspect_sql(inspect.action);
830                let startup_args = with_connection_sources(
831                    startup_args(mode_name, Some(&sql), None, params.len()),
832                    &connection_sources,
833                );
834                Ok(Mode::Cli(CliRequest {
835                    sql,
836                    params,
837                    options: QueryOptions::default(),
838                    session,
839                    output,
840                    log,
841                    startup_args,
842                    startup_env,
843                    startup_requested,
844                    dry_run: false,
845                    psql_mode: false,
846                }))
847            }
848        };
849    }
850
851    match cli.mode {
852        RuntimeMode::Pipe => {
853            return Ok(Mode::Pipe(PipeInit {
854                output,
855                session,
856                log: log.clone(),
857                startup_args: with_connection_sources(
858                    startup_args(mode_name, None, None, 0),
859                    &connection_sources,
860                ),
861                startup_env,
862                startup_requested,
863            }));
864        }
865        RuntimeMode::Cli | RuntimeMode::Psql => {}
866    }
867
868    let startup_sql_file = cli.sql_file.clone();
869    let user_sql = load_sql(cli.sql, cli.sql_file)?;
870    let params = parse_params(&cli.param)?;
871    let sql = if cli.explain {
872        wrap_explain_sql(&user_sql, false)
873    } else if cli.explain_analyze {
874        wrap_explain_sql(&user_sql, true)
875    } else {
876        user_sql
877    };
878    let startup_args = with_connection_sources(
879        startup_args(
880            mode_name,
881            Some(&sql),
882            startup_sql_file.as_deref(),
883            params.len(),
884        ),
885        &connection_sources,
886    );
887
888    let options = QueryOptions {
889        stream_rows: cli.stream_rows,
890        batch_rows: cli.batch_rows,
891        batch_bytes: cli.batch_bytes,
892        statement_timeout_ms: cli.statement_timeout_ms,
893        lock_timeout_ms: cli.lock_timeout_ms,
894        permission: cli.permission,
895        inline_max_rows: cli.inline_max_rows,
896        inline_max_bytes: cli.inline_max_bytes,
897    };
898
899    Ok(Mode::Cli(CliRequest {
900        sql,
901        params,
902        options,
903        session,
904        output,
905        log,
906        startup_args,
907        startup_env,
908        startup_requested,
909        dry_run: cli.dry_run,
910        psql_mode: false,
911    }))
912}
913
914fn command_for_bin(bin_name: &str) -> clap::Command {
915    match bin_name {
916        "afpsql-readonly" => AfdCli::command()
917            .name("afpsql-readonly")
918            .bin_name("afpsql-readonly"),
919        _ => AfdCli::command().name("afpsql").bin_name("afpsql"),
920    }
921}
922
923fn parse_psql_mode(raw: &[String]) -> Result<Mode, String> {
924    let startup_requested = startup_requested_from_raw(raw);
925    let mut state = PsqlModeState::default();
926
927    let mut i = 1usize;
928    while i < raw.len() {
929        let arg = raw[i].as_str();
930        if arg == "--" {
931            i += 1;
932            while i < raw.len() {
933                state.positionals.push(raw[i].clone());
934                i += 1;
935            }
936            break;
937        }
938        if arg.starts_with("--") {
939            parse_psql_long_arg(raw, &mut i, &mut state)?;
940            continue;
941        }
942        if arg.starts_with('-') && arg.len() > 1 {
943            parse_psql_short_arg(raw, &mut i, &mut state)?;
944            continue;
945        }
946        state.positionals.push(raw[i].clone());
947        i += 1;
948    }
949
950    if let Some(reason) = state.interactive_reason {
951        return Ok(Mode::PsqlUnsupported(PsqlUnsupportedRequest { reason }));
952    }
953
954    apply_psql_positionals(&mut state)?;
955    if state.list_databases {
956        state.sql = Some(psql_list_databases_sql());
957        state.sql_file = None;
958    }
959    if state.sql.is_none() && state.sql_file.is_none() {
960        return Ok(Mode::PsqlUnsupported(PsqlUnsupportedRequest {
961            reason: "no -c/--command, -f/--file, or -l/--list was provided".to_string(),
962        }));
963    }
964
965    let connection_sources = connection_source_metadata([
966        (
967            "dsn",
968            state.dsn_secret.is_some(),
969            state.dsn_secret_env.as_deref(),
970            state.dsn_secret_config.as_ref(),
971        ),
972        (
973            "conninfo",
974            state.conninfo_secret.is_some(),
975            state.conninfo_secret_env.as_deref(),
976            state.conninfo_secret_config.as_ref(),
977        ),
978        (
979            "password",
980            state.password_secret.is_some(),
981            state.password_secret_env.as_deref(),
982            state.password_secret_config.as_ref(),
983        ),
984    ]);
985    let dsn_secret = resolve_secret_value(
986        "--dsn-secret",
987        state.dsn_secret,
988        state.dsn_secret_env.as_deref(),
989        state.dsn_secret_config.as_ref(),
990    )?;
991    let password_secret = resolve_secret_value(
992        "--password-secret",
993        state.password_secret,
994        state.password_secret_env.as_deref(),
995        state.password_secret_config.as_ref(),
996    )?;
997    let conninfo_secret = resolve_secret_value(
998        "--conninfo-secret",
999        state.conninfo_secret,
1000        state.conninfo_secret_env.as_deref(),
1001        state.conninfo_secret_config.as_ref(),
1002    )?;
1003    let session = SessionConfig {
1004        dsn_secret,
1005        conninfo_secret,
1006        host: state.host,
1007        port: state.port,
1008        user: state.user,
1009        dbname: state.dbname,
1010        password_secret,
1011        ssh: SshConfig::default(),
1012        container: ContainerConfig {
1013            target: state
1014                .container
1015                .or_else(|| std::env::var("AFPSQL_CONTAINER").ok()),
1016            driver: state
1017                .container_driver
1018                .or_else(|| std::env::var("AFPSQL_CONTAINER_DRIVER").ok()),
1019            runtime: state
1020                .container_runtime
1021                .or_else(|| std::env::var("AFPSQL_CONTAINER_RUNTIME").ok()),
1022            user: state
1023                .container_user
1024                .or_else(|| std::env::var("AFPSQL_CONTAINER_USER").ok()),
1025            namespace: state
1026                .container_namespace
1027                .or_else(|| std::env::var("AFPSQL_CONTAINER_NAMESPACE").ok()),
1028            context: state
1029                .container_context
1030                .or_else(|| std::env::var("AFPSQL_CONTAINER_CONTEXT").ok()),
1031            compose_files: resolve_container_compose_files(state.container_compose_files),
1032            compose_project: state
1033                .container_compose_project
1034                .or_else(|| std::env::var("AFPSQL_CONTAINER_COMPOSE_PROJECT").ok()),
1035            pod_container: state
1036                .container_pod_container
1037                .or_else(|| std::env::var("AFPSQL_CONTAINER_POD_CONTAINER").ok()),
1038        },
1039    };
1040
1041    let startup_sql_file = state.sql_file.clone();
1042    let sql = load_sql(state.sql, state.sql_file)?;
1043    let params = parse_params(&state.params_kv)?;
1044    let startup_args = with_connection_sources(
1045        psql_startup_args(PsqlStartupArgs {
1046            mode: "psql",
1047            sql: Some(&sql),
1048            sql_file: startup_sql_file,
1049            param_count: params.len(),
1050        }),
1051        &connection_sources,
1052    );
1053    Ok(Mode::Cli(CliRequest {
1054        sql,
1055        params,
1056        options: QueryOptions {
1057            permission: Some(if session.uses_container_transport() {
1058                Permission::ContainerWrite
1059            } else {
1060                Permission::Write
1061            }),
1062            ..Default::default()
1063        },
1064        session,
1065        output: state.output,
1066        log: parse_log_categories(&state.log_entries),
1067        startup_args,
1068        startup_env: startup_env_snapshot(),
1069        startup_requested,
1070        dry_run: false,
1071        psql_mode: true,
1072    }))
1073}
1074
1075struct PsqlModeState {
1076    sql: Option<String>,
1077    sql_file: Option<String>,
1078    host: Option<String>,
1079    port: Option<u16>,
1080    user: Option<String>,
1081    dbname: Option<String>,
1082    dsn_secret: Option<String>,
1083    dsn_secret_env: Option<String>,
1084    dsn_secret_config: Option<SecretConfigRef>,
1085    conninfo_secret: Option<String>,
1086    conninfo_secret_env: Option<String>,
1087    conninfo_secret_config: Option<SecretConfigRef>,
1088    password_secret: Option<String>,
1089    password_secret_env: Option<String>,
1090    password_secret_config: Option<SecretConfigRef>,
1091    container: Option<String>,
1092    container_driver: Option<String>,
1093    container_runtime: Option<String>,
1094    container_user: Option<String>,
1095    container_namespace: Option<String>,
1096    container_context: Option<String>,
1097    container_compose_files: Vec<String>,
1098    container_compose_project: Option<String>,
1099    container_pod_container: Option<String>,
1100    params_kv: Vec<String>,
1101    output: OutputFormat,
1102    log_entries: Vec<String>,
1103    list_databases: bool,
1104    positionals: Vec<String>,
1105    interactive_reason: Option<String>,
1106}
1107
1108impl Default for PsqlModeState {
1109    fn default() -> Self {
1110        Self {
1111            sql: None,
1112            sql_file: None,
1113            host: None,
1114            port: None,
1115            user: None,
1116            dbname: None,
1117            dsn_secret: None,
1118            dsn_secret_env: None,
1119            dsn_secret_config: None,
1120            conninfo_secret: None,
1121            conninfo_secret_env: None,
1122            conninfo_secret_config: None,
1123            password_secret: None,
1124            password_secret_env: None,
1125            password_secret_config: None,
1126            container: None,
1127            container_driver: None,
1128            container_runtime: None,
1129            container_user: None,
1130            container_namespace: None,
1131            container_context: None,
1132            container_compose_files: vec![],
1133            container_compose_project: None,
1134            container_pod_container: None,
1135            params_kv: vec![],
1136            output: OutputFormat::Json,
1137            log_entries: vec![],
1138            list_databases: false,
1139            positionals: vec![],
1140            interactive_reason: None,
1141        }
1142    }
1143}
1144
1145impl PsqlModeState {
1146    fn set_sql(&mut self, sql: String, flag: &str) -> Result<(), String> {
1147        if self.sql.is_some() || self.sql_file.is_some() {
1148            return Err(format!(
1149                "psql mode currently supports only one -c/--command or -f/--file source; repeated source at {flag}"
1150            ));
1151        }
1152        self.sql = Some(sql);
1153        Ok(())
1154    }
1155
1156    fn set_sql_file(&mut self, path: String, flag: &str) -> Result<(), String> {
1157        if self.sql.is_some() || self.sql_file.is_some() {
1158            return Err(format!(
1159                "psql mode currently supports only one -c/--command or -f/--file source; repeated source at {flag}"
1160            ));
1161        }
1162        self.sql_file = Some(path);
1163        Ok(())
1164    }
1165}
1166
1167fn parse_psql_long_arg(
1168    raw: &[String],
1169    i: &mut usize,
1170    state: &mut PsqlModeState,
1171) -> Result<(), String> {
1172    let arg = raw[*i].as_str();
1173    if arg == "--mode" {
1174        let value = take_arg_value(raw, i, "--mode")?;
1175        if value != "psql" {
1176            return Err(format!(
1177                "unsupported psql-mode argument: --mode {value}; only --mode psql is allowed with psql translation"
1178            ));
1179        }
1180        return Ok(());
1181    }
1182    if let Some(value) = arg.strip_prefix("--mode=") {
1183        if value != "psql" {
1184            return Err(format!(
1185                "unsupported psql-mode argument: {arg}; only --mode=psql is allowed with psql translation"
1186            ));
1187        }
1188        *i += 1;
1189        return Ok(());
1190    }
1191
1192    if arg == "--help" || arg.starts_with("--help=") {
1193        emit_psql_mode_help();
1194        std::process::exit(0);
1195    }
1196    if arg == "--version" {
1197        emit_psql_mode_version();
1198        std::process::exit(0);
1199    }
1200
1201    match long_name(arg) {
1202        "--command" => {
1203            let value = take_long_arg_value(raw, i, "--command")?;
1204            state.set_sql(value, "--command")
1205        }
1206        "--file" => {
1207            let value = take_long_arg_value(raw, i, "--file")?;
1208            state.set_sql_file(value, "--file")
1209        }
1210        "--host" => {
1211            state.host = Some(take_long_arg_value(raw, i, "--host")?);
1212            Ok(())
1213        }
1214        "--port" => {
1215            state.port = Some(parse_port(
1216                &take_long_arg_value(raw, i, "--port")?,
1217                "--port",
1218            )?);
1219            Ok(())
1220        }
1221        "--username" | "--user" => {
1222            state.user = Some(take_long_arg_value(raw, i, long_name(arg))?);
1223            Ok(())
1224        }
1225        "--dbname" => {
1226            apply_dbname_value(state, take_long_arg_value(raw, i, "--dbname")?);
1227            Ok(())
1228        }
1229        "--set" | "--variable" => {
1230            let value = take_long_arg_value(raw, i, long_name(arg))?;
1231            add_psql_variable(state, value)
1232        }
1233        "--list" => {
1234            state.list_databases = true;
1235            *i += 1;
1236            Ok(())
1237        }
1238        "--no-password"
1239        | "--no-psqlrc"
1240        | "--no-readline"
1241        | "--quiet"
1242        | "--echo-all"
1243        | "--echo-errors"
1244        | "--echo-queries"
1245        | "--echo-hidden"
1246        | "--no-align"
1247        | "--csv"
1248        | "--html"
1249        | "--tuples-only"
1250        | "--expanded"
1251        | "--field-separator-zero"
1252        | "--record-separator-zero"
1253        | "--single-transaction" => {
1254            *i += 1;
1255            Ok(())
1256        }
1257        "--field-separator" | "--record-separator" | "--pset" | "--table-attr" => {
1258            let _ = take_long_arg_value(raw, i, long_name(arg))?;
1259            Ok(())
1260        }
1261        "--password" => {
1262            state.interactive_reason =
1263                Some("--password/-W requests an interactive password prompt".to_string());
1264            *i += 1;
1265            Ok(())
1266        }
1267        "--single-step" => {
1268            state.interactive_reason =
1269                Some("--single-step/-s requires interactive command confirmation".to_string());
1270            *i += 1;
1271            Ok(())
1272        }
1273        "--single-line" => {
1274            state.interactive_reason =
1275                Some("--single-line/-S is a human-interactive input mode".to_string());
1276            *i += 1;
1277            Ok(())
1278        }
1279        "--dsn-secret" => {
1280            state.dsn_secret = Some(take_long_arg_value(raw, i, "--dsn-secret")?);
1281            Ok(())
1282        }
1283        "--dsn-secret-env" => {
1284            state.dsn_secret_env = Some(take_long_arg_value(raw, i, "--dsn-secret-env")?);
1285            Ok(())
1286        }
1287        "--dsn-secret-config" => {
1288            state.dsn_secret_config = Some(take_secret_config_ref(raw, i, "--dsn-secret-config")?);
1289            Ok(())
1290        }
1291        "--conninfo-secret" => {
1292            state.conninfo_secret = Some(take_long_arg_value(raw, i, "--conninfo-secret")?);
1293            Ok(())
1294        }
1295        "--conninfo-secret-env" => {
1296            state.conninfo_secret_env = Some(take_long_arg_value(raw, i, "--conninfo-secret-env")?);
1297            Ok(())
1298        }
1299        "--conninfo-secret-config" => {
1300            state.conninfo_secret_config =
1301                Some(take_secret_config_ref(raw, i, "--conninfo-secret-config")?);
1302            Ok(())
1303        }
1304        "--password-secret" => {
1305            state.password_secret = Some(take_long_arg_value(raw, i, "--password-secret")?);
1306            Ok(())
1307        }
1308        "--password-secret-env" => {
1309            state.password_secret_env = Some(take_long_arg_value(raw, i, "--password-secret-env")?);
1310            Ok(())
1311        }
1312        "--password-secret-config" => {
1313            state.password_secret_config =
1314                Some(take_secret_config_ref(raw, i, "--password-secret-config")?);
1315            Ok(())
1316        }
1317        "--container" => {
1318            state.container = Some(take_long_arg_value(raw, i, "--container")?);
1319            Ok(())
1320        }
1321        "--container-driver" => {
1322            state.container_driver = Some(take_long_arg_value(raw, i, "--container-driver")?);
1323            Ok(())
1324        }
1325        "--container-runtime" => {
1326            state.container_runtime = Some(take_long_arg_value(raw, i, "--container-runtime")?);
1327            Ok(())
1328        }
1329        "--container-user" => {
1330            state.container_user = Some(take_long_arg_value(raw, i, "--container-user")?);
1331            Ok(())
1332        }
1333        "--container-namespace" => {
1334            state.container_namespace = Some(take_long_arg_value(raw, i, "--container-namespace")?);
1335            Ok(())
1336        }
1337        "--container-context" => {
1338            state.container_context = Some(take_long_arg_value(raw, i, "--container-context")?);
1339            Ok(())
1340        }
1341        "--container-compose-file" => {
1342            state.container_compose_files.push(take_long_arg_value(
1343                raw,
1344                i,
1345                "--container-compose-file",
1346            )?);
1347            Ok(())
1348        }
1349        "--container-compose-project" => {
1350            state.container_compose_project =
1351                Some(take_long_arg_value(raw, i, "--container-compose-project")?);
1352            Ok(())
1353        }
1354        "--container-pod-container" => {
1355            state.container_pod_container =
1356                Some(take_long_arg_value(raw, i, "--container-pod-container")?);
1357            Ok(())
1358        }
1359        "--output" => {
1360            let value = take_long_arg_value(raw, i, "--output")?;
1361            state.output = parse_output(&value)?;
1362            Ok(())
1363        }
1364        "--stdout-file" | "--stderr-file" => {
1365            let _ = take_long_arg_value(raw, i, long_name(arg))?;
1366            Ok(())
1367        }
1368        "--log" => {
1369            let values = take_long_arg_value(raw, i, "--log")?;
1370            add_log_entries(state, &values);
1371            Ok(())
1372        }
1373        _ => Err(format!("unsupported psql-mode argument: {arg}")),
1374    }
1375}
1376
1377fn parse_psql_short_arg(
1378    raw: &[String],
1379    i: &mut usize,
1380    state: &mut PsqlModeState,
1381) -> Result<(), String> {
1382    let arg = raw[*i].as_str();
1383    let mut offset = 1usize;
1384    while offset < arg.len() {
1385        let flag = arg.as_bytes()[offset] as char;
1386        offset += 1;
1387        match flag {
1388            '?' => {
1389                emit_psql_mode_help();
1390                std::process::exit(0);
1391            }
1392            'V' => {
1393                emit_psql_mode_version();
1394                std::process::exit(0);
1395            }
1396            'c' => {
1397                let value = take_short_arg_value(raw, i, arg, offset, "-c")?;
1398                return state.set_sql(value, "-c");
1399            }
1400            'f' => {
1401                let value = take_short_arg_value(raw, i, arg, offset, "-f")?;
1402                return state.set_sql_file(value, "-f");
1403            }
1404            'h' => {
1405                state.host = Some(take_short_arg_value(raw, i, arg, offset, "-h")?);
1406                return Ok(());
1407            }
1408            'p' => {
1409                let value = take_short_arg_value(raw, i, arg, offset, "-p")?;
1410                state.port = Some(parse_port(&value, "-p")?);
1411                return Ok(());
1412            }
1413            'U' => {
1414                state.user = Some(take_short_arg_value(raw, i, arg, offset, "-U")?);
1415                return Ok(());
1416            }
1417            'd' => {
1418                apply_dbname_value(state, take_short_arg_value(raw, i, arg, offset, "-d")?);
1419                return Ok(());
1420            }
1421            'v' => {
1422                let value = take_short_arg_value(raw, i, arg, offset, "-v")?;
1423                return add_psql_variable(state, value);
1424            }
1425            'F' | 'P' | 'R' | 'T' => {
1426                let _ = take_short_arg_value(raw, i, arg, offset, &format!("-{flag}"))?;
1427                return Ok(());
1428            }
1429            'o' => {
1430                let value = take_short_arg_value(raw, i, arg, offset, "-o")?;
1431                state.output = parse_output(&value)?;
1432                return Ok(());
1433            }
1434            'l' => state.list_databases = true,
1435            'W' => {
1436                state.interactive_reason =
1437                    Some("--password/-W requests an interactive password prompt".to_string());
1438            }
1439            's' => {
1440                state.interactive_reason =
1441                    Some("--single-step/-s requires interactive command confirmation".to_string());
1442            }
1443            'S' => {
1444                state.interactive_reason =
1445                    Some("--single-line/-S is a human-interactive input mode".to_string());
1446            }
1447            'a' | 'A' | 'b' | 'e' | 'E' | 'H' | 'n' | 'q' | 't' | 'w' | 'x' | 'X' | 'z' | '0'
1448            | '1' => {}
1449            _ => return Err(format!("unsupported psql-mode argument: -{flag}")),
1450        }
1451    }
1452    *i += 1;
1453    Ok(())
1454}
1455
1456fn long_name(arg: &str) -> &str {
1457    arg.split_once('=').map(|(name, _)| name).unwrap_or(arg)
1458}
1459
1460fn take_arg_value(raw: &[String], i: &mut usize, flag: &str) -> Result<String, String> {
1461    *i += 1;
1462    let value = raw
1463        .get(*i)
1464        .ok_or_else(|| format!("{flag} requires value"))?
1465        .clone();
1466    *i += 1;
1467    Ok(value)
1468}
1469
1470fn take_long_arg_value(raw: &[String], i: &mut usize, flag: &str) -> Result<String, String> {
1471    let arg = raw[*i].as_str();
1472    if let Some((_, value)) = arg.split_once('=') {
1473        *i += 1;
1474        return Ok(value.to_string());
1475    }
1476    take_arg_value(raw, i, flag)
1477}
1478
1479fn take_secret_config_ref(
1480    raw: &[String],
1481    i: &mut usize,
1482    flag: &str,
1483) -> Result<SecretConfigRef, String> {
1484    if raw[*i].contains('=') {
1485        return Err(format!(
1486            "{flag} requires space-separated values: {flag} <FILE> <DOT_PATH>"
1487        ));
1488    }
1489    let file = take_long_arg_value(raw, i, flag)?;
1490    let path = raw
1491        .get(*i)
1492        .filter(|value| !value.starts_with('-'))
1493        .ok_or_else(|| format!("{flag} requires exactly two values: <FILE> <DOT_PATH>"))?
1494        .clone();
1495    *i += 1;
1496    if raw.get(*i).is_some_and(|value| !value.starts_with('-')) {
1497        return Err(format!(
1498            "{flag} accepts exactly two values: <FILE> <DOT_PATH>"
1499        ));
1500    }
1501    if file.is_empty() || path.is_empty() {
1502        return Err(format!(
1503            "{flag} requires exactly two non-empty values: <FILE> <DOT_PATH>"
1504        ));
1505    }
1506    Ok(SecretConfigRef {
1507        file: file.into(),
1508        path,
1509    })
1510}
1511
1512fn take_short_arg_value(
1513    raw: &[String],
1514    i: &mut usize,
1515    arg: &str,
1516    offset: usize,
1517    flag: &str,
1518) -> Result<String, String> {
1519    if offset < arg.len() {
1520        let value = arg[offset..].to_string();
1521        *i += 1;
1522        return Ok(value);
1523    }
1524    take_arg_value(raw, i, flag)
1525}
1526
1527fn parse_port(value: &str, flag: &str) -> Result<u16, String> {
1528    value.parse().map_err(|_| format!("invalid {flag} port"))
1529}
1530
1531fn add_log_entries(state: &mut PsqlModeState, values: &str) {
1532    for part in values.split(',') {
1533        let trimmed = part.trim();
1534        if !trimmed.is_empty() {
1535            state.log_entries.push(trimmed.to_string());
1536        }
1537    }
1538}
1539
1540fn add_psql_variable(state: &mut PsqlModeState, value: String) -> Result<(), String> {
1541    let name = value
1542        .split_once('=')
1543        .map(|(name, _)| name)
1544        .unwrap_or(value.as_str());
1545    if name.parse::<usize>().is_ok() {
1546        if value.contains('=') {
1547            state.params_kv.push(value);
1548            return Ok(());
1549        }
1550        return Err(format!("invalid param '{value}', expected N=value"));
1551    }
1552    if is_psql_behavior_variable(name) {
1553        return Ok(());
1554    }
1555    Err(format!(
1556        "invalid or unsupported psql variable '{name}'; afpsql supports numeric -v N=value bind parameters, not client-side :name interpolation"
1557    ))
1558}
1559
1560fn is_psql_behavior_variable(name: &str) -> bool {
1561    matches!(
1562        name.to_ascii_uppercase().as_str(),
1563        "ON_ERROR_STOP"
1564            | "ON_ERROR_ROLLBACK"
1565            | "QUIET"
1566            | "ECHO"
1567            | "ECHO_HIDDEN"
1568            | "FETCH_COUNT"
1569            | "VERBOSITY"
1570            | "SHOW_CONTEXT"
1571            | "HISTCONTROL"
1572            | "HISTFILE"
1573            | "HISTSIZE"
1574            | "IGNOREEOF"
1575            | "PAGER"
1576            | "COLUMNS"
1577    )
1578}
1579
1580fn apply_psql_positionals(state: &mut PsqlModeState) -> Result<(), String> {
1581    let positionals = std::mem::take(&mut state.positionals);
1582    for value in positionals {
1583        if is_postgres_uri(&value) {
1584            state.dsn_secret = Some(value);
1585            continue;
1586        }
1587        if looks_like_conninfo(&value) {
1588            state.conninfo_secret = Some(value);
1589            continue;
1590        }
1591        if state.dbname.is_none() {
1592            state.dbname = Some(value);
1593            continue;
1594        }
1595        if state.user.is_none() {
1596            state.user = Some(value);
1597            continue;
1598        }
1599        return Err(format!("too many positional psql arguments: {value}"));
1600    }
1601    Ok(())
1602}
1603
1604fn apply_dbname_value(state: &mut PsqlModeState, value: String) {
1605    if is_postgres_uri(&value) {
1606        state.dsn_secret = Some(value);
1607    } else if looks_like_conninfo(&value) {
1608        state.conninfo_secret = Some(value);
1609    } else {
1610        state.dbname = Some(value);
1611    }
1612}
1613
1614fn is_postgres_uri(value: &str) -> bool {
1615    value.starts_with("postgresql://") || value.starts_with("postgres://")
1616}
1617
1618fn looks_like_conninfo(value: &str) -> bool {
1619    value.contains('=')
1620}
1621
1622fn psql_list_databases_sql() -> String {
1623    "select datname as name from pg_catalog.pg_database where datallowconn order by datname"
1624        .to_string()
1625}
1626
1627fn emit_psql_mode_version() {
1628    let _ = writeln!(
1629        std::io::stdout(),
1630        "psql (afpsql wrapper) {}",
1631        env!("CARGO_PKG_VERSION")
1632    );
1633}
1634
1635fn emit_psql_mode_help() {
1636    let _ = writeln!(
1637        std::io::stdout(),
1638        "psql (afpsql wrapper) {}\n\
1639Usage:\n  psql [OPTION]... [DBNAME [USERNAME]]\n\n\
1640Supported 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\
1641Output:\n  -o, --output=json|yaml|plain changes afpsql rendering\n  --stdout-file=FILE redirects stdout bytes to FILE\n  --stderr-file=FILE redirects stderr bytes to FILE\n\n\
1642Human-interactive psql modes and psql meta-commands are not supported by this wrapper.",
1643        env!("CARGO_PKG_VERSION")
1644    );
1645}
1646
1647fn psql_admin_action(action: PsqlCliAction) -> PsqlAdminAction {
1648    match action {
1649        PsqlCliAction::Status(args) => PsqlAdminAction::Status {
1650            bin_dir: args.bin_dir,
1651        },
1652        PsqlCliAction::Install(args) => PsqlAdminAction::Install {
1653            bin_dir: args.bin_dir,
1654        },
1655        PsqlCliAction::Uninstall(args) => PsqlAdminAction::Uninstall {
1656            bin_dir: args.bin_dir,
1657        },
1658    }
1659}
1660
1661fn skill_admin_action(action: SkillCliAction) -> SkillAdminAction {
1662    match action {
1663        SkillCliAction::Status(args) => SkillAdminAction::Status(skill_options(args, false)),
1664        SkillCliAction::Install(args) => {
1665            SkillAdminAction::Install(skill_options(args.target, args.force))
1666        }
1667        SkillCliAction::Uninstall(args) => {
1668            SkillAdminAction::Uninstall(skill_options(args.target, args.force))
1669        }
1670    }
1671}
1672
1673fn skill_options(args: SkillTargetArgs, force: bool) -> SkillAdminOptions {
1674    SkillAdminOptions {
1675        agent: args.agent,
1676        scope: args.scope,
1677        skills_dir: args.skills_dir,
1678        force,
1679    }
1680}
1681
1682fn is_psql_mode_requested(raw: &[String]) -> bool {
1683    let mut i = 1usize;
1684    while i < raw.len() {
1685        let arg = raw[i].as_str();
1686        if arg == "--" {
1687            break;
1688        }
1689        if arg == "--mode" {
1690            if let Some(v) = raw.get(i + 1) {
1691                return v == "psql";
1692            }
1693            return false;
1694        }
1695        if arg == "--mode=psql" {
1696            return true;
1697        }
1698        if top_level_arg_consumes_two_values(arg) {
1699            i += if arg.contains('=') { 2 } else { 3 };
1700            continue;
1701        }
1702        if top_level_arg_consumes_value(arg) {
1703            i += if arg.contains('=') { 1 } else { 2 };
1704            continue;
1705        }
1706        if arg.starts_with('-') {
1707            i += 1;
1708            continue;
1709        }
1710        break;
1711    }
1712    false
1713}
1714
1715fn top_level_arg_consumes_two_values(arg: &str) -> bool {
1716    let name = arg.split_once('=').map(|(name, _)| name).unwrap_or(arg);
1717    matches!(
1718        name,
1719        "--dsn-secret-config" | "--conninfo-secret-config" | "--password-secret-config"
1720    )
1721}
1722
1723fn top_level_arg_consumes_value(arg: &str) -> bool {
1724    let name = arg.split_once('=').map(|(name, _)| name).unwrap_or(arg);
1725    matches!(
1726        name,
1727        "--sql"
1728            | "--sql-file"
1729            | "--param"
1730            | "--batch-rows"
1731            | "--batch-bytes"
1732            | "--statement-timeout-ms"
1733            | "--lock-timeout-ms"
1734            | "--inline-max-rows"
1735            | "--inline-max-bytes"
1736            | "--permission"
1737            | "--dsn-secret"
1738            | "--dsn-secret-env"
1739            | "--conninfo-secret"
1740            | "--conninfo-secret-env"
1741            | "--host"
1742            | "--port"
1743            | "--user"
1744            | "--dbname"
1745            | "--password-secret"
1746            | "--password-secret-env"
1747            | "--ssh"
1748            | "--ssh-via"
1749            | "--ssh-option"
1750            | "--ssh-local-host"
1751            | "--ssh-local-port"
1752            | "--ssh-remote-socket"
1753            | "--ssh-sudo-user"
1754            | "--container"
1755            | "--container-driver"
1756            | "--container-runtime"
1757            | "--container-user"
1758            | "--container-namespace"
1759            | "--container-context"
1760            | "--container-compose-file"
1761            | "--container-compose-project"
1762            | "--container-pod-container"
1763            | "--output"
1764            | "--stdout-file"
1765            | "--stderr-file"
1766            | "--log"
1767    )
1768}
1769
1770fn resolve_container_compose_files(cli_files: Vec<String>) -> Vec<String> {
1771    if !cli_files.is_empty() {
1772        return cli_files;
1773    }
1774    std::env::var("AFPSQL_CONTAINER_COMPOSE_FILE")
1775        .ok()
1776        .map(|value| {
1777            value
1778                .split(':')
1779                .filter(|part| !part.is_empty())
1780                .map(std::string::ToString::to_string)
1781                .collect()
1782        })
1783        .unwrap_or_default()
1784}
1785
1786fn load_sql(sql: Option<String>, sql_file: Option<String>) -> Result<String, String> {
1787    match (sql, sql_file) {
1788        (Some(s), None) => validate_sql_size(s),
1789        (None, Some(path)) if path == "-" => {
1790            let stdin = std::io::stdin();
1791            read_limited_sql(stdin.lock(), "read --sql-file -")
1792        }
1793        (None, Some(path)) => {
1794            let metadata =
1795                std::fs::metadata(&path).map_err(|e| format!("read --sql-file failed: {e}"))?;
1796            if metadata.is_file() && metadata.len() > MAX_SQL_BYTES as u64 {
1797                return Err(sql_size_error());
1798            }
1799            let file =
1800                std::fs::File::open(&path).map_err(|e| format!("read --sql-file failed: {e}"))?;
1801            read_limited_sql(file, "read --sql-file")
1802        }
1803        (Some(_), Some(_)) => Err("--sql and --sql-file are mutually exclusive".to_string()),
1804        (None, None) => Err("one of --sql or --sql-file is required".to_string()),
1805    }
1806}
1807
1808fn read_limited_sql<R: Read>(reader: R, context: &str) -> Result<String, String> {
1809    let mut buf = Vec::new();
1810    let mut limited = reader.take(MAX_SQL_BYTES as u64 + 1);
1811    limited
1812        .read_to_end(&mut buf)
1813        .map_err(|e| format!("{context} failed: {e}"))?;
1814    if buf.len() > MAX_SQL_BYTES {
1815        return Err(sql_size_error());
1816    }
1817    String::from_utf8(buf).map_err(|e| format!("{context} failed: {e}"))
1818}
1819
1820fn validate_sql_size(sql: String) -> Result<String, String> {
1821    if sql.len() > MAX_SQL_BYTES {
1822        return Err(sql_size_error());
1823    }
1824    Ok(sql)
1825}
1826
1827fn sql_size_error() -> String {
1828    format!("sql exceeds maximum size; maximum SQL size is {MAX_SQL_BYTES} bytes")
1829}
1830
1831fn parse_output(v: &str) -> Result<OutputFormat, String> {
1832    cli_parse_output(v)
1833}
1834
1835fn parse_log_categories(entries: &[String]) -> LogFilters {
1836    cli_parse_log_filters(entries)
1837}
1838
1839fn parse_csv_env(name: &str) -> Vec<String> {
1840    std::env::var(name)
1841        .ok()
1842        .into_iter()
1843        .flat_map(|value| {
1844            value
1845                .split(',')
1846                .map(str::trim)
1847                .filter(|part| !part.is_empty())
1848                .map(std::string::ToString::to_string)
1849                .collect::<Vec<_>>()
1850        })
1851        .collect()
1852}
1853
1854fn startup_requested_from_raw(raw: &[String]) -> bool {
1855    let mut i = 1usize;
1856    while i < raw.len() {
1857        if raw[i] == "--log" {
1858            if let Some(values) = raw.get(i + 1) {
1859                for part in values.split(',') {
1860                    let v = part.trim().to_ascii_lowercase();
1861                    if matches!(v.as_str(), "startup" | "all" | "*") {
1862                        return true;
1863                    }
1864                }
1865            }
1866            i += 2;
1867            continue;
1868        }
1869        if let Some(values) = raw[i].strip_prefix("--log=") {
1870            for part in values.split(',') {
1871                let v = part.trim().to_ascii_lowercase();
1872                if matches!(v.as_str(), "startup" | "all" | "*") {
1873                    return true;
1874                }
1875            }
1876        }
1877        i += 1;
1878    }
1879    false
1880}
1881
1882fn startup_env_snapshot() -> Value {
1883    Value::Array(
1884        STARTUP_ENV_KEYS
1885            .iter()
1886            .map(|key| {
1887                json!({
1888                    "key": key,
1889                    "present": std::env::var_os(key).is_some(),
1890                })
1891            })
1892            .collect(),
1893    )
1894}
1895
1896fn startup_args(
1897    mode: &str,
1898    sql: Option<&str>,
1899    sql_file: Option<&str>,
1900    param_count: usize,
1901) -> Value {
1902    json!({
1903        "mode": mode,
1904        "sql": startup_sql_summary(sql, sql_file),
1905        "param_count": param_count,
1906    })
1907}
1908
1909fn with_connection_sources(mut args: Value, sources: &Value) -> Value {
1910    if let (Some(args), Some(sources)) = (args.as_object_mut(), sources.as_object())
1911        && !sources.is_empty()
1912    {
1913        args.insert(
1914            "connection_sources".to_string(),
1915            Value::Object(sources.clone()),
1916        );
1917    }
1918    args
1919}
1920
1921fn connection_source_metadata<const N: usize>(
1922    sources: [(&str, bool, Option<&str>, Option<&SecretConfigRef>); N],
1923) -> Value {
1924    let mut metadata = serde_json::Map::new();
1925    for (slot, direct, env_name, config) in sources {
1926        let value = if let Some(reference) = config {
1927            Some(reference.safe_metadata())
1928        } else if let Some(env_name) = env_name {
1929            Some(json!({"kind": "env", "name": env_name}))
1930        } else if direct {
1931            Some(json!({"kind": "direct"}))
1932        } else {
1933            None
1934        };
1935        if let Some(value) = value {
1936            metadata.insert(slot.to_string(), value);
1937        }
1938    }
1939    Value::Object(metadata)
1940}
1941
1942fn startup_sql_summary(sql: Option<&str>, sql_file: Option<&str>) -> Value {
1943    let Some(sql) = sql else {
1944        return json!({
1945            "present": false,
1946            "source": "none",
1947            "bytes": 0,
1948            "chars": 0,
1949            "operation": null,
1950        });
1951    };
1952    json!({
1953        "present": true,
1954        "source": if sql_file.is_some() { "file" } else { "inline" },
1955        "bytes": sql.len(),
1956        "chars": sql.chars().count(),
1957        "operation": sql_operation(sql),
1958    })
1959}
1960
1961fn sql_operation(sql: &str) -> Option<String> {
1962    let sql = trim_leading_sql_comments(sql);
1963    let token: String = sql
1964        .chars()
1965        .skip_while(|c| c.is_whitespace())
1966        .take_while(|c| c.is_ascii_alphabetic() || *c == '_')
1967        .collect();
1968    if token.is_empty() {
1969        None
1970    } else {
1971        Some(token.to_ascii_lowercase())
1972    }
1973}
1974
1975fn trim_leading_sql_comments(mut sql: &str) -> &str {
1976    loop {
1977        sql = sql.trim_start();
1978        if let Some(rest) = sql.strip_prefix("--") {
1979            sql = rest.split_once('\n').map(|(_, rest)| rest).unwrap_or("");
1980            continue;
1981        }
1982        if let Some(rest) = sql.strip_prefix("/*") {
1983            let Some((_, after)) = rest.split_once("*/") else {
1984                return "";
1985            };
1986            sql = after;
1987            continue;
1988        }
1989        return sql;
1990    }
1991}
1992
1993struct PsqlStartupArgs<'a> {
1994    mode: &'a str,
1995    sql: Option<&'a str>,
1996    sql_file: Option<String>,
1997    param_count: usize,
1998}
1999
2000fn psql_startup_args(args: PsqlStartupArgs<'_>) -> Value {
2001    startup_args(
2002        args.mode,
2003        args.sql,
2004        args.sql_file.as_deref(),
2005        args.param_count,
2006    )
2007}
2008
2009fn resolve_secret_value(
2010    flag_name: &str,
2011    direct: Option<String>,
2012    env_name: Option<&str>,
2013    config: Option<&SecretConfigRef>,
2014) -> Result<Option<String>, String> {
2015    let source_count = usize::from(direct.is_some())
2016        + usize::from(env_name.is_some())
2017        + usize::from(config.is_some());
2018    if source_count > 1 {
2019        return Err(format!(
2020            "{flag_name}, {flag_name}-env, and {flag_name}-config are mutually exclusive"
2021        ));
2022    }
2023    match (direct, env_name, config) {
2024        (Some(value), None, None) => Ok(Some(value)),
2025        (None, Some(name), None) => {
2026            if name.is_empty() {
2027                return Err(format!(
2028                    "{flag_name}-env requires a non-empty variable name"
2029                ));
2030            }
2031            std::env::var(name).map(Some).map_err(|_| {
2032                format!("{flag_name}-env references unset environment variable: {name}")
2033            })
2034        }
2035        (None, None, Some(reference)) => {
2036            resolve_config_secret(&format!("{flag_name}-config"), reference).map(Some)
2037        }
2038        (None, None, None) => Ok(None),
2039        _ => Err(format!(
2040            "{flag_name}, {flag_name}-env, and {flag_name}-config are mutually exclusive"
2041        )),
2042    }
2043}
2044
2045pub fn parse_params(entries: &[String]) -> Result<Vec<Value>, String> {
2046    if entries.len() > MAX_PARAMS {
2047        return Err(format!("too many params; maximum params is {MAX_PARAMS}"));
2048    }
2049
2050    let mut by_index: BTreeMap<usize, Value> = BTreeMap::new();
2051    for entry in entries {
2052        let (idx, raw) = split_index_value(entry)?;
2053        if idx == 0 {
2054            return Err("param index must start at 1".to_string());
2055        }
2056        if idx > MAX_PARAMS {
2057            return Err(format!(
2058                "parameter index {idx} exceeds maximum params {MAX_PARAMS}"
2059            ));
2060        }
2061        match by_index.entry(idx) {
2062            Entry::Vacant(slot) => {
2063                slot.insert(parse_param_value(raw));
2064            }
2065            Entry::Occupied(_) => return Err(format!("duplicate parameter index {idx}")),
2066        }
2067    }
2068    if by_index.is_empty() {
2069        return Ok(vec![]);
2070    }
2071    let max = by_index.keys().max().copied().unwrap_or(0);
2072    for i in 1..=max {
2073        if !by_index.contains_key(&i) {
2074            return Err(format!("missing parameter index {i}"));
2075        }
2076    }
2077    Ok(by_index.into_values().collect())
2078}
2079
2080fn split_index_value(entry: &str) -> Result<(usize, &str), String> {
2081    let mut parts = entry.splitn(2, '=');
2082    let left = parts.next().unwrap_or_default();
2083    let right = parts
2084        .next()
2085        .ok_or_else(|| format!("invalid param '{entry}', expected N=value"))?;
2086    let idx = left
2087        .parse::<usize>()
2088        .map_err(|_| format!("invalid param index in '{entry}'"))?;
2089    Ok((idx, right))
2090}
2091
2092fn parse_param_value(v: &str) -> Value {
2093    if v == "null" {
2094        return Value::Null;
2095    }
2096    if v == "true" {
2097        return Value::Bool(true);
2098    }
2099    if v == "false" {
2100        return Value::Bool(false);
2101    }
2102    // Strings are passed verbatim to PostgreSQL via the text bind path so
2103    // that values like "00123" or "1.0" preserve their original form. The
2104    // server coerces them based on the prepared statement's parameter type.
2105    Value::String(v.to_string())
2106}
2107
2108fn wrap_explain_sql(user_sql: &str, analyze: bool) -> String {
2109    let body = user_sql.trim_end_matches([';', ' ', '\n', '\t', '\r']);
2110    if analyze {
2111        format!("explain (analyze true, format json, buffers true) {body}")
2112    } else {
2113        format!("explain (format json) {body}")
2114    }
2115}
2116
2117fn optional_string_value(value: Option<String>) -> Value {
2118    value.map(Value::String).unwrap_or(Value::Null)
2119}
2120
2121fn split_table_name(default_schema: String, name: String) -> (String, String) {
2122    match name.split_once('.') {
2123        Some((schema, table)) => (schema.to_string(), table.to_string()),
2124        None => (default_schema, name),
2125    }
2126}
2127
2128fn split_optional_table(default_schema: String, table: Option<String>) -> (String, Option<String>) {
2129    match table {
2130        Some(name) => {
2131            let (schema, table_name) = split_table_name(default_schema, name);
2132            (schema, Some(table_name))
2133        }
2134        None => (default_schema, None),
2135    }
2136}
2137
2138fn full_schema_snapshot_sql(relation_filter: &str, schema_only_filter: &str) -> String {
2139    format!(
2140        "with relation_filter as ( \
2141             select c.oid, c.relname, c.relkind, c.relpersistence, c.reltuples, c.relowner, \
2142                    n.nspname, pg_catalog.obj_description(c.oid, 'pg_class') as comment \
2143             from pg_catalog.pg_class c \
2144             join pg_catalog.pg_namespace n on n.oid = c.relnamespace \
2145             where n.nspname = $1 \
2146               and c.relkind in ('r', 'p', 'f', 'v', 'm', 'S') \
2147               and ({relation_filter}) \
2148         ), snapshot as ( \
2149             select 'extension'::text as kind, \
2150                    n.nspname::text as schema, \
2151                    null::text as relation, \
2152                    e.extname::text as name, \
2153                    'extension'::text as object_type, \
2154                    null::integer as position, \
2155                    null::text as definition, \
2156                    null::bigint as size_bytes, \
2157                    null::text as size, \
2158                    null::bigint as estimated_rows, \
2159                    pg_catalog.jsonb_build_object('version', e.extversion) as payload \
2160             from pg_catalog.pg_extension e \
2161             join pg_catalog.pg_namespace n on n.oid = e.extnamespace \
2162             where n.nspname = $1 and ({schema_only_filter}) \
2163             union all \
2164             select 'relation'::text as kind, \
2165                    rf.nspname::text as schema, \
2166                    rf.relname::text as relation, \
2167                    rf.relname::text as name, \
2168                    case rf.relkind \
2169                        when 'r' then 'table' \
2170                        when 'p' then 'partitioned table' \
2171                        when 'f' then 'foreign table' \
2172                        when 'v' then 'view' \
2173                        when 'm' then 'materialized view' \
2174                        else rf.relkind::text \
2175                    end as object_type, \
2176                    null::integer as position, \
2177                    case when rf.relkind in ('v', 'm') \
2178                         then pg_catalog.pg_get_viewdef(rf.oid, true) end as definition, \
2179                    case when rf.relkind in ('r', 'p', 'm') \
2180                         then pg_catalog.pg_total_relation_size(rf.oid) end as size_bytes, \
2181                    case when rf.relkind in ('r', 'p', 'm') \
2182                         then pg_catalog.pg_size_pretty(pg_catalog.pg_total_relation_size(rf.oid)) end as size, \
2183                    rf.reltuples::bigint as estimated_rows, \
2184                    pg_catalog.jsonb_build_object( \
2185                        'owner', pg_catalog.pg_get_userbyid(rf.relowner), \
2186                        'persistence', rf.relpersistence, \
2187                        'comment', rf.comment \
2188                    ) as payload \
2189             from relation_filter rf \
2190             where rf.relkind in ('r', 'p', 'f', 'v', 'm') \
2191             union all \
2192             select 'sequence'::text as kind, \
2193                    rf.nspname::text as schema, \
2194                    rf.relname::text as relation, \
2195                    rf.relname::text as name, \
2196                    'sequence'::text as object_type, \
2197                    null::integer as position, \
2198                    null::text as definition, \
2199                    pg_catalog.pg_relation_size(rf.oid) as size_bytes, \
2200                    pg_catalog.pg_size_pretty(pg_catalog.pg_relation_size(rf.oid)) as size, \
2201                    null::bigint as estimated_rows, \
2202                    pg_catalog.jsonb_build_object( \
2203                        'owner', pg_catalog.pg_get_userbyid(rf.relowner), \
2204                        'comment', rf.comment \
2205                    ) as payload \
2206             from relation_filter rf \
2207             where rf.relkind = 'S' \
2208             union all \
2209             select 'column'::text as kind, \
2210                    rf.nspname::text as schema, \
2211                    rf.relname::text as relation, \
2212                    a.attname::text as name, \
2213                    pg_catalog.format_type(a.atttypid, a.atttypmod)::text as object_type, \
2214                    a.attnum::integer as position, \
2215                    pg_catalog.pg_get_expr(ad.adbin, ad.adrelid)::text as definition, \
2216                    null::bigint as size_bytes, \
2217                    null::text as size, \
2218                    null::bigint as estimated_rows, \
2219                    pg_catalog.jsonb_build_object( \
2220                        'nullable', not a.attnotnull, \
2221                        'primary_key', coalesce(pk.is_primary, false), \
2222                        'identity', a.attidentity::text, \
2223                        'generated', a.attgenerated::text, \
2224                        'serial_sequence', pg_catalog.pg_get_serial_sequence( \
2225                            pg_catalog.format('%I.%I', rf.nspname, rf.relname), a.attname), \
2226                        'comment', pg_catalog.col_description(rf.oid, a.attnum) \
2227                    ) as payload \
2228             from pg_catalog.pg_attribute a \
2229             join relation_filter rf on rf.oid = a.attrelid \
2230             left join pg_catalog.pg_attrdef ad on ad.adrelid = a.attrelid and ad.adnum = a.attnum \
2231             left join lateral ( \
2232                 select true as is_primary \
2233                 from pg_catalog.pg_index i \
2234                 where i.indrelid = a.attrelid and i.indisprimary \
2235                   and a.attnum = any(i.indkey) \
2236             ) pk on true \
2237             where rf.relkind in ('r', 'p', 'f', 'v', 'm') \
2238               and a.attnum > 0 and not a.attisdropped \
2239             union all \
2240             select 'constraint'::text as kind, \
2241                    rf.nspname::text as schema, \
2242                    rf.relname::text as relation, \
2243                    con.conname::text as name, \
2244                    case con.contype \
2245                        when 'p' then 'primary key' \
2246                        when 'u' then 'unique' \
2247                        when 'f' then 'foreign key' \
2248                        when 'c' then 'check' \
2249                        when 'x' then 'exclusion' \
2250                        else con.contype::text \
2251                    end as object_type, \
2252                    null::integer as position, \
2253                    pg_catalog.pg_get_constraintdef(con.oid, true)::text as definition, \
2254                    null::bigint as size_bytes, \
2255                    null::text as size, \
2256                    null::bigint as estimated_rows, \
2257                    pg_catalog.jsonb_build_object( \
2258                        'type', con.contype::text, \
2259                        'deferrable', con.condeferrable, \
2260                        'deferred_by_default', con.condeferred, \
2261                        'validated', con.convalidated \
2262                    ) as payload \
2263             from pg_catalog.pg_constraint con \
2264             join relation_filter rf on rf.oid = con.conrelid \
2265             union all \
2266             select 'index'::text as kind, \
2267                    rf.nspname::text as schema, \
2268                    rf.relname::text as relation, \
2269                    ic.relname::text as name, \
2270                    am.amname::text as object_type, \
2271                    null::integer as position, \
2272                    pg_catalog.pg_get_indexdef(i.indexrelid)::text as definition, \
2273                    pg_catalog.pg_relation_size(i.indexrelid) as size_bytes, \
2274                    pg_catalog.pg_size_pretty(pg_catalog.pg_relation_size(i.indexrelid)) as size, \
2275                    null::bigint as estimated_rows, \
2276                    pg_catalog.jsonb_build_object( \
2277                        'unique', i.indisunique, \
2278                        'primary', i.indisprimary, \
2279                        'valid', i.indisvalid, \
2280                        'ready', i.indisready \
2281                    ) as payload \
2282             from pg_catalog.pg_index i \
2283             join pg_catalog.pg_class ic on ic.oid = i.indexrelid \
2284             join relation_filter rf on rf.oid = i.indrelid \
2285             join pg_catalog.pg_am am on am.oid = ic.relam \
2286             union all \
2287             select 'trigger'::text as kind, \
2288                    rf.nspname::text as schema, \
2289                    rf.relname::text as relation, \
2290                    tg.tgname::text as name, \
2291                    'trigger'::text as object_type, \
2292                    null::integer as position, \
2293                    pg_catalog.pg_get_triggerdef(tg.oid, true)::text as definition, \
2294                    null::bigint as size_bytes, \
2295                    null::text as size, \
2296                    null::bigint as estimated_rows, \
2297                    pg_catalog.jsonb_build_object( \
2298                        'enabled', tg.tgenabled::text, \
2299                        'function_schema', fn_ns.nspname, \
2300                        'function_name', fn.proname \
2301                    ) as payload \
2302             from pg_catalog.pg_trigger tg \
2303             join relation_filter rf on rf.oid = tg.tgrelid \
2304             join pg_catalog.pg_proc fn on fn.oid = tg.tgfoid \
2305             join pg_catalog.pg_namespace fn_ns on fn_ns.oid = fn.pronamespace \
2306             where not tg.tgisinternal \
2307             union all \
2308             select 'function'::text as kind, \
2309                    n.nspname::text as schema, \
2310                    null::text as relation, \
2311                    (p.proname || '(' || pg_catalog.pg_get_function_identity_arguments(p.oid) || ')')::text as name, \
2312                    'function'::text as object_type, \
2313                    null::integer as position, \
2314                    pg_catalog.pg_get_functiondef(p.oid)::text as definition, \
2315                    null::bigint as size_bytes, \
2316                    null::text as size, \
2317                    null::bigint as estimated_rows, \
2318                    pg_catalog.jsonb_build_object( \
2319                        'language', l.lanname, \
2320                        'result', pg_catalog.pg_get_function_result(p.oid), \
2321                        'identity_args', pg_catalog.pg_get_function_identity_arguments(p.oid) \
2322                    ) as payload \
2323             from pg_catalog.pg_proc p \
2324             join pg_catalog.pg_namespace n on n.oid = p.pronamespace \
2325             join pg_catalog.pg_language l on l.oid = p.prolang \
2326             where n.nspname = $1 \
2327               and p.prokind = 'f' \
2328               and ({schema_only_filter}) \
2329               and not exists ( \
2330                   select 1 \
2331                   from pg_catalog.pg_depend d \
2332                   where d.classid = 'pg_catalog.pg_proc'::regclass \
2333                     and d.objid = p.oid \
2334                     and d.deptype = 'e' \
2335               ) \
2336         ) \
2337         select * from snapshot \
2338         order by case kind \
2339                    when 'extension' then 0 \
2340                    when 'relation' then 1 \
2341                    when 'sequence' then 2 \
2342                    when 'column' then 3 \
2343                    when 'constraint' then 4 \
2344                    when 'index' then 5 \
2345                    when 'trigger' then 6 \
2346                    when 'function' then 7 \
2347                    else 99 end, \
2348                  schema, relation nulls first, position nulls last, name"
2349    )
2350}
2351
2352fn build_schema_snapshot_sql(args: InspectSchemaArgs) -> (String, Vec<Value>) {
2353    (
2354        full_schema_snapshot_sql("$2::text is null or c.relname like $2", "$2::text is null"),
2355        vec![Value::String(args.schema), optional_string_value(args.like)],
2356    )
2357}
2358
2359fn build_table_full_sql(schema: String, name: String) -> (String, Vec<Value>) {
2360    (
2361        full_schema_snapshot_sql("c.relname = $2", "false"),
2362        vec![Value::String(schema), Value::String(name)],
2363    )
2364}
2365
2366fn build_inspect_indexes_sql(args: InspectIndexesArgs) -> (String, Vec<Value>) {
2367    let (schema, table) = split_optional_table(args.schema, args.table);
2368    let mut sql = String::from(
2369        "select n.nspname as schema, \
2370                tc.relname as table, \
2371                ic.relname as name, \
2372                am.amname as method, \
2373                i.indisunique as unique, \
2374                i.indisprimary as primary, \
2375                i.indisvalid as valid, \
2376                i.indisready as ready, \
2377                pg_catalog.pg_get_indexdef(i.indexrelid) as definition, \
2378                pg_catalog.pg_relation_size(i.indexrelid) as size_bytes, \
2379                pg_catalog.pg_size_pretty(pg_catalog.pg_relation_size(i.indexrelid)) as size",
2380    );
2381    if args.stats {
2382        sql.push_str(
2383            ", s.idx_scan as index_scan_count, \
2384             s.idx_tup_read as index_tuple_read_count, \
2385             s.idx_tup_fetch as index_tuple_fetch_count",
2386        );
2387    }
2388    sql.push_str(
2389        " from pg_catalog.pg_index i \
2390          join pg_catalog.pg_class ic on ic.oid = i.indexrelid \
2391          join pg_catalog.pg_class tc on tc.oid = i.indrelid \
2392          join pg_catalog.pg_namespace n on n.oid = tc.relnamespace \
2393          join pg_catalog.pg_am am on am.oid = ic.relam",
2394    );
2395    if args.stats {
2396        sql.push_str(" left join pg_catalog.pg_stat_user_indexes s on s.indexrelid = i.indexrelid");
2397    }
2398    sql.push_str(" where n.nspname = $1");
2399
2400    let mut params = vec![Value::String(schema)];
2401    if let Some(table_name) = table {
2402        sql.push_str(" and tc.relname = $2");
2403        params.push(Value::String(table_name));
2404    }
2405    sql.push_str(" order by tc.relname, ic.relname");
2406    (sql, params)
2407}
2408
2409fn build_inspect_sql(action: InspectAction) -> (String, Vec<Value>) {
2410    match action {
2411        InspectAction::Databases(args) => {
2412            let mut sql = String::from(
2413                "select d.datname as database, \
2414                        pg_catalog.pg_get_userbyid(d.datdba) as owner, \
2415                        pg_catalog.pg_encoding_to_char(d.encoding) as encoding, \
2416                        d.datcollate as collate, \
2417                        d.datctype as ctype, \
2418                        d.datistemplate as is_template, \
2419                        d.datallowconn as allow_connections, \
2420                        d.datconnlimit as connection_limit, \
2421                        case when has_database_privilege(d.datname, 'CONNECT') \
2422                             then pg_catalog.pg_database_size(d.oid) end as size_bytes, \
2423                        case when has_database_privilege(d.datname, 'CONNECT') \
2424                             then pg_catalog.pg_size_pretty(pg_catalog.pg_database_size(d.oid)) end as size, \
2425                        s.numbackends as active_connections \
2426                 from pg_catalog.pg_database d \
2427                 left join pg_catalog.pg_stat_database s on s.datid = d.oid",
2428            );
2429            if !args.all {
2430                sql.push_str(" where not d.datistemplate");
2431            }
2432            sql.push_str(" order by d.datname");
2433            (sql, vec![])
2434        }
2435        InspectAction::Database => (
2436            "with rels as ( \
2437                 select c.relkind \
2438                 from pg_catalog.pg_class c \
2439                 join pg_catalog.pg_namespace n on n.oid = c.relnamespace \
2440                 where n.nspname not in ('pg_catalog', 'information_schema') \
2441                   and n.nspname not like 'pg_toast%' \
2442                   and n.nspname not like 'pg_temp_%' \
2443             ) \
2444             select current_database() as database, \
2445                    ( select count(*) from pg_catalog.pg_namespace n \
2446                       where n.nspname not in ('pg_catalog', 'information_schema') \
2447                         and n.nspname not like 'pg_toast%' \
2448                         and n.nspname not like 'pg_temp_%' ) as schemas, \
2449                    count(*) filter (where relkind in ('r', 'p')) as tables, \
2450                    count(*) filter (where relkind = 'v') as views, \
2451                    count(*) filter (where relkind = 'm') as materialized_views, \
2452                    count(*) filter (where relkind = 'S') as sequences, \
2453                    pg_catalog.pg_database_size(current_database()) as size_bytes, \
2454                    pg_catalog.pg_size_pretty(pg_catalog.pg_database_size(current_database())) as size \
2455             from rels"
2456                .to_string(),
2457            vec![],
2458        ),
2459        InspectAction::Schemas => (
2460            "select n.nspname as schema, \
2461                    pg_catalog.pg_get_userbyid(n.nspowner) as owner, \
2462                    count(*) filter (where c.relkind in ('r', 'p')) as tables, \
2463                    count(*) filter (where c.relkind = 'v') as views, \
2464                    count(*) filter (where c.relkind = 'm') as materialized_views, \
2465                    count(*) filter (where c.relkind = 'S') as sequences, \
2466                    pg_catalog.pg_size_pretty(coalesce( \
2467                        sum(pg_catalog.pg_total_relation_size(c.oid)) \
2468                            filter (where c.relkind in ('r', 'p', 'm')), 0)) as size \
2469             from pg_catalog.pg_namespace n \
2470             left join pg_catalog.pg_class c on c.relnamespace = n.oid \
2471             where n.nspname not in ('pg_catalog', 'information_schema') \
2472               and n.nspname not like 'pg_toast%' \
2473               and n.nspname not like 'pg_temp_%' \
2474             group by n.nspname, n.nspowner \
2475             order by n.nspname"
2476                .to_string(),
2477            vec![],
2478        ),
2479        InspectAction::Schema(args) | InspectAction::Snapshot(args) => build_schema_snapshot_sql(args),
2480        InspectAction::Tables(args) => {
2481            let mut sql = String::from(
2482                "select n.nspname as schema, \
2483                        c.relname as name, \
2484                        case c.relkind when 'r' then 'table' \
2485                                       when 'p' then 'partitioned table' \
2486                                       when 'f' then 'foreign table' end as kind, \
2487                        pg_catalog.pg_get_userbyid(c.relowner) as owner, \
2488                        c.reltuples::bigint as estimated_rows, \
2489                        pg_catalog.pg_size_pretty(pg_catalog.pg_total_relation_size(c.oid)) as size, \
2490                        pg_catalog.pg_total_relation_size(c.oid) as size_bytes \
2491                 from pg_catalog.pg_class c \
2492                 join pg_catalog.pg_namespace n on n.oid = c.relnamespace \
2493                 where n.nspname = $1 and c.relkind in ('r', 'p', 'f')",
2494            );
2495            let mut params = vec![Value::String(args.schema)];
2496            if let Some(pattern) = args.like {
2497                sql.push_str(" and c.relname like $2");
2498                params.push(Value::String(pattern));
2499            }
2500            sql.push_str(" order by c.relname");
2501            (sql, params)
2502        }
2503        InspectAction::Views(args) => {
2504            let mut sql = String::from(
2505                "select n.nspname as schema, \
2506                        c.relname as name, \
2507                        case c.relkind when 'm' then true else false end as materialized, \
2508                        pg_catalog.pg_get_userbyid(c.relowner) as owner \
2509                 from pg_catalog.pg_class c \
2510                 join pg_catalog.pg_namespace n on n.oid = c.relnamespace \
2511                 where n.nspname = $1 and c.relkind in ('v', 'm')",
2512            );
2513            let mut params = vec![Value::String(args.schema)];
2514            if let Some(pattern) = args.like {
2515                sql.push_str(" and c.relname like $2");
2516                params.push(Value::String(pattern));
2517            }
2518            sql.push_str(" order by c.relname");
2519            (sql, params)
2520        }
2521        InspectAction::Indexes(args) => build_inspect_indexes_sql(args),
2522        InspectAction::Table(args) => {
2523            let (schema, name) = split_table_name("public".to_string(), args.name);
2524            if args.full {
2525                return build_table_full_sql(schema, name);
2526            }
2527            (
2528                "select a.attname as name, \
2529                        pg_catalog.format_type(a.atttypid, a.atttypmod) as type, \
2530                        not a.attnotnull as nullable, \
2531                        pg_catalog.pg_get_expr(ad.adbin, ad.adrelid) as default, \
2532                        a.attnum as position, \
2533                        coalesce(pk.is_primary, false) as primary_key, \
2534                        pg_catalog.col_description(c.oid, a.attnum) as comment \
2535                 from pg_catalog.pg_attribute a \
2536                 join pg_catalog.pg_class c on c.oid = a.attrelid \
2537                 join pg_catalog.pg_namespace n on n.oid = c.relnamespace \
2538                 left join pg_catalog.pg_attrdef ad \
2539                     on ad.adrelid = a.attrelid and ad.adnum = a.attnum \
2540                 left join lateral ( \
2541                     select true as is_primary \
2542                     from pg_catalog.pg_index i \
2543                     where i.indrelid = a.attrelid and i.indisprimary \
2544                       and a.attnum = any(i.indkey) \
2545                 ) pk on true \
2546                 where n.nspname = $1 and c.relname = $2 \
2547                   and a.attnum > 0 and not a.attisdropped \
2548                 order by a.attnum"
2549                    .to_string(),
2550                vec![Value::String(schema), Value::String(name)],
2551            )
2552        }
2553    }
2554}
2555
2556#[cfg(test)]
2557#[path = "../tests/support/unit_cli.rs"]
2558mod tests;