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