1use crate::conn::resolve_session_name;
2use crate::db::{
3 DbExecutor, ExecError, ExecOutcome, ExecRequest, PostgresExecutor, RowSink, StreamOutcome,
4 TransportLogContext,
5};
6use crate::protocol::{command_tag, error_code, log_event};
7use crate::types::*;
8use serde_json::Value;
9use std::sync::Arc;
10use std::sync::atomic::{AtomicU8, Ordering};
11use std::time::Instant;
12use tokio::sync::{Mutex, RwLock, mpsc};
13
14const QUERY_QUEUED: u8 = 0;
15const QUERY_RUNNING: u8 = 1;
16const QUERY_FINISHED: u8 = 2;
17const QUERY_CANCELLED: u8 = 3;
18
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub enum QueryPhase {
21 Queued,
22 Running,
23 Finished,
24 Cancelled,
25}
26
27pub struct QueryState {
28 phase: AtomicU8,
29}
30
31impl QueryState {
32 pub fn queued() -> Self {
33 Self {
34 phase: AtomicU8::new(QUERY_QUEUED),
35 }
36 }
37
38 pub fn phase(&self) -> QueryPhase {
39 match self.phase.load(Ordering::SeqCst) {
40 QUERY_RUNNING => QueryPhase::Running,
41 QUERY_FINISHED => QueryPhase::Finished,
42 QUERY_CANCELLED => QueryPhase::Cancelled,
43 _ => QueryPhase::Queued,
44 }
45 }
46
47 pub fn set_phase(&self, phase: QueryPhase) {
48 let value = match phase {
49 QueryPhase::Queued => QUERY_QUEUED,
50 QueryPhase::Running => QUERY_RUNNING,
51 QueryPhase::Finished => QUERY_FINISHED,
52 QueryPhase::Cancelled => QUERY_CANCELLED,
53 };
54 self.phase.store(value, Ordering::SeqCst);
55 }
56
57 pub fn try_start(&self) -> bool {
58 self.phase
59 .compare_exchange(
60 QUERY_QUEUED,
61 QUERY_RUNNING,
62 Ordering::SeqCst,
63 Ordering::SeqCst,
64 )
65 .is_ok()
66 }
67
68 pub fn is_finished(&self) -> bool {
69 matches!(self.phase(), QueryPhase::Finished | QueryPhase::Cancelled)
70 }
71}
72
73#[derive(Clone)]
74pub struct InFlightQuery {
75 pub cancel_slot: crate::db::CancelSlot,
76 pub state: Arc<QueryState>,
77}
78
79impl InFlightQuery {
80 pub async fn cancel_server_query(&self) -> Result<bool, String> {
81 crate::db::cancel_query(&self.cancel_slot).await
82 }
83}
84
85pub struct App {
86 pub capability: crate::Capability,
87 pub locked_readonly_profile: std::sync::atomic::AtomicBool,
88 pub config: RwLock<RuntimeConfig>,
89 pub executor: Arc<dyn DbExecutor>,
90 pub writer: mpsc::Sender<Output>,
91 pub in_flight: Mutex<std::collections::HashMap<String, InFlightQuery>>,
92 pub requests_total: std::sync::atomic::AtomicU64,
93 pub start_time: Instant,
94}
95
96impl App {
97 pub fn new(
98 config: RuntimeConfig,
99 writer: mpsc::Sender<Output>,
100 capability: crate::Capability,
101 ) -> Self {
102 Self {
103 capability,
104 locked_readonly_profile: std::sync::atomic::AtomicBool::new(false),
105 config: RwLock::new(config),
106 executor: Arc::new(PostgresExecutor::new()),
107 writer,
108 in_flight: Mutex::new(std::collections::HashMap::new()),
109 requests_total: std::sync::atomic::AtomicU64::new(0),
110 start_time: Instant::now(),
111 }
112 }
113}
114
115pub async fn execute_query(
116 app: &Arc<App>,
117 id: Option<String>,
118 session: Option<String>,
119 sql: String,
120 params: Vec<Value>,
121 options: QueryOptions,
122 cancel_slot: Option<crate::db::CancelSlot>,
123) {
124 let start = Instant::now();
125 let cfg = app.config.read().await.clone();
126 let resolved_session = resolve_session_name(&cfg, session.as_deref());
127
128 let Some(session_cfg) = cfg.sessions.get(&resolved_session).cloned() else {
129 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
130 let _ = app
131 .writer
132 .send(Output::Error {
133 id: id.clone(),
134 error_code: error_code::CONNECT_FAILED.to_string(),
135 error: format!("unknown session: {resolved_session}"),
136 sqlstate: None,
137 message: None,
138 detail: None,
139 hint: Some(
140 "check --host/--port or PGHOST/PGPORT environment variables".to_string(),
141 ),
142 retryable: true,
143 trace: trace.clone(),
144 })
145 .await;
146 emit_log(
147 app,
148 log_event::QUERY_ERROR,
149 id.as_deref(),
150 Some(&resolved_session),
151 Some(error_code::CONNECT_FAILED),
152 None,
153 &trace,
154 )
155 .await;
156 return;
157 };
158
159 let resolved_opts = match cfg.resolve_options_for_session(&options, &session_cfg) {
160 Ok(opts) => opts,
161 Err(message) => {
162 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
163 let readonly_write = app.capability == crate::Capability::ReadOnly
164 && options
165 .permission
166 .is_some_and(|value| !value.is_read_only());
167 let (message, hint) = if readonly_write {
168 (
169 "write permission is unavailable in afpsql-readonly".to_string(),
170 crate::readonly_hint().to_string(),
171 )
172 } else {
173 (message, permission_error_hint(&options, &session_cfg))
174 };
175 let _ = app
176 .writer
177 .send(Output::Error {
178 id: id.clone(),
179 error_code: error_code::INVALID_REQUEST.to_string(),
180 error: message,
181 sqlstate: None,
182 message: None,
183 detail: None,
184 hint: Some(hint),
185 retryable: false,
186 trace: trace.clone(),
187 })
188 .await;
189 emit_log(
190 app,
191 log_event::QUERY_ERROR,
192 id.as_deref(),
193 Some(&resolved_session),
194 Some(error_code::INVALID_REQUEST),
195 None,
196 &trace,
197 )
198 .await;
199 return;
200 }
201 };
202
203 let readonly_policy_error = if app.capability == crate::Capability::ReadOnly {
204 crate::readonly_policy::validate_session_with_trust(
205 &session_cfg,
206 app.locked_readonly_profile.load(Ordering::Relaxed),
207 )
208 .err()
209 .map(|error| (error, crate::readonly_local_capability_hint()))
210 .or_else(|| {
211 crate::readonly_policy::validate_sql(&sql)
212 .err()
213 .map(|error| (error, crate::readonly_hint()))
214 })
215 } else {
216 None
217 };
218 if readonly_policy_error.is_some()
219 || (app.capability == crate::Capability::ReadOnly && !resolved_opts.read_only)
220 {
221 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
222 let (error, hint) = readonly_policy_error.unwrap_or_else(|| {
223 (
224 "write permission is unavailable in afpsql-readonly".to_string(),
225 crate::readonly_hint(),
226 )
227 });
228 let _ = app
229 .writer
230 .send(Output::Error {
231 id: id.clone(),
232 error_code: error_code::INVALID_REQUEST.to_string(),
233 error,
234 sqlstate: None,
235 message: None,
236 detail: None,
237 hint: Some(hint.to_string()),
238 retryable: false,
239 trace: trace.clone(),
240 })
241 .await;
242 emit_log(
243 app,
244 log_event::QUERY_ERROR,
245 id.as_deref(),
246 Some(&resolved_session),
247 Some(error_code::INVALID_REQUEST),
248 None,
249 &trace,
250 )
251 .await;
252 return;
253 }
254
255 let cancel_slot_for_suppression = cancel_slot.clone();
256
257 if resolved_opts.stream_rows {
258 let mut sink = OutputRowSink::new(
259 app.clone(),
260 id.clone().unwrap_or_else(|| "cli".to_string()),
261 Some(resolved_session.clone()),
262 resolved_opts.batch_rows,
263 resolved_opts.batch_bytes,
264 );
265 let result = app
266 .executor
267 .execute_streaming(
268 ExecRequest {
269 session_name: &resolved_session,
270 session_cfg: &session_cfg,
271 sql: &sql,
272 params: ¶ms,
273 opts: &resolved_opts,
274 cancel_slot: cancel_slot.clone(),
275 transport_log: Some(TransportLogContext {
276 session: resolved_session.clone(),
277 log: cfg.log.clone(),
278 writer: app.writer.clone(),
279 }),
280 },
281 &mut sink,
282 )
283 .await;
284 if cancel_requested(&cancel_slot_for_suppression) {
285 return;
286 }
287 if !try_claim_terminal_emit(&cancel_slot_for_suppression) {
288 return;
289 }
290 handle_streaming_result(app, id, resolved_session, result, sink, start).await;
291 return;
292 }
293
294 let result = app
295 .executor
296 .execute(ExecRequest {
297 session_name: &resolved_session,
298 session_cfg: &session_cfg,
299 sql: &sql,
300 params: ¶ms,
301 opts: &resolved_opts,
302 cancel_slot: cancel_slot.clone(),
303 transport_log: Some(TransportLogContext {
304 session: resolved_session.clone(),
305 log: cfg.log.clone(),
306 writer: app.writer.clone(),
307 }),
308 })
309 .await;
310
311 if cancel_requested(&cancel_slot) {
312 return;
313 }
314 if !try_claim_terminal_emit(&cancel_slot) {
315 return;
316 }
317
318 match result {
319 Ok(ExecOutcome::Rows {
320 columns,
321 rows,
322 truncated,
323 truncated_at_rows,
324 truncated_at_bytes,
325 }) => {
326 let status = emit_rows_result(
327 app,
328 id.clone(),
329 Some(resolved_session.clone()),
330 columns,
331 rows,
332 InlineTruncation {
333 truncated,
334 at_rows: truncated_at_rows,
335 at_bytes: truncated_at_bytes,
336 },
337 start,
338 &resolved_opts,
339 )
340 .await;
341 let RowEmitStatus::Sent { trace } = status;
342 emit_log(
343 app,
344 log_event::QUERY_RESULT,
345 id.as_deref(),
346 Some(&resolved_session),
347 None,
348 Some(command_tag::SELECT),
349 &trace,
350 )
351 .await;
352 }
353 Ok(ExecOutcome::Command { affected }) => {
354 emit_command_result(app, id, &resolved_session, affected, start).await;
355 }
356 Err(err) => emit_exec_error(app, id, &resolved_session, err, start).await,
357 }
358}
359
360fn permission_error_hint(options: &QueryOptions, session: &SessionConfig) -> String {
361 match (session.transport_kind(), options.permission) {
362 (Ok(TransportKind::Ssh), Some(permission)) if !permission.allows_ssh() => format!(
363 "this session uses afpsql SSH transport, so permission `{}` is invalid; use `ssh-read` for reads or `ssh-write` for writes",
364 permission.as_str()
365 ),
366 (Ok(TransportKind::Container), Some(permission)) if !permission.allows_container() => format!(
367 "this session uses afpsql container transport, so permission `{}` is invalid; use `container-read` for reads or `container-write` for writes",
368 permission.as_str()
369 ),
370 (Ok(TransportKind::Direct), Some(permission)) if permission.allows_ssh() => format!(
371 "this session does not use afpsql SSH transport, so permission `{}` is invalid; use `read` for reads or `write` for writes",
372 permission.as_str()
373 ),
374 (Ok(TransportKind::Direct), Some(permission)) if permission.allows_container() => format!(
375 "this session does not use afpsql container transport, so permission `{}` is invalid; use `read` for reads or `write` for writes",
376 permission.as_str()
377 ),
378 (Err(message), _) => message,
379 _ => {
380 "use read/write for direct connections, ssh-read/ssh-write for afpsql SSH transport, and container-read/container-write for afpsql container transport"
381 .to_string()
382 }
383 }
384}
385
386fn try_claim_terminal_emit(cancel_slot: &Option<crate::db::CancelSlot>) -> bool {
391 cancel_slot
392 .as_ref()
393 .map(|slot| slot.claim_terminal_emit())
394 .unwrap_or(true)
395}
396
397fn cancel_requested(cancel_slot: &Option<crate::db::CancelSlot>) -> bool {
398 cancel_slot
399 .as_ref()
400 .map(|slot| slot.is_cancelled())
401 .unwrap_or(false)
402}
403
404pub async fn handle_session_info(app: &Arc<App>, id: Option<String>, session: Option<String>) {
405 let start = Instant::now();
406 let cfg = app.config.read().await.clone();
407 let resolved_session = resolve_session_name(&cfg, session.as_deref());
408
409 let Some(session_cfg) = cfg.sessions.get(&resolved_session).cloned() else {
410 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
411 let _ = app
412 .writer
413 .send(Output::Error {
414 id: id.clone(),
415 error_code: error_code::INVALID_REQUEST.to_string(),
416 error: format!("unknown session: {resolved_session}"),
417 sqlstate: None,
418 message: None,
419 detail: None,
420 hint: Some(
421 "list active sessions with a `config` request, or pick the default session by omitting `session`"
422 .to_string(),
423 ),
424 retryable: false,
425 trace,
426 })
427 .await;
428 return;
429 };
430
431 let transport_kind = match session_cfg.transport_kind() {
432 Ok(kind) => kind,
433 Err(message) => {
434 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
435 let _ = app
436 .writer
437 .send(Output::Error {
438 id: id.clone(),
439 error_code: error_code::INVALID_REQUEST.to_string(),
440 error: message,
441 sqlstate: None,
442 message: None,
443 detail: None,
444 hint: Some(
445 "this session's transport flags are inconsistent; update the session via a `config` request before requesting `session_info`"
446 .to_string(),
447 ),
448 retryable: false,
449 trace,
450 })
451 .await;
452 return;
453 }
454 };
455
456 let permission_default = match transport_kind {
457 TransportKind::Direct => Permission::Read,
458 TransportKind::Ssh => Permission::SshRead,
459 TransportKind::Container => Permission::ContainerRead,
460 };
461
462 let resolved_opts = match cfg
463 .resolve_options_for_session(&QueryOptions::default(), &session_cfg)
464 {
465 Ok(opts) => opts,
466 Err(message) => {
467 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
468 let _ = app
469 .writer
470 .send(Output::Error {
471 id: id.clone(),
472 error_code: error_code::INVALID_REQUEST.to_string(),
473 error: message,
474 sqlstate: None,
475 message: None,
476 detail: None,
477 hint: Some(
478 "the runtime config could not resolve query defaults for this session; update inline_max_rows/inline_max_bytes via `config` and retry"
479 .to_string(),
480 ),
481 retryable: false,
482 trace,
483 })
484 .await;
485 return;
486 }
487 };
488
489 let (database, user, host, port, server_version) =
490 probe_session_identity(app, &resolved_session, &session_cfg, &resolved_opts).await;
491
492 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
493 let _ = app
494 .writer
495 .send(Output::SessionInfo {
496 id,
497 session: resolved_session,
498 transport_kind: transport_kind.as_str().to_string(),
499 permission_default: permission_default.as_str().to_string(),
500 stream_rows_default: resolved_opts.stream_rows,
501 batch_rows: resolved_opts.batch_rows,
502 batch_bytes: resolved_opts.batch_bytes,
503 inline_max_rows: resolved_opts.inline_max_rows,
504 inline_max_bytes: resolved_opts.inline_max_bytes,
505 statement_timeout_ms: resolved_opts.statement_timeout_ms,
506 lock_timeout_ms: resolved_opts.lock_timeout_ms,
507 database,
508 user,
509 host,
510 port,
511 server_version,
512 trace,
513 })
514 .await;
515}
516
517async fn probe_session_identity(
518 app: &Arc<App>,
519 session_name: &str,
520 session_cfg: &SessionConfig,
521 resolved_opts: &ResolvedOptions,
522) -> (
523 Option<String>,
524 Option<String>,
525 Option<String>,
526 Option<u16>,
527 Option<String>,
528) {
529 let probe = app
530 .executor
531 .execute(ExecRequest {
532 session_name,
533 session_cfg,
534 sql: "select current_database()::text as database, \
535 current_user::text as user, \
536 inet_server_addr()::text as host, \
537 inet_server_port() as port, \
538 current_setting('server_version') as server_version",
539 params: &[],
540 opts: resolved_opts,
541 cancel_slot: None,
542 transport_log: None,
543 })
544 .await;
545
546 if let Ok(ExecOutcome::Rows { rows, .. }) = probe
547 && let Some(row) = rows.first().and_then(|v| v.as_object())
548 {
549 let s = |key: &str| -> Option<String> {
550 row.get(key).and_then(|v| v.as_str().map(|s| s.to_string()))
551 };
552 let port = row
553 .get("port")
554 .and_then(|v| v.as_i64())
555 .and_then(|n| u16::try_from(n).ok())
556 .or_else(|| {
557 row.get("port")
558 .and_then(|v| v.as_str())
559 .and_then(|s| s.parse().ok())
560 });
561 return (
562 s("database").or_else(|| session_cfg.dbname.clone()),
563 s("user").or_else(|| session_cfg.user.clone()),
564 s("host").or_else(|| session_cfg.host.clone()),
565 port.or(session_cfg.port),
566 s("server_version"),
567 );
568 }
569
570 (
571 session_cfg.dbname.clone(),
572 session_cfg.user.clone(),
573 session_cfg.host.clone(),
574 session_cfg.port,
575 None,
576 )
577}
578
579async fn handle_streaming_result(
580 app: &Arc<App>,
581 id: Option<String>,
582 resolved_session: String,
583 result: Result<StreamOutcome, ExecError>,
584 mut sink: OutputRowSink,
585 start: Instant,
586) {
587 match result {
588 Ok(StreamOutcome::Rows {
589 row_count,
590 payload_bytes,
591 }) => {
592 let _ = sink.flush_batch().await;
593 let trace = Trace {
594 duration_ms: start.elapsed().as_millis() as u64,
595 row_count: Some(row_count),
596 payload_bytes: Some(payload_bytes),
597 };
598 let _ = app
599 .writer
600 .send(Output::ResultEnd {
601 id: sink.id.clone(),
602 session: Some(resolved_session.clone()),
603 command_tag: command_tag::rows(row_count),
604 trace: trace.clone(),
605 })
606 .await;
607 emit_log(
608 app,
609 log_event::QUERY_RESULT,
610 id.as_deref(),
611 Some(&resolved_session),
612 None,
613 Some(command_tag::SELECT),
614 &trace,
615 )
616 .await;
617 }
618 Ok(StreamOutcome::Command { affected }) => {
619 emit_command_result(app, id, &resolved_session, affected, start).await;
620 }
621 Err(err) => {
622 emit_exec_error(app, id, &resolved_session, err, start).await;
623 }
624 }
625}
626
627async fn emit_command_result(
628 app: &Arc<App>,
629 id: Option<String>,
630 resolved_session: &str,
631 affected: usize,
632 start: Instant,
633) {
634 let command_tag = command_tag::execute(affected);
635 let trace = Trace {
636 duration_ms: start.elapsed().as_millis() as u64,
637 row_count: Some(0),
638 payload_bytes: Some(0),
639 };
640 let _ = app
641 .writer
642 .send(Output::Result {
643 id: id.clone(),
644 session: Some(resolved_session.to_string()),
645 command_tag: command_tag.clone(),
646 columns: vec![],
647 rows: vec![],
648 row_count: 0,
649 truncated: false,
650 truncated_at_rows: None,
651 truncated_at_bytes: None,
652 trace: trace.clone(),
653 })
654 .await;
655 emit_log(
656 app,
657 log_event::QUERY_RESULT,
658 id.as_deref(),
659 Some(resolved_session),
660 None,
661 Some(command_tag::EXECUTE),
662 &trace,
663 )
664 .await;
665}
666
667pub(crate) async fn emit_exec_error(
668 app: &Arc<App>,
669 id: Option<String>,
670 resolved_session: &str,
671 err: ExecError,
672 start: Instant,
673) {
674 match err {
675 ExecError::Cancelled => {
676 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
677 let _ = app
678 .writer
679 .send(Output::Error {
680 id: id.clone(),
681 error_code: error_code::CANCELLED.to_string(),
682 error: "query cancelled".to_string(),
683 sqlstate: None,
684 message: None,
685 detail: None,
686 hint: Some(
687 "cancellation is final; submit a new query with a fresh id to retry"
688 .to_string(),
689 ),
690 retryable: false,
691 trace: trace.clone(),
692 })
693 .await;
694 emit_log(
695 app,
696 log_event::QUERY_ERROR,
697 id.as_deref(),
698 Some(resolved_session),
699 Some(error_code::CANCELLED),
700 None,
701 &trace,
702 )
703 .await;
704 }
705 ExecError::Connect(connect) => {
706 let connect = *connect;
707 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
708 let _ = app
709 .writer
710 .send(Output::Error {
711 id: id.clone(),
712 error_code: error_code::CONNECT_FAILED.to_string(),
713 error: connect.error,
714 sqlstate: connect.sqlstate,
715 message: connect.message,
716 detail: connect.detail,
717 hint: connect.hint,
718 retryable: connect.retryable,
719 trace: trace.clone(),
720 })
721 .await;
722 emit_log(
723 app,
724 log_event::QUERY_ERROR,
725 id.as_deref(),
726 Some(resolved_session),
727 Some(error_code::CONNECT_FAILED),
728 None,
729 &trace,
730 )
731 .await;
732 }
733 ExecError::Config { message, hint } => {
734 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
735 let _ = app
736 .writer
737 .send(Output::Error {
738 id: id.clone(),
739 error_code: error_code::INVALID_REQUEST.to_string(),
740 error: message,
741 sqlstate: None,
742 message: None,
743 detail: None,
744 hint,
745 retryable: false,
746 trace: trace.clone(),
747 })
748 .await;
749 emit_log(
750 app,
751 log_event::QUERY_ERROR,
752 id.as_deref(),
753 Some(resolved_session),
754 Some(error_code::INVALID_REQUEST),
755 None,
756 &trace,
757 )
758 .await;
759 }
760 ExecError::InvalidParams(message) => {
761 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
762 let _ = app
763 .writer
764 .send(Output::Error {
765 id: id.clone(),
766 error_code: error_code::INVALID_PARAMS.to_string(),
767 error: message,
768 sqlstate: None,
769 message: None,
770 detail: None,
771 hint: Some(
772 "check that `params` count and types match the $1, $2, ... placeholders in `sql`"
773 .to_string(),
774 ),
775 retryable: false,
776 trace: trace.clone(),
777 })
778 .await;
779 emit_log(
780 app,
781 log_event::QUERY_ERROR,
782 id.as_deref(),
783 Some(resolved_session),
784 Some(error_code::INVALID_PARAMS),
785 None,
786 &trace,
787 )
788 .await;
789 }
790 ExecError::ResultTooLarge {
791 row_count,
792 payload_bytes,
793 } => {
794 let trace = Trace {
795 duration_ms: start.elapsed().as_millis() as u64,
796 row_count: Some(row_count),
797 payload_bytes: Some(payload_bytes),
798 };
799 let _ = app
800 .writer
801 .send(Output::Error {
802 id: id.clone(),
803 error_code: error_code::RESULT_TOO_LARGE.to_string(),
804 error: "result exceeds inline limits".to_string(),
805 sqlstate: None,
806 message: None,
807 detail: None,
808 hint: Some(result_too_large_hint()),
809 retryable: false,
810 trace: trace.clone(),
811 })
812 .await;
813 emit_log(
814 app,
815 log_event::QUERY_ERROR,
816 id.as_deref(),
817 Some(resolved_session),
818 Some(error_code::RESULT_TOO_LARGE),
819 None,
820 &trace,
821 )
822 .await;
823 }
824 ExecError::Sql {
825 sqlstate,
826 message,
827 detail,
828 hint,
829 position,
830 } => {
831 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
832 let _ = app
833 .writer
834 .send(Output::SqlError {
835 id: id.clone(),
836 session: Some(resolved_session.to_string()),
837 sqlstate: sqlstate.clone(),
838 message,
839 detail,
840 hint,
841 position,
842 trace: trace.clone(),
843 })
844 .await;
845 emit_log(
846 app,
847 log_event::QUERY_SQL_ERROR,
848 id.as_deref(),
849 Some(resolved_session),
850 Some(&sqlstate),
851 None,
852 &trace,
853 )
854 .await;
855 }
856 ExecError::Internal(message) => {
857 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
858 let _ = app
859 .writer
860 .send(Output::Error {
861 id: id.clone(),
862 error_code: error_code::INVALID_REQUEST.to_string(),
863 error: message,
864 sqlstate: None,
865 message: None,
866 detail: None,
867 hint: Some(
868 "afpsql hit an internal error; retry the query, then restart the session if it persists"
869 .to_string(),
870 ),
871 retryable: false,
872 trace: trace.clone(),
873 })
874 .await;
875 emit_log(
876 app,
877 log_event::QUERY_ERROR,
878 id.as_deref(),
879 Some(resolved_session),
880 Some(error_code::INVALID_REQUEST),
881 None,
882 &trace,
883 )
884 .await;
885 }
886 }
887}
888
889struct OutputRowSink {
890 app: Arc<App>,
891 id: String,
892 session: Option<String>,
893 batch: Vec<Value>,
894 batch_bytes: usize,
895 batch_rows_limit: usize,
896 batch_bytes_limit: usize,
897}
898
899impl OutputRowSink {
900 fn new(
901 app: Arc<App>,
902 id: String,
903 session: Option<String>,
904 batch_rows_limit: usize,
905 batch_bytes_limit: usize,
906 ) -> Self {
907 Self {
908 app,
909 id,
910 session,
911 batch: vec![],
912 batch_bytes: 0,
913 batch_rows_limit,
914 batch_bytes_limit,
915 }
916 }
917
918 async fn flush_batch(&mut self) -> Result<(), ExecError> {
919 if self.batch.is_empty() {
920 return Ok(());
921 }
922 let n = self.batch.len();
923 let rows = std::mem::take(&mut self.batch);
924 self.batch_bytes = 0;
925 self.app
926 .writer
927 .send(Output::ResultRows {
928 id: self.id.clone(),
929 rows,
930 rows_batch_count: n,
931 })
932 .await
933 .map_err(|_| ExecError::Internal("output channel closed".to_string()))
934 }
935}
936
937#[async_trait::async_trait]
938impl RowSink for OutputRowSink {
939 async fn start(&mut self, columns: Vec<ColumnInfo>) -> Result<(), ExecError> {
940 self.app
941 .writer
942 .send(Output::ResultStart {
943 id: self.id.clone(),
944 session: self.session.clone(),
945 columns,
946 })
947 .await
948 .map_err(|_| ExecError::Internal("output channel closed".to_string()))
949 }
950
951 async fn row(&mut self, row: Value, row_bytes: usize) -> Result<(), ExecError> {
952 self.batch_bytes += row_bytes;
953 self.batch.push(row);
954 if self.batch.len() >= self.batch_rows_limit || self.batch_bytes >= self.batch_bytes_limit {
955 self.flush_batch().await?;
956 }
957 Ok(())
958 }
959}
960
961#[derive(Clone)]
962enum RowEmitStatus {
963 Sent { trace: Trace },
964}
965
966#[derive(Clone, Copy, Default)]
969pub(crate) struct InlineTruncation {
970 pub truncated: bool,
971 pub at_rows: Option<usize>,
972 pub at_bytes: Option<usize>,
973}
974
975#[allow(clippy::too_many_arguments)]
976async fn emit_rows_result(
977 app: &Arc<App>,
978 id: Option<String>,
979 session: Option<String>,
980 columns: Vec<ColumnInfo>,
981 rows: Vec<Value>,
982 truncation: InlineTruncation,
983 start: Instant,
984 opts: &ResolvedOptions,
985) -> RowEmitStatus {
986 if opts.stream_rows {
987 let req_id = id.clone().unwrap_or_else(|| "cli".to_string());
988 let _ = app
989 .writer
990 .send(Output::ResultStart {
991 id: req_id.clone(),
992 session: session.clone(),
993 columns: columns.clone(),
994 })
995 .await;
996
997 let mut batch: Vec<Value> = vec![];
998 let mut batch_bytes = 0usize;
999 let mut total_bytes = 0usize;
1000 let mut row_count = 0usize;
1001
1002 for row in rows {
1003 let sz = serde_json::to_vec(&row).map(|b| b.len()).unwrap_or(0);
1004 batch_bytes += sz;
1005 total_bytes += sz;
1006 row_count += 1;
1007 batch.push(row);
1008
1009 if batch.len() >= opts.batch_rows || batch_bytes >= opts.batch_bytes {
1010 let n = batch.len();
1011 let _ = app
1012 .writer
1013 .send(Output::ResultRows {
1014 id: req_id.clone(),
1015 rows: std::mem::take(&mut batch),
1016 rows_batch_count: n,
1017 })
1018 .await;
1019 batch_bytes = 0;
1020 }
1021 }
1022
1023 for tail in std::iter::once(batch).filter(|r| !r.is_empty()) {
1024 let n = tail.len();
1025 let _ = app
1026 .writer
1027 .send(Output::ResultRows {
1028 id: req_id.clone(),
1029 rows: tail,
1030 rows_batch_count: n,
1031 })
1032 .await;
1033 }
1034
1035 let trace = Trace {
1036 duration_ms: start.elapsed().as_millis() as u64,
1037 row_count: Some(row_count),
1038 payload_bytes: Some(total_bytes),
1039 };
1040 let _ = app
1041 .writer
1042 .send(Output::ResultEnd {
1043 id: req_id,
1044 session,
1045 command_tag: command_tag::rows(row_count),
1046 trace: trace.clone(),
1047 })
1048 .await;
1049
1050 return RowEmitStatus::Sent { trace };
1051 }
1052
1053 let mut payload_bytes = 0usize;
1054 for row in &rows {
1055 payload_bytes += serde_json::to_vec(row).map(|b| b.len()).unwrap_or(0);
1056 }
1057
1058 let row_count = rows.len();
1059 let trace = Trace {
1060 duration_ms: start.elapsed().as_millis() as u64,
1061 row_count: Some(row_count),
1062 payload_bytes: Some(payload_bytes),
1063 };
1064 let _ = app
1065 .writer
1066 .send(Output::Result {
1067 id,
1068 session,
1069 command_tag: command_tag::rows(row_count),
1070 columns,
1071 rows,
1072 row_count,
1073 truncated: truncation.truncated,
1074 truncated_at_rows: truncation.at_rows,
1075 truncated_at_bytes: truncation.at_bytes,
1076 trace: trace.clone(),
1077 })
1078 .await;
1079
1080 RowEmitStatus::Sent { trace }
1081}
1082
1083fn result_too_large_hint() -> String {
1084 "retry with stream_rows=true, or increase --inline-max-rows/--inline-max-bytes".to_string()
1085}
1086
1087async fn emit_log(
1088 app: &Arc<App>,
1089 event: &str,
1090 request_id: Option<&str>,
1091 session: Option<&str>,
1092 error_code: Option<&str>,
1093 command_tag: Option<&str>,
1094 trace: &Trace,
1095) {
1096 let enabled = {
1097 let cfg = app.config.read().await;
1098 cfg.log.enabled(event)
1099 };
1100 if !enabled {
1101 return;
1102 }
1103
1104 let _ = app
1105 .writer
1106 .send(Output::Log {
1107 event: event.to_string(),
1108 request_id: request_id.map(std::string::ToString::to_string),
1109 session: session.map(std::string::ToString::to_string),
1110 error_code: error_code.map(std::string::ToString::to_string),
1111 command_tag: command_tag.map(std::string::ToString::to_string),
1112 version: None,
1113 config: None,
1114 args: None,
1115 env: None,
1116 chain: None,
1117 trace: trace.clone(),
1118 })
1119 .await;
1120}
1121
1122#[cfg(test)]
1123#[path = "../tests/support/unit_handler.rs"]
1124mod tests;