1use std::collections::HashMap;
9use std::sync::Arc;
10
11use tokio::io::{AsyncRead, AsyncWrite};
12use tokio::net::TcpListener;
13
14use super::catalog_views::translate_pg_catalog_query;
15use super::protocol::{
16 read_frame, read_startup, write_frame, write_raw_byte, BackendMessage, ColumnDescriptor,
17 DescribeTarget, FrontendMessage, PgWireError, TransactionStatus,
18};
19use super::types::{pg_param_to_param_value, value_to_pg_wire_bytes, PgOid};
20use crate::auth::Role;
21use crate::runtime::query_request::{
22 ParamValue, PreparedId, PreparedRegistry, QueryRequest, QueryRequestExecutor,
23};
24use crate::runtime::RedDBRuntime;
25use crate::storage::query::unified::{UnifiedRecord, UnifiedResult};
26use reddb_rql::ast::QueryExpr;
27use reddb_types::Value;
28
29#[derive(Debug, Clone)]
31pub struct PgWireConfig {
32 pub bind_addr: String,
35 pub server_version: String,
39 pub tls: Option<crate::wire::tls::WireTlsConfig>,
48}
49
50#[derive(Debug, Clone)]
57struct PgAuthContext {
58 username: String,
59 role: Role,
60 tenant: Option<String>,
61}
62
63#[derive(Debug, Clone)]
64struct PgPreparedStatement {
65 sql: String,
66 prepared_id: PreparedId,
67 param_type_oids: Vec<u32>,
68}
69
70#[derive(Debug, Clone)]
71struct PgPortal {
72 sql: String,
73 prepared_id: PreparedId,
74 params: Vec<ParamValue>,
75 #[allow(dead_code)]
76 result_format_codes: Vec<i16>,
77 row_description_sent: bool,
78 result: Option<crate::runtime::OwnedQueryResult>,
79 row_offset: usize,
80}
81
82impl Default for PgWireConfig {
83 fn default() -> Self {
84 Self {
85 bind_addr: "127.0.0.1:5432".to_string(),
86 server_version: "15.0 (RedDB 3.1)".to_string(),
87 tls: None,
88 }
89 }
90}
91
92fn run_runtime_blocking<T>(f: impl FnOnce() -> T) -> T {
112 use tokio::runtime::{Handle, RuntimeFlavor};
113 match Handle::try_current().map(|h| h.runtime_flavor()) {
114 Ok(RuntimeFlavor::MultiThread) => tokio::task::block_in_place(f),
115 _ => f(),
116 }
117}
118
119fn pg_frame_requires_execution(frame: &FrontendMessage) -> bool {
120 matches!(
121 frame,
122 FrontendMessage::Query(_)
123 | FrontendMessage::Parse(_)
124 | FrontendMessage::Describe(_)
125 | FrontendMessage::Execute(_)
126 )
127}
128
129async fn acquire_pg_execution(
130 runtime: &RedDBRuntime,
131) -> Result<crate::transport::execution_admission::ExecutionPermit, &'static str> {
132 use crate::transport::execution_admission::ExecutionAdmissionError;
133
134 let admission = runtime.execution_admission();
135 match tokio::time::timeout(runtime.execution_queue_timeout(), admission.acquire()).await {
136 Ok(Ok(permit)) => Ok(permit),
137 Ok(Err(ExecutionAdmissionError::QueueFull)) => {
138 Err("execution_queue_full: engine execution queue is full; retryable=true")
139 }
140 Ok(Err(ExecutionAdmissionError::Closed)) => {
141 Err("execution_admission_closed: engine execution admission is unavailable")
142 }
143 Err(_) => Err(
144 "execution_queue_timeout: deadline expired while waiting to execute; retryable=true",
145 ),
146 }
147}
148
149struct ExecutionResponseWriter<'a, S> {
154 inner: &'a mut S,
155 permit: Option<crate::transport::execution_admission::ExecutionPermit>,
156}
157
158impl<'a, S> ExecutionResponseWriter<'a, S> {
159 fn new(
160 inner: &'a mut S,
161 permit: Option<crate::transport::execution_admission::ExecutionPermit>,
162 ) -> Self {
163 Self { inner, permit }
164 }
165}
166
167impl<S: AsyncWrite + Unpin> AsyncWrite for ExecutionResponseWriter<'_, S> {
168 fn poll_write(
169 mut self: std::pin::Pin<&mut Self>,
170 cx: &mut std::task::Context<'_>,
171 buf: &[u8],
172 ) -> std::task::Poll<Result<usize, std::io::Error>> {
173 drop(self.permit.take());
174 std::pin::Pin::new(&mut *self.inner).poll_write(cx, buf)
175 }
176
177 fn poll_flush(
178 mut self: std::pin::Pin<&mut Self>,
179 cx: &mut std::task::Context<'_>,
180 ) -> std::task::Poll<Result<(), std::io::Error>> {
181 drop(self.permit.take());
182 std::pin::Pin::new(&mut *self.inner).poll_flush(cx)
183 }
184
185 fn poll_shutdown(
186 mut self: std::pin::Pin<&mut Self>,
187 cx: &mut std::task::Context<'_>,
188 ) -> std::task::Poll<Result<(), std::io::Error>> {
189 drop(self.permit.take());
190 std::pin::Pin::new(&mut *self.inner).poll_shutdown(cx)
191 }
192}
193
194fn execute_with_pg_auth_context<T>(
202 auth_context: Option<&PgAuthContext>,
203 f: impl FnOnce() -> T,
204) -> T {
205 let Some(auth_context) = auth_context else {
206 return f();
207 };
208 let _guard = PgAuthContextGuard::install(auth_context);
209 f()
210}
211
212struct PgAuthContextGuard {
216 previous_tenant: Option<String>,
217}
218
219impl PgAuthContextGuard {
220 fn install(auth_context: &PgAuthContext) -> Self {
221 let previous_tenant = crate::runtime::execution_context::current_tenant();
222 match &auth_context.tenant {
223 Some(tenant) => crate::runtime::execution_context::set_current_tenant(tenant.clone()),
224 None => crate::runtime::execution_context::clear_current_tenant(),
225 }
226 crate::runtime::execution_context::set_current_auth_identity(
227 auth_context.username.clone(),
228 auth_context.role,
229 );
230 Self { previous_tenant }
231 }
232}
233
234impl Drop for PgAuthContextGuard {
235 fn drop(&mut self) {
236 crate::runtime::execution_context::clear_current_auth_identity();
237 match self.previous_tenant.take() {
238 Some(tenant) => crate::runtime::execution_context::set_current_tenant(tenant),
239 None => crate::runtime::execution_context::clear_current_tenant(),
240 }
241 }
242}
243
244fn is_loopback_bind(bind_addr: &str) -> bool {
248 crate::transport::is_loopback_bind(bind_addr)
249}
250
251fn password_from_message(payload: &[u8]) -> Option<String> {
254 let nul = payload.iter().position(|&b| b == 0)?;
255 std::str::from_utf8(&payload[..nul])
256 .ok()
257 .map(|password| password.to_string())
258}
259
260pub async fn start_pg_wire_listener(
263 config: PgWireConfig,
264 runtime: Arc<RedDBRuntime>,
265) -> Result<(), Box<dyn std::error::Error>> {
266 let listener = TcpListener::bind(&config.bind_addr).await?;
267 start_pg_wire_listener_on(listener, config, runtime).await
268}
269
270pub async fn start_pg_wire_listener_on(
272 listener: TcpListener,
273 config: PgWireConfig,
274 runtime: Arc<RedDBRuntime>,
275) -> Result<(), Box<dyn std::error::Error>> {
276 let trust_active = runtime
281 .auth_store()
282 .map(|store| !store.is_enabled())
283 .unwrap_or(true)
284 && is_loopback_bind(&config.bind_addr);
285 if trust_active {
286 tracing::warn!(
287 transport = "pg-wire",
288 bind = %config.bind_addr,
289 "PG-Wire trust authentication is active (no auth store configured); \
290 loopback-only. Configure auth before exposing this listener."
291 );
292 }
293 let tls_acceptor = match config.tls.as_ref() {
299 Some(tls_cfg) => Some(crate::wire::tls::build_tls_acceptor(tls_cfg)?),
300 None => None,
301 };
302 tracing::info!(
303 transport = "pg-wire",
304 bind = %config.bind_addr,
305 tls = tls_acceptor.is_some(),
306 "listener online"
307 );
308 let cfg = Arc::new(config);
309 let admission = crate::wire::admission::ConnectionAdmission::new("pg-wire");
310 loop {
311 let (stream, peer) = listener.accept().await?;
312 let Some(permit) = admission.try_admit() else {
316 continue;
317 };
318 if let Err(error) = stream.set_nodelay(true) {
321 tracing::warn!(transport = "pg-wire", peer = %peer, err = %error,
322 "failed to configure connection");
323 continue;
324 }
325 let rt = Arc::clone(&runtime);
326 let cfg = Arc::clone(&cfg);
327 let acceptor = tls_acceptor.clone();
328 let peer_str = peer.to_string();
329 tokio::spawn(async move {
330 let _permit = permit;
331 if let Err(e) = handle_connection(stream, rt, cfg, acceptor).await {
332 tracing::warn!(
333 transport = "pg-wire",
334 peer = %peer_str,
335 err = %e,
336 "connection failed"
337 );
338 }
339 });
340 }
341}
342
343pub(crate) async fn handle_connection<S>(
355 mut stream: S,
356 runtime: Arc<RedDBRuntime>,
357 config: Arc<PgWireConfig>,
358 tls_acceptor: Option<tokio_rustls::TlsAcceptor>,
359) -> Result<(), PgWireError>
360where
361 S: AsyncRead + AsyncWrite + Unpin + Send,
362{
363 loop {
371 let startup = crate::wire::admission::with_handshake_deadline(read_startup(&mut stream))
372 .await
373 .map_err(PgWireError::from)??;
374 match startup {
375 FrontendMessage::SslRequest => {
376 if let Some(acceptor) = tls_acceptor.as_ref() {
377 write_raw_byte(&mut stream, b'S').await?;
380 let tls_stream =
381 crate::wire::admission::with_handshake_deadline(acceptor.accept(stream))
382 .await??;
383 return serve_session_after_tls(tls_stream, runtime, config).await;
384 }
385 write_raw_byte(&mut stream, b'N').await?;
388 continue;
389 }
390 FrontendMessage::GssEncRequest => {
391 write_raw_byte(&mut stream, b'N').await?;
394 continue;
395 }
396 FrontendMessage::Startup(params) => {
397 return run_pg_session(stream, runtime, config, false, params).await;
398 }
399 FrontendMessage::Unknown { .. } => {
400 return Ok(());
402 }
403 other => {
404 return Err(PgWireError::Protocol(format!(
405 "unexpected startup frame: {other:?}"
406 )));
407 }
408 }
409 }
410}
411
412async fn serve_session_after_tls<S>(
417 mut stream: tokio_rustls::server::TlsStream<S>,
418 runtime: Arc<RedDBRuntime>,
419 config: Arc<PgWireConfig>,
420) -> Result<(), PgWireError>
421where
422 S: AsyncRead + AsyncWrite + Unpin + Send,
423{
424 loop {
425 match read_startup(&mut stream).await? {
426 FrontendMessage::SslRequest | FrontendMessage::GssEncRequest => {
427 write_raw_byte(&mut stream, b'N').await?;
428 continue;
429 }
430 FrontendMessage::Startup(params) => {
431 return run_pg_session(stream, runtime, config, true, params).await;
432 }
433 FrontendMessage::Unknown { .. } => return Ok(()),
434 other => {
435 return Err(PgWireError::Protocol(format!(
436 "unexpected startup frame: {other:?}"
437 )));
438 }
439 }
440 }
441}
442
443async fn run_pg_session<S>(
447 mut stream: S,
448 runtime: Arc<RedDBRuntime>,
449 config: Arc<PgWireConfig>,
450 over_tls: bool,
451 params: super::protocol::StartupParams,
452) -> Result<(), PgWireError>
453where
454 S: AsyncRead + AsyncWrite + Unpin + Send,
455{
456 let auth_context =
457 match authenticate_startup(&mut stream, &runtime, &config, over_tls, ¶ms).await? {
458 Some(context) => context,
459 None => return Ok(()),
462 };
463
464 let mut prepared: HashMap<String, PgPreparedStatement> = HashMap::new();
465 let mut portals: HashMap<String, PgPortal> = HashMap::new();
466 let prepared_registry = PreparedRegistry::new();
467
468 loop {
470 let frame = match read_frame(&mut stream).await {
471 Ok(f) => f,
472 Err(PgWireError::Eof) => return Ok(()),
473 Err(e) => return Err(e),
474 };
475
476 let mut execution_permit = if pg_frame_requires_execution(&frame) {
482 match acquire_pg_execution(&runtime).await {
483 Ok(permit) => Some(permit),
484 Err(message) => {
485 send_error(&mut stream, "53000", message).await?;
489 write_frame(
490 &mut stream,
491 &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
492 )
493 .await?;
494 continue;
495 }
496 }
497 } else {
498 None
499 };
500
501 match frame {
502 FrontendMessage::Query(sql) => {
503 let mut response =
504 ExecutionResponseWriter::new(&mut stream, execution_permit.take());
505 handle_simple_query(&mut response, &runtime, auth_context.as_ref(), &sql).await?;
506 }
507 FrontendMessage::Parse(msg) => {
508 let mut response =
509 ExecutionResponseWriter::new(&mut stream, execution_permit.take());
510 handle_parse(
511 &mut response,
512 &runtime,
513 &prepared_registry,
514 &mut prepared,
515 &portals,
516 msg,
517 )
518 .await?;
519 }
520 FrontendMessage::Bind(msg) => {
521 handle_bind(
522 &mut stream,
523 &prepared_registry,
524 &prepared,
525 &mut portals,
526 msg,
527 )
528 .await?;
529 }
530 FrontendMessage::Describe(msg) => {
531 let mut response =
532 ExecutionResponseWriter::new(&mut stream, execution_permit.take());
533 handle_describe(
534 &mut response,
535 &runtime,
536 auth_context.as_ref(),
537 &prepared_registry,
538 &prepared,
539 &mut portals,
540 msg,
541 )
542 .await?;
543 }
544 FrontendMessage::Execute(msg) => {
545 let mut response =
546 ExecutionResponseWriter::new(&mut stream, execution_permit.take());
547 handle_execute(
548 &mut response,
549 &runtime,
550 auth_context.as_ref(),
551 &prepared_registry,
552 &mut portals,
553 msg,
554 )
555 .await?;
556 }
557 FrontendMessage::Close(msg) => {
558 handle_close(
559 &mut stream,
560 &prepared_registry,
561 &mut prepared,
562 &mut portals,
563 msg,
564 )
565 .await?;
566 }
567 FrontendMessage::Terminate => return Ok(()),
568 FrontendMessage::Flush => {
569 continue;
572 }
573 FrontendMessage::Sync => {
574 write_frame(
575 &mut stream,
576 &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
577 )
578 .await?;
579 }
580 FrontendMessage::PasswordMessage(_) => {
581 continue;
583 }
584 FrontendMessage::Unknown { tag, .. } => {
585 send_error(
586 &mut stream,
587 "0A000",
588 &format!("unsupported frame tag 0x{tag:02x}"),
589 )
590 .await?;
591 write_frame(
592 &mut stream,
593 &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
594 )
595 .await?;
596 }
597 other => {
598 send_error(
599 &mut stream,
600 "0A000",
601 &format!("unsupported frame {other:?}"),
602 )
603 .await?;
604 write_frame(
605 &mut stream,
606 &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
607 )
608 .await?;
609 }
610 }
611 }
612}
613
614async fn handle_parse<S>(
615 stream: &mut S,
616 runtime: &RedDBRuntime,
617 prepared_registry: &PreparedRegistry,
618 prepared: &mut HashMap<String, PgPreparedStatement>,
619 portals: &HashMap<String, PgPortal>,
620 msg: super::protocol::ParseMessage,
621) -> Result<(), PgWireError>
622where
623 S: AsyncWrite + Unpin,
624{
625 let inferred_param_type_oids = infer_pg_cast_param_type_oids(&msg.query);
626 let sql = rewrite_pg_parameter_casts(&msg.query);
627 let prepared_query = match prepared_registry.prepare(runtime, &sql) {
628 Ok(prepared_query) => prepared_query,
629 Err(err) => {
630 send_error(stream, "42601", &err.to_string()).await?;
631 return Ok(());
632 }
633 };
634 let mut param_type_oids = msg.param_type_oids;
635 if param_type_oids.is_empty() {
636 param_type_oids.resize(prepared_query.parameter_count, PgOid::Unknown.as_u32());
637 }
638 for (idx, oid) in inferred_param_type_oids {
639 if idx >= param_type_oids.len() {
640 param_type_oids.resize(idx + 1, PgOid::Unknown.as_u32());
641 }
642 if param_type_oids[idx] == PgOid::Unknown.as_u32() {
643 param_type_oids[idx] = oid;
644 }
645 }
646 let replaced = prepared.insert(
647 msg.statement,
648 PgPreparedStatement {
649 sql,
650 prepared_id: prepared_query.id,
651 param_type_oids,
652 },
653 );
654 if let Some(replaced) = replaced {
655 release_prepared_if_unreferenced(
656 prepared_registry,
657 replaced.prepared_id,
658 prepared,
659 portals,
660 );
661 }
662 write_frame(stream, &BackendMessage::ParseComplete).await
663}
664
665async fn handle_bind<S>(
666 stream: &mut S,
667 prepared_registry: &PreparedRegistry,
668 prepared: &HashMap<String, PgPreparedStatement>,
669 portals: &mut HashMap<String, PgPortal>,
670 msg: super::protocol::BindMessage,
671) -> Result<(), PgWireError>
672where
673 S: AsyncWrite + Unpin,
674{
675 let Some(stmt) = prepared.get(&msg.statement) else {
676 send_error(
677 stream,
678 "26000",
679 &format!("prepared statement {:?} does not exist", msg.statement),
680 )
681 .await?;
682 return Ok(());
683 };
684 let params = match bind_pg_params(stmt, &msg) {
685 Ok(params) => params,
686 Err(err) => {
687 send_error(stream, "22023", &err).await?;
688 return Ok(());
689 }
690 };
691 let replaced = portals.insert(
692 msg.portal,
693 PgPortal {
694 sql: stmt.sql.clone(),
695 prepared_id: stmt.prepared_id,
696 params,
697 result_format_codes: msg.result_format_codes,
698 row_description_sent: false,
699 result: None,
700 row_offset: 0,
701 },
702 );
703 if let Some(replaced) = replaced {
704 release_prepared_if_unreferenced(
705 prepared_registry,
706 replaced.prepared_id,
707 prepared,
708 portals,
709 );
710 }
711 write_frame(stream, &BackendMessage::BindComplete).await
712}
713
714async fn handle_describe<S>(
715 stream: &mut S,
716 runtime: &RedDBRuntime,
717 auth_context: Option<&PgAuthContext>,
718 prepared_registry: &PreparedRegistry,
719 prepared: &HashMap<String, PgPreparedStatement>,
720 portals: &mut HashMap<String, PgPortal>,
721 msg: super::protocol::DescribeMessage,
722) -> Result<(), PgWireError>
723where
724 S: AsyncWrite + Unpin,
725{
726 match msg.target {
727 DescribeTarget::Statement => {
728 let Some(stmt) = prepared.get(&msg.name) else {
729 send_error(
730 stream,
731 "26000",
732 &format!("prepared statement {:?} does not exist", msg.name),
733 )
734 .await?;
735 return Ok(());
736 };
737 write_frame(
738 stream,
739 &BackendMessage::ParameterDescription(stmt.param_type_oids.clone()),
740 )
741 .await?;
742 if is_ask_query(&stmt.sql) {
743 emit_ask_row_description(stream).await
744 } else {
745 write_frame(stream, &BackendMessage::NoData).await
746 }
747 }
748 DescribeTarget::Portal => {
749 let Some(portal) = portals.get_mut(&msg.name) else {
750 send_error(
751 stream,
752 "34000",
753 &format!("portal {:?} does not exist", msg.name),
754 )
755 .await?;
756 return Ok(());
757 };
758 if is_ask_query(&portal.sql) {
759 emit_ask_row_description(stream).await?;
760 portal.row_description_sent = true;
761 Ok(())
762 } else if is_row_returning_query(&portal.sql) {
763 let result = match execute_pg_query_result(
764 runtime,
765 auth_context,
766 prepared_registry,
767 portal.prepared_id,
768 &portal.sql,
769 &portal.params,
770 ) {
771 Ok(result) => result,
772 Err(err) => {
773 let code = classify_sqlstate(&err);
774 send_error(stream, code, &err).await?;
775 return Ok(());
776 }
777 };
778 emit_row_description_for_result(stream, &result).await?;
779 portal.row_description_sent = true;
780 portal.result = Some(result);
781 portal.row_offset = 0;
782 Ok(())
783 } else {
784 write_frame(stream, &BackendMessage::NoData).await
785 }
786 }
787 }
788}
789
790async fn handle_execute<S>(
791 stream: &mut S,
792 runtime: &RedDBRuntime,
793 auth_context: Option<&PgAuthContext>,
794 prepared_registry: &PreparedRegistry,
795 portals: &mut HashMap<String, PgPortal>,
796 msg: super::protocol::ExecuteMessage,
797) -> Result<(), PgWireError>
798where
799 S: AsyncWrite + Unpin,
800{
801 let Some(portal) = portals.get_mut(&msg.portal) else {
802 send_error(
803 stream,
804 "34000",
805 &format!("portal {:?} does not exist", msg.portal),
806 )
807 .await?;
808 return Ok(());
809 };
810 match prepared_registry.ddl_epoch(portal.prepared_id) {
811 Some(epoch) if epoch != runtime.ddl_epoch() => {
812 send_error(stream, "XX000", "prepared_needs_replan").await?;
813 return Ok(());
814 }
815 Some(_) => {}
816 None => {
817 send_error(stream, "26000", "prepared statement does not exist").await?;
818 return Ok(());
819 }
820 }
821 if portal.result.is_none() {
822 match execute_pg_query_result(
823 runtime,
824 auth_context,
825 prepared_registry,
826 portal.prepared_id,
827 &portal.sql,
828 &portal.params,
829 ) {
830 Ok(result) => {
831 portal.result = Some(result);
832 portal.row_offset = 0;
833 }
834 Err(err) => {
835 let code = classify_sqlstate(&err);
836 send_error(stream, code, &err).await?;
837 return Ok(());
838 }
839 }
840 }
841
842 let result = portal.result.as_ref().expect("portal result populated");
843 match emit_portal_execution(
844 stream,
845 result,
846 portal.row_description_sent,
847 portal.row_offset,
848 msg.max_rows,
849 )
850 .await
851 {
852 Ok(PortalExecution::Complete) => {
853 portal.result = None;
854 portal.row_offset = 0;
855 portal.row_description_sent = false;
856 Ok(())
857 }
858 Ok(PortalExecution::Suspended { next_row_offset }) => {
859 portal.row_offset = next_row_offset;
860 portal.row_description_sent = true;
861 Ok(())
862 }
863 Err(err) => {
864 let code = classify_sqlstate(&err);
865 send_error(stream, code, &err).await
866 }
867 }
868}
869
870async fn handle_close<S>(
871 stream: &mut S,
872 prepared_registry: &PreparedRegistry,
873 prepared: &mut HashMap<String, PgPreparedStatement>,
874 portals: &mut HashMap<String, PgPortal>,
875 msg: super::protocol::CloseMessage,
876) -> Result<(), PgWireError>
877where
878 S: AsyncWrite + Unpin,
879{
880 let prepared_id = match msg.target {
881 DescribeTarget::Statement => prepared.remove(&msg.name).map(|stmt| stmt.prepared_id),
882 DescribeTarget::Portal => portals.remove(&msg.name).map(|portal| portal.prepared_id),
883 };
884 if let Some(prepared_id) = prepared_id {
885 release_prepared_if_unreferenced(prepared_registry, prepared_id, prepared, portals);
886 }
887 write_frame(stream, &BackendMessage::CloseComplete).await
888}
889
890fn release_prepared_if_unreferenced(
891 prepared_registry: &PreparedRegistry,
892 prepared_id: PreparedId,
893 prepared: &HashMap<String, PgPreparedStatement>,
894 portals: &HashMap<String, PgPortal>,
895) {
896 let statement_references_id = prepared
897 .values()
898 .any(|statement| statement.prepared_id == prepared_id);
899 let portal_references_id = portals
900 .values()
901 .any(|portal| portal.prepared_id == prepared_id);
902 if !statement_references_id && !portal_references_id {
903 prepared_registry.release(prepared_id);
904 }
905}
906
907fn bind_pg_params(
908 stmt: &PgPreparedStatement,
909 msg: &super::protocol::BindMessage,
910) -> Result<Vec<ParamValue>, String> {
911 if !matches!(msg.param_format_codes.len(), 0 | 1)
912 && msg.param_format_codes.len() != msg.params.len()
913 {
914 return Err("Bind format count must be 0, 1, or match parameter count".to_string());
915 }
916 msg.params
917 .iter()
918 .enumerate()
919 .map(|(idx, param)| {
920 let oid = stmt
921 .param_type_oids
922 .get(idx)
923 .copied()
924 .unwrap_or(PgOid::Unknown.as_u32());
925 let format_code = match msg.param_format_codes.as_slice() {
926 [] => 0,
927 [format] => *format,
928 formats => formats[idx],
929 };
930 pg_param_to_param_value(oid, format_code, param.as_deref())
931 })
932 .collect()
933}
934
935fn execute_pg_query_result(
936 runtime: &RedDBRuntime,
937 auth_context: Option<&PgAuthContext>,
938 prepared_registry: &PreparedRegistry,
939 prepared_id: PreparedId,
940 sql: &str,
941 params: &[ParamValue],
942) -> Result<crate::runtime::OwnedQueryResult, String> {
943 let executor = QueryRequestExecutor::new(runtime, prepared_registry);
944 let bound = executor
945 .bind(QueryRequest::prepared(prepared_id, params.to_vec()))
946 .map_err(|err| err.to_string())?;
947 if pg_scalar_select_param_index(sql).is_some() {
948 let values = params.iter().cloned().map(Value::from).collect::<Vec<_>>();
949 if let Some(result) = try_execute_pg_scalar_select(sql, &values) {
950 return runtime
951 .own_query_result(result)
952 .map_err(|error| error.to_string());
953 }
954 }
955 if params.is_empty() {
956 return match translate_pg_catalog_query(runtime, sql) {
957 Ok(Some(result)) => runtime
958 .own_query_result(crate::runtime::RuntimeQueryResult {
959 query: sql.to_string(),
960 mode: reddb_rql::modes::QueryMode::Sql,
961 statement: "select",
962 engine: "pg-catalog",
963 result,
964 affected_rows: 0,
965 statement_type: "select",
966 bookmark: None,
967 notice: None,
968 })
969 .map_err(|error| error.to_string()),
970 Ok(None) => run_runtime_blocking(|| {
971 execute_with_pg_auth_context(auth_context, || {
972 if matches!(
976 bound.prepared_expr(),
977 Some(
978 QueryExpr::Function(_)
979 | QueryExpr::Table(_)
980 | QueryExpr::Graph(_)
981 | QueryExpr::Path(_)
982 | QueryExpr::Join(_)
983 | QueryExpr::Vector(_)
984 | QueryExpr::Hybrid(_)
985 | QueryExpr::Insert(_)
986 | QueryExpr::Update(_)
987 | QueryExpr::Delete(_)
988 | QueryExpr::QueueCommand(_)
989 | QueryExpr::KvCommand(_)
990 | QueryExpr::SearchCommand(_)
991 | QueryExpr::Ask(_)
992 )
993 ) {
994 executor.execute_bound(bound)
995 } else {
996 runtime
997 .execute_query(sql)
998 .and_then(|result| runtime.own_query_result(result))
999 }
1000 })
1001 })
1002 .map_err(|err| err.to_string()),
1003 Err(err) => Err(err.to_string()),
1004 };
1005 }
1006
1007 run_runtime_blocking(|| {
1008 execute_with_pg_auth_context(auth_context, || executor.execute_bound(bound))
1009 })
1010 .map_err(|err| err.to_string())
1011}
1012
1013fn try_execute_pg_scalar_select(
1014 sql: &str,
1015 params: &[Value],
1016) -> Option<crate::runtime::RuntimeQueryResult> {
1017 let index = pg_scalar_select_param_index(sql)?;
1018 let value = params.get(index)?.clone();
1019 let mut result = UnifiedResult::with_columns(vec!["?column?".to_string()]);
1020 let mut record = UnifiedRecord::new();
1021 record.set("?column?", value);
1022 result.push(record);
1023 Some(crate::runtime::RuntimeQueryResult {
1024 query: sql.to_string(),
1025 mode: reddb_rql::modes::QueryMode::Sql,
1026 statement: "select",
1027 engine: "pg-wire",
1028 result,
1029 affected_rows: 0,
1030 statement_type: "select",
1031 bookmark: None,
1032 notice: None,
1033 })
1034}
1035
1036fn pg_scalar_select_param_index(sql: &str) -> Option<usize> {
1037 let trimmed = sql.trim().trim_end_matches(';').trim();
1038 let lower = trimmed.to_ascii_lowercase();
1039 let body = lower.strip_prefix("select ")?;
1040 let param = if let Some(inner) = body.strip_prefix("cast(") {
1041 let end = inner.find(" as ")?;
1042 &inner[..end]
1043 } else {
1044 body.split_whitespace().next()?
1045 };
1046 let digits = param.strip_prefix('$')?;
1047 let n = digits.parse::<usize>().ok()?;
1048 n.checked_sub(1)
1049}
1050
1051fn rewrite_pg_parameter_casts(sql: &str) -> String {
1052 let mut out = String::with_capacity(sql.len());
1053 let bytes = sql.as_bytes();
1054 let mut cursor = 0;
1055 let mut pos = 0;
1056 while pos < bytes.len() {
1057 if bytes[pos] != b'$' {
1058 pos += 1;
1059 continue;
1060 }
1061 let param_start = pos;
1062 pos += 1;
1063 let digits_start = pos;
1064 while pos < bytes.len() && bytes[pos].is_ascii_digit() {
1065 pos += 1;
1066 }
1067 if digits_start == pos {
1068 continue;
1069 }
1070 if pos + 2 <= bytes.len() && &bytes[pos..pos + 2] == b"::" {
1071 let param_end = pos;
1072 pos += 2;
1073 let type_start = pos;
1074 while pos < bytes.len()
1075 && (bytes[pos].is_ascii_alphanumeric() || matches!(bytes[pos], b'_' | b'.'))
1076 {
1077 pos += 1;
1078 }
1079 if type_start != pos {
1080 out.push_str(&sql[cursor..param_start]);
1081 out.push_str(&sql[param_start..param_end]);
1082 cursor = pos;
1083 continue;
1084 }
1085 }
1086 }
1087 out.push_str(&sql[cursor..]);
1088 out
1089}
1090
1091fn infer_pg_cast_param_type_oids(sql: &str) -> Vec<(usize, u32)> {
1092 let mut out = Vec::new();
1093 let bytes = sql.as_bytes();
1094 let mut pos = 0;
1095 while pos < bytes.len() {
1096 if bytes[pos] != b'$' {
1097 pos += 1;
1098 continue;
1099 }
1100 pos += 1;
1101 let digits_start = pos;
1102 while pos < bytes.len() && bytes[pos].is_ascii_digit() {
1103 pos += 1;
1104 }
1105 if digits_start == pos {
1106 continue;
1107 }
1108 let Some(param_index) = sql[digits_start..pos]
1109 .parse::<usize>()
1110 .ok()
1111 .and_then(|idx| idx.checked_sub(1))
1112 else {
1113 continue;
1114 };
1115 if pos + 2 > bytes.len() || &bytes[pos..pos + 2] != b"::" {
1116 continue;
1117 }
1118 pos += 2;
1119 let type_start = pos;
1120 while pos < bytes.len()
1121 && (bytes[pos].is_ascii_alphanumeric() || matches!(bytes[pos], b'_' | b'.'))
1122 {
1123 pos += 1;
1124 }
1125 if type_start == pos {
1126 continue;
1127 }
1128 if let Some(oid) = pg_cast_type_oid(&sql[type_start..pos]) {
1129 out.push((param_index, oid));
1130 }
1131 }
1132 out
1133}
1134
1135fn pg_cast_type_oid(ty: &str) -> Option<u32> {
1136 let lower = ty.to_ascii_lowercase();
1137 let short = lower.rsplit('.').next().unwrap_or(lower.as_str());
1138 let oid = match short {
1139 "bool" | "boolean" => PgOid::Bool,
1140 "int2" | "smallint" => PgOid::Int2,
1141 "int" | "int4" | "integer" => PgOid::Int4,
1142 "int8" | "bigint" => PgOid::Int8,
1143 "float4" | "real" => PgOid::Float4,
1144 "float8" | "double" | "doubleprecision" => PgOid::Float8,
1145 "numeric" | "decimal" => PgOid::Numeric,
1146 "bytea" => PgOid::Bytea,
1147 "json" => PgOid::Json,
1148 "jsonb" => PgOid::Jsonb,
1149 "text" => PgOid::Text,
1150 "varchar" | "character varying" => PgOid::Varchar,
1151 "uuid" => PgOid::Uuid,
1152 "timestamp" => PgOid::Timestamp,
1153 "timestamptz" | "timestampz" => PgOid::TimestampTz,
1154 "vector" => PgOid::Vector,
1155 _ => return None,
1156 };
1157 Some(oid.as_u32())
1158}
1159
1160fn is_row_returning_query(sql: &str) -> bool {
1161 let trimmed = sql.trim_start().to_ascii_lowercase();
1162 trimmed.starts_with("select")
1163 || trimmed.starts_with("with")
1164 || trimmed.starts_with("ask")
1165 || trimmed.starts_with("search")
1166 || trimmed.starts_with("vector")
1167 || trimmed.starts_with("hybrid")
1168}
1169
1170fn is_ask_query(sql: &str) -> bool {
1171 sql.trim_start().to_ascii_lowercase().starts_with("ask")
1172}
1173
1174async fn authenticate_startup<S>(
1201 stream: &mut S,
1202 runtime: &RedDBRuntime,
1203 config: &PgWireConfig,
1204 over_tls: bool,
1205 params: &super::protocol::StartupParams,
1206) -> Result<Option<Option<PgAuthContext>>, PgWireError>
1207where
1208 S: AsyncRead + AsyncWrite + Unpin,
1209{
1210 let auth_store = runtime.auth_store().filter(|store| store.is_enabled());
1211 let Some(auth_store) = auth_store else {
1212 if is_loopback_bind(&config.bind_addr) {
1214 send_startup_ready(stream, config, params).await?;
1215 return Ok(Some(None));
1216 }
1217 send_error(stream, "28P01", "password authentication failed").await?;
1218 return Ok(None);
1219 };
1220
1221 if !over_tls && !is_loopback_bind(&config.bind_addr) {
1226 send_error(
1227 stream,
1228 "28000",
1229 "cleartext password authentication requires TLS on non-loopback binds",
1230 )
1231 .await?;
1232 return Ok(None);
1233 }
1234
1235 write_frame(stream, &BackendMessage::AuthenticationCleartextPassword).await?;
1236 let password = match read_frame(stream).await? {
1237 FrontendMessage::PasswordMessage(payload) => password_from_message(&payload),
1238 _ => None,
1241 };
1242 let username = params.get("user").unwrap_or("");
1243 let tenant = params.get("tenant");
1244 let Some(password) = password else {
1245 send_error(stream, "28P01", "password authentication failed").await?;
1246 return Ok(None);
1247 };
1248 let session = match auth_store.authenticate_in_tenant(tenant, username, &password) {
1251 Ok(session) => session,
1252 Err(_) => {
1253 send_error(stream, "28P01", "password authentication failed").await?;
1254 return Ok(None);
1255 }
1256 };
1257
1258 send_startup_ready(stream, config, params).await?;
1259 Ok(Some(Some(PgAuthContext {
1260 username: session.username,
1261 role: session.role,
1262 tenant: session.tenant_id,
1263 })))
1264}
1265
1266async fn send_startup_ready<S>(
1267 stream: &mut S,
1268 config: &PgWireConfig,
1269 params: &super::protocol::StartupParams,
1270) -> Result<(), PgWireError>
1271where
1272 S: AsyncWrite + Unpin,
1273{
1274 write_frame(stream, &BackendMessage::AuthenticationOk).await?;
1276
1277 for (name, value) in [
1279 ("server_version", config.server_version.as_str()),
1280 ("server_encoding", "UTF8"),
1281 ("client_encoding", "UTF8"),
1282 ("DateStyle", "ISO, MDY"),
1283 ("TimeZone", "UTC"),
1284 ("integer_datetimes", "on"),
1285 ("standard_conforming_strings", "on"),
1286 (
1287 "application_name",
1288 params.get("application_name").unwrap_or(""),
1289 ),
1290 ] {
1291 write_frame(
1292 stream,
1293 &BackendMessage::ParameterStatus {
1294 name: name.to_string(),
1295 value: value.to_string(),
1296 },
1297 )
1298 .await?;
1299 }
1300
1301 write_frame(
1306 stream,
1307 &BackendMessage::BackendKeyData {
1308 pid: std::process::id(),
1309 key: random_cancel_key(),
1310 },
1311 )
1312 .await?;
1313
1314 write_frame(
1315 stream,
1316 &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
1317 )
1318 .await?;
1319 Ok(())
1320}
1321
1322async fn handle_simple_query<S>(
1323 stream: &mut S,
1324 runtime: &RedDBRuntime,
1325 auth_context: Option<&PgAuthContext>,
1326 sql: &str,
1327) -> Result<(), PgWireError>
1328where
1329 S: AsyncWrite + Unpin,
1330{
1331 if sql.trim().is_empty() {
1334 write_frame(stream, &BackendMessage::EmptyQueryResponse).await?;
1335 write_frame(
1336 stream,
1337 &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
1338 )
1339 .await?;
1340 return Ok(());
1341 }
1342
1343 if let Some(tag) = pg_session_compat_command_tag(sql) {
1344 write_frame(stream, &BackendMessage::CommandComplete(tag.to_string())).await?;
1345 write_frame(
1346 stream,
1347 &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
1348 )
1349 .await?;
1350 return Ok(());
1351 }
1352
1353 let query_result = match translate_pg_catalog_query(runtime, sql) {
1354 Ok(Some(result)) => runtime.own_query_result(crate::runtime::RuntimeQueryResult {
1355 query: sql.to_string(),
1356 mode: reddb_rql::modes::QueryMode::Sql,
1357 statement: "select",
1358 engine: "pg-catalog",
1359 result,
1360 affected_rows: 0,
1361 statement_type: "select",
1362 bookmark: None,
1363 notice: None,
1364 }),
1365 Ok(None) => run_runtime_blocking(|| {
1366 execute_with_pg_auth_context(auth_context, || {
1367 runtime
1368 .execute_query(sql)
1369 .and_then(|result| runtime.own_query_result(result))
1370 })
1371 }),
1372 Err(err) => Err(err),
1373 };
1374
1375 match query_result {
1376 Ok(result) => {
1377 emit_success_result(stream, &result).await?;
1378 }
1379 Err(err) => {
1380 let code = classify_sqlstate(&err.to_string());
1384 send_error(stream, code, &crate::api::client_facing_message(&err)).await?;
1385 }
1386 }
1387
1388 write_frame(
1389 stream,
1390 &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
1391 )
1392 .await?;
1393 Ok(())
1394}
1395
1396fn pg_session_compat_command_tag(sql: &str) -> Option<&'static str> {
1397 let lower = sql.trim().trim_end_matches(';').to_ascii_lowercase();
1398 if lower.starts_with("set ") {
1399 return Some("SET");
1400 }
1401 None
1402}
1403
1404async fn emit_success_result<S>(
1405 stream: &mut S,
1406 result: &crate::runtime::RuntimeQueryResult,
1407) -> Result<(), PgWireError>
1408where
1409 S: AsyncWrite + Unpin,
1410{
1411 if result.statement == "ask" {
1412 emit_ask_result_row(stream, result).await?;
1413 write_frame(
1414 stream,
1415 &BackendMessage::CommandComplete("SELECT 1".to_string()),
1416 )
1417 .await?;
1418 } else if result_returns_rows(result) {
1419 emit_result_rows(stream, &result.result).await?;
1420 write_frame(
1421 stream,
1422 &BackendMessage::CommandComplete(format!("SELECT {}", result.result.records.len())),
1423 )
1424 .await?;
1425 } else {
1426 let tag = match result.statement_type {
1431 "insert" => format!("INSERT 0 {}", result.affected_rows),
1432 "update" => format!("UPDATE {}", result.affected_rows),
1433 "delete" => format!("DELETE {}", result.affected_rows),
1434 other => other.to_uppercase(),
1435 };
1436 write_frame(stream, &BackendMessage::CommandComplete(tag)).await?;
1437 }
1438 Ok(())
1439}
1440
1441async fn emit_success_result_without_row_description<S>(
1442 stream: &mut S,
1443 result: &crate::runtime::RuntimeQueryResult,
1444) -> Result<(), PgWireError>
1445where
1446 S: AsyncWrite + Unpin,
1447{
1448 if result.statement == "ask" {
1449 let row = ask_query_result_to_pg_wire_row(result)
1450 .ok_or_else(|| PgWireError::Protocol("ASK result missing row body".to_string()))?;
1451 write_frame(stream, &BackendMessage::DataRow(row.cells)).await?;
1452 write_frame(
1453 stream,
1454 &BackendMessage::CommandComplete("SELECT 1".to_string()),
1455 )
1456 .await?;
1457 } else if result_returns_rows(result) {
1458 emit_result_data_rows(stream, &result.result).await?;
1459 write_frame(
1460 stream,
1461 &BackendMessage::CommandComplete(format!("SELECT {}", result.result.records.len())),
1462 )
1463 .await?;
1464 } else {
1465 let tag = match result.statement_type {
1466 "insert" => format!("INSERT 0 {}", result.affected_rows),
1467 "update" => format!("UPDATE {}", result.affected_rows),
1468 "delete" => format!("DELETE {}", result.affected_rows),
1469 other => other.to_uppercase(),
1470 };
1471 write_frame(stream, &BackendMessage::CommandComplete(tag)).await?;
1472 }
1473 Ok(())
1474}
1475
1476enum PortalExecution {
1477 Complete,
1478 Suspended { next_row_offset: usize },
1479}
1480
1481async fn emit_portal_execution<S>(
1482 stream: &mut S,
1483 result: &crate::runtime::RuntimeQueryResult,
1484 row_description_sent: bool,
1485 row_offset: usize,
1486 max_rows: u32,
1487) -> Result<PortalExecution, String>
1488where
1489 S: AsyncWrite + Unpin,
1490{
1491 if result.statement == "ask" && row_description_sent {
1492 emit_success_result_without_row_description(stream, result)
1493 .await
1494 .map_err(|err| err.to_string())?;
1495 Ok(PortalExecution::Complete)
1496 } else if result.statement == "ask" {
1497 emit_success_result(stream, result)
1498 .await
1499 .map_err(|err| err.to_string())?;
1500 Ok(PortalExecution::Complete)
1501 } else if result_returns_rows(result) {
1502 emit_row_returning_portal_execution(
1503 stream,
1504 result,
1505 row_description_sent,
1506 row_offset,
1507 max_rows,
1508 )
1509 .await
1510 } else if row_description_sent {
1511 emit_success_result_without_row_description(stream, result)
1512 .await
1513 .map_err(|err| err.to_string())?;
1514 Ok(PortalExecution::Complete)
1515 } else {
1516 emit_success_result(stream, result)
1517 .await
1518 .map_err(|err| err.to_string())?;
1519 Ok(PortalExecution::Complete)
1520 }
1521}
1522
1523async fn emit_row_returning_portal_execution<S>(
1524 stream: &mut S,
1525 result: &crate::runtime::RuntimeQueryResult,
1526 row_description_sent: bool,
1527 row_offset: usize,
1528 max_rows: u32,
1529) -> Result<PortalExecution, String>
1530where
1531 S: AsyncWrite + Unpin,
1532{
1533 if !row_description_sent {
1534 emit_result_row_description(stream, &result.result)
1535 .await
1536 .map_err(|err| err.to_string())?;
1537 }
1538
1539 let total_rows = result.result.records.len();
1540 let remaining_rows = total_rows.saturating_sub(row_offset);
1541 let rows_to_emit = if max_rows == 0 {
1542 remaining_rows
1543 } else {
1544 remaining_rows.min(max_rows as usize)
1545 };
1546 emit_result_data_rows_range(stream, &result.result, row_offset, rows_to_emit)
1547 .await
1548 .map_err(|err| err.to_string())?;
1549
1550 let next_row_offset = row_offset + rows_to_emit;
1551 if next_row_offset < total_rows {
1552 write_frame(stream, &BackendMessage::PortalSuspended)
1553 .await
1554 .map_err(|err| err.to_string())?;
1555 Ok(PortalExecution::Suspended { next_row_offset })
1556 } else {
1557 write_frame(
1558 stream,
1559 &BackendMessage::CommandComplete(format!("SELECT {}", total_rows)),
1560 )
1561 .await
1562 .map_err(|err| err.to_string())?;
1563 Ok(PortalExecution::Complete)
1564 }
1565}
1566
1567async fn emit_row_description_for_result<S>(
1568 stream: &mut S,
1569 result: &crate::runtime::RuntimeQueryResult,
1570) -> Result<(), PgWireError>
1571where
1572 S: AsyncWrite + Unpin,
1573{
1574 if result.statement == "ask" {
1575 emit_ask_row_description(stream).await
1576 } else if result_returns_rows(result) {
1577 emit_result_row_description(stream, &result.result).await
1578 } else {
1579 write_frame(stream, &BackendMessage::NoData).await
1580 }
1581}
1582
1583fn result_returns_rows(result: &crate::runtime::RuntimeQueryResult) -> bool {
1584 result.statement_type == "select"
1585}
1586
1587async fn emit_result_rows<S>(
1588 stream: &mut S,
1589 result: &crate::storage::query::unified::UnifiedResult,
1590) -> Result<(), PgWireError>
1591where
1592 S: AsyncWrite + Unpin,
1593{
1594 emit_result_row_description(stream, result).await?;
1595 emit_result_data_rows(stream, result).await
1596}
1597
1598async fn emit_result_row_description<S>(
1599 stream: &mut S,
1600 result: &crate::storage::query::unified::UnifiedResult,
1601) -> Result<(), PgWireError>
1602where
1603 S: AsyncWrite + Unpin,
1604{
1605 let columns: Vec<String> = if !result.columns.is_empty() {
1609 result.columns.clone()
1610 } else if let Some(first) = result.records.first() {
1611 record_field_names(first)
1612 } else {
1613 Vec::new()
1614 };
1615
1616 let type_oids: Vec<PgOid> = columns
1620 .iter()
1621 .map(|col| {
1622 result
1623 .records
1624 .first()
1625 .and_then(|r| record_get(r, col))
1626 .map(PgOid::from_value)
1627 .unwrap_or(PgOid::Text)
1628 })
1629 .collect();
1630
1631 let descriptors: Vec<ColumnDescriptor> = columns
1632 .iter()
1633 .zip(type_oids.iter())
1634 .map(|(name, oid)| ColumnDescriptor {
1635 name: name.clone(),
1636 table_oid: 0,
1637 column_attr: 0,
1638 type_oid: oid.as_u32(),
1639 type_size: -1,
1640 type_mod: -1,
1641 format: 0,
1642 })
1643 .collect();
1644
1645 write_frame(stream, &BackendMessage::RowDescription(descriptors)).await
1646}
1647
1648async fn emit_result_data_rows<S>(
1649 stream: &mut S,
1650 result: &crate::storage::query::unified::UnifiedResult,
1651) -> Result<(), PgWireError>
1652where
1653 S: AsyncWrite + Unpin,
1654{
1655 emit_result_data_rows_range(stream, result, 0, result.records.len()).await
1656}
1657
1658async fn emit_result_data_rows_range<S>(
1659 stream: &mut S,
1660 result: &crate::storage::query::unified::UnifiedResult,
1661 row_offset: usize,
1662 row_count: usize,
1663) -> Result<(), PgWireError>
1664where
1665 S: AsyncWrite + Unpin,
1666{
1667 let columns: Vec<String> = if !result.columns.is_empty() {
1668 result.columns.clone()
1669 } else if let Some(first) = result.records.first() {
1670 record_field_names(first)
1671 } else {
1672 Vec::new()
1673 };
1674 for record in result.records.iter().skip(row_offset).take(row_count) {
1675 let fields: Vec<Option<Vec<u8>>> = columns
1676 .iter()
1677 .map(|col| record_get(record, col).and_then(value_to_pg_wire_bytes))
1678 .collect();
1679 write_frame(stream, &BackendMessage::DataRow(fields)).await?;
1680 }
1681
1682 Ok(())
1683}
1684
1685async fn emit_ask_result_row<S>(
1686 stream: &mut S,
1687 result: &crate::runtime::RuntimeQueryResult,
1688) -> Result<(), PgWireError>
1689where
1690 S: AsyncWrite + Unpin,
1691{
1692 let row = ask_query_result_to_pg_wire_row(result)
1693 .ok_or_else(|| PgWireError::Protocol("ASK result missing row body".to_string()))?;
1694
1695 emit_ask_row_description(stream).await?;
1696 write_frame(stream, &BackendMessage::DataRow(row.cells)).await?;
1697 Ok(())
1698}
1699
1700async fn emit_ask_row_description<S>(stream: &mut S) -> Result<(), PgWireError>
1701where
1702 S: AsyncWrite + Unpin,
1703{
1704 let descriptors: Vec<ColumnDescriptor> = crate::runtime::ai::pg_wire_ask_row_encoder::columns()
1705 .iter()
1706 .map(|col| ColumnDescriptor {
1707 name: col.name.to_string(),
1708 table_oid: 0,
1709 column_attr: 0,
1710 type_oid: col.oid.as_u32(),
1711 type_size: -1,
1712 type_mod: -1,
1713 format: 0,
1714 })
1715 .collect();
1716 write_frame(stream, &BackendMessage::RowDescription(descriptors)).await
1717}
1718
1719fn ask_query_result_to_pg_wire_row(
1720 result: &crate::runtime::RuntimeQueryResult,
1721) -> Option<crate::runtime::ai::pg_wire_ask_row_encoder::AskRow> {
1722 if result.statement != "ask" {
1723 return None;
1724 }
1725 let ask = crate::presentation::query_result::ask_result_from_unified_result(&result.result)?;
1726
1727 Some(crate::runtime::ai::pg_wire_ask_row_encoder::encode(&ask))
1728}
1729
1730fn record_get<'a>(record: &'a UnifiedRecord, key: &str) -> Option<&'a Value> {
1734 record.get(key)
1735}
1736
1737fn record_field_names(record: &UnifiedRecord) -> Vec<String> {
1746 record
1750 .column_names()
1751 .into_iter()
1752 .map(|k| k.to_string())
1753 .collect()
1754}
1755
1756async fn send_error<S>(stream: &mut S, code: &str, message: &str) -> Result<(), PgWireError>
1757where
1758 S: AsyncWrite + Unpin,
1759{
1760 write_frame(
1761 stream,
1762 &BackendMessage::ErrorResponse {
1763 severity: "ERROR".to_string(),
1764 code: code.to_string(),
1765 message: message.to_string(),
1766 },
1767 )
1768 .await
1769}
1770
1771fn classify_sqlstate(msg: &str) -> &'static str {
1775 let lower = msg.to_ascii_lowercase();
1776 if lower.contains("not found") || lower.contains("does not exist") {
1777 "42P01"
1779 } else if lower.contains("parse") || lower.contains("expected") || lower.contains("syntax") {
1780 "42601"
1781 } else if lower.contains("already exists") {
1782 "42P07"
1783 } else if lower.contains("permission") || lower.contains("auth") {
1784 "28000"
1785 } else {
1786 "XX000"
1787 }
1788}
1789
1790#[cfg(test)]
1791mod tests {
1792 use super::*;
1793 use crate::api::RedDBOptions;
1794 use crate::auth::{AuthConfig, AuthStore};
1795 use crate::runtime::RuntimeQueryResult;
1796 use crate::storage::query::unified::UnifiedResult;
1797 use reddb_rql::modes::QueryMode;
1798 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
1799
1800 #[tokio::test]
1801 async fn response_writer_releases_execution_slot_before_socket_write() {
1802 let runtime = RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap();
1803 runtime.configure_execution_admission(1, 0, 100);
1804 let admission = runtime.execution_admission();
1805 let permit = admission.acquire().await.expect("execution permit");
1806 assert_eq!(admission.active(), 1);
1807
1808 let (mut server_io, mut client_io) = tokio::io::duplex(64);
1809 let mut writer = ExecutionResponseWriter::new(&mut server_io, Some(permit));
1810 writer.write_all(b"ok").await.unwrap();
1811 assert_eq!(admission.active(), 0);
1812
1813 let mut received = [0_u8; 2];
1814 client_io.read_exact(&mut received).await.unwrap();
1815 assert_eq!(&received, b"ok");
1816 }
1817
1818 #[tokio::test]
1819 async fn enabled_auth_store_rejects_bad_pgwire_password_with_28p01() {
1820 let runtime = Arc::new(RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap());
1821 let auth = Arc::new(AuthStore::new(AuthConfig {
1822 enabled: true,
1823 require_auth: true,
1824 ..AuthConfig::default()
1825 }));
1826 auth.create_user("alice", "secret", Role::Write).unwrap();
1827 runtime.set_auth_store(auth);
1828
1829 let config = Arc::new(PgWireConfig::default());
1830 let (server_io, mut client_io) = tokio::io::duplex(64 * 1024);
1831 let server = tokio::spawn(async move {
1832 handle_connection(server_io, runtime, config, None)
1833 .await
1834 .unwrap();
1835 });
1836
1837 write_startup_params(&mut client_io, &[("user", "alice")]).await;
1838 let (tag, body) = read_backend_frame(&mut client_io).await;
1840 assert_eq!(tag, b'R');
1841 assert_eq!(i32::from_be_bytes(body[..4].try_into().unwrap()), 3);
1842
1843 write_frontend_frame(&mut client_io, b'p', cstring_body("wrong")).await;
1844 let (tag, body) = read_backend_frame(&mut client_io).await;
1845 assert_eq!(tag, b'E');
1846 assert_eq!(error_response_field(&body, b'C'), Some("28P01"));
1847
1848 let mut eof = [0u8; 1];
1850 assert_eq!(client_io.read(&mut eof).await.unwrap(), 0);
1851 server.await.unwrap();
1852 }
1853
1854 #[tokio::test]
1855 async fn pgwire_password_auth_installs_statement_identity() {
1856 let runtime = Arc::new(RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap());
1857 let auth = Arc::new(AuthStore::new(AuthConfig {
1858 enabled: true,
1859 require_auth: true,
1860 ..AuthConfig::default()
1861 }));
1862 auth.create_user_in_tenant(Some("acme"), "alice", "secret", Role::Write)
1863 .unwrap();
1864 runtime.set_auth_store(auth);
1865
1866 let config = Arc::new(PgWireConfig::default());
1867 let (server_io, mut client_io) = tokio::io::duplex(64 * 1024);
1868 let server = tokio::spawn(async move {
1869 handle_connection(server_io, runtime, config, None)
1870 .await
1871 .unwrap();
1872 });
1873
1874 write_startup_params(&mut client_io, &[("user", "alice"), ("tenant", "acme")]).await;
1875 let (tag, body) = read_backend_frame(&mut client_io).await;
1876 assert_eq!(tag, b'R');
1877 assert_eq!(i32::from_be_bytes(body[..4].try_into().unwrap()), 3);
1878 write_frontend_frame(&mut client_io, b'p', cstring_body("secret")).await;
1879 read_until_ready(&mut client_io).await;
1880
1881 write_frontend_frame(
1885 &mut client_io,
1886 b'Q',
1887 cstring_body("SELECT CURRENT_USER(), CURRENT_ROLE(), CURRENT_TENANT()"),
1888 )
1889 .await;
1890 let frames = read_until_ready(&mut client_io).await;
1891 assert_eq!(
1892 frames.iter().map(|(tag, _)| *tag).collect::<Vec<_>>(),
1893 b"TDCZ"
1894 );
1895 let cells = decode_data_row(&frames[1].1);
1896 assert_eq!(cells[0].as_deref(), Some(b"alice".as_slice()));
1897 assert_eq!(cells[1].as_deref(), Some(b"write".as_slice()));
1898 assert_eq!(cells[2].as_deref(), Some(b"acme".as_slice()));
1899
1900 write_frontend_frame(&mut client_io, b'X', Vec::new()).await;
1901 server.await.unwrap();
1902 }
1903
1904 #[tokio::test]
1905 async fn extended_parse_bind_execute_returns_rows() {
1906 let runtime = Arc::new(RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap());
1907 let config = Arc::new(PgWireConfig::default());
1908 let (server_io, mut client_io) = tokio::io::duplex(64 * 1024);
1909 let server = tokio::spawn(async move {
1910 handle_connection(server_io, runtime, config, None)
1911 .await
1912 .unwrap();
1913 });
1914
1915 write_startup(&mut client_io).await;
1916 read_until_ready(&mut client_io).await;
1917
1918 write_frontend_frame(
1919 &mut client_io,
1920 b'P',
1921 parse_body("", "SELECT $1::int", &[PgOid::Int4.as_u32()]),
1922 )
1923 .await;
1924 write_frontend_frame(
1925 &mut client_io,
1926 b'B',
1927 bind_body("", "", &[0], &[Some(b"42".as_slice())], &[]),
1928 )
1929 .await;
1930 write_frontend_frame(&mut client_io, b'D', describe_body(b'P', "")).await;
1931 write_frontend_frame(&mut client_io, b'E', execute_body("", 0)).await;
1932 write_frontend_frame(&mut client_io, b'S', Vec::new()).await;
1933
1934 let frames = read_until_ready(&mut client_io).await;
1935 assert_eq!(
1936 frames.iter().map(|(tag, _)| *tag).collect::<Vec<_>>(),
1937 b"12TDCZ"
1938 );
1939 let columns = decode_row_description(&frames[2].1);
1940 assert_eq!(columns.len(), 1);
1941 let cells = decode_data_row(&frames[3].1);
1942 assert_eq!(cells.len(), 1);
1943 assert_eq!(cells[0].as_deref(), Some(b"42".as_slice()));
1944 assert_eq!(decode_command_complete(&frames[4].1), "SELECT 1");
1945
1946 write_frontend_frame(&mut client_io, b'X', Vec::new()).await;
1947 server.await.unwrap();
1948 }
1949
1950 #[test]
1951 fn prepared_pg_utility_statements_preserve_ddl_and_transactions() {
1952 let runtime = RedDBRuntime::in_memory().expect("runtime");
1953 let registry = PreparedRegistry::new();
1954 let execute = |sql: &str, params: Vec<ParamValue>| {
1955 let prepared = registry
1956 .prepare(&runtime, sql)
1957 .expect("prepare utility or DML");
1958 execute_pg_query_result(&runtime, None, ®istry, prepared.id, sql, ¶ms)
1959 .unwrap_or_else(|error| panic!("{sql}: {error}"))
1960 };
1961 execute("CREATE TABLE pg_utilities (id INT)", Vec::new());
1962 execute("BEGIN", Vec::new());
1963 execute(
1964 "INSERT INTO pg_utilities (id) VALUES ($1)",
1965 vec![ParamValue::Int64(1)],
1966 );
1967 execute("SAVEPOINT before_second", Vec::new());
1968 execute(
1969 "INSERT INTO pg_utilities (id) VALUES ($1)",
1970 vec![ParamValue::Int64(2)],
1971 );
1972 execute("ROLLBACK TO SAVEPOINT before_second", Vec::new());
1973 execute("COMMIT", Vec::new());
1974 let result = execute("SELECT id FROM pg_utilities", Vec::new());
1975 assert_eq!(result.result.records.len(), 1);
1976 assert_eq!(result.result.records[0].get("id"), Some(&Value::Integer(1)));
1977 execute("DROP TABLE pg_utilities", Vec::new());
1978 assert!(runtime
1979 .db()
1980 .store()
1981 .get_collection("pg_utilities")
1982 .is_none());
1983 }
1984
1985 #[tokio::test]
1986 async fn extended_repeated_execute_parses_prepared_shape_once() {
1987 let runtime = RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap();
1988 runtime
1989 .execute_query("CREATE TABLE pg_parse_once (id INT)")
1990 .unwrap();
1991 runtime
1992 .execute_query("INSERT INTO pg_parse_once (id) VALUES (42)")
1993 .unwrap();
1994 let prepared_registry = PreparedRegistry::new();
1995 let mut prepared = HashMap::new();
1996 let mut portals = HashMap::new();
1997 let (mut server_io, _client_io) = tokio::io::duplex(64 * 1024);
1998
1999 handle_parse(
2000 &mut server_io,
2001 &runtime,
2002 &prepared_registry,
2003 &mut prepared,
2004 &portals,
2005 super::super::protocol::ParseMessage {
2006 statement: "parse_once_stmt".to_string(),
2007 query: "SELECT id FROM pg_parse_once WHERE id = $1".to_string(),
2008 param_type_oids: vec![PgOid::Int4.as_u32()],
2009 },
2010 )
2011 .await
2012 .unwrap();
2013 assert_eq!(prepared_registry.parse_count(), 1);
2014
2015 handle_bind(
2016 &mut server_io,
2017 &prepared_registry,
2018 &prepared,
2019 &mut portals,
2020 super::super::protocol::BindMessage {
2021 portal: "parse_once_portal".to_string(),
2022 statement: "parse_once_stmt".to_string(),
2023 param_format_codes: vec![0],
2024 params: vec![Some(b"42".to_vec())],
2025 result_format_codes: Vec::new(),
2026 },
2027 )
2028 .await
2029 .unwrap();
2030
2031 for _ in 0..2 {
2032 handle_execute(
2033 &mut server_io,
2034 &runtime,
2035 None,
2036 &prepared_registry,
2037 &mut portals,
2038 super::super::protocol::ExecuteMessage {
2039 portal: "parse_once_portal".to_string(),
2040 max_rows: 0,
2041 },
2042 )
2043 .await
2044 .unwrap();
2045 }
2046 assert_eq!(prepared_registry.parse_count(), 1);
2047 }
2048
2049 #[tokio::test]
2050 async fn non_loopback_plaintext_refuses_cleartext_password_without_tls() {
2051 let runtime = Arc::new(RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap());
2058 let auth = Arc::new(AuthStore::new(AuthConfig {
2059 enabled: true,
2060 require_auth: true,
2061 ..AuthConfig::default()
2062 }));
2063 auth.create_user("alice", "secret", Role::Write).unwrap();
2064 runtime.set_auth_store(auth);
2065
2066 let config = Arc::new(PgWireConfig {
2068 bind_addr: "0.0.0.0:5432".to_string(),
2069 ..PgWireConfig::default()
2070 });
2071 let (server_io, mut client_io) = tokio::io::duplex(64 * 1024);
2072 let server = tokio::spawn(async move {
2073 handle_connection(server_io, runtime, config, None)
2074 .await
2075 .unwrap();
2076 });
2077
2078 write_startup_params(&mut client_io, &[("user", "alice")]).await;
2079 let (tag, body) = read_backend_frame(&mut client_io).await;
2081 assert_eq!(tag, b'E', "expected an ErrorResponse, not a challenge");
2082 assert_eq!(error_response_field(&body, b'C'), Some("28000"));
2083
2084 let mut eof = [0u8; 1];
2086 assert_eq!(client_io.read(&mut eof).await.unwrap(), 0);
2087 server.await.unwrap();
2088 }
2089
2090 #[test]
2091 fn infer_pg_cast_param_type_oids_from_parameter_casts() {
2092 assert_eq!(
2093 infer_pg_cast_param_type_oids("INSERT INTO t (id, name) VALUES ($1::int, $2::text)"),
2094 vec![(0, PgOid::Int4.as_u32()), (1, PgOid::Text.as_u32())]
2095 );
2096 assert_eq!(
2097 infer_pg_cast_param_type_oids("SEARCH SIMILAR [1.0] COLLECTION v LIMIT $1::int8"),
2098 vec![(0, PgOid::Int8.as_u32())]
2099 );
2100 }
2101
2102 #[test]
2103 fn pg_session_compat_accepts_driver_setup_set_commands() {
2104 assert_eq!(
2105 pg_session_compat_command_tag("SET extra_float_digits = 3"),
2106 Some("SET")
2107 );
2108 assert_eq!(
2109 pg_session_compat_command_tag("SET application_name = 'pgjdbc'"),
2110 Some("SET")
2111 );
2112 assert_eq!(pg_session_compat_command_tag("SELECT 1"), None);
2113 }
2114
2115 #[tokio::test]
2116 async fn ask_success_result_uses_canonical_pg_wire_row_shape() {
2117 let mut result = UnifiedResult::with_columns(vec![
2118 "answer".into(),
2119 "provider".into(),
2120 "model".into(),
2121 "prompt_tokens".into(),
2122 "completion_tokens".into(),
2123 "sources_count".into(),
2124 "sources_flat".into(),
2125 "citations".into(),
2126 "validation".into(),
2127 ]);
2128 let mut record = UnifiedRecord::new();
2129 record.set("answer", Value::text("Deploy failed [^1]."));
2130 record.set("provider", Value::text("openai"));
2131 record.set("model", Value::text("gpt-4o-mini"));
2132 record.set("prompt_tokens", Value::Integer(11));
2133 record.set("completion_tokens", Value::Integer(7));
2134 record.set(
2135 "sources_flat",
2136 Value::Json(
2137 br#"[{"urn":"urn:reddb:row:deployments:1","kind":"row","collection":"deployments","id":"1"}]"#
2138 .to_vec(),
2139 ),
2140 );
2141 record.set(
2142 "citations",
2143 Value::Json(br#"[{"marker":1,"urn":"urn:reddb:row:deployments:1"}]"#.to_vec()),
2144 );
2145 record.set(
2146 "validation",
2147 Value::Json(br#"{"ok":true,"warnings":[],"errors":[]}"#.to_vec()),
2148 );
2149 result.push(record);
2150
2151 let qr = RuntimeQueryResult {
2152 query: "ASK 'why did deploy fail?'".to_string(),
2153 mode: QueryMode::Sql,
2154 statement: "ask",
2155 engine: "runtime-ai",
2156 result,
2157 affected_rows: 0,
2158 statement_type: "select",
2159 bookmark: None,
2160 notice: None,
2161 };
2162
2163 let mut out = Vec::new();
2164 emit_success_result(&mut out, &qr).await.unwrap();
2165 let frames = decode_frames(&out);
2166
2167 assert_eq!(
2168 frames.iter().map(|(tag, _)| *tag).collect::<Vec<_>>(),
2169 b"TDC"
2170 );
2171
2172 let columns = decode_row_description(frames[0].1);
2173 assert_eq!(
2174 columns,
2175 vec![
2176 ("answer".to_string(), PgOid::Text.as_u32()),
2177 ("cache_hit".to_string(), PgOid::Bool.as_u32()),
2178 ("citations".to_string(), PgOid::Jsonb.as_u32()),
2179 ("completion_tokens".to_string(), PgOid::Int8.as_u32()),
2180 ("cost_usd".to_string(), PgOid::Numeric.as_u32()),
2181 ("mode".to_string(), PgOid::Text.as_u32()),
2182 ("model".to_string(), PgOid::Text.as_u32()),
2183 ("prompt_tokens".to_string(), PgOid::Int8.as_u32()),
2184 ("provider".to_string(), PgOid::Text.as_u32()),
2185 ("retry_count".to_string(), PgOid::Int8.as_u32()),
2186 ("sources_flat".to_string(), PgOid::Jsonb.as_u32()),
2187 ("validation".to_string(), PgOid::Jsonb.as_u32()),
2188 ]
2189 );
2190
2191 let cells = decode_data_row(frames[1].1);
2192 assert_eq!(cells.len(), 12);
2193 assert_eq!(cells[0].as_deref(), Some(b"Deploy failed [^1].".as_slice()));
2194 assert_eq!(cells[1].as_deref(), Some(b"f".as_slice()));
2195 assert_eq!(cells[4].as_deref(), Some(b"0".as_slice()));
2196 assert_eq!(cells[5].as_deref(), Some(b"strict".as_slice()));
2197 assert_eq!(cells[9].as_deref(), Some(b"0".as_slice()));
2198 assert!(std::str::from_utf8(cells[10].as_deref().unwrap())
2199 .unwrap()
2200 .contains(r#""payload""#));
2201 assert_eq!(decode_command_complete(frames[2].1), "SELECT 1");
2202 }
2203
2204 fn decode_frames(bytes: &[u8]) -> Vec<(u8, &[u8])> {
2205 let mut pos = 0;
2206 let mut frames = Vec::new();
2207 while pos < bytes.len() {
2208 let tag = bytes[pos];
2209 let len = u32::from_be_bytes([
2210 bytes[pos + 1],
2211 bytes[pos + 2],
2212 bytes[pos + 3],
2213 bytes[pos + 4],
2214 ]) as usize;
2215 let body_start = pos + 5;
2216 let body_end = pos + 1 + len;
2217 frames.push((tag, &bytes[body_start..body_end]));
2218 pos = body_end;
2219 }
2220 frames
2221 }
2222
2223 fn decode_row_description(body: &[u8]) -> Vec<(String, u32)> {
2224 let count = i16::from_be_bytes([body[0], body[1]]) as usize;
2225 let mut pos = 2;
2226 let mut columns = Vec::with_capacity(count);
2227 for _ in 0..count {
2228 let end = body[pos..].iter().position(|&b| b == 0).unwrap() + pos;
2229 let name = std::str::from_utf8(&body[pos..end]).unwrap().to_string();
2230 pos = end + 1;
2231 pos += 4; pos += 2; let oid = u32::from_be_bytes([body[pos], body[pos + 1], body[pos + 2], body[pos + 3]]);
2234 pos += 4;
2235 pos += 2; pos += 4; pos += 2; columns.push((name, oid));
2239 }
2240 columns
2241 }
2242
2243 fn decode_data_row(body: &[u8]) -> Vec<Option<Vec<u8>>> {
2244 let count = i16::from_be_bytes([body[0], body[1]]) as usize;
2245 let mut pos = 2;
2246 let mut cells = Vec::with_capacity(count);
2247 for _ in 0..count {
2248 let len = i32::from_be_bytes([body[pos], body[pos + 1], body[pos + 2], body[pos + 3]]);
2249 pos += 4;
2250 if len < 0 {
2251 cells.push(None);
2252 } else {
2253 let len = len as usize;
2254 cells.push(Some(body[pos..pos + len].to_vec()));
2255 pos += len;
2256 }
2257 }
2258 cells
2259 }
2260
2261 fn decode_command_complete(body: &[u8]) -> &str {
2262 let nul = body.iter().position(|&b| b == 0).unwrap_or(body.len());
2263 std::str::from_utf8(&body[..nul]).unwrap()
2264 }
2265
2266 async fn write_startup<W: AsyncWrite + Unpin>(stream: &mut W) {
2267 write_startup_params(stream, &[("user", "reddb")]).await;
2268 }
2269
2270 async fn write_startup_params<W: AsyncWrite + Unpin>(stream: &mut W, params: &[(&str, &str)]) {
2271 let mut payload = Vec::new();
2272 payload.extend_from_slice(&crate::wire::postgres::protocol::PG_PROTOCOL_V3.to_be_bytes());
2273 for (name, value) in params {
2274 push_pg_cstring(&mut payload, name);
2275 push_pg_cstring(&mut payload, value);
2276 }
2277 payload.push(0);
2278 let len = (payload.len() + 4) as u32;
2279 stream.write_all(&len.to_be_bytes()).await.unwrap();
2280 stream.write_all(&payload).await.unwrap();
2281 }
2282
2283 fn cstring_body(value: &str) -> Vec<u8> {
2284 let mut out = Vec::new();
2285 push_pg_cstring(&mut out, value);
2286 out
2287 }
2288
2289 fn error_response_field(body: &[u8], field_tag: u8) -> Option<&str> {
2290 let mut pos = 0;
2291 while pos < body.len() && body[pos] != 0 {
2292 let tag = body[pos];
2293 pos += 1;
2294 let end = body[pos..].iter().position(|&b| b == 0)? + pos;
2295 if tag == field_tag {
2296 return std::str::from_utf8(&body[pos..end]).ok();
2297 }
2298 pos = end + 1;
2299 }
2300 None
2301 }
2302
2303 async fn write_frontend_frame<W: AsyncWrite + Unpin>(
2304 stream: &mut W,
2305 tag: u8,
2306 payload: Vec<u8>,
2307 ) {
2308 stream.write_all(&[tag]).await.unwrap();
2309 stream
2310 .write_all(&((payload.len() + 4) as u32).to_be_bytes())
2311 .await
2312 .unwrap();
2313 stream.write_all(&payload).await.unwrap();
2314 }
2315
2316 async fn read_backend_frame<R: AsyncRead + Unpin>(stream: &mut R) -> (u8, Vec<u8>) {
2317 let mut tag = [0u8; 1];
2318 stream.read_exact(&mut tag).await.unwrap();
2319 let mut len = [0u8; 4];
2320 stream.read_exact(&mut len).await.unwrap();
2321 let len = u32::from_be_bytes(len) as usize;
2322 let mut body = vec![0u8; len - 4];
2323 stream.read_exact(&mut body).await.unwrap();
2324 (tag[0], body)
2325 }
2326
2327 async fn read_until_ready<R: AsyncRead + Unpin>(stream: &mut R) -> Vec<(u8, Vec<u8>)> {
2328 let mut frames = Vec::new();
2329 loop {
2330 let frame = read_backend_frame(stream).await;
2331 let done = frame.0 == b'Z';
2332 frames.push(frame);
2333 if done {
2334 return frames;
2335 }
2336 }
2337 }
2338
2339 fn parse_body(statement: &str, query: &str, oids: &[u32]) -> Vec<u8> {
2340 let mut out = Vec::new();
2341 push_pg_cstring(&mut out, statement);
2342 push_pg_cstring(&mut out, query);
2343 out.extend_from_slice(&(oids.len() as i16).to_be_bytes());
2344 for oid in oids {
2345 out.extend_from_slice(&oid.to_be_bytes());
2346 }
2347 out
2348 }
2349
2350 fn bind_body(
2351 portal: &str,
2352 statement: &str,
2353 formats: &[i16],
2354 params: &[Option<&[u8]>],
2355 result_formats: &[i16],
2356 ) -> Vec<u8> {
2357 let mut out = Vec::new();
2358 push_pg_cstring(&mut out, portal);
2359 push_pg_cstring(&mut out, statement);
2360 out.extend_from_slice(&(formats.len() as i16).to_be_bytes());
2361 for format in formats {
2362 out.extend_from_slice(&format.to_be_bytes());
2363 }
2364 out.extend_from_slice(&(params.len() as i16).to_be_bytes());
2365 for param in params {
2366 match param {
2367 Some(bytes) => {
2368 out.extend_from_slice(&(bytes.len() as i32).to_be_bytes());
2369 out.extend_from_slice(bytes);
2370 }
2371 None => out.extend_from_slice(&(-1i32).to_be_bytes()),
2372 }
2373 }
2374 out.extend_from_slice(&(result_formats.len() as i16).to_be_bytes());
2375 for format in result_formats {
2376 out.extend_from_slice(&format.to_be_bytes());
2377 }
2378 out
2379 }
2380
2381 fn describe_body(target: u8, name: &str) -> Vec<u8> {
2382 let mut out = vec![target];
2383 push_pg_cstring(&mut out, name);
2384 out
2385 }
2386
2387 fn execute_body(portal: &str, max_rows: u32) -> Vec<u8> {
2388 let mut out = Vec::new();
2389 push_pg_cstring(&mut out, portal);
2390 out.extend_from_slice(&max_rows.to_be_bytes());
2391 out
2392 }
2393
2394 fn push_pg_cstring(out: &mut Vec<u8>, value: &str) {
2395 out.extend_from_slice(value.as_bytes());
2396 out.push(0);
2397 }
2398}
2399
2400fn random_cancel_key() -> u32 {
2405 let mut bytes = [0u8; 4];
2406 if crate::crypto::os_random::fill_bytes(&mut bytes).is_ok() {
2407 return u32::from_le_bytes(bytes);
2408 }
2409 let nanos = std::time::SystemTime::now()
2410 .duration_since(std::time::UNIX_EPOCH)
2411 .map(|d| d.subsec_nanos())
2412 .unwrap_or(0);
2413 std::process::id() ^ nanos
2414}