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::INVALID_REQUEST.to_string(),
135 error: format!("unknown session: {resolved_session}"),
136 sqlstate: None,
137 message: None,
138 detail: None,
139 hint: Some(
140 "list active sessions with a `config` request, or omit `session` to use the default"
141 .to_string(),
142 ),
143 retryable: false,
144 trace: trace.clone(),
145 })
146 .await;
147 emit_log(
148 app,
149 log_event::QUERY_ERROR,
150 id.as_deref(),
151 Some(&resolved_session),
152 Some(error_code::INVALID_REQUEST),
153 None,
154 &trace,
155 )
156 .await;
157 return;
158 };
159
160 let resolved_opts = match cfg.resolve_options_for_session(&options, &session_cfg) {
161 Ok(opts) => opts,
162 Err(message) => {
163 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
164 let readonly_write = app.capability == crate::Capability::ReadOnly
165 && options
166 .permission
167 .is_some_and(|value| !value.is_read_only());
168 let (message, hint) = if readonly_write {
169 (
170 "write permission is unavailable in afpsql-readonly".to_string(),
171 crate::readonly_hint().to_string(),
172 )
173 } else {
174 (message, permission_error_hint(&options, &session_cfg))
175 };
176 let _ = app
177 .writer
178 .send(Output::Error {
179 id: id.clone(),
180 error_code: error_code::INVALID_REQUEST.to_string(),
181 error: message,
182 sqlstate: None,
183 message: None,
184 detail: None,
185 hint: Some(hint),
186 retryable: false,
187 trace: trace.clone(),
188 })
189 .await;
190 emit_log(
191 app,
192 log_event::QUERY_ERROR,
193 id.as_deref(),
194 Some(&resolved_session),
195 Some(error_code::INVALID_REQUEST),
196 None,
197 &trace,
198 )
199 .await;
200 return;
201 }
202 };
203
204 let readonly_policy_error = if app.capability == crate::Capability::ReadOnly {
205 crate::readonly_policy::validate_session_with_trust(
206 &session_cfg,
207 app.locked_readonly_profile.load(Ordering::Relaxed),
208 )
209 .err()
210 .map(|error| (error, crate::readonly_local_capability_hint()))
211 .or_else(|| {
212 crate::readonly_policy::validate_sql(&sql)
213 .err()
214 .map(|error| (error, crate::readonly_hint()))
215 })
216 } else {
217 None
218 };
219 if readonly_policy_error.is_some()
220 || (app.capability == crate::Capability::ReadOnly && !resolved_opts.read_only)
221 {
222 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
223 let (error, hint) = readonly_policy_error.unwrap_or_else(|| {
224 (
225 "write permission is unavailable in afpsql-readonly".to_string(),
226 crate::readonly_hint(),
227 )
228 });
229 let _ = app
230 .writer
231 .send(Output::Error {
232 id: id.clone(),
233 error_code: error_code::INVALID_REQUEST.to_string(),
234 error,
235 sqlstate: None,
236 message: None,
237 detail: None,
238 hint: Some(hint.to_string()),
239 retryable: false,
240 trace: trace.clone(),
241 })
242 .await;
243 emit_log(
244 app,
245 log_event::QUERY_ERROR,
246 id.as_deref(),
247 Some(&resolved_session),
248 Some(error_code::INVALID_REQUEST),
249 None,
250 &trace,
251 )
252 .await;
253 return;
254 }
255
256 let cancel_slot_for_suppression = cancel_slot.clone();
257
258 if resolved_opts.stream_rows {
259 let mut sink = OutputRowSink::new(
260 app.clone(),
261 id.clone().unwrap_or_else(|| "cli".to_string()),
262 Some(resolved_session.clone()),
263 resolved_opts.batch_rows,
264 resolved_opts.batch_bytes,
265 );
266 let result = app
267 .executor
268 .execute_streaming(
269 ExecRequest {
270 session_name: &resolved_session,
271 session_cfg: &session_cfg,
272 sql: &sql,
273 params: ¶ms,
274 opts: &resolved_opts,
275 cancel_slot: cancel_slot.clone(),
276 transport_log: Some(TransportLogContext {
277 session: resolved_session.clone(),
278 log: cfg.log.clone(),
279 writer: app.writer.clone(),
280 }),
281 },
282 &mut sink,
283 )
284 .await;
285 if cancel_requested(&cancel_slot_for_suppression) {
286 return;
287 }
288 if !try_claim_terminal_emit(&cancel_slot_for_suppression) {
289 return;
290 }
291 handle_streaming_result(app, id, resolved_session, result, sink, start).await;
292 return;
293 }
294
295 let result = app
296 .executor
297 .execute(ExecRequest {
298 session_name: &resolved_session,
299 session_cfg: &session_cfg,
300 sql: &sql,
301 params: ¶ms,
302 opts: &resolved_opts,
303 cancel_slot: cancel_slot.clone(),
304 transport_log: Some(TransportLogContext {
305 session: resolved_session.clone(),
306 log: cfg.log.clone(),
307 writer: app.writer.clone(),
308 }),
309 })
310 .await;
311
312 if cancel_requested(&cancel_slot) {
313 return;
314 }
315 if !try_claim_terminal_emit(&cancel_slot) {
316 return;
317 }
318
319 match result {
320 Ok(ExecOutcome::Rows {
321 columns,
322 rows,
323 truncated,
324 truncated_at_rows,
325 truncated_at_bytes,
326 }) => {
327 let trace = emit_rows_result(
328 app,
329 id.clone(),
330 Some(resolved_session.clone()),
331 columns,
332 rows,
333 InlineTruncation {
334 truncated,
335 at_rows: truncated_at_rows,
336 at_bytes: truncated_at_bytes,
337 },
338 start,
339 )
340 .await;
341 emit_log(
342 app,
343 log_event::QUERY_RESULT,
344 id.as_deref(),
345 Some(&resolved_session),
346 None,
347 Some(command_tag::SELECT),
348 &trace,
349 )
350 .await;
351 }
352 Ok(ExecOutcome::Command { affected }) => {
353 emit_command_result(app, id, &resolved_session, affected, start).await;
354 }
355 Err(err) => emit_exec_error(app, id, &resolved_session, err, start).await,
356 }
357}
358
359fn permission_error_hint(options: &QueryOptions, session: &SessionConfig) -> String {
360 match (session.transport_kind(), options.permission) {
361 (Ok(TransportKind::Ssh), Some(permission)) if !permission.allows_ssh() => format!(
362 "this session uses afpsql SSH transport, so permission `{}` is invalid; use `ssh-read` for reads or `ssh-write` for writes",
363 permission.as_str()
364 ),
365 (Ok(TransportKind::Container), Some(permission)) if !permission.allows_container() => format!(
366 "this session uses afpsql container transport, so permission `{}` is invalid; use `container-read` for reads or `container-write` for writes",
367 permission.as_str()
368 ),
369 (Ok(TransportKind::Direct), Some(permission)) if permission.allows_ssh() => format!(
370 "this session does not use afpsql SSH transport, so permission `{}` is invalid; use `read` for reads or `write` for writes",
371 permission.as_str()
372 ),
373 (Ok(TransportKind::Direct), Some(permission)) if permission.allows_container() => format!(
374 "this session does not use afpsql container transport, so permission `{}` is invalid; use `read` for reads or `write` for writes",
375 permission.as_str()
376 ),
377 (Err(_), _) => {
380 "keep this session inside one container driver flag family; fields from two families name two drivers"
381 .to_string()
382 }
383 _ => {
384 "use read/write for direct connections, ssh-read/ssh-write for afpsql SSH transport, and container-read/container-write for afpsql container transport"
385 .to_string()
386 }
387 }
388}
389
390fn try_claim_terminal_emit(cancel_slot: &Option<crate::db::CancelSlot>) -> bool {
395 cancel_slot
396 .as_ref()
397 .map(|slot| slot.claim_terminal_emit())
398 .unwrap_or(true)
399}
400
401fn cancel_requested(cancel_slot: &Option<crate::db::CancelSlot>) -> bool {
402 cancel_slot
403 .as_ref()
404 .map(|slot| slot.is_cancelled())
405 .unwrap_or(false)
406}
407
408pub async fn handle_session_info(app: &Arc<App>, id: Option<String>, session: Option<String>) {
409 let start = Instant::now();
410 let cfg = app.config.read().await.clone();
411 let resolved_session = resolve_session_name(&cfg, session.as_deref());
412
413 let Some(session_cfg) = cfg.sessions.get(&resolved_session).cloned() else {
414 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
415 let _ = app
416 .writer
417 .send(Output::Error {
418 id: id.clone(),
419 error_code: error_code::INVALID_REQUEST.to_string(),
420 error: format!("unknown session: {resolved_session}"),
421 sqlstate: None,
422 message: None,
423 detail: None,
424 hint: Some(
425 "list active sessions with a `config` request, or pick the default session by omitting `session`"
426 .to_string(),
427 ),
428 retryable: false,
429 trace,
430 })
431 .await;
432 return;
433 };
434
435 let transport_kind = match session_cfg.transport_kind() {
436 Ok(kind) => kind,
437 Err(message) => {
438 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
439 let _ = app
440 .writer
441 .send(Output::Error {
442 id: id.clone(),
443 error_code: error_code::INVALID_REQUEST.to_string(),
444 error: message,
445 sqlstate: None,
446 message: None,
447 detail: None,
448 hint: Some(
449 "this session's transport fields are inconsistent; update the session via a `config` request before requesting `session_info`"
450 .to_string(),
451 ),
452 retryable: false,
453 trace,
454 })
455 .await;
456 return;
457 }
458 };
459
460 let permission_default = match transport_kind {
461 TransportKind::Direct => Permission::Read,
462 TransportKind::Ssh => Permission::SshRead,
463 TransportKind::Container => Permission::ContainerRead,
464 };
465
466 let resolved_opts = match cfg
467 .resolve_options_for_session(&QueryOptions::default(), &session_cfg)
468 {
469 Ok(opts) => opts,
470 Err(message) => {
471 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
472 let _ = app
473 .writer
474 .send(Output::Error {
475 id: id.clone(),
476 error_code: error_code::INVALID_REQUEST.to_string(),
477 error: message,
478 sqlstate: None,
479 message: None,
480 detail: None,
481 hint: Some(
482 "the runtime config could not resolve query defaults for this session; update inline_max_rows/inline_max_bytes via `config` and retry"
483 .to_string(),
484 ),
485 retryable: false,
486 trace,
487 })
488 .await;
489 return;
490 }
491 };
492
493 let (database, user, host, port, server_version) =
494 if app.executor.explicit_tx_open(&resolved_session).await {
495 (None, None, None, None, None)
496 } else {
497 probe_session_identity(app, &resolved_session, &session_cfg, &resolved_opts).await
498 };
499
500 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
501 let _ = app
502 .writer
503 .send(Output::SessionInfo {
504 id,
505 session: resolved_session,
506 transport_kind: transport_kind.as_str().to_string(),
507 permission_default: permission_default.as_str().to_string(),
508 stream_rows_default: resolved_opts.stream_rows,
509 batch_rows: resolved_opts.batch_rows,
510 batch_bytes: resolved_opts.batch_bytes,
511 inline_max_rows: resolved_opts.inline_max_rows,
512 inline_max_bytes: resolved_opts.inline_max_bytes,
513 statement_timeout_ms: resolved_opts.statement_timeout_ms,
514 lock_timeout_ms: resolved_opts.lock_timeout_ms,
515 database,
516 user,
517 host,
518 port,
519 server_version,
520 trace,
521 })
522 .await;
523}
524
525async fn probe_session_identity(
526 app: &Arc<App>,
527 session_name: &str,
528 session_cfg: &SessionConfig,
529 resolved_opts: &ResolvedOptions,
530) -> (
531 Option<String>,
532 Option<String>,
533 Option<String>,
534 Option<u16>,
535 Option<String>,
536) {
537 let probe = app
538 .executor
539 .execute(ExecRequest {
540 session_name,
541 session_cfg,
542 sql: "select current_database()::text as database, \
543 current_user::text as user, \
544 inet_server_addr()::text as host, \
545 inet_server_port() as port, \
546 current_setting('server_version') as server_version",
547 params: &[],
548 opts: resolved_opts,
549 cancel_slot: None,
550 transport_log: None,
551 })
552 .await;
553
554 if let Ok(ExecOutcome::Rows { rows, .. }) = probe
555 && let Some(row) = rows.first().and_then(|v| v.as_object())
556 {
557 let s = |key: &str| -> Option<String> {
558 row.get(key).and_then(|v| v.as_str().map(|s| s.to_string()))
559 };
560 let port = row
561 .get("port")
562 .and_then(|v| v.as_i64())
563 .and_then(|n| u16::try_from(n).ok())
564 .or_else(|| {
565 row.get("port")
566 .and_then(|v| v.as_str())
567 .and_then(|s| s.parse().ok())
568 });
569 return (
570 s("database").or_else(|| session_cfg.dbname.clone()),
571 s("user").or_else(|| session_cfg.user.clone()),
572 s("host").or_else(|| session_cfg.host.clone()),
573 port.or(session_cfg.port),
574 s("server_version"),
575 );
576 }
577
578 (
579 session_cfg.dbname.clone(),
580 session_cfg.user.clone(),
581 session_cfg.host.clone(),
582 session_cfg.port,
583 None,
584 )
585}
586
587async fn handle_streaming_result(
588 app: &Arc<App>,
589 id: Option<String>,
590 resolved_session: String,
591 result: Result<StreamOutcome, ExecError>,
592 mut sink: OutputRowSink,
593 start: Instant,
594) {
595 match result {
596 Ok(StreamOutcome::Rows {
597 row_count,
598 payload_bytes,
599 }) => {
600 let _ = sink.flush_batch().await;
601 let trace = Trace {
602 duration_ms: start.elapsed().as_millis() as u64,
603 row_count: Some(row_count),
604 payload_bytes: Some(payload_bytes),
605 };
606 let _ = app
607 .writer
608 .send(Output::ResultEnd {
609 id: sink.id.clone(),
610 session: Some(resolved_session.clone()),
611 command_tag: command_tag::rows(row_count),
612 trace: trace.clone(),
613 })
614 .await;
615 emit_log(
616 app,
617 log_event::QUERY_RESULT,
618 id.as_deref(),
619 Some(&resolved_session),
620 None,
621 Some(command_tag::SELECT),
622 &trace,
623 )
624 .await;
625 }
626 Ok(StreamOutcome::Command { affected }) => {
627 emit_command_result(app, id, &resolved_session, affected, start).await;
628 }
629 Err(err) => {
630 emit_exec_error(app, id, &resolved_session, err, start).await;
631 }
632 }
633}
634
635async fn emit_command_result(
636 app: &Arc<App>,
637 id: Option<String>,
638 resolved_session: &str,
639 affected: usize,
640 start: Instant,
641) {
642 let command_tag = command_tag::execute(affected);
643 let trace = Trace {
644 duration_ms: start.elapsed().as_millis() as u64,
645 row_count: Some(0),
646 payload_bytes: Some(0),
647 };
648 let _ = app
649 .writer
650 .send(Output::Result {
651 id: id.clone(),
652 session: Some(resolved_session.to_string()),
653 command_tag: command_tag.clone(),
654 columns: vec![],
655 rows: vec![],
656 row_count: 0,
657 truncated: false,
658 truncated_at_rows: None,
659 truncated_at_bytes: None,
660 trace: trace.clone(),
661 })
662 .await;
663 emit_log(
664 app,
665 log_event::QUERY_RESULT,
666 id.as_deref(),
667 Some(resolved_session),
668 None,
669 Some(command_tag::EXECUTE),
670 &trace,
671 )
672 .await;
673}
674
675pub(crate) async fn emit_exec_error(
676 app: &Arc<App>,
677 id: Option<String>,
678 resolved_session: &str,
679 err: ExecError,
680 start: Instant,
681) {
682 match err {
683 ExecError::Cancelled => {
684 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
685 let _ = app
686 .writer
687 .send(Output::Error {
688 id: id.clone(),
689 error_code: error_code::CANCELLED.to_string(),
690 error: "query cancelled".to_string(),
691 sqlstate: None,
692 message: None,
693 detail: None,
694 hint: Some(
695 "cancellation is final; submit a new query with a fresh id to retry"
696 .to_string(),
697 ),
698 retryable: false,
699 trace: trace.clone(),
700 })
701 .await;
702 emit_log(
703 app,
704 log_event::QUERY_ERROR,
705 id.as_deref(),
706 Some(resolved_session),
707 Some(error_code::CANCELLED),
708 None,
709 &trace,
710 )
711 .await;
712 }
713 ExecError::Connect(connect) => {
714 let connect = *connect;
715 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
716 let _ = app
717 .writer
718 .send(Output::Error {
719 id: id.clone(),
720 error_code: error_code::CONNECT_FAILED.to_string(),
721 error: connect.error,
722 sqlstate: connect.sqlstate,
723 message: connect.message,
724 detail: connect.detail,
725 hint: connect.hint,
726 retryable: connect.retryable,
727 trace: trace.clone(),
728 })
729 .await;
730 emit_log(
731 app,
732 log_event::QUERY_ERROR,
733 id.as_deref(),
734 Some(resolved_session),
735 Some(error_code::CONNECT_FAILED),
736 None,
737 &trace,
738 )
739 .await;
740 }
741 ExecError::Config { message, hint } | ExecError::InvalidRequest { message, hint } => {
742 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
743 let _ = app
744 .writer
745 .send(Output::Error {
746 id: id.clone(),
747 error_code: error_code::INVALID_REQUEST.to_string(),
748 error: message,
749 sqlstate: None,
750 message: None,
751 detail: None,
752 hint,
753 retryable: false,
754 trace: trace.clone(),
755 })
756 .await;
757 emit_log(
758 app,
759 log_event::QUERY_ERROR,
760 id.as_deref(),
761 Some(resolved_session),
762 Some(error_code::INVALID_REQUEST),
763 None,
764 &trace,
765 )
766 .await;
767 }
768 ExecError::InvalidParams(message) => {
769 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
770 let _ = app
771 .writer
772 .send(Output::Error {
773 id: id.clone(),
774 error_code: error_code::INVALID_PARAMS.to_string(),
775 error: message,
776 sqlstate: None,
777 message: None,
778 detail: None,
779 hint: Some(
780 "check that `params` count and types match the $1, $2, ... placeholders in `sql`"
781 .to_string(),
782 ),
783 retryable: false,
784 trace: trace.clone(),
785 })
786 .await;
787 emit_log(
788 app,
789 log_event::QUERY_ERROR,
790 id.as_deref(),
791 Some(resolved_session),
792 Some(error_code::INVALID_PARAMS),
793 None,
794 &trace,
795 )
796 .await;
797 }
798 ExecError::Sql {
799 sqlstate,
800 message,
801 detail,
802 hint,
803 position,
804 } => {
805 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
806 let _ = app
807 .writer
808 .send(Output::SqlError {
809 id: id.clone(),
810 session: Some(resolved_session.to_string()),
811 sqlstate: sqlstate.clone(),
812 message,
813 detail,
814 hint,
815 position,
816 retryable: crate::protocol::sqlstate_retryable(&sqlstate),
817 trace: trace.clone(),
818 })
819 .await;
820 emit_log(
821 app,
822 log_event::QUERY_SQL_ERROR,
823 id.as_deref(),
824 Some(resolved_session),
825 Some(&sqlstate),
826 None,
827 &trace,
828 )
829 .await;
830 }
831 ExecError::Internal(message) => {
832 let trace = Trace::only_duration(start.elapsed().as_millis() as u64);
833 let _ = app
834 .writer
835 .send(Output::Error {
836 id: id.clone(),
837 error_code: error_code::INVALID_REQUEST.to_string(),
838 error: message,
839 sqlstate: None,
840 message: None,
841 detail: None,
842 hint: Some(
843 "afpsql hit an internal error; retry the query, then restart the session if it persists"
844 .to_string(),
845 ),
846 retryable: false,
847 trace: trace.clone(),
848 })
849 .await;
850 emit_log(
851 app,
852 log_event::QUERY_ERROR,
853 id.as_deref(),
854 Some(resolved_session),
855 Some(error_code::INVALID_REQUEST),
856 None,
857 &trace,
858 )
859 .await;
860 }
861 }
862}
863
864struct OutputRowSink {
865 app: Arc<App>,
866 id: String,
867 session: Option<String>,
868 batch: Vec<Value>,
869 batch_bytes: usize,
870 batch_rows_limit: usize,
871 batch_bytes_limit: usize,
872}
873
874impl OutputRowSink {
875 fn new(
876 app: Arc<App>,
877 id: String,
878 session: Option<String>,
879 batch_rows_limit: usize,
880 batch_bytes_limit: usize,
881 ) -> Self {
882 Self {
883 app,
884 id,
885 session,
886 batch: vec![],
887 batch_bytes: 0,
888 batch_rows_limit,
889 batch_bytes_limit,
890 }
891 }
892
893 async fn flush_batch(&mut self) -> Result<(), ExecError> {
894 if self.batch.is_empty() {
895 return Ok(());
896 }
897 let n = self.batch.len();
898 let rows = std::mem::take(&mut self.batch);
899 self.batch_bytes = 0;
900 self.app
901 .writer
902 .send(Output::ResultRows {
903 id: self.id.clone(),
904 rows,
905 rows_batch_count: n,
906 })
907 .await
908 .map_err(|_| ExecError::Internal("output channel closed".to_string()))
909 }
910}
911
912#[async_trait::async_trait]
913impl RowSink for OutputRowSink {
914 async fn start(&mut self, columns: Vec<ColumnInfo>) -> Result<(), ExecError> {
915 self.app
916 .writer
917 .send(Output::ResultStart {
918 id: self.id.clone(),
919 session: self.session.clone(),
920 columns,
921 })
922 .await
923 .map_err(|_| ExecError::Internal("output channel closed".to_string()))
924 }
925
926 async fn row(&mut self, row: Value, row_bytes: usize) -> Result<(), ExecError> {
927 self.batch_bytes += row_bytes;
928 self.batch.push(row);
929 if self.batch.len() >= self.batch_rows_limit || self.batch_bytes >= self.batch_bytes_limit {
930 self.flush_batch().await?;
931 }
932 Ok(())
933 }
934}
935
936#[derive(Clone, Copy, Default)]
939pub(crate) struct InlineTruncation {
940 pub truncated: bool,
941 pub at_rows: Option<usize>,
942 pub at_bytes: Option<usize>,
943}
944
945#[allow(clippy::too_many_arguments)]
946async fn emit_rows_result(
947 app: &Arc<App>,
948 id: Option<String>,
949 session: Option<String>,
950 columns: Vec<ColumnInfo>,
951 rows: Vec<Value>,
952 truncation: InlineTruncation,
953 start: Instant,
954) -> Trace {
955 let mut payload_bytes = 0usize;
956 for row in &rows {
957 payload_bytes += serde_json::to_vec(row).map(|b| b.len()).unwrap_or(0);
958 }
959
960 let row_count = rows.len();
961 let trace = Trace {
962 duration_ms: start.elapsed().as_millis() as u64,
963 row_count: Some(row_count),
964 payload_bytes: Some(payload_bytes),
965 };
966 let _ = app
967 .writer
968 .send(Output::Result {
969 id,
970 session,
971 command_tag: command_tag::rows(row_count),
972 columns,
973 rows,
974 row_count,
975 truncated: truncation.truncated,
976 truncated_at_rows: truncation.at_rows,
977 truncated_at_bytes: truncation.at_bytes,
978 trace: trace.clone(),
979 })
980 .await;
981
982 trace
983}
984
985async fn emit_log(
986 app: &Arc<App>,
987 event: &str,
988 request_id: Option<&str>,
989 session: Option<&str>,
990 error_code: Option<&str>,
991 command_tag: Option<&str>,
992 trace: &Trace,
993) {
994 let enabled = {
995 let cfg = app.config.read().await;
996 cfg.log.enabled(event)
997 };
998 if !enabled {
999 return;
1000 }
1001
1002 let _ = app
1003 .writer
1004 .send(Output::Log {
1005 event: event.to_string(),
1006 request_id: request_id.map(std::string::ToString::to_string),
1007 session: session.map(std::string::ToString::to_string),
1008 error_code: error_code.map(std::string::ToString::to_string),
1009 command_tag: command_tag.map(std::string::ToString::to_string),
1010 version: None,
1011 config: None,
1012 args: None,
1013 env: None,
1014 chain: None,
1015 trace: trace.clone(),
1016 })
1017 .await;
1018}
1019
1020#[cfg(test)]
1021#[path = "../tests/support/unit_handler.rs"]
1022mod tests;