agent_first_psql/
cli_runner.rs

1use crate::db::ExecRequest;
2use crate::handler;
3use crate::handler::App;
4use crate::limits::OUTPUT_CHANNEL_CAPACITY;
5use crate::logutil::build_startup_log;
6use crate::protocol::log_event;
7use crate::types::{Output, QueryOptions, RuntimeConfig, Trace};
8use agent_first_data::OutputFormat;
9use std::io::Write;
10use std::sync::Arc;
11use std::sync::atomic::Ordering;
12use tokio::sync::mpsc;
13
14pub async fn run(
15    req: crate::cli::CliRequest,
16    capability: crate::Capability,
17    locked_readonly_profile: bool,
18) {
19    let crate::cli::CliRequest {
20        sql,
21        params,
22        options,
23        session,
24        output: output_format,
25        log,
26        startup_args,
27        startup_env,
28        startup_requested,
29        dry_run,
30        psql_mode,
31    } = req;
32
33    let stdout = std::io::stdout();
34    let mut sink = CliOutputSink::new(stdout.lock(), output_format);
35
36    let config = RuntimeConfig::default();
37    let (tx, mut rx) = mpsc::channel::<Output>(OUTPUT_CHANNEL_CAPACITY);
38    let app = Arc::new(App::new(config, tx, capability));
39    app.locked_readonly_profile
40        .store(locked_readonly_profile, Ordering::Relaxed);
41
42    {
43        let mut cfg = app.config.write().await;
44        cfg.sessions.insert("default".to_string(), session.clone());
45        if !log.is_empty() {
46            cfg.log = log.clone();
47        }
48    }
49
50    if dry_run {
51        let cfg = app.config.read().await.clone();
52        let session_cfg = cfg.sessions.get("default").cloned().unwrap_or_default();
53        let resolved_opts = match cfg.resolve_options_for_session(&options, &session_cfg) {
54            Ok(options) => options,
55            Err(error) => {
56                if crate::emit::emit_cli_error(&error, None, output_format).is_err() {
57                    std::process::exit(4);
58                }
59                std::process::exit(2);
60            }
61        };
62        if capability == crate::Capability::ReadOnly
63            && let Err(error) = crate::readonly_policy::validate_sql(&sql)
64        {
65            if crate::emit::emit_cli_error(&error, Some(crate::readonly_hint()), output_format)
66                .is_err()
67            {
68                std::process::exit(4);
69            }
70            std::process::exit(2);
71        }
72        if capability == crate::Capability::ReadOnly && !resolved_opts.read_only {
73            if crate::emit::emit_cli_error(
74                "write permission is unavailable in afpsql-readonly",
75                Some(crate::readonly_hint()),
76                output_format,
77            )
78            .is_err()
79            {
80                std::process::exit(4);
81            }
82            std::process::exit(2);
83        }
84        let start = std::time::Instant::now();
85        let outcome = app
86            .executor
87            .prepare_only(ExecRequest {
88                session_name: "default",
89                session_cfg: &session_cfg,
90                sql: &sql,
91                params: &params,
92                opts: &resolved_opts,
93                cancel_slot: None,
94                transport_log: None,
95            })
96            .await;
97        let mut had_error = false;
98        match outcome {
99            Ok(info) => {
100                let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
101                if sink
102                    .emit(&Output::DryRun {
103                        id: None,
104                        sql: sql.clone(),
105                        params: params.iter().map(|v| v.to_string()).collect(),
106                        session: Some("default".to_string()),
107                        param_types: info.param_types,
108                        columns: info.columns,
109                        trace,
110                    })
111                    .is_err()
112                {
113                    std::process::exit(4);
114                }
115            }
116            Err(err) => {
117                had_error = true;
118                handler::emit_exec_error(&app, None, "default", err, start).await;
119            }
120        }
121        app.executor.shutdown().await;
122        drop(app);
123        while let Some(event) = rx.recv().await {
124            if sink.emit(&event).is_err() {
125                std::process::exit(4);
126            }
127        }
128        std::process::exit(if had_error { 1 } else { 0 });
129    }
130
131    if startup_requested {
132        let event = build_startup_log(Some("default"), &startup_args, &startup_env);
133        if sink.emit(&event).is_err() {
134            std::process::exit(4);
135        }
136    }
137
138    if psql_mode
139        && log.enabled(log_event::MODE_PERMISSION_DEFAULT_CHANGED)
140        && sink.emit(&psql_mode_permission_event(&options)).is_err()
141    {
142        std::process::exit(4);
143    }
144
145    app.requests_total.fetch_add(1, Ordering::Relaxed);
146    handler::execute_query(
147        &app,
148        None,
149        Some("default".to_string()),
150        sql,
151        params,
152        options,
153        None,
154    )
155    .await;
156
157    app.executor.shutdown().await;
158    drop(app);
159
160    let mut had_error = false;
161    while let Some(event) = rx.recv().await {
162        if matches!(event, Output::Error { .. } | Output::SqlError { .. }) {
163            had_error = true;
164        }
165        if sink.emit(&event).is_err() {
166            std::process::exit(4);
167        }
168    }
169
170    std::process::exit(if had_error { 1 } else { 0 });
171}
172
173struct CliOutputSink<W: Write> {
174    writer: W,
175    format: OutputFormat,
176}
177
178impl<W: Write> CliOutputSink<W> {
179    fn new(writer: W, format: OutputFormat) -> Self {
180        Self { writer, format }
181    }
182
183    fn emit(&mut self, out: &Output) -> Result<(), agent_first_data::CliEmitterError> {
184        crate::output_fmt::emit_output(&mut self.writer, out, self.format)
185    }
186}
187
188fn psql_mode_permission_event(options: &QueryOptions) -> Output {
189    let permission = options.permission.map(|p| p.as_str()).unwrap_or("write");
190    let mut config = serde_json::Map::new();
191    config.insert("mode".to_string(), serde_json::Value::from("psql"));
192    config.insert(
193        "permission".to_string(),
194        serde_json::Value::from(permission),
195    );
196    config.insert(
197        "note".to_string(),
198        serde_json::Value::from(
199            "psql mode inherits psql's writable default; native mode defaults to read",
200        ),
201    );
202    Output::Log {
203        event: log_event::MODE_PERMISSION_DEFAULT_CHANGED.to_string(),
204        request_id: None,
205        session: Some("default".to_string()),
206        error_code: None,
207        command_tag: None,
208        version: None,
209        config: Some(serde_json::Value::Object(config)),
210        args: None,
211        env: None,
212        chain: None,
213        trace: Trace::only_duration(0),
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220    use crate::types::Permission;
221
222    #[test]
223    fn psql_mode_event_reports_resolved_permission_and_filter_prefix() {
224        let opts = QueryOptions {
225            permission: Some(Permission::ContainerWrite),
226            ..Default::default()
227        };
228        let emitted = psql_mode_permission_event(&opts);
229        assert!(matches!(emitted, Output::Log { .. }));
230        let Output::Log { event, config, .. } = emitted else {
231            return;
232        };
233        assert_eq!(event, "mode.permission_default_changed");
234        let cfg = config.unwrap_or_default();
235        assert_eq!(cfg.get("mode").and_then(|v| v.as_str()), Some("psql"));
236        assert_eq!(
237            cfg.get("permission").and_then(|v| v.as_str()),
238            Some("container-write")
239        );
240        assert!(cfg.get("note").is_some());
241        assert!(
242            agent_first_data::LogFilters::new(["mode"])
243                .enabled(log_event::MODE_PERMISSION_DEFAULT_CHANGED)
244        );
245        assert!(
246            !agent_first_data::LogFilters::default()
247                .enabled(log_event::MODE_PERMISSION_DEFAULT_CHANGED)
248        );
249    }
250}
agent_first_psql/cli_runner.rs

agent_first_psql/
cli_runner.rs
