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