Skip to main content

reddb_server/wire/postgres/
server.rs

1//! PostgreSQL wire-protocol listener (Phase 3.1 PG parity).
2//!
3//! Accepts TCP connections from PG-compatible clients, drives the startup
4//! handshake, and routes simple/extended-query frames into the existing
5//! `RedDBRuntime::execute_query` path. Results are adapted back into PG
6//! `RowDescription` + `DataRow` frames via `types::value_to_pg_wire_bytes`.
7
8use 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/// Startup-tuned configuration for the PG wire listener.
30#[derive(Debug, Clone)]
31pub struct PgWireConfig {
32    /// TCP bind address ("host:port"). The caller is responsible for
33    /// keeping this disjoint from the native wire / gRPC / HTTP listeners.
34    pub bind_addr: String,
35    /// PG version string sent back in `ParameterStatus`. Many drivers
36    /// sniff this to enable/disable features. RedDB advertises a
37    /// recent-enough version to get the broadest client support.
38    pub server_version: String,
39    /// Optional TLS material. When set, an `SSLRequest` is answered with
40    /// `'S'` and the connection is upgraded via a rustls handshake before
41    /// the startup flow continues over the encrypted stream. The cert/key
42    /// surface is the *same* [`WireTlsConfig`] the native RedWire listener
43    /// uses — the two transports share one cert config rather than
44    /// inventing a parallel one. `None` leaves the historical behaviour:
45    /// `SSLRequest` is declined with `'N'` and clients continue in
46    /// plaintext.
47    pub tls: Option<crate::wire::tls::WireTlsConfig>,
48}
49
50/// The authenticated identity resolved during the PG-Wire startup
51/// handshake. Installed as the thread-local execution identity for the
52/// duration of each statement so every query dispatched over this
53/// connection flows through the same IAM/policy gates as native
54/// transports. `None` connections (loopback trust) run under the ambient
55/// embedded identity, preserving the pre-auth behaviour.
56#[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
92/// Run a synchronous, possibly long-parking runtime call without
93/// monopolising the async worker thread.
94///
95/// `QUEUE READ … WAIT <budget>` parks the calling thread on a condvar
96/// for up to the wait budget. Doing that directly on a tokio worker
97/// holds the worker for the whole budget, which serialises every other
98/// connection's query behind the parked waiter — a producer's
99/// `QUEUE PUSH` could not run to notify the waiter, so an enqueue during
100/// a WAIT would never be delivered — and it stalls the runtime's timer
101/// driver. `block_in_place` hands the remaining tasks on this worker to
102/// a sibling worker for the duration while keeping the current OS thread
103/// — and therefore the per-connection thread-locals `execute_query`
104/// relies on (current tenant, connection id, transaction context) —
105/// intact, which `spawn_blocking` would not.
106///
107/// `block_in_place` panics on a current-thread runtime, so we fall back
108/// to a direct call there: a single-threaded runtime has no sibling
109/// worker to hand work to and nothing else to starve (e.g. the
110/// `#[tokio::test]` harnesses that drive one connection at a time).
111fn 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
119/// Run `f` with the authenticated PG-Wire identity installed as the
120/// thread-local execution identity, restoring the prior state afterwards.
121///
122/// The runtime's query dispatch resolves `CURRENT_USER/ROLE/TENANT` and
123/// folds RLS/IAM policies from these thread-locals, so installing them
124/// here is what makes PG-Wire statements transport-invariant with the
125/// native wire. A `None` context leaves the ambient identity untouched.
126fn 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
137/// RAII guard that installs the PG-Wire auth identity + tenant on the
138/// current thread and restores the previous state on drop, so pooled
139/// runtime threads never leak an identity across statements.
140struct 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
169/// Whether `bind_addr` ("host:port") resolves to a loopback host. Trust
170/// authentication survives only for loopback binds with no credentials
171/// configured; every other bind requires a verified password.
172fn is_loopback_bind(bind_addr: &str) -> bool {
173    crate::transport::is_loopback_bind(bind_addr)
174}
175
176/// Extract the cleartext password from a `PasswordMessage` payload (a
177/// single NUL-terminated C-string), returning `None` on malformed input.
178fn 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
185/// Spawn the PG wire listener. Blocks until the listener errors out.
186/// Each connection is handled in its own tokio task.
187pub 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
195/// Serve PostgreSQL wire traffic on a listener bound during node bootstrap.
196pub 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    // Posture: with no auth store enabled, PG-Wire falls back to trust
202    // auth for loopback binds only (see `authenticate_startup`). Surface a
203    // startup warning whenever that fallback is active so an operator never
204    // exposes an unauthenticated listener without noticing.
205    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    // Build the shared TLS acceptor once, up front, so a bad cert/key
219    // fails the whole listener rather than every individual connection.
220    // The acceptor is cheap to clone (an `Arc` internally) and is handed
221    // to each connection so `SSLRequest` can be answered with `'S'` and
222    // the rustls handshake performed before the startup flow resumes.
223    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        // Bound concurrent sessions. Without this the accept loop spawned a
238        // task per connection with no ceiling, so an unauthenticated peer
239        // could exhaust descriptors before presenting a credential.
240        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
261/// Drive one connection's lifetime: SSL negotiation → startup →
262/// authentication → query loop.
263///
264/// The first frame may be `SSLRequest` / `GSSENCRequest` (pre-auth
265/// negotiation) or a plain `Startup`. When a TLS acceptor is present an
266/// `SSLRequest` is answered with `'S'` and the socket is upgraded via a
267/// rustls handshake; the rest of the session (a fresh `Startup`,
268/// authentication, and the query loop) then runs over the encrypted
269/// stream. Without an acceptor — or for `GSSENCRequest`, which we never
270/// support — the request is declined with `'N'` and the client continues
271/// in plaintext on the same socket.
272pub(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    // Bound the whole pre-session negotiation. `read_startup` and the rustls
282    // handshake both block on the peer, and neither had a deadline: a client
283    // that connected and sent nothing (or stalled mid-handshake) pinned a
284    // task and a connection slot indefinitely, before any credential was
285    // presented. The loop below can iterate (SSLRequest/GSSENCRequest are
286    // declined and the client retries), so each read is bounded rather than
287    // the loop as a whole.
288    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                    // 'S' = SSL supported. Perform the rustls handshake and
296                    // hand the encrypted stream to the session driver.
297                    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                // 'N' = not supported — client continues in plaintext and
304                // re-sends a normal Startup on the same socket.
305                write_raw_byte(&mut stream, b'N').await?;
306                continue;
307            }
308            FrontendMessage::GssEncRequest => {
309                // GSSAPI encryption is never offered; decline and let the
310                // client retry with SSLRequest or plaintext Startup.
311                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                // CancelRequest: no response expected; drop the socket.
319                return Ok(());
320            }
321            other => {
322                return Err(PgWireError::Protocol(format!(
323                    "unexpected startup frame: {other:?}"
324                )));
325            }
326        }
327    }
328}
329
330/// Continue the startup flow after a successful TLS upgrade. The client
331/// now sends a fresh `Startup` over the encrypted stream; a second
332/// `SSLRequest` / `GSSENCRequest` cannot be honoured (no nested
333/// encryption) and is declined with `'N'`.
334async 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
361/// Authenticate a `Startup` and drive the query loop over `stream`.
362/// `over_tls` records whether the transport is encrypted, which gates the
363/// cleartext-password posture in [`authenticate_startup`].
364async 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, &params).await? {
376            Some(context) => context,
377            // Authentication failed: the error frame was already sent and
378            // the socket must close.
379            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    // Main query loop.
387    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                // Frames are written immediately; no additional marker is
455                // needed. ReadyForQuery belongs to Sync, not Flush.
456                continue;
457            }
458            FrontendMessage::Sync => {
459                write_frame(
460                    &mut stream,
461                    &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
462                )
463                .await?;
464            }
465            FrontendMessage::PasswordMessage(_) => {
466                // Should only arrive during auth. Ignore post-auth.
467                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                    // Utility statements use the full dispatcher after registry validation.
854                    // Their zero-parameter SQL can be reparsed without losing bindings;
855                    // data queries retain the cached prepared shape on repeated Execute.
856                    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
1053/// Drive the PG-Wire authentication handshake after a `StartupMessage`.
1054///
1055/// Posture:
1056///   * Auth store enabled → request `AuthenticationCleartextPassword`,
1057///     read the client's `PasswordMessage`, and verify it against the same
1058///     auth store the native transports use. Success installs the resolved
1059///     identity; failure emits `28P01` and closes the socket.
1060///   * No auth store enabled + loopback bind → legacy trust: send
1061///     `AuthenticationOk` with no credentials required (`None` context).
1062///   * No auth store enabled + non-loopback bind → refuse with `28P01`;
1063///     an unconfigured listener must never grant unauthenticated access
1064///     off loopback.
1065///
1066/// TLS posture: a cleartext password may only cross an *unencrypted* link
1067/// on a loopback bind. On a non-loopback bind we refuse to even challenge
1068/// for a password unless the connection is already over TLS (`over_tls`),
1069/// so a client using `sslmode=disable` against a remote listener never
1070/// leaks its password in the clear. Loopback and TLS connections are
1071/// challenged as before.
1072///
1073/// Returns `Ok(Some(context))` once the connection is ready for queries
1074/// (`context` is `None` for trust connections), or `Ok(None)` when
1075/// authentication failed and the error frame has already been sent.
1076///
1077/// Failures give the same `28P01` "password authentication failed" for a
1078/// wrong password and a missing user, so there is no user-existence oracle.
1079async 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        // No credentials configured: trust survives for loopback only.
1092        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    // Refuse to challenge for a cleartext password over an unencrypted
1101    // non-loopback link: it would traverse the network in the clear.
1102    // Clients must connect with `sslmode=require` (TLS) to authenticate
1103    // against a remote bind.
1104    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        // Anything other than a PasswordMessage is a protocol violation
1118        // here; treat it as an auth failure rather than leaking detail.
1119        _ => 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    // A missing user and a wrong password both surface as the same error
1128    // class below — no user-existence oracle leaks over the wire.
1129    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    // Authentication has succeeded (or trust applies): signal AuthenticationOk.
1154    write_frame(stream, &BackendMessage::AuthenticationOk).await?;
1155
1156    // Standard ParameterStatus frames. Drivers gate capabilities on these.
1157    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    // BackendKeyData: (pid, secret_key). CancelRequest is not honoured
1181    // yet, but the key is what a future cancel path authenticates on, so
1182    // it is random per connection rather than a constant a client could
1183    // learn once and reuse against every session.
1184    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    // Empty query convention: PG emits EmptyQueryResponse instead of a
1211    // CommandComplete. Some clients (psql `\;`) rely on this.
1212    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            // PG SQLSTATE class 42 covers syntax / binding errors; we use
1256            // 42P01 (undefined_table) and 42601 (syntax_error) when we can
1257            // detect; otherwise fall back to XX000 (internal error).
1258            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        // DDL / DML / config statements: echo the runtime's
1302        // high-level statement tag back. PG format is
1303        // "<CMD> [<OID>] <COUNT>"; we keep the count where
1304        // applicable and fall back to the runtime's message.
1305        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    // RowDescription: derived from the first record's column ordering.
1481    // When `result.columns` is non-empty we honour that order; otherwise
1482    // we synthesise one from the record's field order.
1483    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    // Peek at the first record for per-column type OIDs. When there's no
1492    // data row we fall back to TEXT for every column — clients render
1493    // empty result sets happily.
1494    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
1605/// Best-effort field lookup on a `UnifiedRecord`. The record API lives in
1606/// `storage::query::unified` and today uses `HashMap<String, Value>` under
1607/// the hood — we use `get` if it exists, else fall back to serialised map.
1608fn record_get<'a>(record: &'a UnifiedRecord, key: &str) -> Option<&'a Value> {
1609    record.get(key)
1610}
1611
1612/// Extract column names in iteration order from a single record. When
1613/// the caller didn't supply an explicit `columns` projection we use the
1614/// first record's field ordering as the canonical tuple shape.
1615///
1616/// HashMap iteration order is non-deterministic — for Phase 3.1 we
1617/// accept the shuffle since PG clients receive the ordered header via
1618/// RowDescription and match cells positionally. A stable ordering
1619/// would require keeping an insertion-order index alongside `values`.
1620fn record_field_names(record: &UnifiedRecord) -> Vec<String> {
1621    // `column_names()` merges the columnar scan side-channel with
1622    // the HashMap so scan rows (which populate only columnar) still
1623    // surface their field names in PG wire output.
1624    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
1646/// Heuristically map a runtime error message onto a PG SQLSTATE. Full
1647/// coverage would map every `RedDBError` variant; this is enough for the
1648/// common psql / JDBC paths.
1649fn 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 undefined_table; close enough for collection-not-found.
1653        "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        // Server must challenge for a cleartext password first.
1696        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        // The connection is closed after an auth failure.
1706        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        // The authenticated identity must drive CURRENT_USER/ROLE/TENANT so
1739        // the simple-query path yields real rows (T,D,C,Z), transport-
1740        // invariant with the native wire.
1741        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, &registry, prepared.id, sql, &params)
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        // Posture (issue #1653): a cleartext password may not be challenged
1909        // over an unencrypted non-loopback link. The server must refuse
1910        // *before* sending AuthenticationCleartextPassword, so a driver
1911        // using sslmode=disable against a remote bind never leaks its
1912        // password in the clear. `tls_acceptor = None` + no SSLRequest is
1913        // exactly that plaintext, no-TLS path.
1914        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        // Non-loopback bind: the wildcard address is not a loopback host.
1924        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        // The refusal arrives directly — never an 'R' cleartext challenge.
1937        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        // The connection is closed after the refusal.
1942        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; // table oid
2089            pos += 2; // column attr
2090            let oid = u32::from_be_bytes([body[pos], body[pos + 1], body[pos + 2], body[pos + 3]]);
2091            pos += 4;
2092            pos += 2; // type size
2093            pos += 4; // type mod
2094            pos += 2; // format
2095            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
2257/// A fresh 32-bit cancel key from the OS CSPRNG. Falls back to a
2258/// process/time-derived value only if the CSPRNG is unavailable — the
2259/// key guards nothing yet, so refusing the connection would be the
2260/// wrong trade.
2261fn 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}