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

agent_first_psql/
cli_runner.rs
