agent_first_psql/
pipe.rs

1use crate::config::sessions_to_invalidate;
2use crate::conn::resolve_session_name;
3use crate::emit::emit_output;
4use crate::handler::{self, App, QueryPhase, QueryState};
5use crate::limits::{
6    MAX_ACTIVE_QUERIES, MAX_PARAMS, MAX_PIPE_LINE_BYTES, MAX_SQL_BYTES, OUTPUT_CHANNEL_CAPACITY,
7};
8use crate::logutil::build_startup_log;
9use crate::protocol::error_code;
10use crate::types::*;
11use std::collections::HashMap;
12use std::sync::Arc;
13use std::sync::atomic::Ordering;
14use std::time::Instant;
15use tokio::io::{AsyncBufRead, AsyncBufReadExt};
16use tokio::sync::mpsc;
17
18pub async fn run(
19    init: crate::cli::PipeInit,
20    capability: crate::Capability,
21    locked_readonly_profile: bool,
22) {
23    let crate::cli::PipeInit {
24        output,
25        session,
26        log,
27        startup_args,
28        startup_env,
29        startup_requested,
30    } = init;
31
32    let mut config = RuntimeConfig::default();
33    if has_session_override(&session) {
34        config
35            .sessions
36            .insert(config.default_session.clone(), session.clone());
37    }
38    if !log.is_empty() {
39        config.log = log.clone();
40    }
41    if startup_requested {
42        let event = build_startup_log(None, &startup_args, &startup_env);
43        if emit_output(&event, output).is_err() {
44            std::process::exit(4);
45        }
46    }
47
48    let (tx, rx) = mpsc::channel::<Output>(OUTPUT_CHANNEL_CAPACITY);
49    let writer = tokio::spawn(crate::writer::writer_task(rx, output));
50
51    let app = Arc::new(App::new(config, tx, capability));
52    app.locked_readonly_profile
53        .store(locked_readonly_profile, Ordering::Relaxed);
54    let runtime = Arc::new(PipeRuntime::new(app));
55
56    let stdin = tokio::io::stdin();
57    let reader = tokio::io::BufReader::new(stdin);
58    let mut reader = reader;
59
60    loop {
61        let line = match read_limited_line(&mut reader, MAX_PIPE_LINE_BYTES).await {
62            Ok(Some(Ok(line))) => line,
63            Ok(Some(Err(()))) => {
64                send_protocol_error(
65                    &runtime.app,
66                    None,
67                    error_code::INVALID_REQUEST,
68                    "input line exceeds maximum size",
69                    Some("split large requests or reduce SQL/params payload size"),
70                    false,
71                )
72                .await;
73                continue;
74            }
75            Ok(None) => break,
76            Err(e) => {
77                send_protocol_error(
78                    &runtime.app,
79                    None,
80                    error_code::INVALID_REQUEST,
81                    &format!("read error: {e}"),
82                    None,
83                    false,
84                )
85                .await;
86                break;
87            }
88        };
89        let trimmed = line.trim();
90        if trimmed.is_empty() {
91            continue;
92        }
93
94        let input: Input = match serde_json::from_str(trimmed) {
95            Ok(v) => v,
96            Err(e) => {
97                send_protocol_error(
98                    &runtime.app,
99                    None,
100                    error_code::INVALID_REQUEST,
101                    &format!("parse error: {e}"),
102                    None,
103                    false,
104                )
105                .await;
106                continue;
107            }
108        };
109
110        if dispatch_input(&runtime, input).await {
111            break;
112        }
113    }
114
115    wait_for_workers_shutdown(&runtime).await;
116    runtime.app.executor.shutdown().await;
117    send_close_event(&runtime.app).await;
118    drop(runtime);
119    match writer.await {
120        Ok(Ok(())) => {}
121        Ok(Err(_)) | Err(_) => std::process::exit(4),
122    }
123}
124
125struct PipeRuntime {
126    app: Arc<App>,
127    workers: tokio::sync::Mutex<HashMap<String, SessionWorker>>,
128}
129
130struct SessionWorker {
131    tx: mpsc::Sender<SessionOp>,
132    handle: tokio::task::JoinHandle<()>,
133}
134
135/// One operation queued for a session's worker. Tx control (`begin`/
136/// `commit`/`rollback`) flows through the same FIFO as `query` so that
137/// the agent's input order is the order PostgreSQL sees, even though
138/// queries normally run on a separate task.
139enum SessionOp {
140    Query(QueuedQuery),
141    Begin {
142        id: Option<String>,
143        session: String,
144        read_only: bool,
145    },
146    Commit {
147        id: Option<String>,
148        session: String,
149    },
150    Rollback {
151        id: Option<String>,
152        session: String,
153    },
154}
155
156struct QueuedQuery {
157    id: String,
158    session: String,
159    sql: String,
160    params: Vec<serde_json::Value>,
161    options: QueryOptions,
162    cancel_slot: crate::db::CancelSlot,
163    state: Arc<QueryState>,
164}
165
166impl PipeRuntime {
167    fn new(app: Arc<App>) -> Self {
168        Self {
169            app,
170            workers: tokio::sync::Mutex::new(HashMap::new()),
171        }
172    }
173}
174
175async fn dispatch_input(runtime: &Arc<PipeRuntime>, input: Input) -> bool {
176    match input {
177        Input::Query {
178            id,
179            session,
180            sql,
181            params,
182            options,
183        } => dispatch_query(runtime, id, session, sql, params, options).await,
184        Input::Config(patch) => {
185            if runtime.app.capability == crate::Capability::ReadOnly
186                && runtime.app.locked_readonly_profile.load(Ordering::Relaxed)
187                && patch.sessions.is_some()
188            {
189                send_protocol_error(
190                    &runtime.app,
191                    None,
192                    error_code::INVALID_REQUEST,
193                    "pipe session config cannot override an administrator-locked afpsql-readonly profile",
194                    Some(crate::readonly_local_capability_hint()),
195                    false,
196                )
197                .await;
198                return false;
199            }
200            let sessions = sessions_to_invalidate(&patch);
201            let mut cfg_snapshot = runtime.app.config.read().await.clone();
202            cfg_snapshot.apply_update(patch);
203            if runtime.app.capability == crate::Capability::ReadOnly
204                && let Some(error) = cfg_snapshot
205                    .sessions
206                    .values()
207                    .find_map(|session| crate::readonly_policy::validate_session(session).err())
208            {
209                send_protocol_error(
210                    &runtime.app,
211                    None,
212                    error_code::INVALID_REQUEST,
213                    &error,
214                    Some(crate::readonly_local_capability_hint()),
215                    false,
216                )
217                .await;
218                return false;
219            }
220            *runtime.app.config.write().await = cfg_snapshot.clone();
221            runtime.app.executor.invalidate_sessions(&sessions).await;
222            let _ = runtime.app.writer.send(Output::Config(cfg_snapshot)).await;
223        }
224        Input::Cancel { id } => dispatch_cancel(&runtime.app, id).await,
225        Input::SessionInfo { id, session } => {
226            handler::handle_session_info(&runtime.app, id, session).await;
227        }
228        Input::Ping => {
229            cleanup_finished_queries(&runtime.app).await;
230            let _ = runtime
231                .app
232                .writer
233                .send(Output::Pong {
234                    trace: PongTrace {
235                        uptime_s: runtime.app.start_time.elapsed().as_secs(),
236                        requests_total: runtime.app.requests_total.load(Ordering::Relaxed),
237                        in_flight: runtime.app.in_flight.lock().await.len(),
238                    },
239                })
240                .await;
241        }
242        Input::Begin {
243            id,
244            session,
245            read_only,
246            permission,
247        } => dispatch_begin(runtime, id, session, read_only, permission).await,
248        Input::Commit { id, session } => dispatch_commit(runtime, id, session).await,
249        Input::Rollback { id, session } => dispatch_rollback(runtime, id, session).await,
250        Input::Close => return true,
251    }
252    false
253}
254
255async fn dispatch_begin(
256    runtime: &Arc<PipeRuntime>,
257    id: Option<String>,
258    session: Option<String>,
259    read_only: bool,
260    permission: Option<crate::types::Permission>,
261) {
262    let app = &runtime.app;
263    if app.capability == crate::Capability::ReadOnly
264        && (!read_only || permission.is_some_and(|value| !value.is_read_only()))
265    {
266        send_protocol_error(
267            app,
268            id,
269            error_code::INVALID_REQUEST,
270            "read-write transactions are unavailable in afpsql-readonly",
271            Some(crate::readonly_hint()),
272            false,
273        )
274        .await;
275        return;
276    }
277    let cfg = app.config.read().await.clone();
278    let resolved_session = crate::conn::resolve_session_name(&cfg, session.as_deref());
279    let Some(session_cfg) = cfg.sessions.get(&resolved_session).cloned() else {
280        let _ = app
281            .writer
282            .send(Output::Error {
283                id,
284                error_code: error_code::CONNECT_FAILED.to_string(),
285                error: format!("unknown session: {resolved_session}"),
286                sqlstate: None,
287                message: None,
288                detail: None,
289                hint: Some(
290                    "check --host/--port or PGHOST/PGPORT environment variables".to_string(),
291                ),
292                retryable: true,
293                trace: Trace::only_duration(0),
294            })
295            .await;
296        return;
297    };
298    if app.capability == crate::Capability::ReadOnly
299        && let Err(error) = crate::readonly_policy::validate_session_with_trust(
300            &session_cfg,
301            app.locked_readonly_profile.load(Ordering::Relaxed),
302        )
303    {
304        send_protocol_error(
305            app,
306            id,
307            error_code::INVALID_REQUEST,
308            &error,
309            Some(crate::readonly_local_capability_hint()),
310            false,
311        )
312        .await;
313        return;
314    }
315
316    // A read-write begin needs a write permission matching the session's
317    // transport. Read-only begin always passes (it's strictly less than the
318    // session default).
319    if !read_only {
320        let probe = QueryOptions {
321            permission,
322            ..Default::default()
323        };
324        if let Err(message) = cfg.resolve_options_for_session(&probe, &session_cfg) {
325            let _ = app
326                .writer
327                .send(Output::Error {
328                    id,
329                    error_code: error_code::INVALID_REQUEST.to_string(),
330                    error: message,
331                    sqlstate: None,
332                    message: None,
333                    detail: None,
334                    hint: Some(
335                        "pass `permission` matching this session's transport (write / ssh-write / container-write), or set read_only:true"
336                            .to_string(),
337                    ),
338                    retryable: false,
339                    trace: Trace::only_duration(0),
340                })
341                .await;
342            return;
343        }
344    }
345
346    let tx = get_session_worker(runtime, &resolved_session).await;
347    if tx
348        .send(SessionOp::Begin {
349            id: id.clone(),
350            session: resolved_session,
351            read_only,
352        })
353        .await
354        .is_err()
355    {
356        let _ = app
357            .writer
358            .send(Output::Error {
359                id,
360                error_code: error_code::INVALID_REQUEST.to_string(),
361                error: "session worker is unavailable".to_string(),
362                sqlstate: None,
363                message: None,
364                detail: None,
365                hint: Some("retry; the session worker will be restarted".to_string()),
366                retryable: true,
367                trace: Trace::only_duration(0),
368            })
369            .await;
370    }
371}
372
373async fn dispatch_commit(runtime: &Arc<PipeRuntime>, id: Option<String>, session: Option<String>) {
374    enqueue_tx_finish(runtime, id, session, true).await;
375}
376
377async fn dispatch_rollback(
378    runtime: &Arc<PipeRuntime>,
379    id: Option<String>,
380    session: Option<String>,
381) {
382    enqueue_tx_finish(runtime, id, session, false).await;
383}
384
385async fn enqueue_tx_finish(
386    runtime: &Arc<PipeRuntime>,
387    id: Option<String>,
388    session: Option<String>,
389    commit: bool,
390) {
391    let app = &runtime.app;
392    let cfg = app.config.read().await.clone();
393    let resolved_session = crate::conn::resolve_session_name(&cfg, session.as_deref());
394    if !cfg.sessions.contains_key(&resolved_session) {
395        let _ = app
396            .writer
397            .send(Output::Error {
398                id,
399                error_code: error_code::CONNECT_FAILED.to_string(),
400                error: format!("unknown session: {resolved_session}"),
401                sqlstate: None,
402                message: None,
403                detail: None,
404                hint: None,
405                retryable: true,
406                trace: Trace::only_duration(0),
407            })
408            .await;
409        return;
410    }
411    let tx = get_session_worker(runtime, &resolved_session).await;
412    let op = if commit {
413        SessionOp::Commit {
414            id: id.clone(),
415            session: resolved_session,
416        }
417    } else {
418        SessionOp::Rollback {
419            id: id.clone(),
420            session: resolved_session,
421        }
422    };
423    if tx.send(op).await.is_err() {
424        let _ = app
425            .writer
426            .send(Output::Error {
427                id,
428                error_code: error_code::INVALID_REQUEST.to_string(),
429                error: "session worker is unavailable".to_string(),
430                sqlstate: None,
431                message: None,
432                detail: None,
433                hint: Some("retry; the session worker will be restarted".to_string()),
434                retryable: true,
435                trace: Trace::only_duration(0),
436            })
437            .await;
438    }
439}
440
441async fn emit_tx_error(
442    app: &Arc<App>,
443    id: Option<String>,
444    resolved_session: &str,
445    err: crate::db::ExecError,
446    start: Instant,
447) {
448    use crate::db::ExecError;
449    let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
450    let output = match err {
451        ExecError::Sql {
452            sqlstate,
453            message,
454            detail,
455            hint,
456            position,
457        } => Output::SqlError {
458            id,
459            session: Some(resolved_session.to_string()),
460            sqlstate,
461            message,
462            detail,
463            hint,
464            position,
465            trace,
466        },
467        ExecError::Connect(connect) => {
468            let c = *connect;
469            Output::Error {
470                id,
471                error_code: error_code::CONNECT_FAILED.to_string(),
472                error: c.error,
473                sqlstate: c.sqlstate,
474                message: c.message,
475                detail: c.detail,
476                hint: c.hint,
477                retryable: c.retryable,
478                trace,
479            }
480        }
481        ExecError::InvalidParams(message) => Output::Error {
482            id,
483            error_code: error_code::INVALID_REQUEST.to_string(),
484            error: message,
485            sqlstate: None,
486            message: None,
487            detail: None,
488            hint: None,
489            retryable: false,
490            trace,
491        },
492        other => Output::Error {
493            id,
494            error_code: error_code::INVALID_REQUEST.to_string(),
495            error: format!("{other:?}"),
496            sqlstate: None,
497            message: None,
498            detail: None,
499            hint: None,
500            retryable: false,
501            trace,
502        },
503    };
504    let _ = app.writer.send(output).await;
505}
506
507async fn dispatch_query(
508    runtime: &Arc<PipeRuntime>,
509    id: String,
510    session: Option<String>,
511    sql: String,
512    params: Vec<serde_json::Value>,
513    options: QueryOptions,
514) {
515    let app = &runtime.app;
516    if let Some(error) = validate_query_request(&id, &sql, &params) {
517        let _ = app.writer.send(error).await;
518        return;
519    }
520
521    cleanup_finished_queries(app).await;
522
523    let key = id.clone();
524    let mut rejection: Option<Output> = None;
525    let cancel_slot = crate::db::new_cancel_slot();
526    let state = Arc::new(QueryState::queued());
527    {
528        let mut in_flight = app.in_flight.lock().await;
529        if let Some(existing) = in_flight.get(&key) {
530            if existing.state.is_finished() {
531                let _ = in_flight.remove(&key);
532            } else {
533                rejection = Some(Output::Error {
534                    id: Some(key.clone()),
535                    error_code: error_code::INVALID_REQUEST.to_string(),
536                    error: "duplicate active query id".to_string(),
537                    sqlstate: None,
538                    message: None,
539                    detail: None,
540                    hint: Some(
541                        "pick a unique `id` per in-flight query, or cancel the prior one first"
542                            .to_string(),
543                    ),
544                    retryable: false,
545                    trace: Trace::only_duration(0),
546                });
547            }
548        }
549
550        if rejection.is_none() && in_flight.len() >= MAX_ACTIVE_QUERIES {
551            rejection = Some(Output::Error {
552                id: Some(key.clone()),
553                error_code: error_code::INVALID_REQUEST.to_string(),
554                error: "too many queued or running queries".to_string(),
555                sqlstate: None,
556                message: None,
557                detail: None,
558                hint: Some(format!(
559                    "maximum queued or running queries is {MAX_ACTIVE_QUERIES}"
560                )),
561                retryable: true,
562                trace: Trace::only_duration(0),
563            });
564        }
565
566        if rejection.is_none() {
567            in_flight.insert(
568                key.clone(),
569                handler::InFlightQuery {
570                    cancel_slot: cancel_slot.clone(),
571                    state: state.clone(),
572                },
573            );
574        }
575    }
576
577    if let Some(output) = rejection {
578        let _ = app.writer.send(output).await;
579        return;
580    }
581
582    let resolved_session = {
583        let cfg = app.config.read().await;
584        resolve_session_name(&cfg, session.as_deref())
585    };
586    let tx = get_session_worker(runtime, &resolved_session).await;
587    app.requests_total.fetch_add(1, Ordering::Relaxed);
588    let queued = QueuedQuery {
589        id: key.clone(),
590        session: resolved_session,
591        sql,
592        params,
593        options,
594        cancel_slot,
595        state,
596    };
597    if tx.send(SessionOp::Query(queued)).await.is_err() {
598        let _ = app.in_flight.lock().await.remove(&key);
599        let _ = app
600            .writer
601            .send(Output::Error {
602                id: Some(key),
603                error_code: error_code::INVALID_REQUEST.to_string(),
604                error: "session worker is unavailable".to_string(),
605                sqlstate: None,
606                message: None,
607                detail: None,
608                hint: Some("retry the query; the session worker will be restarted".to_string()),
609                retryable: true,
610                trace: Trace::only_duration(0),
611            })
612            .await;
613    }
614}
615
616async fn dispatch_cancel(app: &Arc<App>, id: String) {
617    let query = {
618        let in_flight = app.in_flight.lock().await;
619        in_flight.get(&id).cloned()
620    };
621
622    if let Some(query) = query {
623        // Try to claim the terminal emit slot. If the handler already won
624        // (result/sql_error has been sent), report `already finished` and
625        // do not emit a second terminal event.
626        let won_emit = query.cancel_slot.claim_terminal_emit();
627        if !won_emit || query.state.phase() == QueryPhase::Finished {
628            let _ = app.in_flight.lock().await.remove(&id);
629            let _ = app
630                .writer
631                .send(Output::Error {
632                    id: Some(id),
633                    error_code: error_code::INVALID_REQUEST.to_string(),
634                    error: "query already finished".to_string(),
635                    sqlstate: None,
636                    message: None,
637                    detail: None,
638                    hint: Some(
639                        "cancel raced completion; the prior result/error event holds the outcome"
640                            .to_string(),
641                    ),
642                    retryable: false,
643                    trace: Trace::only_duration(0),
644                })
645                .await;
646        } else {
647            query.state.set_phase(QueryPhase::Cancelled);
648            let hint = match query.cancel_server_query().await {
649                Ok(true) => Some("server-side cancel requested".to_string()),
650                Ok(false) => {
651                    Some("query cancelled before execution reached the database".to_string())
652                }
653                Err(e) => Some(e),
654            };
655            let _ = app.in_flight.lock().await.remove(&id);
656            let _ = app
657                .writer
658                .send(Output::Error {
659                    id: Some(id),
660                    error_code: error_code::CANCELLED.to_string(),
661                    error: "query cancelled".to_string(),
662                    sqlstate: None,
663                    message: None,
664                    detail: None,
665                    hint,
666                    retryable: false,
667                    trace: Trace::only_duration(0),
668                })
669                .await;
670        }
671    } else {
672        let _ = app
673            .writer
674            .send(Output::Error {
675                id: Some(id),
676                error_code: error_code::INVALID_REQUEST.to_string(),
677                error: "no queued or running query with this id".to_string(),
678                sqlstate: None,
679                message: None,
680                detail: None,
681                hint: Some(
682                    "no matching in-flight query for this id; it may have already completed or never been submitted"
683                        .to_string(),
684                ),
685                retryable: false,
686                trace: Trace::only_duration(0),
687            })
688            .await;
689    }
690}
691
692async fn get_session_worker(runtime: &Arc<PipeRuntime>, session: &str) -> mpsc::Sender<SessionOp> {
693    let mut workers = runtime.workers.lock().await;
694    if let Some(worker) = workers.get(session)
695        && !worker.handle.is_finished()
696    {
697        return worker.tx.clone();
698    }
699
700    workers.remove(session);
701    let (tx, rx) = mpsc::channel(MAX_ACTIVE_QUERIES);
702    let app = runtime.app.clone();
703    let handle = tokio::spawn(async move {
704        session_worker_loop(app, rx).await;
705    });
706    workers.insert(
707        session.to_string(),
708        SessionWorker {
709            tx: tx.clone(),
710            handle,
711        },
712    );
713    tx
714}
715
716async fn session_worker_loop(app: Arc<App>, mut rx: mpsc::Receiver<SessionOp>) {
717    while let Some(op) = rx.recv().await {
718        match op {
719            SessionOp::Query(query) => {
720                if query.cancel_slot.is_cancelled() || !query.state.try_start() {
721                    query.state.set_phase(QueryPhase::Cancelled);
722                    continue;
723                }
724
725                handler::execute_query(
726                    &app,
727                    Some(query.id),
728                    Some(query.session),
729                    query.sql,
730                    query.params,
731                    query.options,
732                    Some(query.cancel_slot.clone()),
733                )
734                .await;
735
736                if query.cancel_slot.is_cancelled() || query.state.phase() == QueryPhase::Cancelled
737                {
738                    query.state.set_phase(QueryPhase::Cancelled);
739                } else {
740                    query.state.set_phase(QueryPhase::Finished);
741                }
742            }
743            SessionOp::Begin {
744                id,
745                session,
746                read_only,
747            } => exec_tx_begin_on_worker(&app, id, session, read_only).await,
748            SessionOp::Commit { id, session } => {
749                exec_tx_finish_on_worker(&app, id, session, true).await
750            }
751            SessionOp::Rollback { id, session } => {
752                exec_tx_finish_on_worker(&app, id, session, false).await
753            }
754        }
755    }
756}
757
758async fn exec_tx_begin_on_worker(
759    app: &Arc<App>,
760    id: Option<String>,
761    session: String,
762    read_only: bool,
763) {
764    let cfg = app.config.read().await.clone();
765    let Some(session_cfg) = cfg.sessions.get(&session).cloned() else {
766        let _ = app
767            .writer
768            .send(Output::Error {
769                id,
770                error_code: error_code::CONNECT_FAILED.to_string(),
771                error: format!("unknown session: {session}"),
772                sqlstate: None,
773                message: None,
774                detail: None,
775                hint: None,
776                retryable: true,
777                trace: Trace::only_duration(0),
778            })
779            .await;
780        return;
781    };
782    let start = Instant::now();
783    match app
784        .executor
785        .tx_begin(&session, &session_cfg, read_only)
786        .await
787    {
788        Ok(()) => {
789            let _ = app
790                .writer
791                .send(Output::Result {
792                    id,
793                    session: Some(session),
794                    command_tag: crate::protocol::command_tag::BEGIN.to_string(),
795                    columns: vec![],
796                    rows: vec![],
797                    row_count: 0,
798                    truncated: false,
799                    truncated_at_rows: None,
800                    truncated_at_bytes: None,
801                    trace: Trace::only_duration(start.elapsed().as_millis() as u64),
802                })
803                .await;
804        }
805        Err(err) => emit_tx_error(app, id, &session, err, start).await,
806    }
807}
808
809async fn exec_tx_finish_on_worker(
810    app: &Arc<App>,
811    id: Option<String>,
812    session: String,
813    commit: bool,
814) {
815    let cfg = app.config.read().await.clone();
816    let Some(session_cfg) = cfg.sessions.get(&session).cloned() else {
817        let _ = app
818            .writer
819            .send(Output::Error {
820                id,
821                error_code: error_code::CONNECT_FAILED.to_string(),
822                error: format!("unknown session: {session}"),
823                sqlstate: None,
824                message: None,
825                detail: None,
826                hint: None,
827                retryable: true,
828                trace: Trace::only_duration(0),
829            })
830            .await;
831        return;
832    };
833    let start = Instant::now();
834    let result = if commit {
835        app.executor.tx_commit(&session, &session_cfg).await
836    } else {
837        app.executor.tx_rollback(&session, &session_cfg).await
838    };
839    match result {
840        Ok(()) => {
841            let _ = app
842                .writer
843                .send(Output::Result {
844                    id,
845                    session: Some(session),
846                    command_tag: if commit {
847                        crate::protocol::command_tag::COMMIT.to_string()
848                    } else {
849                        crate::protocol::command_tag::ROLLBACK.to_string()
850                    },
851                    columns: vec![],
852                    rows: vec![],
853                    row_count: 0,
854                    truncated: false,
855                    truncated_at_rows: None,
856                    truncated_at_bytes: None,
857                    trace: Trace::only_duration(start.elapsed().as_millis() as u64),
858                })
859                .await;
860        }
861        Err(err) => emit_tx_error(app, id, &session, err, start).await,
862    }
863}
864
865async fn cleanup_finished_queries(app: &Arc<App>) {
866    app.in_flight
867        .lock()
868        .await
869        .retain(|_, query| !query.state.is_finished());
870}
871
872async fn wait_for_workers_shutdown(runtime: &Arc<PipeRuntime>) {
873    let handles: Vec<tokio::task::JoinHandle<()>> = runtime
874        .workers
875        .lock()
876        .await
877        .drain()
878        .map(|(_, worker)| worker.handle)
879        .collect();
880    let deadline = Instant::now() + std::time::Duration::from_secs(5);
881    for handle in handles {
882        let now = Instant::now();
883        let remain = deadline.saturating_duration_since(now);
884        if tokio::time::timeout(remain, handle).await.is_err() {
885            // timeout waiting this task; move on
886        }
887    }
888}
889
890async fn send_close_event(app: &Arc<App>) {
891    let _ = app
892        .writer
893        .send(Output::Close {
894            message: "shutdown".to_string(),
895            trace: CloseTrace {
896                uptime_s: app.start_time.elapsed().as_secs(),
897                requests_total: app.requests_total.load(Ordering::Relaxed),
898            },
899        })
900        .await;
901}
902
903pub(crate) async fn read_limited_line<R>(
904    reader: &mut R,
905    max_bytes: usize,
906) -> std::io::Result<Option<Result<String, ()>>>
907where
908    R: AsyncBufRead + Unpin,
909{
910    let mut out = Vec::new();
911    let mut too_long = false;
912
913    loop {
914        let available = reader.fill_buf().await?;
915        if available.is_empty() {
916            if out.is_empty() && !too_long {
917                return Ok(None);
918            }
919            break;
920        }
921
922        let take = available
923            .iter()
924            .position(|b| *b == b'\n')
925            .map_or(available.len(), |pos| pos + 1);
926
927        if !too_long {
928            if out.len() + take <= max_bytes {
929                out.extend_from_slice(&available[..take]);
930            } else {
931                too_long = true;
932            }
933        }
934
935        let ended = available.get(take.saturating_sub(1)) == Some(&b'\n');
936        reader.consume(take);
937        if ended {
938            break;
939        }
940    }
941
942    if too_long {
943        return Ok(Some(Err(())));
944    }
945
946    Ok(Some(Ok(String::from_utf8_lossy(&out).to_string())))
947}
948
949pub(crate) fn validate_query_request(
950    id: &str,
951    sql: &str,
952    params: &[serde_json::Value],
953) -> Option<Output> {
954    if sql.len() > MAX_SQL_BYTES {
955        return Some(Output::Error {
956            id: Some(id.to_string()),
957            error_code: error_code::INVALID_REQUEST.to_string(),
958            error: "sql exceeds maximum size".to_string(),
959            sqlstate: None,
960            message: None,
961            detail: None,
962            hint: Some(format!("maximum SQL size is {MAX_SQL_BYTES} bytes")),
963            retryable: false,
964            trace: Trace::only_duration(0),
965        });
966    }
967    if params.len() > MAX_PARAMS {
968        return Some(Output::Error {
969            id: Some(id.to_string()),
970            error_code: error_code::INVALID_REQUEST.to_string(),
971            error: "too many params".to_string(),
972            sqlstate: None,
973            message: None,
974            detail: None,
975            hint: Some(format!("maximum params is {MAX_PARAMS}")),
976            retryable: false,
977            trace: Trace::only_duration(0),
978        });
979    }
980    None
981}
982
983async fn send_protocol_error(
984    app: &Arc<App>,
985    id: Option<String>,
986    error_code: &str,
987    error: &str,
988    hint: Option<&str>,
989    retryable: bool,
990) {
991    let _ = app
992        .writer
993        .send(Output::Error {
994            id,
995            error_code: error_code.to_string(),
996            error: error.to_string(),
997            sqlstate: None,
998            message: None,
999            detail: None,
1000            hint: hint.map(std::string::ToString::to_string),
1001            retryable,
1002            trace: Trace::only_duration(0),
1003        })
1004        .await;
1005}
1006
1007pub(crate) fn has_session_override(session: &SessionConfig) -> bool {
1008    session.dsn_secret.is_some()
1009        || session.conninfo_secret.is_some()
1010        || session.host.is_some()
1011        || session.port.is_some()
1012        || session.user.is_some()
1013        || session.dbname.is_some()
1014        || session.password_secret.is_some()
1015        || session.ssh.has_transport_fields()
1016        || session.container.has_transport_fields()
1017}
agent_first_psql/pipe.rs

agent_first_psql/
pipe.rs
