Skip to main content

powdb_server/
handler.rs

1use crate::metrics::{Metrics, QueryOutcome, SyncOperation, SyncOutcome, SyncRepairLabel};
2use crate::protocol::{
3    ErrorClass, Message, WireParam, WireRetainedUnit, WireSyncRepairAction, WireSyncStatus,
4};
5use powdb_auth::{Permission, Role, UserStore};
6use powdb_query::executor::{is_read_only_statement, Engine, WalDurabilityTicket};
7use powdb_query::parser;
8use powdb_query::result::{QueryError, QueryResult};
9use powdb_query::sql;
10use powdb_storage::types::Value;
11use powdb_sync::{
12    acknowledge_replica_apply, read_identity, read_units_through, replica_sync_status,
13    retained_segments_dir, validate_retained_tail_available, validate_v1_retained_units_applyable,
14    ReplicaSyncStatus, RetainedUnit, SegmentIdentity, SyncRepairAction,
15    RETAINED_SEGMENT_FORMAT_VERSION,
16};
17use std::collections::{HashMap, VecDeque};
18use std::net::IpAddr;
19use std::path::{Path, PathBuf};
20use std::sync::{Arc, Mutex, RwLock};
21use std::time::{Duration, Instant};
22use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
23use tokio::sync::{watch, OwnedSemaphorePermit, Semaphore};
24use tracing::{debug, error, info, warn};
25use zeroize::Zeroizing;
26
27/// Tracks per-IP authentication failure counts for rate limiting.
28pub type AuthRateLimiter = Arc<Mutex<HashMap<IpAddr, (u32, Instant)>>>;
29
30/// Fixed reader-permit pool. Read-only autocommit statements take one permit;
31/// writers, sync operations, and explicit transactions take the entire pool.
32/// Tokio's fair semaphore queue prevents a waiting writer from being starved by
33/// later readers.
34#[derive(Clone)]
35pub struct TxGate {
36    semaphore: Arc<Semaphore>,
37    permit_count: u32,
38}
39
40pub const DEFAULT_TX_GATE_READER_PERMITS: u32 = 1024;
41
42/// Create a transaction gate for a shared engine.
43pub fn new_tx_gate() -> TxGate {
44    new_tx_gate_with_permits(DEFAULT_TX_GATE_READER_PERMITS)
45}
46
47/// Create a transaction gate with an explicit reader capacity.
48///
49/// The configurable constructor exists so benchmark and compatibility tests can
50/// reproduce the former single-permit admission policy using the exact same
51/// handler code. Production uses [`new_tx_gate`].
52pub fn new_tx_gate_with_permits(permit_count: u32) -> TxGate {
53    assert!(permit_count > 0, "transaction gate requires a permit");
54    TxGate {
55        semaphore: Arc::new(Semaphore::new(permit_count as usize)),
56        permit_count,
57    }
58}
59
60impl TxGate {
61    pub fn permit_count(&self) -> u32 {
62        self.permit_count
63    }
64
65    fn available_permits(&self) -> usize {
66        self.semaphore.available_permits()
67    }
68
69    async fn acquire_many_owned(
70        self,
71        permits: u32,
72    ) -> Result<OwnedSemaphorePermit, tokio::sync::AcquireError> {
73        self.semaphore.acquire_many_owned(permits).await
74    }
75
76    fn try_acquire_many_owned(
77        self,
78        permits: u32,
79    ) -> Result<OwnedSemaphorePermit, tokio::sync::TryAcquireError> {
80        self.semaphore.try_acquire_many_owned(permits)
81    }
82}
83
84/// Maximum query text length accepted from the wire (1 MB).
85const MAX_QUERY_LENGTH: usize = 1024 * 1024;
86
87/// Maximum payload accepted by the post-auth cancellation-safe frame reader.
88/// Keep this equal to the protocol reader's public wire limit.
89const MAX_WIRE_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
90
91/// Frames received while a query is executing are retained for normal
92/// pipelined processing. Both limits are deliberately much smaller than the
93/// ordinary 64 MiB per-frame protocol limit: in-flight read-ahead is merely a
94/// liveness aid, not a second request buffer. Reaching either cap cancels the
95/// query and closes the connection; socket monitoring is never disabled.
96const MAX_IN_FLIGHT_READ_AHEAD_FRAMES: usize = 128;
97const MAX_IN_FLIGHT_READ_AHEAD_BYTES: usize = 1024 * 1024;
98
99/// Server-side cap for private sync pull batches.
100const MAX_SYNC_PULL_UNITS: u32 = 4096;
101
102/// Server-side cap for retained-unit payload bytes in a private sync pull.
103const MAX_SYNC_PULL_BYTES: u64 = 16 * 1024 * 1024;
104
105/// Maximum encoded response payload size (64 MB). The wire format is still a
106/// single frame today, so oversized result sets must fail cleanly instead of
107/// building an unbounded `Vec<Vec<String>>` and frame in memory.
108#[cfg(not(test))]
109const MAX_RESPONSE_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
110#[cfg(test)]
111const MAX_RESPONSE_PAYLOAD_SIZE: usize = 1024;
112
113/// Timeout for writing a response to the client. Prevents slow-drain
114/// clients from blocking the handler indefinitely.
115const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
116
117/// Maximum number of auth failures per IP within the rate-limit window.
118const MAX_AUTH_FAILURES: u32 = 5;
119
120/// Window during which auth failures are counted (60 seconds).
121const AUTH_FAILURE_WINDOW: Duration = Duration::from_secs(60);
122
123/// Create a new shared rate limiter.
124pub fn new_rate_limiter() -> AuthRateLimiter {
125    Arc::new(Mutex::new(HashMap::new()))
126}
127
128/// Check whether an IP is rate-limited and record a failure if requested.
129/// Returns `true` if the IP should be rejected.
130fn is_rate_limited(limiter: &AuthRateLimiter, ip: IpAddr) -> bool {
131    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
132    // Clean up stale entries while we have the lock.
133    let now = Instant::now();
134    map.retain(|_, (_, ts)| now.duration_since(*ts) < AUTH_FAILURE_WINDOW);
135
136    if let Some((count, _)) = map.get(&ip) {
137        *count >= MAX_AUTH_FAILURES
138    } else {
139        false
140    }
141}
142
143/// Record an auth failure for the given IP.
144fn record_auth_failure(limiter: &AuthRateLimiter, ip: IpAddr) {
145    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
146    let now = Instant::now();
147    let entry = map.entry(ip).or_insert((0, now));
148    // Reset counter if the window has elapsed.
149    if now.duration_since(entry.1) >= AUTH_FAILURE_WINDOW {
150        *entry = (1, now);
151    } else {
152        entry.0 += 1;
153    }
154}
155
156/// Clear the failure counter on successful auth.
157fn clear_auth_failures(limiter: &AuthRateLimiter, ip: IpAddr) {
158    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
159    map.remove(&ip);
160}
161
162/// Constant-time password comparison. Hashes both inputs to fixed-size
163/// SHA-256 digests so neither length nor content leaks through timing.
164fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
165    use sha2::{Digest, Sha256};
166    let ha = Sha256::digest(a);
167    let hb = Sha256::digest(b);
168    let mut diff = 0u8;
169    for (x, y) in ha.iter().zip(hb.iter()) {
170        diff |= x ^ y;
171    }
172    diff == 0
173}
174
175/// An authenticated connection's identity. Bound at connect time and consulted
176/// on every query by `dispatch_query` to enforce the user's role: a
177/// `readonly` principal may only execute read statements.
178#[derive(Debug, Clone, PartialEq, Eq)]
179pub struct Principal {
180    pub name: String,
181    pub role: String,
182}
183
184/// Whether a parsed statement is data-definition (schema) work: creating,
185/// altering, or dropping a type or view. `explain <ddl>` is classified by its
186/// inner statement so `explain drop User` needs the same permission as
187/// `drop User`. Mutations that change *rows* (insert/update/delete/upsert/
188/// refresh) and transaction control are NOT DDL — they fall under `Write`.
189fn is_ddl_statement(stmt: &powdb_query::ast::Statement) -> bool {
190    use powdb_query::ast::Statement;
191    let inner = match stmt {
192        Statement::Explain(inner) => inner.as_ref(),
193        other => other,
194    };
195    matches!(
196        inner,
197        Statement::CreateType(_)
198            | Statement::AlterTable(_)
199            | Statement::DropTable(_)
200            | Statement::CreateView(_)
201            | Statement::DropView(_)
202    )
203}
204
205/// The capability a parsed statement requires under the RBAC lattice
206/// (`crates/auth/src/role.rs`). Reads need [`Permission::Read`]; schema
207/// definition needs [`Permission::Ddl`]; every other mutation needs
208/// [`Permission::Write`]. [`Permission::Admin`] is reserved for user/role
209/// management, which is CLI-only today and never reaches this wire path.
210fn required_permission(stmt: &powdb_query::ast::Statement) -> Permission {
211    if is_read_only_statement(stmt) {
212        Permission::Read
213    } else if is_ddl_statement(stmt) {
214        Permission::Ddl
215    } else {
216        Permission::Write
217    }
218}
219
220/// Enforce the principal's role against a parsed statement using the full
221/// permission lattice. Reads are always permitted (any authenticated role can
222/// read — unknown role names still read but fail closed on any mutation).
223/// Mutations require the specific capability the statement maps to: row
224/// mutations need `Write`, schema changes need `Ddl`. Unknown role names
225/// resolve to no builtin and therefore grant nothing beyond reads.
226///
227/// Classification uses the parsed AST via
228/// [`powdb_query::executor::is_read_only_statement`] — the exact same
229/// classifier the RwLock read/write split relies on — so the permission
230/// boundary and the concurrency boundary can never disagree.
231fn check_statement_permitted(
232    principal: Option<&Principal>,
233    stmt: &powdb_query::ast::Statement,
234) -> Result<(), QueryError> {
235    let Some(p) = principal else {
236        // No per-user identity (shared-password or open mode): full access,
237        // byte-identical to the pre-RBAC behavior.
238        return Ok(());
239    };
240    // Reads are permitted for every authenticated principal (preserves the
241    // pre-lattice contract that any connected role may run read-only queries).
242    if is_read_only_statement(stmt) {
243        return Ok(());
244    }
245    let needed = required_permission(stmt);
246    if Role::builtin(&p.role).is_some_and(|r| r.allows(needed)) {
247        return Ok(());
248    }
249    let kind = if needed == Permission::Ddl {
250        "schema-definition"
251    } else {
252        "write"
253    };
254    Err(QueryError::Execution(format!(
255        "permission denied: role '{}' cannot execute {kind} statements",
256        p.role
257    )))
258}
259
260/// Result of the connect-time authentication decision.
261#[derive(Debug, Clone, PartialEq, Eq)]
262pub enum AuthOutcome {
263    /// Authenticated. `principal` is `Some` when a named user authenticated via
264    /// the UserStore, and `None` for the legacy shared-password / open paths
265    /// where there is no per-user identity.
266    Authenticated { principal: Option<Principal> },
267    /// Rejected. The caller sends a generic "authentication failed" error and
268    /// records a rate-limit failure — it must not reveal which check failed.
269    Rejected,
270}
271
272/// Pure, exhaustively-testable authentication decision for a CONNECT handshake.
273///
274/// Policy:
275/// - If `users` has at least one user, multi-user auth is in force: a
276///   `username` is required and `users.authenticate(username, password)` must
277///   succeed. Unknown user, wrong password, or a missing username all reject
278///   with an indistinguishable `Rejected` (no user-vs-password leak).
279/// - If `users` is empty, fall back verbatim to the legacy behavior: when
280///   `expected_password` is `Some`, the candidate must match it (constant time);
281///   when `None`, no auth is required (open). The `username` is ignored here so
282///   that a new client talking to a shared-password server still connects.
283pub fn authenticate_connect(
284    users: &UserStore,
285    expected_password: Option<&str>,
286    username: Option<&str>,
287    password: Option<&str>,
288) -> AuthOutcome {
289    if !users.is_empty() {
290        // Multi-user mode: a username is mandatory.
291        let Some(name) = username else {
292            return AuthOutcome::Rejected;
293        };
294        let Some(candidate) = password else {
295            return AuthOutcome::Rejected;
296        };
297        match users.authenticate(name, candidate) {
298            Some(user) => AuthOutcome::Authenticated {
299                principal: Some(Principal {
300                    name: user.name.clone(),
301                    role: user.role.clone(),
302                }),
303            },
304            None => AuthOutcome::Rejected,
305        }
306    } else {
307        // Legacy shared-password fallback (byte-identical to prior behavior).
308        match expected_password {
309            Some(expected) => {
310                if password.is_some_and(|p| constant_time_eq(p.as_bytes(), expected.as_bytes())) {
311                    AuthOutcome::Authenticated { principal: None }
312                } else {
313                    AuthOutcome::Rejected
314                }
315            }
316            None => AuthOutcome::Authenticated { principal: None },
317        }
318    }
319}
320
321/// The sentinel database name clients send when the user selected none. Both
322/// the CLI and the TS client default to this, so it means "no specific
323/// database" and is always accepted — even when the server is pinned to a name.
324const DEFAULT_DB_NAME: &str = "default";
325
326/// Decide whether a CONNECT's requested `db_name` is served by this process.
327///
328/// One server process serves exactly one global database. When it is pinned to
329/// a name (`configured = Some`), a request that *explicitly* names a different
330/// database is rejected so a client can never silently read/write the wrong
331/// store. An empty name or the client default sentinel (`"default"`) means "no
332/// specific database selected" and is always accepted. When unpinned (`None`)
333/// every name is accepted (0.9.x back-compat); the caller warns on a non-default
334/// name so the silent-mismatch footgun is at least visible in the logs.
335fn check_db_name(configured: Option<&str>, requested: &str) -> Result<(), String> {
336    if requested.is_empty() || requested == DEFAULT_DB_NAME {
337        return Ok(());
338    }
339    match configured {
340        None => Ok(()),
341        Some(name) if requested == name => Ok(()),
342        Some(name) => Err(format!(
343            "unknown database '{requested}'; this server serves '{name}'"
344        )),
345    }
346}
347
348/// Error messages that are safe to forward to the client verbatim.
349const SAFE_ERROR_PREFIXES: &[&str] = &[
350    "table not found",
351    // The executor's actual phrasing is `table 'X' not found`, which the
352    // bare prefix above never matches — keep both so the real message
353    // reaches clients.
354    "table '",
355    "type '",
356    "column not found",
357    // Lexer diagnostics (`at position N: unterminated quoted identifier`)
358    // are derived purely from the client's own query text.
359    "at position",
360    "parse error",
361    "type mismatch",
362    "unknown table",
363    "unknown column",
364    "unknown function",
365    "syntax error",
366    "expected",
367    "unexpected",
368    "missing",
369    "duplicate",
370    "invalid",
371    "cannot",
372    "no such",
373    "already exists",
374    "permission denied",
375    "row too large",
376    "unique constraint violation",
377    // Resource-limit errors carry actionable guidance (e.g. "add a LIMIT
378    // clause") and leak no internal state, so surface them verbatim instead
379    // of masking them to the generic message. See QueryError::{SortLimit,
380    // JoinLimit,MemoryLimit}Exceeded in crates/query/src/result.rs.
381    "sort input exceeds",
382    "join result exceeds",
383    "query exceeded memory budget",
384    "result too large",
385    // A failed covering fsync means the statement executed in memory but was
386    // never made durable — the client MUST be able to distinguish this from
387    // an ordinary failed query (the statement may still be visible until the
388    // server restarts). The io::Error detail leaks no internal state.
389    "wal durability sync failed",
390    // Cooperative query cancellation. Both messages are derived purely from
391    // the configured timeout / a client disconnect and leak no internal state.
392    // See QueryError::{Timeout,Cancelled} in crates/query/src/result.rs.
393    "query timeout after",
394    "query cancelled",
395    // Read-only snapshot serving refuses mutations (and reads needing a writer,
396    // e.g. a stale materialized view) with an operator-facing message that names
397    // the mode and the fix. It leaks no internal state.
398    // See QueryError::ReadonlyMode in crates/query/src/result.rs.
399    "readonly mode",
400];
401
402/// Build the client-facing error frame: sanitized message plus the stable
403/// 1-byte [`ErrorClass`]. The class is orthogonal to the message text: it is
404/// derived from the typed error (or the call site), never from the message,
405/// so sanitization to a generic string does not degrade it.
406fn error_response(message: impl Into<String>, class: ErrorClass) -> Message {
407    Message::ErrorWithClass {
408        message: message.into(),
409        class,
410    }
411}
412
413/// Map a [`QueryError`] to its stable wire [`ErrorClass`].
414///
415/// [`QueryError::ReadonlyNeedsWrite`] is an internal retry sentinel the
416/// server intercepts before Display; if it ever reaches classification it is
417/// reported as [`ErrorClass::Internal`], matching the generic message the
418/// caller sends for that path.
419fn classify_query_error(e: &QueryError) -> ErrorClass {
420    match e {
421        QueryError::Parse(_) => ErrorClass::Parse,
422        QueryError::Timeout { .. } => ErrorClass::Timeout,
423        QueryError::Cancelled => ErrorClass::Cancelled,
424        QueryError::ReadonlyMode => ErrorClass::ReadonlyRefused,
425        QueryError::ReadonlyNeedsWrite => ErrorClass::Internal,
426        QueryError::JoinLimitExceeded
427        | QueryError::NestedLoopPairLimitExceeded { .. }
428        | QueryError::SortLimitExceeded
429        | QueryError::MemoryLimitExceeded { .. } => ErrorClass::LimitExceeded,
430        QueryError::TableNotFound(_)
431        | QueryError::ColumnNotFound { .. }
432        | QueryError::TypeError(_)
433        | QueryError::IndexError(_)
434        | QueryError::ViewError(_) => ErrorClass::Execution,
435        QueryError::StorageError(_) => ErrorClass::Internal,
436        QueryError::Execution(msg) => {
437            if msg.starts_with("unique constraint violation") {
438                ErrorClass::ConstraintViolation
439            } else if msg.starts_with("result too large") {
440                ErrorClass::LimitExceeded
441            } else {
442                ErrorClass::Execution
443            }
444        }
445    }
446}
447
448/// Sanitize an error message before sending it to the client.
449/// Known safe errors are passed through; everything else is replaced
450/// with a generic message to avoid leaking internal details.
451fn sanitize_error(e: &str) -> String {
452    let lower = e.to_lowercase();
453    for prefix in SAFE_ERROR_PREFIXES {
454        if lower.starts_with(prefix) {
455            return e.to_string();
456        }
457    }
458    "query execution error".into()
459}
460
461/// Write a message to the client with a timeout. Returns false if the
462/// write failed or timed out (caller should close the connection).
463async fn write_msg<W: AsyncWrite + Unpin>(writer: &mut BufWriter<W>, msg: &Message) -> bool {
464    let write_fut = async {
465        if msg.write_to(writer).await.is_err() {
466            return false;
467        }
468        writer.flush().await.is_ok()
469    };
470    tokio::time::timeout(WRITE_TIMEOUT, write_fut)
471        .await
472        .unwrap_or_default()
473}
474
475/// Options for a single connection, bundled to keep `handle_connection`'s
476/// argument list short.
477pub struct ConnOpts<'a> {
478    pub engine: Arc<RwLock<Engine>>,
479    pub tx_gate: TxGate,
480    /// Expected client password. Wrapped in `Zeroizing` so the secret is wiped
481    /// from memory on drop (defends against leaking via a core dump).
482    pub expected_password: Option<Zeroizing<String>>,
483    /// Multi-user store loaded from the data dir at startup. When it has users,
484    /// the handshake authenticates `(username, password)` against it; when empty
485    /// the server falls back to `expected_password`. Shared across connections.
486    pub users: Arc<UserStore>,
487    pub shutdown_rx: &'a mut watch::Receiver<bool>,
488    pub idle_timeout: Duration,
489    pub query_timeout: Duration,
490    pub rate_limiter: Option<&'a AuthRateLimiter>,
491    pub peer_addr: Option<std::net::SocketAddr>,
492    /// Shared server metrics. Always present; tests pass `Arc::new(Metrics::new())`.
493    pub metrics: Arc<Metrics>,
494    /// How long an explicit `begin` waits to acquire the transaction gate while
495    /// another connection holds an open explicit transaction, before giving up
496    /// with a clear timeout error instead of queueing indefinitely.
497    pub tx_wait_timeout: Duration,
498    /// When `Some`, the single database name this server serves. A CONNECT that
499    /// explicitly names a *different* database is rejected at connect time.
500    /// `None` accepts any name (0.9.x behavior) and only warns.
501    pub db_name: Option<String>,
502}
503
504/// Execute a query against the engine under the RwLock. Read-only
505/// statements acquire `.read()` so concurrent SELECTs can scan in
506/// parallel; mutations acquire `.write()`.
507///
508/// When `principal` is `Some`, the user's role is enforced first: a role
509/// without the `Write` permission (i.e. `readonly`) gets a clean
510/// "permission denied" error for any non-read statement, before any lock
511/// is taken or any engine state is touched.
512/// A statement's execution result plus its (not-yet-waited) WAL durability
513/// obligation. The ticket is settled by [`finalize_durability`] AFTER the
514/// TxGate permit is dropped, so overlapping committers on other connections
515/// can share a single fsync (group commit).
516type DispatchOutcome = (Result<QueryResult, QueryError>, Option<WalDurabilityTicket>);
517
518/// Execute a mutating statement under the engine write lock with WAL group
519/// commit: the commit's fsync obligation is registered (not performed) while
520/// the lock is held, then the lock is dropped and the un-waited ticket is
521/// returned. The caller must settle the ticket (via [`finalize_durability`])
522/// before acknowledging the statement — and must do so with the TxGate
523/// permit already released, or committers can never overlap.
524fn execute_write_deferred(
525    engine: &Arc<RwLock<Engine>>,
526    f: impl FnOnce(&mut Engine) -> Result<QueryResult, QueryError>,
527) -> DispatchOutcome {
528    let mut eng = match engine.write() {
529        Ok(eng) => eng,
530        Err(e) => {
531            return (
532                Err(QueryError::Execution(format!("lock poisoned: {e}"))),
533                None,
534            )
535        }
536    };
537    eng.run_with_deferred_durability(f)
538}
539
540fn dispatch_query(
541    engine: &Arc<RwLock<Engine>>,
542    query: &str,
543    principal: Option<&Principal>,
544    allow_readonly_escalation: bool,
545) -> DispatchOutcome {
546    let stmt_result = parser::parse(query).map_err(|e| e.to_string());
547    dispatch_query_parsed(
548        engine,
549        query,
550        &stmt_result,
551        principal,
552        allow_readonly_escalation,
553    )
554}
555
556fn dispatch_query_parsed(
557    engine: &Arc<RwLock<Engine>>,
558    query: &str,
559    stmt_result: &Result<powdb_query::ast::Statement, String>,
560    principal: Option<&Principal>,
561    allow_readonly_escalation: bool,
562) -> DispatchOutcome {
563    // Role enforcement happens on the parsed AST. Statements that fail to
564    // parse fall through — the engine returns the parse error itself and
565    // can never execute anything for them.
566    if let Ok(stmt) = &stmt_result {
567        if let Err(e) = check_statement_permitted(principal, stmt) {
568            return (Err(e), None);
569        }
570    }
571
572    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
573    if can_try_read {
574        let res = {
575            let eng = match engine.read() {
576                Ok(eng) => eng,
577                Err(e) => {
578                    return (
579                        Err(QueryError::Execution(format!("lock poisoned: {e}"))),
580                        None,
581                    )
582                }
583            };
584            eng.execute_powql_readonly(query)
585        };
586        match res {
587            Ok(r) => return (Ok(r), None),
588            Err(QueryError::ReadonlyNeedsWrite) => {
589                if !allow_readonly_escalation {
590                    return (Err(QueryError::ReadonlyNeedsWrite), None);
591                }
592                // The caller already owns writer admission, so it is safe to
593                // fall through and retry under the engine write lock.
594            }
595            Err(e) => return (Err(e), None),
596        }
597    }
598
599    if matches!(
600        parsed_transaction_control(stmt_result),
601        Some(TransactionControl::Rollback)
602    ) {
603        let mut eng = match engine.write() {
604            Ok(eng) => eng,
605            Err(e) => {
606                return (
607                    Err(QueryError::Execution(format!("lock poisoned: {e}"))),
608                    None,
609                )
610            }
611        };
612        return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
613    }
614    execute_write_deferred(engine, |eng| eng.execute_powql(query))
615}
616
617#[cfg(test)]
618fn dispatch_sql_query(
619    engine: &Arc<RwLock<Engine>>,
620    query: &str,
621    principal: Option<&Principal>,
622    allow_readonly_escalation: bool,
623) -> DispatchOutcome {
624    let stmt_result = sql::parse_sql(query).map_err(|e| e.to_string());
625    dispatch_sql_query_parsed(
626        engine,
627        query,
628        &stmt_result,
629        principal,
630        allow_readonly_escalation,
631    )
632}
633
634fn dispatch_sql_query_parsed(
635    engine: &Arc<RwLock<Engine>>,
636    query: &str,
637    stmt_result: &Result<powdb_query::ast::Statement, String>,
638    principal: Option<&Principal>,
639    allow_readonly_escalation: bool,
640) -> DispatchOutcome {
641    if let Ok(stmt) = &stmt_result {
642        if let Err(e) = check_statement_permitted(principal, stmt) {
643            return (Err(e), None);
644        }
645    }
646
647    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
648    if can_try_read {
649        let res = {
650            let eng = match engine.read() {
651                Ok(eng) => eng,
652                Err(e) => {
653                    return (
654                        Err(QueryError::Execution(format!("lock poisoned: {e}"))),
655                        None,
656                    )
657                }
658            };
659            eng.execute_sql_readonly(query)
660        };
661        match res {
662            Ok(r) => return (Ok(r), None),
663            Err(QueryError::ReadonlyNeedsWrite) => {
664                if !allow_readonly_escalation {
665                    return (Err(QueryError::ReadonlyNeedsWrite), None);
666                }
667            }
668            Err(e) => return (Err(e), None),
669        }
670    }
671
672    if matches!(
673        parsed_transaction_control(stmt_result),
674        Some(TransactionControl::Rollback)
675    ) {
676        let mut eng = match engine.write() {
677            Ok(eng) => eng,
678            Err(e) => {
679                return (
680                    Err(QueryError::Execution(format!("lock poisoned: {e}"))),
681                    None,
682                )
683            }
684        };
685        return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
686    }
687    execute_write_deferred(engine, |eng| eng.execute_sql(query))
688}
689
690/// Convert a wire parameter into the query-crate [`ParamValue`] used for
691/// token-level binding.
692fn wire_param_to_value(p: &WireParam) -> powdb_query::ast::ParamValue {
693    use powdb_query::ast::ParamValue;
694    match p {
695        WireParam::Null => ParamValue::Null,
696        WireParam::Int(v) => ParamValue::Int(*v),
697        WireParam::Float(v) => ParamValue::Float(*v),
698        WireParam::Bool(v) => ParamValue::Bool(*v),
699        WireParam::Str(s) => ParamValue::Str(s.clone()),
700    }
701}
702
703/// Parameterized counterpart of [`dispatch_query`]. Routes through the exact
704/// same role-enforcement and read/write escalation logic, but binds the
705/// `$N` placeholders at the token level via the query crate's
706/// `parse_with_params` path. A string parameter can never change the query's
707/// shape — it is substituted as a literal token, not interpolated text.
708#[derive(Debug, Clone, Copy, PartialEq, Eq)]
709enum TransactionControl {
710    Begin,
711    Commit,
712    Rollback,
713}
714
715#[derive(Debug, Clone, Copy, PartialEq, Eq)]
716enum AdmissionMode {
717    Reader,
718    Writer,
719}
720
721#[derive(Clone, Copy, Debug, PartialEq, Eq)]
722enum WireResultMode {
723    LegacyText,
724    Native,
725}
726
727fn statement_admission(stmt: &powdb_query::ast::Statement) -> AdmissionMode {
728    if is_read_only_statement(stmt) {
729        AdmissionMode::Reader
730    } else {
731        AdmissionMode::Writer
732    }
733}
734
735/// The terminal error a read-only (snapshot-serving) server returns for a
736/// mutating statement, transaction-control frame, or a read that needs a writer
737/// (a stale materialized view). The connection stays usable afterward.
738fn readonly_terminal_message() -> Message {
739    error_response(
740        sanitize_error(&QueryError::ReadonlyMode.to_string()),
741        ErrorClass::ReadonlyRefused,
742    )
743}
744
745#[cfg(test)]
746fn classify_query_admission(query: &str) -> AdmissionMode {
747    parser::parse(query)
748        .map(|stmt| statement_admission(&stmt))
749        .unwrap_or(AdmissionMode::Writer)
750}
751
752#[cfg(test)]
753fn classify_sql_admission(query: &str) -> AdmissionMode {
754    sql::parse_sql(query)
755        .map(|stmt| statement_admission(&stmt))
756        .unwrap_or(AdmissionMode::Writer)
757}
758
759#[cfg(test)]
760fn classify_params_admission(query: &str, params: &[WireParam]) -> AdmissionMode {
761    let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
762    parser::parse_with_params(query, &bound)
763        .map(|stmt| statement_admission(&stmt))
764        .unwrap_or(AdmissionMode::Writer)
765}
766
767fn transaction_control(stmt: &powdb_query::ast::Statement) -> Option<TransactionControl> {
768    use powdb_query::ast::Statement;
769    match stmt {
770        Statement::Begin => Some(TransactionControl::Begin),
771        Statement::Commit => Some(TransactionControl::Commit),
772        Statement::Rollback => Some(TransactionControl::Rollback),
773        _ => None,
774    }
775}
776
777fn parsed_transaction_control(
778    stmt_result: &Result<powdb_query::ast::Statement, String>,
779) -> Option<TransactionControl> {
780    stmt_result.as_ref().ok().and_then(transaction_control)
781}
782
783fn execute_rollback_preserving_sync_if_needed(
784    engine: &mut Engine,
785) -> Result<QueryResult, QueryError> {
786    engine.rollback_transaction_preserving_wal_archive()
787}
788
789#[cfg(test)]
790fn dispatch_query_with_params(
791    engine: &Arc<RwLock<Engine>>,
792    query: &str,
793    params: &[WireParam],
794    principal: Option<&Principal>,
795    allow_readonly_escalation: bool,
796) -> DispatchOutcome {
797    let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
798
799    // Parse once (with params bound) so role enforcement and read/write
800    // classification see exactly the statement that will execute.
801    let stmt_result = parser::parse_with_params(query, &bound).map_err(|e| e.to_string());
802    dispatch_query_with_bound_params_parsed(
803        engine,
804        query,
805        &bound,
806        &stmt_result,
807        principal,
808        allow_readonly_escalation,
809    )
810}
811
812fn dispatch_query_with_bound_params_parsed(
813    engine: &Arc<RwLock<Engine>>,
814    query: &str,
815    bound: &[powdb_query::ast::ParamValue],
816    stmt_result: &Result<powdb_query::ast::Statement, String>,
817    principal: Option<&Principal>,
818    allow_readonly_escalation: bool,
819) -> DispatchOutcome {
820    if let Ok(stmt) = &stmt_result {
821        if let Err(e) = check_statement_permitted(principal, stmt) {
822            return (Err(e), None);
823        }
824    }
825
826    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
827    if can_try_read {
828        let res = {
829            let eng = match engine.read() {
830                Ok(eng) => eng,
831                Err(e) => {
832                    return (
833                        Err(QueryError::Execution(format!("lock poisoned: {e}"))),
834                        None,
835                    )
836                }
837            };
838            eng.execute_powql_readonly_with_params(query, bound)
839        };
840        match res {
841            Ok(r) => return (Ok(r), None),
842            Err(QueryError::ReadonlyNeedsWrite) => {
843                if !allow_readonly_escalation {
844                    return (Err(QueryError::ReadonlyNeedsWrite), None);
845                }
846            }
847            Err(e) => return (Err(e), None),
848        }
849    }
850
851    if matches!(
852        parsed_transaction_control(stmt_result),
853        Some(TransactionControl::Rollback)
854    ) {
855        let mut eng = match engine.write() {
856            Ok(eng) => eng,
857            Err(e) => {
858                return (
859                    Err(QueryError::Execution(format!("lock poisoned: {e}"))),
860                    None,
861                )
862            }
863        };
864        return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
865    }
866    execute_write_deferred(engine, |eng| eng.execute_powql_with_params(query, bound))
867}
868
869#[derive(Debug, Clone)]
870struct SyncPullRequest {
871    replica_id: String,
872    since_lsn: u64,
873    max_units: u32,
874    max_bytes: u64,
875    database_id: [u8; 16],
876    primary_generation: u64,
877    wal_format_version: u16,
878    catalog_version: u16,
879    segment_format_version: u16,
880}
881
882#[derive(Debug, Clone, Copy, PartialEq, Eq)]
883enum SyncErrorClass {
884    AuthRequired,
885    PermissionDenied,
886    InvalidReplicaId,
887    ActiveTransaction,
888    QueryExecution,
889    InvalidMaxUnits,
890    InvalidMaxBytes,
891    SyncContext,
892    StatusRead,
893    CursorLsnMismatch,
894    IdentityRead,
895    IdentityOrFormatMismatch,
896    RetainedRead,
897    RetainedUnitEncoding,
898    RetainedChunkNotApplyable,
899    LsnAheadOfRemote,
900    AckValidation,
901    AckUpdate,
902    Internal,
903}
904
905impl SyncErrorClass {
906    const fn as_label(self) -> &'static str {
907        match self {
908            Self::AuthRequired => "auth_required",
909            Self::PermissionDenied => "permission_denied",
910            Self::InvalidReplicaId => "invalid_replica_id",
911            Self::ActiveTransaction => "active_transaction",
912            Self::QueryExecution => "query_execution",
913            Self::InvalidMaxUnits => "invalid_max_units",
914            Self::InvalidMaxBytes => "invalid_max_bytes",
915            Self::SyncContext => "sync_context",
916            Self::StatusRead => "status_read",
917            Self::CursorLsnMismatch => "cursor_lsn_mismatch",
918            Self::IdentityRead => "identity_read",
919            Self::IdentityOrFormatMismatch => "identity_or_format_mismatch",
920            Self::RetainedRead => "retained_read",
921            Self::RetainedUnitEncoding => "retained_unit_encoding",
922            Self::RetainedChunkNotApplyable => "retained_chunk_not_applyable",
923            Self::LsnAheadOfRemote => "lsn_ahead_of_remote",
924            Self::AckValidation => "ack_validation",
925            Self::AckUpdate => "ack_update",
926            Self::Internal => "internal",
927        }
928    }
929}
930
931#[derive(Debug, Clone)]
932struct SyncDecision {
933    message: Message,
934    error_class: Option<SyncErrorClass>,
935}
936
937impl SyncDecision {
938    fn ok(message: Message) -> Self {
939        Self {
940            message,
941            error_class: None,
942        }
943    }
944
945    fn error(class: SyncErrorClass, message: impl Into<String>) -> Self {
946        Self {
947            message: Message::Error {
948                message: message.into(),
949            },
950            error_class: Some(class),
951        }
952    }
953}
954
955fn check_sync_protocol_permitted(
956    credential_authenticated: bool,
957    principal: Option<&Principal>,
958) -> Result<(), (SyncErrorClass, String)> {
959    if !credential_authenticated {
960        return Err((
961            SyncErrorClass::AuthRequired,
962            "sync protocol requires authentication".to_string(),
963        ));
964    }
965    if let Some(principal) = principal {
966        let allowed =
967            Role::builtin(&principal.role).is_some_and(|role| role.allows(Permission::Write));
968        if !allowed {
969            return Err((
970                SyncErrorClass::PermissionDenied,
971                format!(
972                    "permission denied: role '{}' cannot use sync protocol",
973                    principal.role
974                ),
975            ));
976        }
977    }
978    Ok(())
979}
980
981fn validate_wire_replica_id(replica_id: &str) -> Result<(), String> {
982    if replica_id.is_empty() {
983        return Err("replica id must be non-empty".to_string());
984    }
985    if replica_id.len() > 128 {
986        return Err("replica id must be at most 128 bytes".to_string());
987    }
988    if !replica_id
989        .bytes()
990        .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b':'))
991    {
992        return Err("replica id contains unsupported characters".to_string());
993    }
994    Ok(())
995}
996
997fn sync_context(engine: &Arc<RwLock<Engine>>) -> Result<SyncContext, String> {
998    let engine = engine
999        .read()
1000        .map_err(|e| format!("lock poisoned while reading sync context: {e}"))?;
1001    let catalog = engine.catalog();
1002    Ok(SyncContext {
1003        data_dir: catalog.data_dir().to_path_buf(),
1004        remote_lsn: catalog.max_lsn(),
1005        active_catalog_version: catalog.active_catalog_version(),
1006    })
1007}
1008
1009/// Snapshot of the primary's sync-relevant state captured under the engine lock.
1010/// `active_catalog_version` is the database's *active* on-disk catalog format,
1011/// used to stamp the expected segment identity a pulling replica is checked
1012/// against (a database that never activated v6 keeps expecting v5).
1013struct SyncContext {
1014    data_dir: PathBuf,
1015    remote_lsn: u64,
1016    active_catalog_version: u16,
1017}
1018
1019fn wire_repair_action(action: SyncRepairAction) -> WireSyncRepairAction {
1020    match action {
1021        SyncRepairAction::None => WireSyncRepairAction::None,
1022        SyncRepairAction::Pull => WireSyncRepairAction::Pull,
1023        SyncRepairAction::AwaitArchive => WireSyncRepairAction::AwaitArchive,
1024        SyncRepairAction::Rebootstrap => WireSyncRepairAction::Rebootstrap,
1025    }
1026}
1027
1028fn wire_sync_status(status: ReplicaSyncStatus) -> WireSyncStatus {
1029    WireSyncStatus {
1030        replica_id: status.replica_id,
1031        active: status.active,
1032        last_applied_lsn: status.last_applied_lsn,
1033        remote_lsn: status.remote_lsn,
1034        servable_lsn: status.servable_lsn,
1035        unarchived_lsn: status.unarchived_lsn,
1036        lag_lsn: status.lag_lsn,
1037        lag_bytes: status.lag_bytes,
1038        lag_ms: status.lag_ms,
1039        stale: status.stale,
1040        repair_action: wire_repair_action(status.repair_action),
1041        last_sync_error: status.last_sync_error,
1042    }
1043}
1044
1045fn wire_retained_unit(unit: RetainedUnit) -> WireRetainedUnit {
1046    WireRetainedUnit {
1047        tx_id: unit.tx_id,
1048        record_type: unit.record_type,
1049        lsn: unit.lsn,
1050        data: unit.data,
1051    }
1052}
1053
1054fn sync_operation_outcome(message: &Message) -> SyncOutcome {
1055    match message {
1056        Message::SyncStatusResult { .. }
1057        | Message::SyncPullResult { .. }
1058        | Message::SyncAckResult { .. } => SyncOutcome::Ok,
1059        _ => SyncOutcome::Error,
1060    }
1061}
1062
1063fn sync_operation_label(operation: SyncOperation) -> &'static str {
1064    match operation {
1065        SyncOperation::Status => "status",
1066        SyncOperation::Pull => "pull",
1067        SyncOperation::Ack => "ack",
1068    }
1069}
1070
1071fn sync_pull_payload_bytes(units: &[WireRetainedUnit]) -> u64 {
1072    units.iter().fold(0, |total, unit| {
1073        total.saturating_add(unit.encoded_len().unwrap_or(0))
1074    })
1075}
1076
1077fn sync_repair_label(action: WireSyncRepairAction) -> SyncRepairLabel {
1078    match action {
1079        WireSyncRepairAction::None => SyncRepairLabel::None,
1080        WireSyncRepairAction::Pull => SyncRepairLabel::Pull,
1081        WireSyncRepairAction::AwaitArchive => SyncRepairLabel::AwaitArchive,
1082        WireSyncRepairAction::Rebootstrap => SyncRepairLabel::Rebootstrap,
1083    }
1084}
1085
1086fn wire_repair_action_label(action: WireSyncRepairAction) -> &'static str {
1087    match action {
1088        WireSyncRepairAction::None => "none",
1089        WireSyncRepairAction::Pull => "pull",
1090        WireSyncRepairAction::AwaitArchive => "await_archive",
1091        WireSyncRepairAction::Rebootstrap => "rebootstrap",
1092    }
1093}
1094
1095const FNV1A64_OFFSET: u64 = 0xcbf29ce484222325;
1096const FNV1A64_PRIME: u64 = 0x100000001b3;
1097const INVALID_REPLICA_FINGERPRINT: &str = "invalid";
1098
1099fn replica_fingerprint(replica_id: &str) -> String {
1100    let mut hash = FNV1A64_OFFSET;
1101    for byte in replica_id.as_bytes() {
1102        hash ^= u64::from(*byte);
1103        hash = hash.wrapping_mul(FNV1A64_PRIME);
1104    }
1105    format!("{hash:016x}")
1106}
1107
1108fn log_replica_fingerprint(replica_id: &str) -> String {
1109    if validate_wire_replica_id(replica_id).is_ok() {
1110        replica_fingerprint(replica_id)
1111    } else {
1112        INVALID_REPLICA_FINGERPRINT.to_string()
1113    }
1114}
1115
1116#[derive(Debug, Clone)]
1117struct SyncLogContext {
1118    replica_fingerprint: String,
1119    since_lsn: Option<u64>,
1120    applied_lsn: Option<u64>,
1121    observed_remote_lsn: Option<u64>,
1122    max_units: Option<u32>,
1123    max_bytes: Option<u64>,
1124}
1125
1126impl SyncLogContext {
1127    fn status(replica_id: &str) -> Self {
1128        Self::base(replica_id)
1129    }
1130
1131    fn pull(request: &SyncPullRequest) -> Self {
1132        Self {
1133            replica_fingerprint: log_replica_fingerprint(&request.replica_id),
1134            since_lsn: Some(request.since_lsn),
1135            applied_lsn: None,
1136            observed_remote_lsn: None,
1137            max_units: Some(request.max_units),
1138            max_bytes: Some(request.max_bytes),
1139        }
1140    }
1141
1142    fn ack(replica_id: &str, applied_lsn: u64, observed_remote_lsn: u64) -> Self {
1143        Self {
1144            replica_fingerprint: log_replica_fingerprint(replica_id),
1145            since_lsn: None,
1146            applied_lsn: Some(applied_lsn),
1147            observed_remote_lsn: Some(observed_remote_lsn),
1148            max_units: None,
1149            max_bytes: None,
1150        }
1151    }
1152
1153    fn base(replica_id: &str) -> Self {
1154        Self {
1155            replica_fingerprint: log_replica_fingerprint(replica_id),
1156            since_lsn: None,
1157            applied_lsn: None,
1158            observed_remote_lsn: None,
1159            max_units: None,
1160            max_bytes: None,
1161        }
1162    }
1163}
1164
1165struct SyncExecutionContext<'a> {
1166    tx_gate: TxGate,
1167    connection_has_transaction: bool,
1168    operation: SyncOperation,
1169    log_context: SyncLogContext,
1170    metrics: &'a Arc<Metrics>,
1171    query_timeout: Duration,
1172}
1173
1174fn log_sync_decision(
1175    operation: SyncOperation,
1176    context: &SyncLogContext,
1177    elapsed: Duration,
1178    decision: &SyncDecision,
1179) {
1180    let operation = sync_operation_label(operation);
1181    let elapsed_ms = elapsed.as_secs_f64() * 1000.0;
1182    match &decision.message {
1183        Message::SyncStatusResult { status } => {
1184            let repair_action = wire_repair_action_label(status.repair_action);
1185            if status.stale || status.repair_action != WireSyncRepairAction::None {
1186                info!(
1187                    operation = operation,
1188                    replica_fingerprint = %context.replica_fingerprint,
1189                    remote_lsn = status.remote_lsn,
1190                    last_applied_lsn = ?status.last_applied_lsn,
1191                    servable_lsn = ?status.servable_lsn,
1192                    unarchived_lsn = ?status.unarchived_lsn,
1193                    lag_lsn = ?status.lag_lsn,
1194                    lag_bytes = ?status.lag_bytes,
1195                    lag_ms = ?status.lag_ms,
1196                    stale = status.stale,
1197                    repair_action,
1198                    elapsed_ms,
1199                    "sync decision"
1200                );
1201            } else {
1202                debug!(
1203                    operation = operation,
1204                    replica_fingerprint = %context.replica_fingerprint,
1205                    remote_lsn = status.remote_lsn,
1206                    last_applied_lsn = ?status.last_applied_lsn,
1207                    repair_action,
1208                    elapsed_ms,
1209                    "sync decision"
1210                );
1211            }
1212        }
1213        Message::SyncPullResult {
1214            status,
1215            units,
1216            has_more,
1217        } => {
1218            let repair_action = wire_repair_action_label(status.repair_action);
1219            let units_len = units.len();
1220            let payload_bytes = sync_pull_payload_bytes(units);
1221            if *has_more || status.stale || status.repair_action != WireSyncRepairAction::None {
1222                info!(
1223                    operation = operation,
1224                    replica_fingerprint = %context.replica_fingerprint,
1225                    since_lsn = ?context.since_lsn,
1226                    max_units = ?context.max_units,
1227                    max_bytes = ?context.max_bytes,
1228                    units = units_len,
1229                    payload_bytes,
1230                    has_more = *has_more,
1231                    remote_lsn = status.remote_lsn,
1232                    last_applied_lsn = ?status.last_applied_lsn,
1233                    servable_lsn = ?status.servable_lsn,
1234                    unarchived_lsn = ?status.unarchived_lsn,
1235                    lag_lsn = ?status.lag_lsn,
1236                    stale = status.stale,
1237                    repair_action,
1238                    elapsed_ms,
1239                    "sync decision"
1240                );
1241            } else {
1242                debug!(
1243                    operation = operation,
1244                    replica_fingerprint = %context.replica_fingerprint,
1245                    since_lsn = ?context.since_lsn,
1246                    units = units_len,
1247                    payload_bytes,
1248                    has_more = *has_more,
1249                    remote_lsn = status.remote_lsn,
1250                    repair_action,
1251                    elapsed_ms,
1252                    "sync decision"
1253                );
1254            }
1255        }
1256        Message::SyncAckResult {
1257            previous_applied_lsn,
1258            applied_lsn,
1259            remote_lsn,
1260            advanced,
1261            status,
1262        } => {
1263            let repair_action = wire_repair_action_label(status.repair_action);
1264            if *advanced || status.stale || status.repair_action != WireSyncRepairAction::None {
1265                info!(
1266                    operation = operation,
1267                    replica_fingerprint = %context.replica_fingerprint,
1268                    requested_applied_lsn = ?context.applied_lsn,
1269                    observed_remote_lsn = ?context.observed_remote_lsn,
1270                    previous_applied_lsn = *previous_applied_lsn,
1271                    applied_lsn = *applied_lsn,
1272                    remote_lsn = *remote_lsn,
1273                    advanced = *advanced,
1274                    stale = status.stale,
1275                    repair_action,
1276                    elapsed_ms,
1277                    "sync decision"
1278                );
1279            } else {
1280                debug!(
1281                    operation = operation,
1282                    replica_fingerprint = %context.replica_fingerprint,
1283                    requested_applied_lsn = ?context.applied_lsn,
1284                    previous_applied_lsn = *previous_applied_lsn,
1285                    applied_lsn = *applied_lsn,
1286                    remote_lsn = *remote_lsn,
1287                    advanced = *advanced,
1288                    repair_action,
1289                    elapsed_ms,
1290                    "sync decision"
1291                );
1292            }
1293        }
1294        Message::Error { .. } => {
1295            warn!(
1296                operation = operation,
1297                replica_fingerprint = %context.replica_fingerprint,
1298                since_lsn = ?context.since_lsn,
1299                applied_lsn = ?context.applied_lsn,
1300                observed_remote_lsn = ?context.observed_remote_lsn,
1301                max_units = ?context.max_units,
1302                max_bytes = ?context.max_bytes,
1303                error_class = decision
1304                    .error_class
1305                    .unwrap_or(SyncErrorClass::Internal)
1306                    .as_label(),
1307                elapsed_ms,
1308                "sync decision rejected"
1309            );
1310        }
1311        _ => {
1312            debug!(
1313                operation = operation,
1314                replica_fingerprint = %context.replica_fingerprint,
1315                elapsed_ms,
1316                "unexpected sync decision response"
1317            );
1318        }
1319    }
1320}
1321
1322fn trim_to_applyable_v1_prefix(
1323    raw_units: &mut Vec<RetainedUnit>,
1324    wire_units: &mut Vec<WireRetainedUnit>,
1325) -> Result<(), String> {
1326    let mut last_error = None;
1327    while !raw_units.is_empty() {
1328        match validate_v1_retained_units_applyable(raw_units) {
1329            Ok(()) => return Ok(()),
1330            Err(err) => {
1331                last_error = Some(err.to_string());
1332                raw_units.pop();
1333                wire_units.pop();
1334            }
1335        }
1336    }
1337    if let Some(error) = last_error {
1338        return Err(format!(
1339            "sync pull cannot serve an applyable V1 retained chunk with current limits: {error}"
1340        ));
1341    }
1342    Ok(())
1343}
1344
1345fn validate_sync_ack_applyable_boundary(
1346    data_dir: &Path,
1347    replica_id: &str,
1348    applied_lsn: u64,
1349    remote_lsn: u64,
1350) -> Result<(), String> {
1351    let status =
1352        replica_sync_status(data_dir, replica_id, remote_lsn).map_err(|err| err.to_string())?;
1353    let Some(previous_lsn) = status.last_applied_lsn else {
1354        return Ok(());
1355    };
1356    if applied_lsn <= previous_lsn {
1357        return Ok(());
1358    }
1359    let range_len = applied_lsn - previous_lsn;
1360    if range_len > u64::from(MAX_SYNC_PULL_UNITS) {
1361        return Err(format!(
1362            "sync ack range contains {range_len} units; acknowledge ranges no larger than {MAX_SYNC_PULL_UNITS}"
1363        ));
1364    }
1365    let max_units =
1366        usize::try_from(range_len).map_err(|_| "sync ack range is too large to validate")?;
1367    let identity = read_identity(data_dir).map_err(|err| err.to_string())?;
1368    let segment_dir = retained_segments_dir(data_dir);
1369    let units = read_units_through(
1370        &segment_dir,
1371        identity.segment_identity(),
1372        previous_lsn,
1373        applied_lsn,
1374        max_units,
1375    )
1376    .map_err(|err| err.to_string())?;
1377    if units.len() != max_units || units.last().map(|unit| unit.lsn) != Some(applied_lsn) {
1378        return Err(
1379            "sync ack does not cover a complete retained-unit range; rebootstrap required".into(),
1380        );
1381    }
1382    validate_v1_retained_units_applyable(&units).map_err(|err| err.to_string())
1383}
1384
1385#[cfg(test)]
1386fn dispatch_sync_status(
1387    engine: &Arc<RwLock<Engine>>,
1388    replica_id: String,
1389    credential_authenticated: bool,
1390    principal: Option<&Principal>,
1391) -> Message {
1392    dispatch_sync_status_decision(engine, replica_id, credential_authenticated, principal).message
1393}
1394
1395fn dispatch_sync_status_decision(
1396    engine: &Arc<RwLock<Engine>>,
1397    replica_id: String,
1398    credential_authenticated: bool,
1399    principal: Option<&Principal>,
1400) -> SyncDecision {
1401    if let Err((class, message)) =
1402        check_sync_protocol_permitted(credential_authenticated, principal)
1403    {
1404        return SyncDecision::error(class, message);
1405    }
1406    if let Err(message) = validate_wire_replica_id(&replica_id) {
1407        return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1408    }
1409    let SyncContext {
1410        data_dir,
1411        remote_lsn,
1412        ..
1413    } = match sync_context(engine) {
1414        Ok(context) => context,
1415        Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1416    };
1417    match replica_sync_status(&data_dir, &replica_id, remote_lsn) {
1418        Ok(status) => SyncDecision::ok(Message::SyncStatusResult {
1419            status: wire_sync_status(status),
1420        }),
1421        Err(err) => SyncDecision::error(SyncErrorClass::StatusRead, err.to_string()),
1422    }
1423}
1424
1425#[cfg(test)]
1426fn dispatch_sync_pull(
1427    engine: &Arc<RwLock<Engine>>,
1428    request: SyncPullRequest,
1429    credential_authenticated: bool,
1430    principal: Option<&Principal>,
1431) -> Message {
1432    dispatch_sync_pull_decision(engine, request, credential_authenticated, principal).message
1433}
1434
1435fn dispatch_sync_pull_decision(
1436    engine: &Arc<RwLock<Engine>>,
1437    request: SyncPullRequest,
1438    credential_authenticated: bool,
1439    principal: Option<&Principal>,
1440) -> SyncDecision {
1441    if let Err((class, message)) =
1442        check_sync_protocol_permitted(credential_authenticated, principal)
1443    {
1444        return SyncDecision::error(class, message);
1445    }
1446    if let Err(message) = validate_wire_replica_id(&request.replica_id) {
1447        return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1448    }
1449    if request.max_units == 0 || request.max_units > MAX_SYNC_PULL_UNITS {
1450        return SyncDecision::error(
1451            SyncErrorClass::InvalidMaxUnits,
1452            format!("sync pull maxUnits must be between 1 and {MAX_SYNC_PULL_UNITS}"),
1453        );
1454    }
1455    if request.max_bytes == 0 || request.max_bytes > MAX_SYNC_PULL_BYTES {
1456        return SyncDecision::error(
1457            SyncErrorClass::InvalidMaxBytes,
1458            format!("sync pull maxBytes must be between 1 and {MAX_SYNC_PULL_BYTES}"),
1459        );
1460    }
1461
1462    let SyncContext {
1463        data_dir,
1464        remote_lsn,
1465        active_catalog_version,
1466    } = match sync_context(engine) {
1467        Ok(context) => context,
1468        Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1469    };
1470    let status = match replica_sync_status(&data_dir, &request.replica_id, remote_lsn) {
1471        Ok(status) => status,
1472        Err(err) => {
1473            return SyncDecision::error(SyncErrorClass::StatusRead, err.to_string());
1474        }
1475    };
1476    let Some(cursor_lsn) = status.last_applied_lsn else {
1477        return SyncDecision::ok(Message::SyncPullResult {
1478            status: wire_sync_status(status),
1479            units: Vec::new(),
1480            has_more: false,
1481        });
1482    };
1483    if status.repair_action != SyncRepairAction::Pull {
1484        return SyncDecision::ok(Message::SyncPullResult {
1485            status: wire_sync_status(status),
1486            units: Vec::new(),
1487            has_more: false,
1488        });
1489    }
1490    if request.since_lsn != cursor_lsn {
1491        return SyncDecision::error(
1492            SyncErrorClass::CursorLsnMismatch,
1493            format!(
1494                "sync pull sinceLsn {} does not match primary cursor LSN {cursor_lsn}",
1495                request.since_lsn
1496            ),
1497        );
1498    }
1499
1500    let identity = match read_identity(&data_dir) {
1501        Ok(identity) => identity,
1502        Err(err) => {
1503            return SyncDecision::error(SyncErrorClass::IdentityRead, err.to_string());
1504        }
1505    };
1506    // Stamp the expected identity with the database's *active* catalog version,
1507    // not this binary's compile-time maximum, so a database that never activated
1508    // v6 still expects v5 and accepts a v0.12 replica.
1509    let expected = SegmentIdentity::with_catalog_version(
1510        identity.database_id,
1511        identity.primary_generation,
1512        active_catalog_version,
1513    );
1514    if request.database_id != expected.database_id
1515        || request.primary_generation != expected.primary_generation
1516        || request.wal_format_version != expected.wal_format_version
1517        || request.segment_format_version != RETAINED_SEGMENT_FORMAT_VERSION
1518    {
1519        return SyncDecision::error(
1520            SyncErrorClass::IdentityOrFormatMismatch,
1521            "sync pull identity or format version mismatch; rebootstrap required",
1522        );
1523    }
1524    // The request states the maximum catalog format the replica can read. Accept
1525    // any replica that can read the active format (>= active); reject a replica
1526    // whose maximum is older than the data it would receive.
1527    if request.catalog_version < expected.catalog_version {
1528        return SyncDecision::error(
1529            SyncErrorClass::IdentityOrFormatMismatch,
1530            format!(
1531                "sync pull replica catalog format v{} cannot read this database's active catalog format v{}; rebootstrap with an upgraded replica required",
1532                request.catalog_version, expected.catalog_version
1533            ),
1534        );
1535    }
1536
1537    let effective_max_units = request.max_units.min(MAX_SYNC_PULL_UNITS) as usize;
1538    let requested_through_lsn = request
1539        .since_lsn
1540        .saturating_add(request.max_units as u64)
1541        .min(remote_lsn)
1542        .min(status.servable_lsn.unwrap_or(request.since_lsn));
1543    let segment_dir = retained_segments_dir(&data_dir);
1544    if requested_through_lsn > request.since_lsn {
1545        if let Err(err) = validate_retained_tail_available(
1546            &segment_dir,
1547            expected,
1548            request.since_lsn,
1549            requested_through_lsn,
1550        ) {
1551            let mut rebootstrap_status = status;
1552            rebootstrap_status.stale = true;
1553            rebootstrap_status.repair_action = SyncRepairAction::Rebootstrap;
1554            rebootstrap_status.last_sync_error = Some(format!(
1555                "retained history is unavailable; rebootstrap required: {err}"
1556            ));
1557            return SyncDecision::ok(Message::SyncPullResult {
1558                status: wire_sync_status(rebootstrap_status),
1559                units: Vec::new(),
1560                has_more: false,
1561            });
1562        }
1563    }
1564
1565    let raw_units = match read_units_through(
1566        &segment_dir,
1567        expected,
1568        request.since_lsn,
1569        requested_through_lsn,
1570        effective_max_units,
1571    ) {
1572        Ok(units) => units,
1573        Err(err) => {
1574            return SyncDecision::error(SyncErrorClass::RetainedRead, err.to_string());
1575        }
1576    };
1577
1578    let mut selected_raw = Vec::new();
1579    let mut selected = Vec::new();
1580    let mut selected_bytes = 0u64;
1581    for unit in raw_units {
1582        let wire_unit = wire_retained_unit(unit.clone());
1583        let unit_bytes = match wire_unit.encoded_len() {
1584            Ok(bytes) => bytes,
1585            Err(message) => {
1586                return SyncDecision::error(SyncErrorClass::RetainedUnitEncoding, message);
1587            }
1588        };
1589        if selected_bytes.saturating_add(unit_bytes) > request.max_bytes {
1590            if selected.is_empty() {
1591                return SyncDecision::error(
1592                    SyncErrorClass::InvalidMaxBytes,
1593                    "sync pull maxBytes is too small for the next retained unit",
1594                );
1595            }
1596            break;
1597        }
1598        selected_bytes += unit_bytes;
1599        selected_raw.push(unit);
1600        selected.push(wire_unit);
1601    }
1602    if let Err(message) = trim_to_applyable_v1_prefix(&mut selected_raw, &mut selected) {
1603        return SyncDecision::error(SyncErrorClass::RetainedChunkNotApplyable, message);
1604    }
1605
1606    let fetchable_through_lsn = status.servable_lsn.unwrap_or(remote_lsn).min(remote_lsn);
1607    let has_more = selected
1608        .last()
1609        .is_some_and(|unit| unit.lsn < fetchable_through_lsn);
1610    SyncDecision::ok(Message::SyncPullResult {
1611        status: wire_sync_status(status),
1612        units: selected,
1613        has_more,
1614    })
1615}
1616
1617#[cfg(test)]
1618fn dispatch_sync_ack(
1619    engine: &Arc<RwLock<Engine>>,
1620    replica_id: String,
1621    applied_lsn: u64,
1622    observed_remote_lsn: u64,
1623    credential_authenticated: bool,
1624    principal: Option<&Principal>,
1625) -> Message {
1626    dispatch_sync_ack_decision(
1627        engine,
1628        replica_id,
1629        applied_lsn,
1630        observed_remote_lsn,
1631        credential_authenticated,
1632        principal,
1633    )
1634    .message
1635}
1636
1637fn dispatch_sync_ack_decision(
1638    engine: &Arc<RwLock<Engine>>,
1639    replica_id: String,
1640    applied_lsn: u64,
1641    observed_remote_lsn: u64,
1642    credential_authenticated: bool,
1643    principal: Option<&Principal>,
1644) -> SyncDecision {
1645    if let Err((class, message)) =
1646        check_sync_protocol_permitted(credential_authenticated, principal)
1647    {
1648        return SyncDecision::error(class, message);
1649    }
1650    if let Err(message) = validate_wire_replica_id(&replica_id) {
1651        return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1652    }
1653    if applied_lsn > observed_remote_lsn {
1654        return SyncDecision::error(
1655            SyncErrorClass::LsnAheadOfRemote,
1656            format!(
1657                "sync ack appliedLsn {applied_lsn} is ahead of observed remoteLsn {observed_remote_lsn}"
1658            ),
1659        );
1660    }
1661    let SyncContext {
1662        data_dir,
1663        remote_lsn,
1664        ..
1665    } = match sync_context(engine) {
1666        Ok(context) => context,
1667        Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1668    };
1669    if observed_remote_lsn > remote_lsn {
1670        return SyncDecision::error(
1671            SyncErrorClass::LsnAheadOfRemote,
1672            format!(
1673                "sync ack remoteLsn {observed_remote_lsn} is ahead of primary LSN {remote_lsn}"
1674            ),
1675        );
1676    }
1677    if let Err(message) =
1678        validate_sync_ack_applyable_boundary(&data_dir, &replica_id, applied_lsn, remote_lsn)
1679    {
1680        return SyncDecision::error(SyncErrorClass::AckValidation, message);
1681    }
1682    match acknowledge_replica_apply(&data_dir, &replica_id, applied_lsn, remote_lsn) {
1683        Ok(summary) => SyncDecision::ok(Message::SyncAckResult {
1684            previous_applied_lsn: summary.previous_applied_lsn,
1685            applied_lsn: summary.applied_lsn,
1686            remote_lsn: summary.remote_lsn,
1687            advanced: summary.advanced,
1688            status: wire_sync_status(summary.status),
1689        }),
1690        Err(err) => SyncDecision::error(SyncErrorClass::AckUpdate, err.to_string()),
1691    }
1692}
1693
1694async fn run_blocking_sync<T, F>(input: T, query_timeout: Duration, f: F) -> SyncDecision
1695where
1696    T: Send + 'static,
1697    F: FnOnce(T) -> SyncDecision + Send + 'static,
1698{
1699    let mut handle = tokio::task::spawn_blocking(move || f(input));
1700    tokio::select! {
1701        result = &mut handle => match result {
1702            Ok(decision) => decision,
1703            Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1704        },
1705        _ = tokio::time::sleep(query_timeout) => match handle.await {
1706            Ok(decision) => decision,
1707            Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1708        },
1709    }
1710}
1711
1712async fn execute_gated_sync<T, F>(context: SyncExecutionContext<'_>, input: T, f: F) -> Message
1713where
1714    T: Send + 'static,
1715    F: FnOnce(T) -> SyncDecision + Send + 'static,
1716{
1717    let SyncExecutionContext {
1718        tx_gate,
1719        connection_has_transaction,
1720        operation,
1721        log_context,
1722        metrics,
1723        query_timeout,
1724    } = context;
1725    let start = Instant::now();
1726    if connection_has_transaction {
1727        let decision = SyncDecision::error(
1728            SyncErrorClass::ActiveTransaction,
1729            "sync protocol is unavailable inside an active transaction",
1730        );
1731        let elapsed = start.elapsed();
1732        log_sync_decision(operation, &log_context, elapsed, &decision);
1733        metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1734        return decision.message;
1735    }
1736
1737    let permit_count = tx_gate.permit_count();
1738    let permit = match tx_gate.acquire_many_owned(permit_count).await {
1739        Ok(permit) => permit,
1740        Err(_) => {
1741            let decision =
1742                SyncDecision::error(SyncErrorClass::QueryExecution, "query execution error");
1743            let elapsed = start.elapsed();
1744            log_sync_decision(operation, &log_context, elapsed, &decision);
1745            metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1746            return decision.message;
1747        }
1748    };
1749    let decision = run_blocking_sync(input, query_timeout, f).await;
1750    drop(permit);
1751    match &decision.message {
1752        Message::SyncStatusResult { status } => {
1753            metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1754        }
1755        Message::SyncPullResult { status, units, .. } => {
1756            metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1757            metrics.record_sync_pull_payload(units.len() as u64, sync_pull_payload_bytes(units));
1758        }
1759        Message::SyncAckResult {
1760            advanced, status, ..
1761        } => {
1762            metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1763            if *advanced {
1764                metrics.inc_sync_ack_advanced();
1765            }
1766        }
1767        _ => {}
1768    }
1769    let elapsed = start.elapsed();
1770    log_sync_decision(operation, &log_context, elapsed, &decision);
1771    metrics.record_sync_operation(
1772        operation,
1773        elapsed,
1774        sync_operation_outcome(&decision.message),
1775    );
1776    decision.message
1777}
1778
1779/// Acquire the TxGate for an explicit `begin`, bounded by `tx_wait_timeout`.
1780/// Overlapping explicit transactions queue behind the permit rather than being
1781/// rejected, but a connection gives up with a clear, client-facing error once
1782/// the wait elapses — so a transaction stalled (or held open) on another
1783/// connection can never block this one indefinitely. A timeout is recorded so
1784/// `powdb_tx_gate_timeouts_total` (and the error total) stay truthful.
1785async fn acquire_begin_permit(
1786    tx_gate: &TxGate,
1787    tx_wait_timeout: Duration,
1788    metrics: &Arc<Metrics>,
1789) -> Result<OwnedSemaphorePermit, Message> {
1790    match tokio::time::timeout(
1791        tx_wait_timeout,
1792        tx_gate.clone().acquire_many_owned(tx_gate.permit_count()),
1793    )
1794    .await
1795    {
1796        Ok(Ok(permit)) => Ok(permit),
1797        Ok(Err(_)) => Err(error_response(
1798            "query execution error",
1799            ErrorClass::Internal,
1800        )),
1801        Err(_) => {
1802            metrics.inc_tx_gate_timeout();
1803            Err(error_response(
1804                format!(
1805                    "transaction gate timeout after {}ms waiting for concurrent transaction to complete",
1806                    tx_wait_timeout.as_millis()
1807                ),
1808                ErrorClass::Timeout,
1809            ))
1810        }
1811    }
1812}
1813
1814/// Acquire the TxGate for a BARE autocommit statement, bounded by
1815/// `tx_wait_timeout` exactly like [`acquire_begin_permit`]. Autocommit writes
1816/// serialize through the same gate as explicit transactions, so a stalled (or
1817/// held-open) transaction on another connection would otherwise block this
1818/// write indefinitely. Bounding the acquire turns that indefinite wait into a
1819/// clear, client-facing timeout error and records the timeout so
1820/// `powdb_tx_gate_timeouts_total` (and the error total) stay truthful. This
1821/// only bounds the ACQUIRE; the permit is still dropped BEFORE the caller's
1822/// durability wait so overlapping committers can share an fsync.
1823async fn acquire_autocommit_permit(
1824    tx_gate: &TxGate,
1825    admission: AdmissionMode,
1826    tx_wait_timeout: Duration,
1827    metrics: &Arc<Metrics>,
1828) -> Result<OwnedSemaphorePermit, Message> {
1829    let permits = match admission {
1830        AdmissionMode::Reader => 1,
1831        AdmissionMode::Writer => tx_gate.permit_count(),
1832    };
1833
1834    // The uncontended path is overwhelmingly common for autocommit work.
1835    // Avoid constructing and polling a timeout-wrapped semaphore future when
1836    // the permits are already available. If a reader or writer is queued, the
1837    // try-acquire fails and we fall back to Tokio's fair semaphore queue, so a
1838    // waiting writer still cannot be bypassed by later readers.
1839    if let Ok(permit) = tx_gate.clone().try_acquire_many_owned(permits) {
1840        return Ok(permit);
1841    }
1842
1843    let acquire = tx_gate.clone().acquire_many_owned(permits);
1844    match tokio::time::timeout(tx_wait_timeout, acquire).await {
1845        Ok(Ok(permit)) => Ok(permit),
1846        Ok(Err(_)) => Err(error_response(
1847            "query execution error",
1848            ErrorClass::Internal,
1849        )),
1850        Err(_) => {
1851            metrics.inc_tx_gate_timeout();
1852            Err(error_response(
1853                format!(
1854                    "transaction gate timeout after {}ms waiting for concurrent transaction to complete",
1855                    tx_wait_timeout.as_millis()
1856                ),
1857                ErrorClass::Timeout,
1858            ))
1859        }
1860    }
1861}
1862
1863/// Run the shared four-arm transaction-routing state machine for one wire
1864/// query frame, returning the response plus its un-waited WAL durability
1865/// ticket. The TxGate permit is managed here and, crucially, is already
1866/// released (bare statements, commit/rollback) by the time this returns, so
1867/// the caller's `finalize_durability` wait happens OUTSIDE the gate and
1868/// overlapping committers can share an fsync.
1869///
1870/// The three wire dialects (PowQL, SQL, parameterized PowQL) differ only in
1871/// how a frame is parsed and dispatched; they share this routing, so a
1872/// behavior fix (e.g. cancellation rollback parity) lands here exactly once.
1873/// `parsed_query` is the dialect's already-parsed frame, `tx_control` its
1874/// transaction-control classification, `autocommit_admission` the admission
1875/// mode for a bare (non-transaction) statement, and `dispatch` the closure
1876/// that executes the parsed frame (its `bool` argument is
1877/// `allow_readonly_escalation`).
1878#[allow(clippy::too_many_arguments)]
1879async fn run_wire_query_state_machine<Inner, D, R>(
1880    engine: Arc<RwLock<Engine>>,
1881    tx_gate: TxGate,
1882    tx_permit: &mut Option<OwnedSemaphorePermit>,
1883    parsed_query: Arc<Inner>,
1884    tx_control: Option<TransactionControl>,
1885    autocommit_admission: AdmissionMode,
1886    result_mode: WireResultMode,
1887    principal: Option<Principal>,
1888    query_timeout: Duration,
1889    query_deadline: Instant,
1890    tx_wait_timeout: Duration,
1891    metrics: &Arc<Metrics>,
1892    reader: &mut BufReader<R>,
1893    wire_read_buffer: &mut Vec<u8>,
1894    pending_messages: &mut InFlightReadAhead,
1895    dispatch: D,
1896) -> (
1897    Message,
1898    Option<PendingDurability>,
1899    Option<ConnectionTermination>,
1900)
1901where
1902    Inner: Send + Sync + 'static,
1903    D: Fn(Arc<RwLock<Engine>>, Arc<Inner>, Option<Principal>, bool) -> DispatchOutcome
1904        + Clone
1905        + Send
1906        + Sync
1907        + 'static,
1908    R: AsyncRead + Unpin,
1909{
1910    // A read-only (snapshot-serving) engine never takes writer admission and
1911    // never runs a mutating statement. The flag is fixed for the engine's
1912    // lifetime; reading it costs one uncontended shared lock per frame.
1913    let read_only = engine.read().map(|eng| eng.is_read_only()).unwrap_or(false);
1914    if read_only {
1915        // Transaction control and any write statement require a writer that does
1916        // not exist in this mode: return the terminal error without touching the
1917        // gate or the engine write lock.
1918        let is_write = tx_control.is_some() || autocommit_admission == AdmissionMode::Writer;
1919        if is_write {
1920            return (readonly_terminal_message(), None, None);
1921        }
1922    }
1923
1924    match tx_control {
1925        Some(TransactionControl::Begin) => {
1926            if tx_permit.is_some() {
1927                return (
1928                    error_response(
1929                        sanitize_error(
1930                            "cannot begin: a transaction is already active on this connection",
1931                        ),
1932                        ErrorClass::Execution,
1933                    ),
1934                    None,
1935                    None,
1936                );
1937            }
1938            let permit = match acquire_begin_permit(&tx_gate, tx_wait_timeout, metrics).await {
1939                Ok(permit) => permit,
1940                Err(response) => return (response, None, None),
1941            };
1942            let dispatch_begin = dispatch.clone();
1943            let (response, ticket, mut termination, _) = run_blocking_query(
1944                engine.clone(),
1945                Arc::clone(&parsed_query),
1946                principal.clone(),
1947                result_mode,
1948                query_timeout,
1949                query_deadline,
1950                metrics,
1951                reader,
1952                wire_read_buffer,
1953                pending_messages,
1954                move |engine, parsed_query, principal| {
1955                    dispatch_begin(engine, parsed_query, principal, true)
1956                },
1957            )
1958            .await;
1959            if is_success_response(&response) {
1960                *tx_permit = Some(permit);
1961            } else if is_query_cancellation_response(&response) {
1962                // Parity with the commit/rollback and in-transaction arms: a
1963                // cancelled begin must not leave the engine's transaction
1964                // state ownerless. Install the just-acquired permit so the
1965                // shared rollback helper can undo any transaction the begin
1966                // opened, then release the permit and close the connection.
1967                *tx_permit = Some(permit);
1968                rollback_connection_transaction(engine, principal, tx_permit).await;
1969                termination = Some(ConnectionTermination::Closed);
1970            }
1971            (response, ticket, termination)
1972        }
1973        Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1974            let standalone_permit = if tx_permit.is_none() {
1975                match acquire_autocommit_permit(
1976                    &tx_gate,
1977                    AdmissionMode::Writer,
1978                    tx_wait_timeout,
1979                    metrics,
1980                )
1981                .await
1982                {
1983                    Ok(permit) => Some(permit),
1984                    Err(response) => return (response, None, None),
1985                }
1986            } else {
1987                None
1988            };
1989            let dispatch_commit = dispatch.clone();
1990            let (response, ticket, mut termination, _) = run_blocking_query(
1991                engine.clone(),
1992                Arc::clone(&parsed_query),
1993                principal.clone(),
1994                result_mode,
1995                query_timeout,
1996                query_deadline,
1997                metrics,
1998                reader,
1999                wire_read_buffer,
2000                pending_messages,
2001                move |engine, parsed_query, principal| {
2002                    dispatch_commit(engine, parsed_query, principal, true)
2003                },
2004            )
2005            .await;
2006            if is_success_response(&response) {
2007                // Release the gate BEFORE the caller waits on the commit's
2008                // ticket: the engine work is done and WAL order is fixed, so
2009                // another connection's commit can start (and share the fsync)
2010                // while this one waits.
2011                tx_permit.take();
2012            } else if is_query_cancellation_response(&response) {
2013                rollback_connection_transaction(engine, principal, tx_permit).await;
2014                termination = Some(ConnectionTermination::Closed);
2015            }
2016            drop(standalone_permit);
2017            (response, ticket, termination)
2018        }
2019        None if tx_permit.is_some() => {
2020            let dispatch_in_tx = dispatch.clone();
2021            let mut out = run_blocking_query(
2022                engine.clone(),
2023                Arc::clone(&parsed_query),
2024                principal.clone(),
2025                result_mode,
2026                query_timeout,
2027                query_deadline,
2028                metrics,
2029                reader,
2030                wire_read_buffer,
2031                pending_messages,
2032                move |engine, parsed_query, principal| {
2033                    dispatch_in_tx(engine, parsed_query, principal, true)
2034                },
2035            )
2036            .await;
2037            if is_query_cancellation_response(&out.0) {
2038                rollback_connection_transaction(engine, principal, tx_permit).await;
2039                out.2 = Some(ConnectionTermination::Closed);
2040            }
2041            (out.0, out.1, out.2)
2042        }
2043        None => {
2044            let admission = autocommit_admission;
2045            let permit = match acquire_autocommit_permit(
2046                &tx_gate,
2047                admission,
2048                tx_wait_timeout,
2049                metrics,
2050            )
2051            .await
2052            {
2053                Ok(permit) => permit,
2054                Err(response) => return (response, None, None),
2055            };
2056            // The parsed AST can be large. Share the first parse with the
2057            // rare dirty-view retry instead of cloning it on every successful
2058            // autocommit read.
2059            let retry_engine = Arc::clone(&engine);
2060            let retry_parsed_query = Arc::clone(&parsed_query);
2061            let retry_principal = principal.clone();
2062            let allow_readonly_escalation = admission == AdmissionMode::Writer;
2063            let dispatch_first = dispatch.clone();
2064            let mut out = run_blocking_query(
2065                engine,
2066                parsed_query,
2067                principal,
2068                result_mode,
2069                query_timeout,
2070                query_deadline,
2071                metrics,
2072                reader,
2073                wire_read_buffer,
2074                pending_messages,
2075                move |engine, parsed_query, principal| {
2076                    dispatch_first(engine, parsed_query, principal, allow_readonly_escalation)
2077                },
2078            )
2079            .await;
2080            drop(permit);
2081            if read_only && out.3 {
2082                // A read reached ReadonlyNeedsWrite (e.g. a stale materialized
2083                // view). There is no writer to escalate to in read-only mode:
2084                // surface the terminal error telling the operator to refresh
2085                // before snapshotting, rather than retrying under a writer.
2086                out.0 = readonly_terminal_message();
2087                out.3 = false;
2088            }
2089            if out.3 {
2090                let writer_permit = match acquire_autocommit_permit(
2091                    &tx_gate,
2092                    AdmissionMode::Writer,
2093                    tx_wait_timeout,
2094                    metrics,
2095                )
2096                .await
2097                {
2098                    Ok(permit) => permit,
2099                    Err(response) => return (response, None, None),
2100                };
2101                let dispatch_retry = dispatch.clone();
2102                out = run_blocking_query(
2103                    retry_engine,
2104                    retry_parsed_query,
2105                    retry_principal,
2106                    result_mode,
2107                    query_timeout,
2108                    query_deadline,
2109                    metrics,
2110                    reader,
2111                    wire_read_buffer,
2112                    pending_messages,
2113                    move |engine, parsed_query, principal| {
2114                        dispatch_retry(engine, parsed_query, principal, true)
2115                    },
2116                )
2117                .await;
2118                drop(writer_permit);
2119            }
2120            (out.0, out.1, out.2)
2121        }
2122    }
2123}
2124
2125/// Execute one PowQL wire query frame. Thin dialect wrapper over
2126/// [`run_wire_query_state_machine`].
2127#[allow(clippy::too_many_arguments)]
2128async fn execute_wire_query<R>(
2129    engine: Arc<RwLock<Engine>>,
2130    tx_gate: TxGate,
2131    tx_permit: &mut Option<OwnedSemaphorePermit>,
2132    query: String,
2133    result_mode: WireResultMode,
2134    principal: Option<Principal>,
2135    query_timeout: Duration,
2136    tx_wait_timeout: Duration,
2137    metrics: &Arc<Metrics>,
2138    reader: &mut BufReader<R>,
2139    wire_read_buffer: &mut Vec<u8>,
2140    pending_messages: &mut InFlightReadAhead,
2141) -> (
2142    Message,
2143    Option<PendingDurability>,
2144    Option<ConnectionTermination>,
2145)
2146where
2147    R: AsyncRead + Unpin,
2148{
2149    let query_deadline = Instant::now() + query_timeout;
2150    // Parse each frame once for transaction routing, admission, and role
2151    // enforcement. The engine still canonicalizes/parses as needed for plan
2152    // cache execution, but the server no longer repeats the same parse in
2153    // three separate routing helpers before reaching it.
2154    let stmt_result = parser::parse(&query).map_err(|e| e.to_string());
2155    let parsed_query = Arc::new((query, stmt_result));
2156    let tx_control = parsed_transaction_control(&parsed_query.1);
2157    let autocommit_admission = parsed_query
2158        .1
2159        .as_ref()
2160        .map(statement_admission)
2161        .unwrap_or(AdmissionMode::Writer);
2162    run_wire_query_state_machine(
2163        engine,
2164        tx_gate,
2165        tx_permit,
2166        parsed_query,
2167        tx_control,
2168        autocommit_admission,
2169        result_mode,
2170        principal,
2171        query_timeout,
2172        query_deadline,
2173        tx_wait_timeout,
2174        metrics,
2175        reader,
2176        wire_read_buffer,
2177        pending_messages,
2178        |engine,
2179         parsed_query: Arc<(String, Result<powdb_query::ast::Statement, String>)>,
2180         principal: Option<Principal>,
2181         allow| {
2182            dispatch_query_parsed(
2183                &engine,
2184                &parsed_query.0,
2185                &parsed_query.1,
2186                principal.as_ref(),
2187                allow,
2188            )
2189        },
2190    )
2191    .await
2192}
2193
2194#[allow(clippy::too_many_arguments)]
2195async fn execute_wire_query_sql<R>(
2196    engine: Arc<RwLock<Engine>>,
2197    tx_gate: TxGate,
2198    tx_permit: &mut Option<OwnedSemaphorePermit>,
2199    query: String,
2200    result_mode: WireResultMode,
2201    principal: Option<Principal>,
2202    query_timeout: Duration,
2203    tx_wait_timeout: Duration,
2204    metrics: &Arc<Metrics>,
2205    reader: &mut BufReader<R>,
2206    wire_read_buffer: &mut Vec<u8>,
2207    pending_messages: &mut InFlightReadAhead,
2208) -> (
2209    Message,
2210    Option<PendingDurability>,
2211    Option<ConnectionTermination>,
2212)
2213where
2214    R: AsyncRead + Unpin,
2215{
2216    let query_deadline = Instant::now() + query_timeout;
2217    let stmt_result = sql::parse_sql(&query).map_err(|e| e.to_string());
2218    let parsed_query = Arc::new((query, stmt_result));
2219    let tx_control = parsed_transaction_control(&parsed_query.1);
2220    let autocommit_admission = parsed_query
2221        .1
2222        .as_ref()
2223        .map(statement_admission)
2224        .unwrap_or(AdmissionMode::Writer);
2225    run_wire_query_state_machine(
2226        engine,
2227        tx_gate,
2228        tx_permit,
2229        parsed_query,
2230        tx_control,
2231        autocommit_admission,
2232        result_mode,
2233        principal,
2234        query_timeout,
2235        query_deadline,
2236        tx_wait_timeout,
2237        metrics,
2238        reader,
2239        wire_read_buffer,
2240        pending_messages,
2241        |engine,
2242         parsed_query: Arc<(String, Result<powdb_query::ast::Statement, String>)>,
2243         principal: Option<Principal>,
2244         allow| {
2245            dispatch_sql_query_parsed(
2246                &engine,
2247                &parsed_query.0,
2248                &parsed_query.1,
2249                principal.as_ref(),
2250                allow,
2251            )
2252        },
2253    )
2254    .await
2255}
2256
2257// One over clippy's default arg limit: the metrics handle was threaded through
2258// to instrument the typed query result. Bundling these into a struct would add
2259// more noise than it removes for an internal dispatcher.
2260#[allow(clippy::too_many_arguments)]
2261async fn execute_wire_query_with_params<R>(
2262    engine: Arc<RwLock<Engine>>,
2263    tx_gate: TxGate,
2264    tx_permit: &mut Option<OwnedSemaphorePermit>,
2265    query: String,
2266    params: Vec<WireParam>,
2267    result_mode: WireResultMode,
2268    principal: Option<Principal>,
2269    query_timeout: Duration,
2270    tx_wait_timeout: Duration,
2271    metrics: &Arc<Metrics>,
2272    reader: &mut BufReader<R>,
2273    wire_read_buffer: &mut Vec<u8>,
2274    pending_messages: &mut InFlightReadAhead,
2275) -> (
2276    Message,
2277    Option<PendingDurability>,
2278    Option<ConnectionTermination>,
2279)
2280where
2281    R: AsyncRead + Unpin,
2282{
2283    let query_deadline = Instant::now() + query_timeout;
2284    let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
2285    let stmt_result = parser::parse_with_params(&query, &bound).map_err(|e| e.to_string());
2286    let parsed_query = Arc::new((query, bound, stmt_result));
2287    let tx_control = parsed_transaction_control(&parsed_query.2);
2288    let autocommit_admission = parsed_query
2289        .2
2290        .as_ref()
2291        .map(statement_admission)
2292        .unwrap_or(AdmissionMode::Writer);
2293    run_wire_query_state_machine(
2294        engine,
2295        tx_gate,
2296        tx_permit,
2297        parsed_query,
2298        tx_control,
2299        autocommit_admission,
2300        result_mode,
2301        principal,
2302        query_timeout,
2303        query_deadline,
2304        tx_wait_timeout,
2305        metrics,
2306        reader,
2307        wire_read_buffer,
2308        pending_messages,
2309        |engine,
2310         parsed_query: Arc<(
2311            String,
2312            Vec<powdb_query::ast::ParamValue>,
2313            Result<powdb_query::ast::Statement, String>,
2314        )>,
2315         principal: Option<Principal>,
2316         allow| {
2317            dispatch_query_with_bound_params_parsed(
2318                &engine,
2319                &parsed_query.0,
2320                &parsed_query.1,
2321                &parsed_query.2,
2322                principal.as_ref(),
2323                allow,
2324            )
2325        },
2326    )
2327    .await
2328}
2329
2330/// A statement's metric sample whose recording is deferred until its WAL
2331/// durability obligation settles: a Full-mode fsync failure downgrades the
2332/// client's success reply to an error, and the metrics must tell the same
2333/// story (and the latency must include the wait the client observed).
2334struct DeferredQueryMetric {
2335    start: Instant,
2336    outcome: QueryOutcome,
2337    exceeded_timeout: bool,
2338}
2339
2340/// Durability ticket + the deferred metric of the statement that produced it.
2341type PendingDurability = (WalDurabilityTicket, DeferredQueryMetric);
2342
2343#[allow(clippy::too_many_arguments)]
2344async fn run_blocking_query<T, F, R>(
2345    engine: Arc<RwLock<Engine>>,
2346    input: T,
2347    principal: Option<Principal>,
2348    result_mode: WireResultMode,
2349    query_timeout: Duration,
2350    query_deadline: Instant,
2351    metrics: &Arc<Metrics>,
2352    reader: &mut BufReader<R>,
2353    wire_read_buffer: &mut Vec<u8>,
2354    pending_messages: &mut InFlightReadAhead,
2355    f: F,
2356) -> (
2357    Message,
2358    Option<PendingDurability>,
2359    Option<ConnectionTermination>,
2360    bool,
2361)
2362where
2363    T: Send + 'static,
2364    F: FnOnce(Arc<RwLock<Engine>>, T, Option<Principal>) -> DispatchOutcome + Send + 'static,
2365    R: AsyncRead + Unpin,
2366{
2367    let _in_flight = metrics.in_flight_guard();
2368    let start = Instant::now();
2369
2370    // Cooperative cancellation. The blocking closure installs this token for the
2371    // executor thread; every unbounded executor loop polls it. We give it a
2372    // deadline of `now + query_timeout` so the query self-terminates even if the
2373    // async timeout arm below is slow to be scheduled — that is what actually
2374    // makes the timeout enforceable (a `spawn_blocking` thread cannot be
2375    // aborted, so before this the timeout arm just awaited the runaway query to
2376    // completion while it held the engine lock / tx-gate permit).
2377    let timeout_ms = query_timeout.as_millis().min(u128::from(u64::MAX)) as u64;
2378    let cancel = Arc::new(powdb_query::cancel::ExecCancel::with_deadline(
2379        query_deadline,
2380        timeout_ms,
2381    ));
2382    let cancel_task = Arc::clone(&cancel);
2383    let mut handle = tokio::task::spawn_blocking(move || {
2384        let _cancel_guard = powdb_query::cancel::install(cancel_task);
2385        f(engine, input, principal)
2386    });
2387    let mut exceeded_timeout = false;
2388    let mut termination = None;
2389    let timeout = tokio::time::sleep(query_deadline.saturating_duration_since(Instant::now()));
2390    tokio::pin!(timeout);
2391    let join_result = loop {
2392        tokio::select! {
2393            result = &mut handle => break result,
2394            _ = &mut timeout => {
2395                exceeded_timeout = true;
2396                // Signal the executor to stop at its next cancellation checkpoint,
2397                // then await the (now promptly returning) handle. The closure
2398                // returns a typed timeout error and releases the engine lock /
2399                // tx-gate permit as it unwinds.
2400                cancel.cancel(powdb_query::cancel::CancelReason::Timeout);
2401                break handle.await;
2402            }
2403            read = read_message_cancel_safe(
2404                reader,
2405                wire_read_buffer,
2406                pending_messages.remaining_bytes(),
2407            ) => {
2408                match read {
2409                    Ok(Some(DecodedWireMessage { message: Message::Disconnect, .. })) => {
2410                        cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2411                        termination = Some(ConnectionTermination::Closed);
2412                        break handle.await;
2413                    }
2414                    Ok(Some(frame)) => {
2415                        if pending_messages.len() + 1 >= MAX_IN_FLIGHT_READ_AHEAD_FRAMES
2416                            || pending_messages.wire_bytes + frame.wire_len
2417                                >= MAX_IN_FLIGHT_READ_AHEAD_BYTES
2418                        {
2419                            // Never stop observing the socket behind a full
2420                            // queue. Reaching either hard cap cancels the query
2421                            // and closes this connection immediately.
2422                            cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2423                            termination = Some(ConnectionTermination::ReadError);
2424                            break handle.await;
2425                        }
2426                        // Preserve frames that arrive while the blocking query
2427                        // runs. The normal batching/main-loop path consumes them
2428                        // in order once execution completes.
2429                        pending_messages.push_back(frame);
2430                    }
2431                    Ok(None) => {
2432                        cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2433                        termination = Some(ConnectionTermination::Closed);
2434                        break handle.await;
2435                    }
2436                    Err(_) => {
2437                        cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2438                        termination = Some(ConnectionTermination::ReadError);
2439                        break handle.await;
2440                    }
2441                }
2442            }
2443        }
2444    };
2445
2446    let (message, ticket, outcome, readonly_needs_write) = match join_result {
2447        Ok((Ok(result), ticket)) => match query_result_to_message(result, result_mode) {
2448            Ok(message) => (message, ticket, QueryOutcome::Ok, false),
2449            Err(e) => (
2450                error_response(sanitize_error(&e.to_string()), classify_query_error(&e)),
2451                ticket,
2452                QueryOutcome::Error,
2453                false,
2454            ),
2455        },
2456        Ok((Err(QueryError::ReadonlyNeedsWrite), ticket)) => {
2457            if exceeded_timeout {
2458                (
2459                    error_response(
2460                        sanitize_error(&QueryError::Timeout { timeout_ms }.to_string()),
2461                        ErrorClass::Timeout,
2462                    ),
2463                    ticket,
2464                    QueryOutcome::Timeout,
2465                    true,
2466                )
2467            } else {
2468                (
2469                    error_response("query execution error", ErrorClass::Internal),
2470                    ticket,
2471                    QueryOutcome::Error,
2472                    true,
2473                )
2474            }
2475        }
2476        Ok((Err(e), ticket)) => {
2477            // A deadline-driven cancellation returns Timeout even when the async
2478            // timeout arm has not fired yet (the executor self-cancels): treat it
2479            // as a timeout for metrics either way.
2480            if matches!(e, QueryError::Timeout { .. }) {
2481                exceeded_timeout = true;
2482            }
2483            let outcome = if matches!(e, QueryError::MemoryLimitExceeded { .. }) {
2484                QueryOutcome::MemoryLimit
2485            } else if matches!(e, QueryError::Timeout { .. }) {
2486                QueryOutcome::Timeout
2487            } else {
2488                QueryOutcome::Error
2489            };
2490            (
2491                error_response(sanitize_error(&e.to_string()), classify_query_error(&e)),
2492                ticket,
2493                outcome,
2494                false,
2495            )
2496        }
2497        Err(e) => (
2498            error_response(format!("internal error: {e}"), ErrorClass::Internal),
2499            None,
2500            QueryOutcome::Error,
2501            false,
2502        ),
2503    };
2504    let readonly_retry = readonly_needs_write && !exceeded_timeout && termination.is_none();
2505    match ticket {
2506        // The statement's durability (and thus its true outcome and the
2507        // latency the client observes) settles at batch end — defer the
2508        // metric to the settlement site instead of recording a success that
2509        // a failed fsync would falsify.
2510        Some(ticket) => (
2511            message,
2512            Some((
2513                ticket,
2514                DeferredQueryMetric {
2515                    start,
2516                    outcome,
2517                    exceeded_timeout,
2518                },
2519            )),
2520            termination,
2521            readonly_retry,
2522        ),
2523        None => {
2524            if !readonly_retry {
2525                if exceeded_timeout {
2526                    metrics.record_query(start.elapsed(), QueryOutcome::Timeout);
2527                } else {
2528                    metrics.record_query(start.elapsed(), outcome);
2529                }
2530            }
2531            (message, None, termination, readonly_retry)
2532        }
2533    }
2534}
2535
2536/// Settle a WAL durability ticket off the async path, AFTER the TxGate
2537/// permit has been dropped — that ordering is what lets committers on other
2538/// connections append (and share the fsync) while this one waits.
2539///
2540/// Returns `None` when the covering fsync succeeded, or `Some(client-facing
2541/// error message)` when it failed — in which case no statement the ticket
2542/// covers may be acknowledged as durable (it executed in memory only).
2543async fn settle_durability_ticket(ticket: WalDurabilityTicket) -> Option<String> {
2544    match tokio::task::spawn_blocking(move || ticket.wait()).await {
2545        Ok(Ok(())) => None,
2546        Ok(Err(e)) => Some(sanitize_error(&format!("WAL durability sync failed: {e}"))),
2547        Err(e) => Some(format!("internal error: {e}")),
2548    }
2549}
2550
2551fn is_success_response(msg: &Message) -> bool {
2552    matches!(
2553        msg,
2554        Message::ResultRows { .. }
2555            | Message::ResultScalar { .. }
2556            | Message::ResultRowsNative { .. }
2557            | Message::ResultScalarNative { .. }
2558            | Message::ResultOk { .. }
2559            | Message::ResultMessage { .. }
2560    )
2561}
2562
2563fn rollback_open_transaction(engine: Arc<RwLock<Engine>>, principal: Option<Principal>) {
2564    let (res, ticket) = dispatch_query(&engine, "rollback", principal.as_ref(), true);
2565    let _ = res;
2566    // Rollback takes the sync-preserving path (no ticket), but settle one
2567    // defensively if it ever appears so the durability watermark stays honest.
2568    if let Some(ticket) = ticket {
2569        let _ = ticket.wait();
2570    }
2571}
2572
2573fn is_query_cancellation_response(message: &Message) -> bool {
2574    matches!(
2575        message,
2576        Message::Error { message } | Message::ErrorWithClass { message, .. }
2577            if message.starts_with("query timeout after")
2578                || message == "query cancelled by client disconnect"
2579    )
2580}
2581
2582/// Roll back this connection's explicit transaction while it still owns the
2583/// transaction-gate permit, then release the permit. A timed-out/cancelled
2584/// statement cannot leave an ambiguous transaction open and block every later
2585/// writer; releasing first would let another connection enter the engine before
2586/// this rollback has restored the prior snapshot.
2587async fn rollback_connection_transaction(
2588    engine: Arc<RwLock<Engine>>,
2589    principal: Option<Principal>,
2590    tx_permit: &mut Option<OwnedSemaphorePermit>,
2591) {
2592    if tx_permit.is_none() {
2593        return;
2594    }
2595    let _ = tokio::task::spawn_blocking(move || rollback_open_transaction(engine, principal)).await;
2596    tx_permit.take();
2597}
2598
2599/// Read one post-auth wire frame without losing partially-read bytes when the
2600/// future is cancelled by `tokio::select!`.
2601///
2602/// `Message::read_from` uses `read_exact`, whose future is not cancellation
2603/// safe: racing it against query completion can consume part of the next frame
2604/// and then drop those bytes. This reader stores every completed `read` in a
2605/// connection-owned buffer before awaiting again, so a query may safely race
2606/// socket EOF / `DISCONNECT` while preserving ordinary pipelined frames.
2607struct DecodedWireMessage {
2608    message: Message,
2609    wire_len: usize,
2610}
2611
2612#[derive(Default)]
2613struct InFlightReadAhead {
2614    frames: VecDeque<DecodedWireMessage>,
2615    wire_bytes: usize,
2616}
2617
2618impl InFlightReadAhead {
2619    fn len(&self) -> usize {
2620        self.frames.len()
2621    }
2622
2623    fn is_empty(&self) -> bool {
2624        self.frames.is_empty()
2625    }
2626
2627    fn remaining_bytes(&self) -> usize {
2628        MAX_IN_FLIGHT_READ_AHEAD_BYTES.saturating_sub(self.wire_bytes)
2629    }
2630
2631    fn push_back(&mut self, frame: DecodedWireMessage) {
2632        self.wire_bytes += frame.wire_len;
2633        self.frames.push_back(frame);
2634    }
2635
2636    fn pop_front(&mut self) -> Option<Message> {
2637        let frame = self.frames.pop_front()?;
2638        self.wire_bytes -= frame.wire_len;
2639        Some(frame.message)
2640    }
2641}
2642
2643async fn read_message_cancel_safe<R>(
2644    reader: &mut BufReader<R>,
2645    buffered: &mut Vec<u8>,
2646    max_frame_len: usize,
2647) -> std::io::Result<Option<DecodedWireMessage>>
2648where
2649    R: AsyncRead + Unpin,
2650{
2651    loop {
2652        if buffered.len() >= 6 {
2653            let payload_len = u32::from_le_bytes(
2654                buffered[2..6]
2655                    .try_into()
2656                    .expect("four-byte wire payload length"),
2657            ) as usize;
2658            if payload_len > MAX_WIRE_PAYLOAD_SIZE {
2659                return Err(std::io::Error::new(
2660                    std::io::ErrorKind::InvalidData,
2661                    format!("payload too large: {payload_len} bytes (max {MAX_WIRE_PAYLOAD_SIZE})"),
2662                ));
2663            }
2664            let frame_len = 6usize.checked_add(payload_len).ok_or_else(|| {
2665                std::io::Error::new(
2666                    std::io::ErrorKind::InvalidData,
2667                    "wire frame length overflow",
2668                )
2669            })?;
2670            if frame_len > max_frame_len {
2671                return Err(std::io::Error::new(
2672                    std::io::ErrorKind::InvalidData,
2673                    format!(
2674                        "wire frame exceeds the available in-flight read-ahead budget: \
2675                         {frame_len} bytes (available {max_frame_len})"
2676                    ),
2677                ));
2678            }
2679            if buffered.len() >= frame_len {
2680                let frame: Vec<u8> = buffered.drain(..frame_len).collect();
2681                return Message::decode(&frame)
2682                    .map(|message| {
2683                        Some(DecodedWireMessage {
2684                            message,
2685                            wire_len: frame_len,
2686                        })
2687                    })
2688                    .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e));
2689            }
2690        }
2691
2692        let mut chunk = [0u8; 8192];
2693        // Read only the bytes needed for this frame stage. Besides preserving
2694        // cancellation safety, this prevents a large pipelined payload from
2695        // overshooting the in-flight byte budget in one buffered read.
2696        let wanted = if buffered.len() < 6 {
2697            6 - buffered.len()
2698        } else {
2699            let payload_len = u32::from_le_bytes(
2700                buffered[2..6]
2701                    .try_into()
2702                    .expect("four-byte wire payload length"),
2703            ) as usize;
2704            6usize
2705                .checked_add(payload_len)
2706                .and_then(|frame_len| frame_len.checked_sub(buffered.len()))
2707                .ok_or_else(|| {
2708                    std::io::Error::new(
2709                        std::io::ErrorKind::InvalidData,
2710                        "invalid buffered wire frame length",
2711                    )
2712                })?
2713        };
2714        let read_limit = wanted.min(chunk.len());
2715        let read = reader.read(&mut chunk[..read_limit]).await?;
2716        if read == 0 {
2717            if buffered.len() < 6 {
2718                // Match the existing protocol behavior: EOF before a complete
2719                // header is a clean connection close.
2720                buffered.clear();
2721                return Ok(None);
2722            }
2723            return Err(std::io::Error::new(
2724                std::io::ErrorKind::UnexpectedEof,
2725                "connection closed in the middle of a wire frame",
2726            ));
2727        }
2728        buffered.extend_from_slice(&chunk[..read]);
2729    }
2730}
2731
2732#[derive(Clone, Copy, Debug)]
2733enum ConnectionTermination {
2734    Closed,
2735    ReadError,
2736}
2737
2738pub async fn handle_connection<S>(stream: S, opts: ConnOpts<'_>)
2739where
2740    S: AsyncRead + AsyncWrite + Unpin,
2741{
2742    let ConnOpts {
2743        engine,
2744        tx_gate,
2745        expected_password,
2746        users,
2747        shutdown_rx,
2748        idle_timeout,
2749        query_timeout,
2750        rate_limiter,
2751        peer_addr,
2752        metrics,
2753        tx_wait_timeout,
2754        db_name: server_db_name,
2755    } = opts;
2756
2757    let peer = peer_addr
2758        .map(|a| a.to_string())
2759        .unwrap_or_else(|| "unknown".into());
2760    let peer_ip = peer_addr.map(|a| a.ip());
2761
2762    let (reader, writer) = tokio::io::split(stream);
2763    let mut reader = BufReader::new(reader);
2764    let mut writer = BufWriter::new(writer);
2765
2766    // Wait for Connect message (with idle timeout).
2767    // Accept Ping messages before authentication so load balancers can
2768    // health-check without completing a full CONNECT handshake.
2769    // Uses the smaller pre-auth payload limit (4 KB) to prevent memory abuse.
2770    let connect_msg = loop {
2771        match tokio::time::timeout(idle_timeout, Message::read_from_preauth(&mut reader)).await {
2772            Ok(Ok(Some(Message::Ping))) => {
2773                debug!(peer = %peer, "pre-auth ping");
2774                if !write_msg(&mut writer, &Message::Pong).await {
2775                    return;
2776                }
2777                continue;
2778            }
2779            Ok(Ok(Some(msg))) => break msg,
2780            Ok(Ok(None)) => {
2781                debug!(peer = %peer, "client closed before CONNECT");
2782                return;
2783            }
2784            Ok(Err(e)) => {
2785                error!(peer = %peer, error = %e, "error reading CONNECT");
2786                return;
2787            }
2788            Err(_) => {
2789                warn!(peer = %peer, "idle timeout waiting for CONNECT");
2790                return;
2791            }
2792        }
2793    };
2794
2795    // The authenticated identity for this connection. Bound at connect time
2796    // and enforced on every query by `dispatch_query`.
2797    let principal: Option<Principal>;
2798    let credential_auth_configured = !users.is_empty() || expected_password.is_some();
2799    match connect_msg {
2800        Message::Connect {
2801            db_name,
2802            password,
2803            username,
2804        } => {
2805            // Check rate limiting before verifying credentials.
2806            if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2807                if is_rate_limited(limiter, ip) {
2808                    warn!(peer = %peer, "rate limited: too many auth failures");
2809                    let err = error_response(
2810                        "too many auth failures, try again later",
2811                        ErrorClass::RateLimited,
2812                    );
2813                    write_msg(&mut writer, &err).await;
2814                    return;
2815                }
2816            }
2817
2818            let outcome = authenticate_connect(
2819                &users,
2820                expected_password.as_ref().map(|p| p.as_str()),
2821                username.as_deref(),
2822                password.as_ref().map(|p| p.as_str()),
2823            );
2824
2825            match outcome {
2826                AuthOutcome::Rejected => {
2827                    warn!(peer = %peer, db = %db_name, "auth rejected");
2828                    metrics.inc_auth_failure();
2829                    // Record the failure for rate limiting.
2830                    if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2831                        record_auth_failure(limiter, ip);
2832                    }
2833                    let err = error_response("authentication failed", ErrorClass::AuthFailed);
2834                    write_msg(&mut writer, &err).await;
2835                    return;
2836                }
2837                AuthOutcome::Authenticated {
2838                    principal: auth_principal,
2839                } => {
2840                    // Auth succeeded — clear any prior failure count.
2841                    if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2842                        clear_auth_failures(limiter, ip);
2843                    }
2844                    match &auth_principal {
2845                        Some(p) => {
2846                            info!(peer = %peer, db = %db_name, user = %p.name, role = %p.role, "authenticated");
2847                        }
2848                        None => {
2849                            info!(peer = %peer, db = %db_name, "client connected");
2850                        }
2851                    }
2852                    principal = auth_principal;
2853                }
2854            }
2855
2856            // One process serves one database. When pinned to a name, reject a
2857            // CONNECT that explicitly asks for a different one (checked after
2858            // auth so db existence never leaks to unauthenticated clients).
2859            // When unpinned, accept anything but warn once per process so the
2860            // silent one-db-per-process mismatch is visible without spamming
2861            // the log on every pooled connection.
2862            match check_db_name(server_db_name.as_deref(), &db_name) {
2863                Ok(()) => {
2864                    if server_db_name.is_none() && !db_name.is_empty() && db_name != DEFAULT_DB_NAME
2865                    {
2866                        static NAMED_DB_WARNED: std::sync::atomic::AtomicBool =
2867                            std::sync::atomic::AtomicBool::new(false);
2868                        if !NAMED_DB_WARNED.swap(true, std::sync::atomic::Ordering::Relaxed) {
2869                            warn!(
2870                                peer = %peer, db = %db_name,
2871                                "client requested a named database but this server serves a single global database; name ignored"
2872                            );
2873                        }
2874                    }
2875                }
2876                Err(msg) => {
2877                    warn!(peer = %peer, db = %db_name, "rejected: unknown database");
2878                    let err = error_response(msg, ErrorClass::AuthFailed);
2879                    write_msg(&mut writer, &err).await;
2880                    return;
2881                }
2882            }
2883
2884            let ok = Message::ConnectOk {
2885                version: env!("CARGO_PKG_VERSION").into(),
2886            };
2887            if !write_msg(&mut writer, &ok).await {
2888                return;
2889            }
2890        }
2891        _ => {
2892            warn!(peer = %peer, "first message was not CONNECT");
2893            let err = error_response("expected CONNECT", ErrorClass::Internal);
2894            write_msg(&mut writer, &err).await;
2895            return;
2896        }
2897    }
2898
2899    let mut tx_permit: Option<OwnedSemaphorePermit> = None;
2900    // Persistent framing state makes reads cancellation-safe while they race
2901    // an in-flight blocking query. Frames decoded during execution retain
2902    // their original order here for normal pipelined processing afterwards.
2903    let mut wire_read_buffer = Vec::new();
2904    let mut pending_messages = InFlightReadAhead::default();
2905    // A non-query frame decoded during read-ahead batching, carried over to
2906    // the next iteration of the main loop.
2907    let mut carry: Option<Message> = None;
2908
2909    // Main query loop with idle timeout and shutdown awareness.
2910    'conn: loop {
2911        let msg = if let Some(m) = carry.take() {
2912            m
2913        } else if let Some(m) = pending_messages.pop_front() {
2914            m
2915        } else {
2916            tokio::select! {
2917                // Read next message with idle timeout.
2918                result = tokio::time::timeout(
2919                    idle_timeout,
2920                    read_message_cancel_safe(
2921                        &mut reader,
2922                        &mut wire_read_buffer,
2923                        MAX_WIRE_PAYLOAD_SIZE + 6,
2924                    ),
2925                ) => {
2926                    match result {
2927                        Ok(Ok(Some(frame))) => frame.message,
2928                        Ok(Ok(None)) => break,
2929                        Ok(Err(e)) => {
2930                            error!(peer = %peer, error = %e, "read error");
2931                            break;
2932                        }
2933                        Err(_) => {
2934                            info!(peer = %peer, "idle timeout, closing connection");
2935                            let err = error_response("idle timeout", ErrorClass::Timeout);
2936                            write_msg(&mut writer, &err).await;
2937                            break;
2938                        }
2939                    }
2940                }
2941                // If server is shutting down, notify client and close.
2942                _ = shutdown_rx.changed() => {
2943                    if *shutdown_rx.borrow() {
2944                        info!(peer = %peer, "server shutting down, closing connection");
2945                        let err = error_response("server shutting down", ErrorClass::Internal);
2946                        write_msg(&mut writer, &err).await;
2947                        break;
2948                    }
2949                    continue;
2950                }
2951            }
2952        };
2953
2954        // Plain query frames take the batched path: a pipelining client's
2955        // whole burst is executed with ONE durability wait at the end
2956        // (durability generations are cumulative, so the newest statement's
2957        // ticket covers every earlier one). Everything else is handled one
2958        // frame at a time exactly as before.
2959        if matches!(
2960            msg,
2961            Message::Query { .. }
2962                | Message::QuerySql { .. }
2963                | Message::QueryWithParams { .. }
2964                | Message::QueryNative { .. }
2965                | Message::QuerySqlNative { .. }
2966                | Message::QueryWithParamsNative { .. }
2967        ) {
2968            /// Read-ahead cap per batch: bounds unflushed responses and keeps
2969            /// the reply latency of the first statement bounded.
2970            const MAX_PIPELINE_BATCH: usize = 128;
2971            /// Byte cap on retained (unflushed) response payloads per batch:
2972            /// large row results stop read-ahead, so one connection can never
2973            /// hold gigabytes of replies hostage to the batch's durability
2974            /// wait.
2975            const MAX_PIPELINE_BATCH_BYTES: usize = 4 << 20;
2976
2977            /// Approximate encoded size of a response, for the batch byte
2978            /// cap. Counts the dominant string payloads; exact per-frame
2979            /// overhead is irrelevant at the 4 MiB cap.
2980            fn approx_response_bytes(msg: &Message) -> usize {
2981                match msg {
2982                    Message::ResultRows { columns, rows } => {
2983                        columns.iter().map(|c| c.len() + 4).sum::<usize>()
2984                            + rows
2985                                .iter()
2986                                .map(|r| r.iter().map(|v| v.len() + 4).sum::<usize>())
2987                                .sum::<usize>()
2988                    }
2989                    Message::ResultScalar { value } => value.len(),
2990                    Message::ResultRowsNative { columns, rows } => {
2991                        columns.iter().map(|c| c.len() + 4).sum::<usize>()
2992                            + rows
2993                                .iter()
2994                                .map(|row| {
2995                                    row.iter()
2996                                        .map(|value| 5 + native_value_body_len(value))
2997                                        .sum::<usize>()
2998                                })
2999                                .sum::<usize>()
3000                    }
3001                    Message::ResultScalarNative { value } => 5 + native_value_body_len(value),
3002                    Message::ResultMessage { message }
3003                    | Message::Error { message }
3004                    | Message::ErrorWithClass { message, .. } => message.len(),
3005                    _ => 16,
3006                }
3007            }
3008
3009            /// Whether the reader's buffered bytes hold at least one COMPLETE
3010            /// frame (6-byte header + payload). Read-ahead must never await
3011            /// the socket: blocking on a partial frame would hold the batch's
3012            /// durability settlement and unflushed replies hostage to a slow
3013            /// (or malicious) client, up to the idle timeout.
3014            fn complete_frame_buffered(buf: &[u8]) -> bool {
3015                buf.len() >= 6 && {
3016                    let payload_len =
3017                        u32::from_le_bytes(buf[2..6].try_into().expect("4-byte slice")) as usize;
3018                    buf.len() - 6 >= payload_len
3019                }
3020            }
3021
3022            let mut responses: Vec<Message> = Vec::new();
3023            let mut response_bytes: usize = 0;
3024            let mut last_ticket: Option<WalDurabilityTicket> = None;
3025            let mut deferred_metrics: Vec<DeferredQueryMetric> = Vec::new();
3026            let mut fatal: Option<ConnectionTermination> = None;
3027            let mut current = msg;
3028            loop {
3029                let (response, ticket, termination) = match current {
3030                    Message::Query { query } => {
3031                        if query.len() > MAX_QUERY_LENGTH {
3032                            (
3033                                error_response(
3034                                    format!(
3035                                        "query too large: {} bytes (max {})",
3036                                        query.len(),
3037                                        MAX_QUERY_LENGTH
3038                                    ),
3039                                    ErrorClass::LimitExceeded,
3040                                ),
3041                                None,
3042                                None,
3043                            )
3044                        } else {
3045                            debug!(peer = %peer, query = %query, "received query");
3046                            execute_wire_query(
3047                                engine.clone(),
3048                                tx_gate.clone(),
3049                                &mut tx_permit,
3050                                query,
3051                                WireResultMode::LegacyText,
3052                                principal.clone(),
3053                                query_timeout,
3054                                tx_wait_timeout,
3055                                &metrics,
3056                                &mut reader,
3057                                &mut wire_read_buffer,
3058                                &mut pending_messages,
3059                            )
3060                            .await
3061                        }
3062                    }
3063                    Message::QuerySql { query } => {
3064                        if query.len() > MAX_QUERY_LENGTH {
3065                            (
3066                                error_response(
3067                                    format!(
3068                                        "query too large: {} bytes (max {})",
3069                                        query.len(),
3070                                        MAX_QUERY_LENGTH
3071                                    ),
3072                                    ErrorClass::LimitExceeded,
3073                                ),
3074                                None,
3075                                None,
3076                            )
3077                        } else {
3078                            debug!(peer = %peer, query = %query, "received SQL query");
3079                            execute_wire_query_sql(
3080                                engine.clone(),
3081                                tx_gate.clone(),
3082                                &mut tx_permit,
3083                                query,
3084                                WireResultMode::LegacyText,
3085                                principal.clone(),
3086                                query_timeout,
3087                                tx_wait_timeout,
3088                                &metrics,
3089                                &mut reader,
3090                                &mut wire_read_buffer,
3091                                &mut pending_messages,
3092                            )
3093                            .await
3094                        }
3095                    }
3096                    Message::QueryWithParams { query, params } => {
3097                        if query.len() > MAX_QUERY_LENGTH {
3098                            (
3099                                error_response(
3100                                    format!(
3101                                        "query too large: {} bytes (max {})",
3102                                        query.len(),
3103                                        MAX_QUERY_LENGTH
3104                                    ),
3105                                    ErrorClass::LimitExceeded,
3106                                ),
3107                                None,
3108                                None,
3109                            )
3110                        } else {
3111                            debug!(peer = %peer, query = %query, n_params = params.len(), "received parameterized query");
3112                            execute_wire_query_with_params(
3113                                engine.clone(),
3114                                tx_gate.clone(),
3115                                &mut tx_permit,
3116                                query,
3117                                params,
3118                                WireResultMode::LegacyText,
3119                                principal.clone(),
3120                                query_timeout,
3121                                tx_wait_timeout,
3122                                &metrics,
3123                                &mut reader,
3124                                &mut wire_read_buffer,
3125                                &mut pending_messages,
3126                            )
3127                            .await
3128                        }
3129                    }
3130                    Message::QueryNative { query } => {
3131                        if query.len() > MAX_QUERY_LENGTH {
3132                            (
3133                                error_response(
3134                                    format!(
3135                                        "query too large: {} bytes (max {})",
3136                                        query.len(),
3137                                        MAX_QUERY_LENGTH
3138                                    ),
3139                                    ErrorClass::LimitExceeded,
3140                                ),
3141                                None,
3142                                None,
3143                            )
3144                        } else {
3145                            debug!(peer = %peer, query = %query, "received native query");
3146                            execute_wire_query(
3147                                engine.clone(),
3148                                tx_gate.clone(),
3149                                &mut tx_permit,
3150                                query,
3151                                WireResultMode::Native,
3152                                principal.clone(),
3153                                query_timeout,
3154                                tx_wait_timeout,
3155                                &metrics,
3156                                &mut reader,
3157                                &mut wire_read_buffer,
3158                                &mut pending_messages,
3159                            )
3160                            .await
3161                        }
3162                    }
3163                    Message::QuerySqlNative { query } => {
3164                        if query.len() > MAX_QUERY_LENGTH {
3165                            (
3166                                error_response(
3167                                    format!(
3168                                        "query too large: {} bytes (max {})",
3169                                        query.len(),
3170                                        MAX_QUERY_LENGTH
3171                                    ),
3172                                    ErrorClass::LimitExceeded,
3173                                ),
3174                                None,
3175                                None,
3176                            )
3177                        } else {
3178                            debug!(peer = %peer, query = %query, "received native SQL query");
3179                            execute_wire_query_sql(
3180                                engine.clone(),
3181                                tx_gate.clone(),
3182                                &mut tx_permit,
3183                                query,
3184                                WireResultMode::Native,
3185                                principal.clone(),
3186                                query_timeout,
3187                                tx_wait_timeout,
3188                                &metrics,
3189                                &mut reader,
3190                                &mut wire_read_buffer,
3191                                &mut pending_messages,
3192                            )
3193                            .await
3194                        }
3195                    }
3196                    Message::QueryWithParamsNative { query, params } => {
3197                        if query.len() > MAX_QUERY_LENGTH {
3198                            (
3199                                error_response(
3200                                    format!(
3201                                        "query too large: {} bytes (max {})",
3202                                        query.len(),
3203                                        MAX_QUERY_LENGTH
3204                                    ),
3205                                    ErrorClass::LimitExceeded,
3206                                ),
3207                                None,
3208                                None,
3209                            )
3210                        } else {
3211                            debug!(peer = %peer, query = %query, n_params = params.len(), "received native parameterized query");
3212                            execute_wire_query_with_params(
3213                                engine.clone(),
3214                                tx_gate.clone(),
3215                                &mut tx_permit,
3216                                query,
3217                                params,
3218                                WireResultMode::Native,
3219                                principal.clone(),
3220                                query_timeout,
3221                                tx_wait_timeout,
3222                                &metrics,
3223                                &mut reader,
3224                                &mut wire_read_buffer,
3225                                &mut pending_messages,
3226                            )
3227                            .await
3228                        }
3229                    }
3230                    _ => unreachable!("batch loop only receives plain query frames"),
3231                };
3232                if let Some((t, m)) = ticket {
3233                    // Later tickets cover earlier generations — keep only the
3234                    // newest; the batch-end wait settles them all. Every
3235                    // deferred metric is kept: each records after settlement.
3236                    last_ticket = Some(t);
3237                    deferred_metrics.push(m);
3238                }
3239                response_bytes += approx_response_bytes(&response);
3240                responses.push(response);
3241                if let Some(reason) = termination {
3242                    fatal = Some(reason);
3243                    break;
3244                }
3245
3246                // Read ahead only when a COMPLETE next frame is already
3247                // buffered (never await the socket mid-batch) and the
3248                // retained replies stay small. While an explicit transaction
3249                // is open the connection holds the TxGate, so batching would
3250                // only extend the exclusive window — flush instead.
3251                if tx_permit.is_some()
3252                    || responses.len() >= MAX_PIPELINE_BATCH
3253                    || response_bytes >= MAX_PIPELINE_BATCH_BYTES
3254                    || (pending_messages.is_empty()
3255                        && !complete_frame_buffered(&wire_read_buffer)
3256                        && !complete_frame_buffered(reader.buffer()))
3257                {
3258                    break;
3259                }
3260                // The full frame is buffered, so this returns without socket
3261                // I/O; the timeout is a defensive backstop only.
3262                let next_message = if let Some(message) = pending_messages.pop_front() {
3263                    Ok(Some(message))
3264                } else {
3265                    tokio::time::timeout(
3266                        idle_timeout,
3267                        read_message_cancel_safe(
3268                            &mut reader,
3269                            &mut wire_read_buffer,
3270                            MAX_WIRE_PAYLOAD_SIZE + 6,
3271                        ),
3272                    )
3273                    .await
3274                    .map(|result| result.map(|frame| frame.map(|frame| frame.message)))
3275                    .unwrap_or_else(|_| {
3276                        Err(std::io::Error::new(
3277                            std::io::ErrorKind::TimedOut,
3278                            "timeout decoding fully-buffered frame",
3279                        ))
3280                    })
3281                };
3282                match next_message {
3283                    Ok(Some(
3284                        next @ (Message::Query { .. }
3285                        | Message::QuerySql { .. }
3286                        | Message::QueryWithParams { .. }
3287                        | Message::QueryNative { .. }
3288                        | Message::QuerySqlNative { .. }
3289                        | Message::QueryWithParamsNative { .. }),
3290                    )) => {
3291                        // If another connection currently holds the TxGate,
3292                        // the next statement would block on the gate with
3293                        // this batch's replies still unflushed (pre-batching,
3294                        // they'd already have been written). Flush first and
3295                        // handle the frame on the next main-loop iteration.
3296                        // Benign TOCTOU: worst case is one early flush or one
3297                        // gate wait with an empty reply queue.
3298                        if tx_gate.available_permits() == 0 {
3299                            carry = Some(next);
3300                            break;
3301                        }
3302                        current = next;
3303                    }
3304                    Ok(Some(other)) => {
3305                        // Not a plain query — flush this batch, then handle
3306                        // the frame on the next main-loop iteration.
3307                        carry = Some(other);
3308                        break;
3309                    }
3310                    Ok(None) => {
3311                        fatal = Some(ConnectionTermination::Closed);
3312                        break;
3313                    }
3314                    Err(e) => {
3315                        error!(peer = %peer, error = %e, "read error");
3316                        fatal = Some(ConnectionTermination::ReadError);
3317                        break;
3318                    }
3319                }
3320            }
3321
3322            // ONE durability wait for the whole batch, then the deferred
3323            // metrics: a durability failure records Ok statements as errors,
3324            // and latency includes the settlement wait the client observed.
3325            let mut durability_failed = false;
3326            if let Some(ticket) = last_ticket {
3327                if let Some(message) = settle_durability_ticket(ticket).await {
3328                    // The covering fsync failed: nothing in this batch may be
3329                    // acknowledged as durable.
3330                    durability_failed = true;
3331                    for r in responses.iter_mut() {
3332                        if is_success_response(r) {
3333                            *r = error_response(message.clone(), ErrorClass::Internal);
3334                        }
3335                    }
3336                }
3337            }
3338            for m in deferred_metrics.drain(..) {
3339                let outcome = if m.exceeded_timeout {
3340                    QueryOutcome::Timeout
3341                } else if durability_failed && matches!(m.outcome, QueryOutcome::Ok) {
3342                    QueryOutcome::Error
3343                } else {
3344                    m.outcome
3345                };
3346                metrics.record_query(m.start.elapsed(), outcome);
3347            }
3348
3349            for r in &responses {
3350                if !write_msg(&mut writer, r).await {
3351                    break 'conn;
3352                }
3353            }
3354            match fatal {
3355                None => continue,
3356                Some(ConnectionTermination::Closed | ConnectionTermination::ReadError) => break,
3357            }
3358        }
3359
3360        let response = match msg {
3361            Message::Ping => {
3362                debug!(peer = %peer, "ping");
3363                Message::Pong
3364            }
3365            Message::SyncStatus { replica_id } => {
3366                let engine = engine.clone();
3367                let principal = principal.clone();
3368                let log_context = SyncLogContext::status(&replica_id);
3369                execute_gated_sync(
3370                    SyncExecutionContext {
3371                        tx_gate: tx_gate.clone(),
3372                        connection_has_transaction: tx_permit.is_some(),
3373                        operation: SyncOperation::Status,
3374                        log_context,
3375                        metrics: &metrics,
3376                        query_timeout,
3377                    },
3378                    (engine, replica_id, credential_auth_configured, principal),
3379                    |(engine, replica_id, credential_authenticated, principal)| {
3380                        dispatch_sync_status_decision(
3381                            &engine,
3382                            replica_id,
3383                            credential_authenticated,
3384                            principal.as_ref(),
3385                        )
3386                    },
3387                )
3388                .await
3389            }
3390            Message::SyncPull {
3391                replica_id,
3392                since_lsn,
3393                max_units,
3394                max_bytes,
3395                database_id,
3396                primary_generation,
3397                wal_format_version,
3398                catalog_version,
3399                segment_format_version,
3400            } => {
3401                let engine = engine.clone();
3402                let principal = principal.clone();
3403                let request = SyncPullRequest {
3404                    replica_id,
3405                    since_lsn,
3406                    max_units,
3407                    max_bytes,
3408                    database_id,
3409                    primary_generation,
3410                    wal_format_version,
3411                    catalog_version,
3412                    segment_format_version,
3413                };
3414                let log_context = SyncLogContext::pull(&request);
3415                execute_gated_sync(
3416                    SyncExecutionContext {
3417                        tx_gate: tx_gate.clone(),
3418                        connection_has_transaction: tx_permit.is_some(),
3419                        operation: SyncOperation::Pull,
3420                        log_context,
3421                        metrics: &metrics,
3422                        query_timeout,
3423                    },
3424                    (engine, request, credential_auth_configured, principal),
3425                    |(engine, request, credential_authenticated, principal)| {
3426                        dispatch_sync_pull_decision(
3427                            &engine,
3428                            request,
3429                            credential_authenticated,
3430                            principal.as_ref(),
3431                        )
3432                    },
3433                )
3434                .await
3435            }
3436            Message::SyncAck {
3437                replica_id,
3438                applied_lsn,
3439                remote_lsn,
3440            } => {
3441                let engine = engine.clone();
3442                let principal = principal.clone();
3443                let log_context = SyncLogContext::ack(&replica_id, applied_lsn, remote_lsn);
3444                execute_gated_sync(
3445                    SyncExecutionContext {
3446                        tx_gate: tx_gate.clone(),
3447                        connection_has_transaction: tx_permit.is_some(),
3448                        operation: SyncOperation::Ack,
3449                        log_context,
3450                        metrics: &metrics,
3451                        query_timeout,
3452                    },
3453                    (
3454                        engine,
3455                        replica_id,
3456                        applied_lsn,
3457                        remote_lsn,
3458                        credential_auth_configured,
3459                        principal,
3460                    ),
3461                    |(
3462                        engine,
3463                        replica_id,
3464                        applied_lsn,
3465                        observed_remote_lsn,
3466                        credential_authenticated,
3467                        principal,
3468                    )| {
3469                        dispatch_sync_ack_decision(
3470                            &engine,
3471                            replica_id,
3472                            applied_lsn,
3473                            observed_remote_lsn,
3474                            credential_authenticated,
3475                            principal.as_ref(),
3476                        )
3477                    },
3478                )
3479                .await
3480            }
3481            Message::Disconnect => {
3482                debug!(peer = %peer, "received DISCONNECT");
3483                break;
3484            }
3485            _ => error_response("unexpected message type", ErrorClass::Internal),
3486        };
3487
3488        if !write_msg(&mut writer, &response).await {
3489            break;
3490        }
3491    }
3492
3493    // Roll back any open transaction the client left behind on disconnect.
3494    // The permit must stay alive in `tx_permit` for the duration of the awaited
3495    // rollback and be released only afterwards — mirroring the query-timeout
3496    // path above. Using `tx_permit.take().is_some()` here would drop the permit
3497    // (freeing the TxGate) *before* the rollback runs, letting another
3498    // connection BEGIN a transaction that this stale rollback would then clobber.
3499    if tx_permit.is_some() {
3500        let engine = engine.clone();
3501        let principal = principal.clone();
3502        let _ =
3503            tokio::task::spawn_blocking(move || rollback_open_transaction(engine, principal)).await;
3504    }
3505    tx_permit.take();
3506
3507    info!(peer = %peer, "client disconnected");
3508}
3509
3510fn charge_response_bytes(total: &mut usize, bytes: usize) -> Result<(), QueryError> {
3511    *total = total.saturating_add(bytes);
3512    if *total > MAX_RESPONSE_PAYLOAD_SIZE {
3513        return Err(QueryError::Execution(format!(
3514            "result too large: encoded response exceeds {} bytes; add a limit or narrower projection",
3515            MAX_RESPONSE_PAYLOAD_SIZE
3516        )));
3517    }
3518    Ok(())
3519}
3520
3521fn native_value_body_len(value: &Value) -> usize {
3522    match value {
3523        Value::Empty => 0,
3524        Value::Int(_) | Value::Float(_) | Value::DateTime(_) => 8,
3525        Value::Bool(_) => 1,
3526        Value::Str(value) => value.len(),
3527        Value::Uuid(_) => 16,
3528        Value::Bytes(value) => value.len(),
3529        Value::Json(value) => value.len(),
3530    }
3531}
3532
3533fn query_result_to_message(
3534    result: QueryResult,
3535    result_mode: WireResultMode,
3536) -> Result<Message, QueryError> {
3537    match result {
3538        QueryResult::Rows { columns, rows } => {
3539            let mut encoded_bytes = 2usize; // column count
3540            for col in &columns {
3541                charge_response_bytes(&mut encoded_bytes, 4 + col.len())?;
3542            }
3543            charge_response_bytes(&mut encoded_bytes, 4)?; // row count
3544
3545            match result_mode {
3546                WireResultMode::Native => {
3547                    for row in &rows {
3548                        for value in row {
3549                            charge_response_bytes(
3550                                &mut encoded_bytes,
3551                                5 + native_value_body_len(value),
3552                            )?;
3553                        }
3554                    }
3555                    Ok(Message::ResultRowsNative { columns, rows })
3556                }
3557                WireResultMode::LegacyText => {
3558                    let mut str_rows = Vec::with_capacity(rows.len());
3559                    for row in rows {
3560                        let mut str_row = Vec::with_capacity(row.len());
3561                        for value in row {
3562                            let display = value_to_display(&value);
3563                            charge_response_bytes(&mut encoded_bytes, 4 + display.len())?;
3564                            str_row.push(display);
3565                        }
3566                        str_rows.push(str_row);
3567                    }
3568                    Ok(Message::ResultRows {
3569                        columns,
3570                        rows: str_rows,
3571                    })
3572                }
3573            }
3574        }
3575        QueryResult::Scalar(value) => match result_mode {
3576            WireResultMode::Native => {
3577                let mut encoded_bytes = 0;
3578                charge_response_bytes(&mut encoded_bytes, 5 + native_value_body_len(&value))?;
3579                Ok(Message::ResultScalarNative { value })
3580            }
3581            WireResultMode::LegacyText => Ok(Message::ResultScalar {
3582                value: value_to_display(&value),
3583            }),
3584        },
3585        QueryResult::Modified(n) => Ok(Message::ResultOk { affected: n }),
3586        QueryResult::Created(name) => Ok(Message::ResultMessage {
3587            message: format!("type {name} created"),
3588        }),
3589        QueryResult::Executed { message } => Ok(Message::ResultMessage { message }),
3590    }
3591}
3592
3593// Canonical wire rendering lives on `Value` (`powdb_storage`) so the server,
3594// CLI, and embedded bindings render results identically. Kept as a thin alias
3595// to minimize churn at the call sites in this module.
3596fn value_to_display(v: &Value) -> String {
3597    v.to_wire_string()
3598}
3599
3600#[cfg(test)]
3601mod tests {
3602    use super::*;
3603    use powdb_storage::wal::WalRecordType;
3604    use powdb_sync::{
3605        write_identity_snapshot, write_segment_atomic, DatabaseIdentity, IdentitySnapshot,
3606        ReplicaCursor, RetainedSegment, RetainedUnit,
3607    };
3608
3609    // ---- Wire NULL rendering (Fix: remote protocol rendered NULL as `{}`) ----
3610
3611    #[test]
3612    fn null_serializes_as_null_bareword_on_wire() {
3613        assert_eq!(value_to_display(&Value::Empty), "null");
3614    }
3615
3616    // ---- Error sanitization allowlist ----
3617
3618    #[test]
3619    fn unique_violation_error_surfaces_to_remote_clients() {
3620        // The storage layer reports the actionable message; the server must
3621        // not replace it with the generic "query execution error".
3622        assert_eq!(
3623            sanitize_error("unique constraint violation on User.email"),
3624            "unique constraint violation on User.email"
3625        );
3626    }
3627
3628    #[test]
3629    fn internal_errors_stay_generic() {
3630        assert_eq!(
3631            sanitize_error("some internal io panic detail"),
3632            "query execution error"
3633        );
3634    }
3635
3636    #[test]
3637    fn cancellation_errors_surface_to_remote_clients() {
3638        // A cancelled/timed-out query must reach the client with its real
3639        // message (both are derived from the configured timeout or a client
3640        // disconnect and leak no internal state) rather than the generic mask.
3641        for msg in [
3642            &QueryError::Timeout { timeout_ms: 2000 }.to_string(),
3643            &QueryError::Cancelled.to_string(),
3644        ] {
3645            assert_eq!(sanitize_error(msg), *msg, "should pass through verbatim");
3646        }
3647        // Sanity-check the exact wording the executor emits.
3648        assert_eq!(
3649            QueryError::Timeout { timeout_ms: 2000 }.to_string(),
3650            "query timeout after 2000ms"
3651        );
3652        assert_eq!(
3653            QueryError::Cancelled.to_string(),
3654            "query cancelled by client disconnect"
3655        );
3656    }
3657
3658    // ---- JSON (v0.12): canonical-text wire rendering + parse-error passthrough ----
3659
3660    #[test]
3661    fn json_cell_renders_canonical_text_on_wire() {
3662        // A Json value flows through the same string-cell path as every other
3663        // value (value_to_display -> Value::to_wire_string). PJ1 is canonical,
3664        // so keys come back sorted bytewise regardless of input order and the
3665        // client receives parseable JSON text with no protocol change.
3666        let pj1 = powdb_storage::pj1::parse_json_text(r#"{"b":2,"a":1,"nested":{"z":true}}"#)
3667            .expect("valid JSON");
3668        let result = QueryResult::Rows {
3669            columns: vec!["doc".into()],
3670            rows: vec![vec![Value::Json(pj1.into())]],
3671        };
3672        match query_result_to_message(result, WireResultMode::LegacyText).expect("encodes") {
3673            Message::ResultRows { columns, rows } => {
3674                assert_eq!(columns, vec!["doc"]);
3675                assert_eq!(
3676                    rows,
3677                    vec![vec![r#"{"a":1,"b":2,"nested":{"z":true}}"#.to_string()]]
3678                );
3679            }
3680            other => panic!("expected ResultRows, got {other:?}"),
3681        }
3682    }
3683
3684    #[test]
3685    fn json_parse_error_surfaces_to_remote_clients() {
3686        // Lane B rejects invalid JSON on insert as QueryError::TypeError, whose
3687        // Display is "type mismatch: <detail>" (crates/query/src/result.rs).
3688        // That prefix is allowlisted, so the actionable detail reaches the
3689        // client instead of the generic "query execution error". The raw
3690        // storage-layer phrasing ("invalid JSON: ...") is also allowlisted as
3691        // defense-in-depth. Internal PJ1 corruption ("malformed PJ1: ...") is
3692        // deliberately NOT allowlisted: it leaks storage internals and never
3693        // occurs on the client-driven insert path.
3694        for msg in [
3695            "type mismatch: invalid JSON: unexpected character 'x' at position 3",
3696            "invalid JSON: nesting exceeds depth cap 128",
3697        ] {
3698            assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
3699        }
3700        assert_eq!(
3701            sanitize_error("malformed PJ1: truncated"),
3702            "query execution error",
3703            "internal storage corruption must stay masked"
3704        );
3705    }
3706
3707    // `describe <Type>` renders a json column's type as the bareword "json"
3708    // over the wire. introspect_describe emits type_id_to_name(TypeId::Json) =
3709    // "json" (crates/query/src/executor/compiled.rs) as a Str cell, which flows
3710    // through value_to_display unchanged; Lane B's DDL keyword makes `type Doc
3711    // { body: json }` accepted, so this runs end to end (v0.12, Lane D).
3712    #[test]
3713    fn describe_shows_json_type_over_the_wire() {
3714        let dir = tempfile::tempdir().unwrap();
3715        let mut engine = Engine::new(dir.path()).unwrap();
3716        engine
3717            .execute_powql("type Doc { required id: int, body: json }")
3718            .expect("json column DDL should be accepted once Lane B lands");
3719        let result = engine.execute_powql("describe Doc").expect("describe runs");
3720        let msg = query_result_to_message(result, WireResultMode::LegacyText).expect("encodes");
3721        match msg {
3722            Message::ResultRows { columns, rows } => {
3723                assert_eq!(columns[1], "type");
3724                // The `body` column's type cell must be the bareword "json".
3725                let body = rows
3726                    .iter()
3727                    .find(|r| r[0] == "body")
3728                    .expect("body column present");
3729                assert_eq!(body[1], "json");
3730            }
3731            other => panic!("expected ResultRows, got {other:?}"),
3732        }
3733    }
3734
3735    // ---- Named-database gate (P-10) ----
3736
3737    #[test]
3738    fn db_name_unpinned_accepts_any_name() {
3739        for requested in ["", "default", "prod", "anything"] {
3740            assert!(
3741                check_db_name(None, requested).is_ok(),
3742                "rejected {requested}"
3743            );
3744        }
3745    }
3746
3747    #[test]
3748    fn db_name_pinned_accepts_match_empty_and_default_sentinel() {
3749        // The configured name, the empty name, and the client default sentinel
3750        // are all "no foreign database explicitly requested".
3751        assert!(check_db_name(Some("prod"), "prod").is_ok());
3752        assert!(check_db_name(Some("prod"), "").is_ok());
3753        assert!(check_db_name(Some("prod"), DEFAULT_DB_NAME).is_ok());
3754    }
3755
3756    #[test]
3757    fn db_name_pinned_rejects_foreign_with_clear_message() {
3758        let err = check_db_name(Some("prod"), "staging").unwrap_err();
3759        assert_eq!(err, "unknown database 'staging'; this server serves 'prod'");
3760    }
3761
3762    // ---- Explicit-transaction gate wait timeout (P-4) ----
3763
3764    #[tokio::test]
3765    async fn begin_permit_acquires_when_gate_is_free() {
3766        let gate = new_tx_gate();
3767        let metrics = Arc::new(Metrics::new());
3768        let permit = acquire_begin_permit(&gate, Duration::from_secs(5), &metrics)
3769            .await
3770            .expect("should acquire a free gate");
3771        assert_eq!(gate.available_permits(), 0, "permit must be held");
3772        drop(permit);
3773        assert_eq!(
3774            gate.available_permits(),
3775            DEFAULT_TX_GATE_READER_PERMITS as usize,
3776            "permit pool must release on drop"
3777        );
3778    }
3779
3780    #[tokio::test]
3781    async fn begin_permit_times_out_with_clear_error_and_truthful_metric() {
3782        let gate = new_tx_gate();
3783        let metrics = Arc::new(Metrics::new());
3784        // Hold the full writer admission so the next acquire must time out.
3785        let _held = gate
3786            .clone()
3787            .acquire_many_owned(DEFAULT_TX_GATE_READER_PERMITS)
3788            .await
3789            .unwrap();
3790        let err = acquire_begin_permit(&gate, Duration::from_millis(25), &metrics)
3791            .await
3792            .expect_err("must time out while the gate is held");
3793        match err {
3794            Message::ErrorWithClass { message, class } => {
3795                assert_eq!(class, ErrorClass::Timeout);
3796                assert!(
3797                    message.contains("transaction gate timeout after 25ms"),
3798                    "unexpected message: {message}"
3799                );
3800                assert!(
3801                    message.contains("waiting for concurrent transaction"),
3802                    "unexpected message: {message}"
3803                );
3804            }
3805            other => panic!("expected Error, got {other:?}"),
3806        }
3807        let rendered = metrics.render();
3808        assert!(rendered.contains("powdb_tx_gate_timeouts_total 1"));
3809        // A timed-out begin is a failed statement from the client's view.
3810        assert!(rendered.contains("powdb_queries_total{result=\"error\"} 1"));
3811    }
3812
3813    #[test]
3814    fn admission_classification_has_query_shape_parity_and_fails_closed() {
3815        assert_eq!(
3816            classify_query_admission("User filter .id = 1"),
3817            AdmissionMode::Reader
3818        );
3819        assert_eq!(
3820            classify_sql_admission("SELECT * FROM User WHERE id = 1"),
3821            AdmissionMode::Reader
3822        );
3823        assert_eq!(
3824            classify_params_admission("User filter .id = $1", &[WireParam::Int(1)]),
3825            AdmissionMode::Reader
3826        );
3827
3828        assert_eq!(
3829            classify_query_admission("insert User { id := 1 }"),
3830            AdmissionMode::Writer
3831        );
3832        assert_eq!(
3833            classify_sql_admission("INSERT INTO User (id) VALUES (1)"),
3834            AdmissionMode::Writer
3835        );
3836        assert_eq!(
3837            classify_params_admission("insert User { id := $1 }", &[WireParam::Int(1)]),
3838            AdmissionMode::Writer
3839        );
3840        assert_eq!(
3841            classify_query_admission("this is not valid PowQL"),
3842            AdmissionMode::Writer,
3843            "uncertain statements must never enter through reader admission"
3844        );
3845    }
3846
3847    #[tokio::test]
3848    async fn writer_admission_excludes_readers() {
3849        let gate = new_tx_gate();
3850        let metrics = Arc::new(Metrics::new());
3851        let writer = acquire_begin_permit(&gate, Duration::from_secs(1), &metrics)
3852            .await
3853            .expect("writer admission");
3854        let blocked = acquire_autocommit_permit(
3855            &gate,
3856            AdmissionMode::Reader,
3857            Duration::from_millis(25),
3858            &metrics,
3859        )
3860        .await;
3861        assert!(blocked.is_err(), "reader must wait behind writer admission");
3862        drop(writer);
3863        let _reader = acquire_autocommit_permit(
3864            &gate,
3865            AdmissionMode::Reader,
3866            Duration::from_secs(1),
3867            &metrics,
3868        )
3869        .await
3870        .expect("reader must proceed after writer releases");
3871    }
3872
3873    #[tokio::test]
3874    async fn queued_writer_is_not_starved_by_later_readers() {
3875        let gate = new_tx_gate();
3876        let metrics = Arc::new(Metrics::new());
3877        let first_reader = acquire_autocommit_permit(
3878            &gate,
3879            AdmissionMode::Reader,
3880            Duration::from_secs(1),
3881            &metrics,
3882        )
3883        .await
3884        .expect("first reader admission");
3885
3886        let writer_gate = gate.clone();
3887        let writer_metrics = metrics.clone();
3888        let mut writer = tokio::spawn(async move {
3889            acquire_begin_permit(&writer_gate, Duration::from_secs(1), &writer_metrics).await
3890        });
3891        tokio::time::sleep(Duration::from_millis(10)).await;
3892
3893        let late_gate = gate.clone();
3894        let late_metrics = metrics.clone();
3895        let mut late_reader = tokio::spawn(async move {
3896            acquire_autocommit_permit(
3897                &late_gate,
3898                AdmissionMode::Reader,
3899                Duration::from_secs(1),
3900                &late_metrics,
3901            )
3902            .await
3903        });
3904        assert!(
3905            tokio::time::timeout(Duration::from_millis(25), &mut late_reader)
3906                .await
3907                .is_err(),
3908            "a later reader must queue behind the waiting writer"
3909        );
3910
3911        drop(first_reader);
3912        let writer_permit = tokio::time::timeout(Duration::from_secs(1), &mut writer)
3913            .await
3914            .expect("writer must acquire once prior readers drain")
3915            .expect("writer task")
3916            .expect("writer admission");
3917        assert!(
3918            tokio::time::timeout(Duration::from_millis(25), &mut late_reader)
3919                .await
3920                .is_err(),
3921            "writer must hold exclusive admission before the late reader"
3922        );
3923        drop(writer_permit);
3924        let _late_reader = tokio::time::timeout(Duration::from_secs(1), late_reader)
3925            .await
3926            .expect("late reader must eventually acquire")
3927            .expect("late reader task")
3928            .expect("late reader admission");
3929    }
3930
3931    fn dirty_view_engine() -> (tempfile::TempDir, Arc<RwLock<Engine>>) {
3932        let dir = tempfile::tempdir().unwrap();
3933        let mut engine = Engine::new(dir.path()).unwrap();
3934        engine
3935            .execute_powql("type Source { required id: int }")
3936            .unwrap();
3937        engine.execute_powql("insert Source { id := 1 }").unwrap();
3938        engine
3939            .execute_powql("materialize Snapshot as Source")
3940            .unwrap();
3941        engine.execute_powql("insert Source { id := 2 }").unwrap();
3942        (dir, Arc::new(RwLock::new(engine)))
3943    }
3944
3945    #[test]
3946    fn dirty_view_requests_explicit_escalation_on_every_frontend() {
3947        let (_dir, engine) = dirty_view_engine();
3948
3949        assert!(matches!(
3950            dispatch_query(&engine, "Snapshot", None, false).0,
3951            Err(QueryError::ReadonlyNeedsWrite)
3952        ));
3953        assert!(matches!(
3954            dispatch_sql_query(&engine, "SELECT * FROM Snapshot", None, false).0,
3955            Err(QueryError::ReadonlyNeedsWrite)
3956        ));
3957        assert!(matches!(
3958            dispatch_query_with_params(
3959                &engine,
3960                "Snapshot filter .id = $1",
3961                &[WireParam::Int(1)],
3962                None,
3963                false,
3964            )
3965            .0,
3966            Err(QueryError::ReadonlyNeedsWrite)
3967        ));
3968    }
3969
3970    #[tokio::test]
3971    async fn dirty_view_upgrade_waits_for_held_reader_then_records_once() {
3972        let (_dir, engine) = dirty_view_engine();
3973        let gate = new_tx_gate_with_permits(2);
3974        let metrics = Arc::new(Metrics::new());
3975        let held_reader = acquire_autocommit_permit(
3976            &gate,
3977            AdmissionMode::Reader,
3978            Duration::from_secs(1),
3979            &metrics,
3980        )
3981        .await
3982        .expect("held reader admission");
3983
3984        // Keep the peer open so the query monitor waits instead of treating
3985        // EOF as a client disconnect while the admission upgrade is blocked.
3986        let (_client, server) = tokio::io::duplex(1024);
3987        let task_gate = gate.clone();
3988        let task_metrics = metrics.clone();
3989        let mut task = tokio::spawn(async move {
3990            let mut reader = BufReader::new(server);
3991            let mut wire_read_buffer = Vec::new();
3992            let mut pending_messages = InFlightReadAhead::default();
3993            let mut tx_permit = None;
3994            execute_wire_query(
3995                engine,
3996                task_gate,
3997                &mut tx_permit,
3998                "Snapshot".into(),
3999                WireResultMode::Native,
4000                None,
4001                Duration::from_secs(2),
4002                Duration::from_secs(1),
4003                &task_metrics,
4004                &mut reader,
4005                &mut wire_read_buffer,
4006                &mut pending_messages,
4007            )
4008            .await
4009        });
4010
4011        assert!(
4012            tokio::time::timeout(Duration::from_millis(100), &mut task)
4013                .await
4014                .is_err(),
4015            "dirty-view retry must wait for exclusive admission while another reader is held"
4016        );
4017        drop(held_reader);
4018
4019        let (message, ticket, termination) = tokio::time::timeout(Duration::from_secs(2), task)
4020            .await
4021            .expect("upgrade must finish after the held reader releases")
4022            .expect("query task");
4023        assert!(matches!(message, Message::ResultRowsNative { .. }));
4024        assert!(termination.is_none());
4025
4026        let (ticket, metric) = ticket.expect("view refresh must defer its WAL metric");
4027        drop(ticket);
4028        metrics.record_query(metric.start.elapsed(), metric.outcome);
4029
4030        let rendered = metrics.render();
4031        assert!(rendered.contains("powdb_queries_total{result=\"ok\"} 1"));
4032        assert!(rendered.contains("powdb_queries_total{result=\"error\"} 0"));
4033    }
4034
4035    #[tokio::test]
4036    async fn timed_out_readonly_escalation_is_not_retried_or_reported_as_generic_error() {
4037        let (_dir, engine) = dirty_view_engine();
4038        let metrics = Arc::new(Metrics::new());
4039        // Keep the peer open so the socket monitor does not turn this into a
4040        // disconnect before the deadline fires.
4041        let (_client, server) = tokio::io::duplex(1024);
4042        let mut reader = BufReader::new(server);
4043        let mut wire_read_buffer = Vec::new();
4044        let mut pending_messages = InFlightReadAhead::default();
4045        let query_timeout = Duration::from_millis(20);
4046        let query_deadline = Instant::now() + query_timeout;
4047
4048        let (message, ticket, termination, retry) = run_blocking_query(
4049            engine,
4050            (),
4051            None,
4052            WireResultMode::Native,
4053            query_timeout,
4054            query_deadline,
4055            &metrics,
4056            &mut reader,
4057            &mut wire_read_buffer,
4058            &mut pending_messages,
4059            |_engine, (), _principal| {
4060                // Ignore the token deliberately: this reproduces the race in
4061                // which the async deadline wins but the joined task's final
4062                // result is the internal dirty-view escalation sentinel.
4063                std::thread::sleep(Duration::from_millis(50));
4064                (Err(QueryError::ReadonlyNeedsWrite), None)
4065            },
4066        )
4067        .await;
4068
4069        assert!(ticket.is_none());
4070        assert!(termination.is_none());
4071        assert!(!retry, "a timed-out query must never be resurrected");
4072        match message {
4073            Message::ErrorWithClass { message, class } => {
4074                assert_eq!(class, ErrorClass::Timeout);
4075                assert!(
4076                    message.contains("query timeout after 20ms"),
4077                    "timeout must remain client-visible, got {message}"
4078                );
4079            }
4080            other => panic!("expected timeout error, got {other:?}"),
4081        }
4082        let rendered = metrics.render();
4083        assert!(rendered.contains("powdb_query_timeouts_total 1"));
4084        assert!(rendered.contains("powdb_queries_total{result=\"error\"} 1"));
4085    }
4086
4087    #[test]
4088    fn resource_limit_errors_surface_actionable_hints() {
4089        // These carry user-actionable guidance and leak no internal state, so
4090        // they must reach the client verbatim — not be masked to the generic
4091        // message. The exact strings come from QueryError's Display impl
4092        // (crates/query/src/result.rs).
4093        for msg in [
4094            "sort input exceeds row limit — add a LIMIT clause",
4095            "join result exceeds row limit",
4096            "query exceeded memory budget: requested 100 bytes, limit 50 bytes",
4097            "result too large: encoded response exceeds 1024 bytes; add a limit or narrower projection",
4098        ] {
4099            assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
4100        }
4101    }
4102
4103    #[test]
4104    fn oversized_result_is_rejected_before_wire_encoding() {
4105        let long = "x".repeat(MAX_RESPONSE_PAYLOAD_SIZE);
4106        let result = QueryResult::Rows {
4107            columns: vec!["payload".into()],
4108            rows: vec![vec![Value::Str(long)]],
4109        };
4110        let err = query_result_to_message(result, WireResultMode::LegacyText).unwrap_err();
4111        assert!(
4112            err.to_string().starts_with("result too large"),
4113            "unexpected error: {err}"
4114        );
4115    }
4116
4117    // ---- Role enforcement (Fix: readonly role was not enforced) ----
4118
4119    fn parsed(q: &str) -> powdb_query::ast::Statement {
4120        parser::parse(q).unwrap()
4121    }
4122
4123    fn principal(role: &str) -> Option<Principal> {
4124        Some(Principal {
4125            name: "u".into(),
4126            role: role.into(),
4127        })
4128    }
4129
4130    #[test]
4131    fn readonly_can_read_but_not_write() {
4132        let p = principal("readonly");
4133        // Reads pass.
4134        assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
4135        assert!(check_statement_permitted(p.as_ref(), &parsed("count(User)")).is_ok());
4136        assert!(check_statement_permitted(p.as_ref(), &parsed("explain User")).is_ok());
4137        // Writes, DDL, and transaction control are denied.
4138        for q in [
4139            r#"insert User { name := "x" }"#,
4140            "User filter .id = 1 update { age := 2 }",
4141            "User filter .id = 1 delete",
4142            "drop User",
4143            "alter User add column c: str",
4144            "type T { required id: int }",
4145            "begin",
4146            "commit",
4147            "rollback",
4148        ] {
4149            let err = check_statement_permitted(p.as_ref(), &parsed(q))
4150                .expect_err(&format!("must deny: {q}"));
4151            assert!(
4152                err.to_string().contains("permission denied"),
4153                "unexpected error for {q}: {err}"
4154            );
4155        }
4156    }
4157
4158    #[test]
4159    fn readwrite_and_admin_have_full_query_access() {
4160        for role in ["readwrite", "admin"] {
4161            let p = principal(role);
4162            assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
4163            assert!(check_statement_permitted(
4164                p.as_ref(),
4165                &parsed(r#"insert User { name := "x" }"#)
4166            )
4167            .is_ok());
4168            assert!(check_statement_permitted(p.as_ref(), &parsed("drop User")).is_ok());
4169        }
4170    }
4171
4172    #[test]
4173    fn unknown_role_fails_closed_for_writes() {
4174        let p = principal("mystery");
4175        assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
4176        assert!(
4177            check_statement_permitted(p.as_ref(), &parsed(r#"insert User { name := "x" }"#))
4178                .is_err()
4179        );
4180    }
4181
4182    #[test]
4183    fn no_principal_means_full_access() {
4184        // Shared-password / open mode: no per-user identity, no restriction.
4185        assert!(check_statement_permitted(None, &parsed("drop User")).is_ok());
4186        assert!(check_statement_permitted(None, &parsed(r#"insert User { name := "x" }"#)).is_ok());
4187    }
4188
4189    fn store_with_alice() -> UserStore {
4190        let mut s = UserStore::new();
4191        s.create_user("alice", "pw", "readwrite").unwrap();
4192        s
4193    }
4194
4195    // ---- Empty store: legacy shared-password fallback ----
4196
4197    #[test]
4198    fn empty_store_no_password_is_open() {
4199        let s = UserStore::new();
4200        assert_eq!(
4201            authenticate_connect(&s, None, None, None),
4202            AuthOutcome::Authenticated { principal: None }
4203        );
4204        // Even a stray username/password is accepted (legacy open behavior).
4205        assert_eq!(
4206            authenticate_connect(&s, None, Some("x"), Some("y")),
4207            AuthOutcome::Authenticated { principal: None }
4208        );
4209    }
4210
4211    #[test]
4212    fn empty_store_correct_shared_password_succeeds() {
4213        let s = UserStore::new();
4214        assert_eq!(
4215            authenticate_connect(&s, Some("pw"), None, Some("pw")),
4216            AuthOutcome::Authenticated { principal: None }
4217        );
4218    }
4219
4220    #[test]
4221    fn empty_store_wrong_shared_password_rejected() {
4222        let s = UserStore::new();
4223        assert_eq!(
4224            authenticate_connect(&s, Some("pw"), None, Some("bad")),
4225            AuthOutcome::Rejected
4226        );
4227    }
4228
4229    #[test]
4230    fn empty_store_missing_password_rejected_when_expected() {
4231        let s = UserStore::new();
4232        assert_eq!(
4233            authenticate_connect(&s, Some("pw"), None, None),
4234            AuthOutcome::Rejected
4235        );
4236    }
4237
4238    #[test]
4239    fn empty_store_ignores_username_for_shared_password() {
4240        // A new client may send a username even against a shared-password
4241        // server; the username is ignored and the password still governs.
4242        let s = UserStore::new();
4243        assert_eq!(
4244            authenticate_connect(&s, Some("pw"), Some("whoever"), Some("pw")),
4245            AuthOutcome::Authenticated { principal: None }
4246        );
4247    }
4248
4249    // ---- Populated store: multi-user auth ----
4250
4251    #[test]
4252    fn user_auth_success_binds_principal() {
4253        let s = store_with_alice();
4254        assert_eq!(
4255            authenticate_connect(&s, None, Some("alice"), Some("pw")),
4256            AuthOutcome::Authenticated {
4257                principal: Some(Principal {
4258                    name: "alice".into(),
4259                    role: "readwrite".into(),
4260                })
4261            }
4262        );
4263    }
4264
4265    #[test]
4266    fn user_auth_wrong_password_rejected() {
4267        let s = store_with_alice();
4268        assert_eq!(
4269            authenticate_connect(&s, None, Some("alice"), Some("bad")),
4270            AuthOutcome::Rejected
4271        );
4272    }
4273
4274    #[test]
4275    fn user_auth_unknown_user_rejected() {
4276        let s = store_with_alice();
4277        assert_eq!(
4278            authenticate_connect(&s, None, Some("mallory"), Some("pw")),
4279            AuthOutcome::Rejected
4280        );
4281    }
4282
4283    #[test]
4284    fn user_auth_missing_username_rejected() {
4285        let s = store_with_alice();
4286        assert_eq!(
4287            authenticate_connect(&s, None, None, Some("pw")),
4288            AuthOutcome::Rejected
4289        );
4290    }
4291
4292    #[test]
4293    fn user_auth_missing_password_rejected() {
4294        let s = store_with_alice();
4295        assert_eq!(
4296            authenticate_connect(&s, Some("pw"), Some("alice"), None),
4297            AuthOutcome::Rejected
4298        );
4299    }
4300
4301    #[test]
4302    fn user_auth_ignores_shared_password_when_users_present() {
4303        // With users present, the shared password is irrelevant: supplying it as
4304        // the password without a valid user must NOT authenticate.
4305        let s = store_with_alice();
4306        assert_eq!(
4307            authenticate_connect(&s, Some("shared"), None, Some("shared")),
4308            AuthOutcome::Rejected
4309        );
4310    }
4311
4312    #[test]
4313    fn replica_fingerprint_is_stable_and_redacted() {
4314        let replica_id = "customer-prod-replica-a";
4315        let fingerprint = replica_fingerprint(replica_id);
4316        assert_eq!(fingerprint, replica_fingerprint(replica_id));
4317        assert_eq!(fingerprint, log_replica_fingerprint(replica_id));
4318        assert_ne!(fingerprint, replica_fingerprint("customer-prod-replica-b"));
4319        assert_eq!(fingerprint.len(), 16);
4320        assert!(fingerprint.chars().all(|c| c.is_ascii_hexdigit()));
4321        assert!(!fingerprint.contains("customer"));
4322        assert!(!fingerprint.contains("replica"));
4323        assert!(!fingerprint.contains(replica_id));
4324    }
4325
4326    #[test]
4327    fn invalid_replica_ids_use_fixed_log_fingerprint() {
4328        assert_eq!(log_replica_fingerprint(""), INVALID_REPLICA_FINGERPRINT);
4329        assert_eq!(
4330            log_replica_fingerprint("customer/prod/replica"),
4331            INVALID_REPLICA_FINGERPRINT
4332        );
4333        assert_eq!(
4334            log_replica_fingerprint(&"a".repeat(4096)),
4335            INVALID_REPLICA_FINGERPRINT
4336        );
4337    }
4338
4339    #[test]
4340    fn sync_error_classes_use_bounded_labels() {
4341        assert_eq!(SyncErrorClass::AuthRequired.as_label(), "auth_required");
4342        assert_eq!(
4343            SyncErrorClass::PermissionDenied.as_label(),
4344            "permission_denied"
4345        );
4346        assert_eq!(
4347            SyncErrorClass::IdentityOrFormatMismatch.as_label(),
4348            "identity_or_format_mismatch"
4349        );
4350        assert_eq!(SyncErrorClass::AckValidation.as_label(), "ack_validation");
4351        assert_eq!(SyncErrorClass::Internal.as_label(), "internal");
4352    }
4353
4354    fn sync_identity() -> DatabaseIdentity {
4355        DatabaseIdentity {
4356            database_id: *b"server-sync-test",
4357            primary_generation: 1,
4358        }
4359    }
4360
4361    fn retained_unit(lsn: u64) -> RetainedUnit {
4362        RetainedUnit {
4363            tx_id: 1,
4364            record_type: 4,
4365            lsn,
4366            data: lsn.to_le_bytes().to_vec(),
4367        }
4368    }
4369
4370    fn retained_unit_with(tx_id: u64, record_type: WalRecordType, lsn: u64) -> RetainedUnit {
4371        RetainedUnit {
4372            tx_id,
4373            record_type: record_type as u8,
4374            lsn,
4375            data: lsn.to_le_bytes().to_vec(),
4376        }
4377    }
4378
4379    fn write_sync_identity_and_tail(data_dir: &std::path::Path, through_lsn: u64) {
4380        let identity = sync_identity();
4381        write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
4382        let units = (1..=through_lsn).map(retained_unit).collect();
4383        let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
4384        write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
4385    }
4386
4387    fn write_sync_identity_and_units(data_dir: &std::path::Path, units: Vec<RetainedUnit>) {
4388        let identity = sync_identity();
4389        write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
4390        let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
4391        write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
4392    }
4393
4394    fn write_sync_identity_only(data_dir: &std::path::Path) {
4395        let identity = sync_identity();
4396        write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
4397    }
4398
4399    fn admin_principal() -> Principal {
4400        Principal {
4401            name: "admin".into(),
4402            role: "admin".into(),
4403        }
4404    }
4405
4406    #[test]
4407    fn sync_protocol_requires_credential_auth_and_rejects_readonly() {
4408        let dir = tempfile::tempdir().unwrap();
4409        let engine = Arc::new(RwLock::new(Engine::new(dir.path()).unwrap()));
4410
4411        match dispatch_sync_status(&engine, "replica-a".into(), false, None) {
4412            Message::Error { message } => {
4413                assert!(message.contains("requires authentication"));
4414            }
4415            other => panic!("expected auth error, got {other:?}"),
4416        }
4417
4418        let readonly = Principal {
4419            name: "reader".into(),
4420            role: "readonly".into(),
4421        };
4422        match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&readonly)) {
4423            Message::Error { message } => {
4424                assert!(message.contains("permission denied"));
4425            }
4426            other => panic!("expected permission error, got {other:?}"),
4427        }
4428    }
4429
4430    #[test]
4431    fn sync_status_pull_and_ack_use_server_remote_lsn() {
4432        let dir = tempfile::tempdir().unwrap();
4433        let mut engine = Engine::new(dir.path()).unwrap();
4434        engine
4435            .execute_powql("type SyncT { required id: int, v: str }")
4436            .unwrap();
4437        engine
4438            .execute_powql(r#"insert SyncT { id := 1, v := "one" }"#)
4439            .unwrap();
4440        let remote_lsn = engine.catalog().max_lsn();
4441        assert!(remote_lsn > 0);
4442        write_sync_identity_and_tail(dir.path(), remote_lsn);
4443        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4444            .unwrap();
4445
4446        let engine = Arc::new(RwLock::new(engine));
4447        let principal = admin_principal();
4448        let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
4449        {
4450            Message::SyncStatusResult { status } => status,
4451            other => panic!("expected sync status, got {other:?}"),
4452        };
4453        assert_eq!(status.remote_lsn, remote_lsn);
4454        assert_eq!(status.servable_lsn, Some(remote_lsn));
4455        assert_eq!(status.unarchived_lsn, Some(0));
4456        assert_eq!(status.last_applied_lsn, Some(0));
4457        assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4458        assert!(status.stale);
4459
4460        let identity = sync_identity().segment_identity();
4461        let pull = SyncPullRequest {
4462            replica_id: "replica-a".into(),
4463            since_lsn: 0,
4464            max_units: MAX_SYNC_PULL_UNITS,
4465            max_bytes: MAX_SYNC_PULL_BYTES,
4466            database_id: identity.database_id,
4467            primary_generation: identity.primary_generation,
4468            wal_format_version: identity.wal_format_version,
4469            catalog_version: identity.catalog_version,
4470            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4471        };
4472        let units = match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4473            Message::SyncPullResult {
4474                status,
4475                units,
4476                has_more,
4477            } => {
4478                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4479                assert!(!has_more);
4480                units
4481            }
4482            other => panic!("expected sync pull result, got {other:?}"),
4483        };
4484        assert_eq!(units.len() as u64, remote_lsn);
4485        assert_eq!(units.last().unwrap().lsn, remote_lsn);
4486
4487        let ack = match dispatch_sync_ack(
4488            &engine,
4489            "replica-a".into(),
4490            remote_lsn,
4491            remote_lsn,
4492            true,
4493            Some(&principal),
4494        ) {
4495            Message::SyncAckResult {
4496                previous_applied_lsn,
4497                applied_lsn,
4498                remote_lsn: ack_remote_lsn,
4499                advanced,
4500                status,
4501            } => {
4502                assert_eq!(previous_applied_lsn, 0);
4503                assert_eq!(applied_lsn, remote_lsn);
4504                assert_eq!(ack_remote_lsn, remote_lsn);
4505                assert!(advanced);
4506                status
4507            }
4508            other => panic!("expected sync ack result, got {other:?}"),
4509        };
4510        assert_eq!(ack.repair_action, WireSyncRepairAction::None);
4511        assert!(!ack.stale);
4512        assert_eq!(ack.lag_lsn, Some(0));
4513    }
4514
4515    fn seed_pullable_replica(engine: &mut Engine) -> u64 {
4516        let data_dir = engine.catalog().data_dir().to_path_buf();
4517        let remote_lsn = engine.catalog().max_lsn();
4518        assert!(remote_lsn > 0);
4519        write_sync_identity_and_tail(&data_dir, remote_lsn);
4520        powdb_sync::upsert_replica_cursor(&data_dir, ReplicaCursor::active("replica-a", 0))
4521            .unwrap();
4522        remote_lsn
4523    }
4524
4525    fn pull_request_with_catalog_version(catalog_version: u16) -> SyncPullRequest {
4526        let identity = sync_identity().segment_identity();
4527        SyncPullRequest {
4528            replica_id: "replica-a".into(),
4529            since_lsn: 0,
4530            max_units: MAX_SYNC_PULL_UNITS,
4531            max_bytes: MAX_SYNC_PULL_BYTES,
4532            database_id: identity.database_id,
4533            primary_generation: identity.primary_generation,
4534            wal_format_version: identity.wal_format_version,
4535            catalog_version,
4536            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4537        }
4538    }
4539
4540    #[test]
4541    fn fresh_database_expects_legacy_catalog_version_and_accepts_v5_replica() {
4542        use powdb_storage::catalog::{CATALOG_VERSION, LEGACY_CATALOG_VERSION};
4543
4544        let dir = tempfile::tempdir().unwrap();
4545        let mut engine = Engine::new(dir.path()).unwrap();
4546        engine
4547            .execute_powql("type Doc { required id: int, data: json }")
4548            .unwrap();
4549        engine
4550            .execute_powql(r#"insert Doc { id := 1, data := "{\"score\":20}" }"#)
4551            .unwrap();
4552        // No expression index created yet: the database stays at the legacy
4553        // catalog format, exactly as a v0.12 database on disk.
4554        assert_eq!(
4555            engine.catalog().active_catalog_version(),
4556            LEGACY_CATALOG_VERSION
4557        );
4558        let remote_lsn = seed_pullable_replica(&mut engine);
4559
4560        let engine = Arc::new(RwLock::new(engine));
4561        let principal = admin_principal();
4562
4563        // A replica whose maximum is the legacy version (as v0.12 clients state)
4564        // is accepted against a legacy-active server.
4565        let pull = pull_request_with_catalog_version(LEGACY_CATALOG_VERSION);
4566        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4567            Message::SyncPullResult { units, .. } => {
4568                assert_eq!(units.len() as u64, remote_lsn);
4569            }
4570            other => panic!("expected sync pull result, got {other:?}"),
4571        }
4572
4573        // A newer replica (states this binary's max) is also accepted.
4574        let pull = pull_request_with_catalog_version(CATALOG_VERSION);
4575        assert!(matches!(
4576            dispatch_sync_pull(&engine, pull, true, Some(&principal)),
4577            Message::SyncPullResult { .. }
4578        ));
4579
4580        // A replica whose maximum is older than the active format is rejected
4581        // with a message naming both versions.
4582        let pull = pull_request_with_catalog_version(LEGACY_CATALOG_VERSION - 1);
4583        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4584            Message::Error { message } => {
4585                assert!(message.contains("v4"), "message: {message}");
4586                assert!(message.contains("v5"), "message: {message}");
4587                assert!(
4588                    message.contains("rebootstrap with an upgraded replica required"),
4589                    "message: {message}"
4590                );
4591            }
4592            other => panic!("expected identity mismatch error, got {other:?}"),
4593        }
4594    }
4595
4596    #[test]
4597    fn activated_database_expects_v6_and_rejects_v5_replica() {
4598        use powdb_storage::catalog::{CATALOG_VERSION, LEGACY_CATALOG_VERSION};
4599
4600        let dir = tempfile::tempdir().unwrap();
4601        let mut engine = Engine::new(dir.path()).unwrap();
4602        engine
4603            .execute_powql("type Doc { required id: int, data: json }")
4604            .unwrap();
4605        engine
4606            .execute_powql(r#"insert Doc { id := 1, data := "{\"score\":20}" }"#)
4607            .unwrap();
4608        // Creating a JSON-path expression index activates the v6 catalog format.
4609        engine
4610            .execute_powql("alter Doc add index (.data->score)")
4611            .unwrap();
4612        assert_eq!(engine.catalog().active_catalog_version(), CATALOG_VERSION);
4613        let remote_lsn = seed_pullable_replica(&mut engine);
4614
4615        let engine = Arc::new(RwLock::new(engine));
4616        let principal = admin_principal();
4617
4618        // A v0.12 replica (states catalog_version 5) genuinely cannot read the
4619        // now-activated v6 data and is rejected with the targeted message.
4620        let pull = pull_request_with_catalog_version(LEGACY_CATALOG_VERSION);
4621        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4622            Message::Error { message } => {
4623                assert!(message.contains("v5"), "message: {message}");
4624                assert!(message.contains("v6"), "message: {message}");
4625                assert!(
4626                    message.contains("rebootstrap with an upgraded replica required"),
4627                    "message: {message}"
4628                );
4629            }
4630            other => panic!("expected identity mismatch error, got {other:?}"),
4631        }
4632
4633        // A v6-capable replica is accepted.
4634        let pull = pull_request_with_catalog_version(CATALOG_VERSION);
4635        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4636            Message::SyncPullResult { units, .. } => {
4637                assert_eq!(units.len() as u64, remote_lsn);
4638            }
4639            other => panic!("expected sync pull result, got {other:?}"),
4640        }
4641    }
4642
4643    #[test]
4644    fn sync_pull_and_ack_reject_transaction_cut_boundaries() {
4645        let dir = tempfile::tempdir().unwrap();
4646        let mut engine = Engine::new(dir.path()).unwrap();
4647        engine
4648            .execute_powql("type SyncT { required id: int }")
4649            .unwrap();
4650        for id in 1..=3 {
4651            engine
4652                .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
4653                .unwrap();
4654        }
4655        let remote_lsn = engine.catalog().max_lsn();
4656        assert!(remote_lsn >= 3);
4657        write_sync_identity_and_units(
4658            dir.path(),
4659            vec![
4660                retained_unit_with(77, WalRecordType::Begin, 1),
4661                retained_unit_with(77, WalRecordType::Insert, 2),
4662                retained_unit_with(77, WalRecordType::Commit, 3),
4663            ],
4664        );
4665        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4666            .unwrap();
4667
4668        let engine = Arc::new(RwLock::new(engine));
4669        let principal = admin_principal();
4670        let identity = sync_identity().segment_identity();
4671        let cut_pull = SyncPullRequest {
4672            replica_id: "replica-a".into(),
4673            since_lsn: 0,
4674            max_units: 2,
4675            max_bytes: MAX_SYNC_PULL_BYTES,
4676            database_id: identity.database_id,
4677            primary_generation: identity.primary_generation,
4678            wal_format_version: identity.wal_format_version,
4679            catalog_version: identity.catalog_version,
4680            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4681        };
4682        match dispatch_sync_pull(&engine, cut_pull, true, Some(&principal)) {
4683            Message::Error { message } => assert!(message.contains("cuts through transaction")),
4684            other => panic!("expected transaction-cut pull error, got {other:?}"),
4685        }
4686
4687        let cut_bytes_pull = SyncPullRequest {
4688            replica_id: "replica-a".into(),
4689            since_lsn: 0,
4690            max_units: 3,
4691            max_bytes: 58,
4692            database_id: identity.database_id,
4693            primary_generation: identity.primary_generation,
4694            wal_format_version: identity.wal_format_version,
4695            catalog_version: identity.catalog_version,
4696            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4697        };
4698        match dispatch_sync_pull(&engine, cut_bytes_pull, true, Some(&principal)) {
4699            Message::Error { message } => assert!(message.contains("cuts through transaction")),
4700            other => panic!("expected byte-capped transaction-cut pull error, got {other:?}"),
4701        }
4702
4703        let full_pull = SyncPullRequest {
4704            replica_id: "replica-a".into(),
4705            since_lsn: 0,
4706            max_units: 3,
4707            max_bytes: MAX_SYNC_PULL_BYTES,
4708            database_id: identity.database_id,
4709            primary_generation: identity.primary_generation,
4710            wal_format_version: identity.wal_format_version,
4711            catalog_version: identity.catalog_version,
4712            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4713        };
4714        match dispatch_sync_pull(&engine, full_pull, true, Some(&principal)) {
4715            Message::SyncPullResult { units, .. } => {
4716                assert_eq!(units.len(), 3);
4717                assert_eq!(units.last().unwrap().lsn, 3);
4718            }
4719            other => panic!("expected complete transaction pull, got {other:?}"),
4720        }
4721
4722        match dispatch_sync_ack(
4723            &engine,
4724            "replica-a".into(),
4725            2,
4726            remote_lsn,
4727            true,
4728            Some(&principal),
4729        ) {
4730            Message::Error { message } => assert!(message.contains("cuts through transaction")),
4731            other => panic!("expected transaction-cut ack error, got {other:?}"),
4732        }
4733        let cursor = powdb_sync::read_replica_cursors(dir.path()).unwrap();
4734        assert_eq!(cursor[0].applied_lsn, 0);
4735
4736        match dispatch_sync_ack(
4737            &engine,
4738            "replica-a".into(),
4739            3,
4740            remote_lsn,
4741            true,
4742            Some(&principal),
4743        ) {
4744            Message::SyncAckResult {
4745                previous_applied_lsn,
4746                applied_lsn,
4747                advanced,
4748                ..
4749            } => {
4750                assert_eq!(previous_applied_lsn, 0);
4751                assert_eq!(applied_lsn, 3);
4752                assert!(advanced);
4753            }
4754            other => panic!("expected complete transaction ack, got {other:?}"),
4755        }
4756    }
4757
4758    #[test]
4759    fn sync_pull_byte_cap_returns_applyable_prefix_with_reused_tx_id() {
4760        let dir = tempfile::tempdir().unwrap();
4761        let mut engine = Engine::new(dir.path()).unwrap();
4762        engine
4763            .execute_powql("type SyncT { required id: int }")
4764            .unwrap();
4765        for id in 1..=6 {
4766            engine
4767                .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
4768                .unwrap();
4769        }
4770        let remote_lsn = engine.catalog().max_lsn();
4771        assert!(remote_lsn >= 6);
4772        write_sync_identity_and_units(
4773            dir.path(),
4774            vec![
4775                retained_unit_with(1, WalRecordType::Begin, 1),
4776                retained_unit_with(1, WalRecordType::Insert, 2),
4777                retained_unit_with(1, WalRecordType::Commit, 3),
4778                retained_unit_with(1, WalRecordType::Begin, 4),
4779                retained_unit_with(1, WalRecordType::Insert, 5),
4780                retained_unit_with(1, WalRecordType::Commit, 6),
4781            ],
4782        );
4783        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4784            .unwrap();
4785
4786        let engine = Arc::new(RwLock::new(engine));
4787        let principal = admin_principal();
4788        let identity = sync_identity().segment_identity();
4789        let pull = SyncPullRequest {
4790            replica_id: "replica-a".into(),
4791            since_lsn: 0,
4792            max_units: 6,
4793            max_bytes: 100,
4794            database_id: identity.database_id,
4795            primary_generation: identity.primary_generation,
4796            wal_format_version: identity.wal_format_version,
4797            catalog_version: identity.catalog_version,
4798            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4799        };
4800
4801        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4802            Message::SyncPullResult {
4803                status,
4804                units,
4805                has_more,
4806            } => {
4807                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4808                assert_eq!(units.len(), 3);
4809                assert_eq!(units.last().unwrap().lsn, 3);
4810                assert!(has_more);
4811            }
4812            other => panic!("expected byte-capped applyable prefix, got {other:?}"),
4813        }
4814    }
4815
4816    #[test]
4817    fn sync_pull_never_serves_units_beyond_server_remote_lsn() {
4818        let dir = tempfile::tempdir().unwrap();
4819        let mut engine = Engine::new(dir.path()).unwrap();
4820        engine
4821            .execute_powql("type SyncT { required id: int }")
4822            .unwrap();
4823        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4824        let remote_lsn = engine.catalog().max_lsn();
4825        assert!(remote_lsn > 0);
4826        write_sync_identity_and_tail(dir.path(), remote_lsn + 2);
4827        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4828            .unwrap();
4829
4830        let engine = Arc::new(RwLock::new(engine));
4831        let principal = admin_principal();
4832        let identity = sync_identity().segment_identity();
4833        let pull = SyncPullRequest {
4834            replica_id: "replica-a".into(),
4835            since_lsn: 0,
4836            max_units: MAX_SYNC_PULL_UNITS,
4837            max_bytes: MAX_SYNC_PULL_BYTES,
4838            database_id: identity.database_id,
4839            primary_generation: identity.primary_generation,
4840            wal_format_version: identity.wal_format_version,
4841            catalog_version: identity.catalog_version,
4842            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4843        };
4844
4845        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4846            Message::SyncPullResult {
4847                status,
4848                units,
4849                has_more,
4850            } => {
4851                assert_eq!(status.remote_lsn, remote_lsn);
4852                assert_eq!(status.servable_lsn, Some(remote_lsn));
4853                assert_eq!(status.unarchived_lsn, Some(0));
4854                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4855                assert!(!has_more);
4856                assert_eq!(units.len() as u64, remote_lsn);
4857                assert_eq!(units.last().unwrap().lsn, remote_lsn);
4858                assert!(units.iter().all(|unit| unit.lsn <= remote_lsn));
4859            }
4860            other => panic!("expected capped sync pull result, got {other:?}"),
4861        }
4862    }
4863
4864    #[test]
4865    fn sync_status_reports_await_archive_when_primary_outruns_retained_tail() {
4866        let dir = tempfile::tempdir().unwrap();
4867        let mut engine = Engine::new(dir.path()).unwrap();
4868        engine
4869            .execute_powql("type SyncT { required id: int }")
4870            .unwrap();
4871        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4872        let remote_lsn = engine.catalog().max_lsn();
4873        assert!(remote_lsn > 0);
4874        write_sync_identity_only(dir.path());
4875        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4876            .unwrap();
4877
4878        let engine = Arc::new(RwLock::new(engine));
4879        let principal = admin_principal();
4880        let identity = sync_identity().segment_identity();
4881        let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
4882        {
4883            Message::SyncStatusResult { status } => status,
4884            other => panic!("expected sync status, got {other:?}"),
4885        };
4886        assert_eq!(status.remote_lsn, remote_lsn);
4887        assert_eq!(status.servable_lsn, Some(0));
4888        assert_eq!(status.unarchived_lsn, Some(remote_lsn));
4889        assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
4890        assert!(status
4891            .last_sync_error
4892            .as_deref()
4893            .unwrap()
4894            .contains("not yet archived"));
4895
4896        let pull = SyncPullRequest {
4897            replica_id: "replica-a".into(),
4898            since_lsn: 0,
4899            max_units: MAX_SYNC_PULL_UNITS,
4900            max_bytes: MAX_SYNC_PULL_BYTES,
4901            database_id: identity.database_id,
4902            primary_generation: identity.primary_generation,
4903            wal_format_version: identity.wal_format_version,
4904            catalog_version: identity.catalog_version,
4905            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4906        };
4907        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4908            Message::SyncPullResult {
4909                status,
4910                units,
4911                has_more,
4912            } => {
4913                assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
4914                assert!(units.is_empty());
4915                assert!(!has_more);
4916            }
4917            other => panic!("expected await-archive sync pull result, got {other:?}"),
4918        }
4919    }
4920
4921    #[test]
4922    fn sync_pull_serves_partial_retained_prefix_when_archive_lags_remote_lsn() {
4923        let dir = tempfile::tempdir().unwrap();
4924        let mut engine = Engine::new(dir.path()).unwrap();
4925        engine
4926            .execute_powql("type SyncT { required id: int }")
4927            .unwrap();
4928        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4929        engine.execute_powql("insert SyncT { id := 2 }").unwrap();
4930        let remote_lsn = engine.catalog().max_lsn();
4931        assert!(remote_lsn > 1);
4932        let servable_lsn = remote_lsn - 1;
4933        write_sync_identity_and_tail(dir.path(), servable_lsn);
4934        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4935            .unwrap();
4936
4937        let engine = Arc::new(RwLock::new(engine));
4938        let principal = admin_principal();
4939        let identity = sync_identity().segment_identity();
4940        let pull = SyncPullRequest {
4941            replica_id: "replica-a".into(),
4942            since_lsn: 0,
4943            max_units: MAX_SYNC_PULL_UNITS,
4944            max_bytes: MAX_SYNC_PULL_BYTES,
4945            database_id: identity.database_id,
4946            primary_generation: identity.primary_generation,
4947            wal_format_version: identity.wal_format_version,
4948            catalog_version: identity.catalog_version,
4949            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4950        };
4951
4952        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4953            Message::SyncPullResult {
4954                status,
4955                units,
4956                has_more,
4957            } => {
4958                assert_eq!(status.remote_lsn, remote_lsn);
4959                assert_eq!(status.servable_lsn, Some(servable_lsn));
4960                assert_eq!(status.unarchived_lsn, Some(1));
4961                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4962                assert!(!has_more);
4963                assert_eq!(units.len() as u64, servable_lsn);
4964                assert_eq!(units.last().unwrap().lsn, servable_lsn);
4965                assert!(units.iter().all(|unit| unit.lsn <= servable_lsn));
4966            }
4967            other => panic!("expected partial sync pull result, got {other:?}"),
4968        }
4969    }
4970
4971    #[test]
4972    fn sync_pull_rejects_cursor_or_format_mismatch() {
4973        let dir = tempfile::tempdir().unwrap();
4974        let mut engine = Engine::new(dir.path()).unwrap();
4975        engine
4976            .execute_powql("type SyncT { required id: int }")
4977            .unwrap();
4978        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4979        let remote_lsn = engine.catalog().max_lsn();
4980        write_sync_identity_and_tail(dir.path(), remote_lsn);
4981        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4982            .unwrap();
4983        let engine = Arc::new(RwLock::new(engine));
4984        let principal = admin_principal();
4985        let identity = sync_identity().segment_identity();
4986
4987        let wrong_cursor = SyncPullRequest {
4988            replica_id: "replica-a".into(),
4989            since_lsn: 1,
4990            max_units: 10,
4991            max_bytes: MAX_SYNC_PULL_BYTES,
4992            database_id: identity.database_id,
4993            primary_generation: identity.primary_generation,
4994            wal_format_version: identity.wal_format_version,
4995            catalog_version: identity.catalog_version,
4996            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4997        };
4998        match dispatch_sync_pull(&engine, wrong_cursor, true, Some(&principal)) {
4999            Message::Error { message } => assert!(message.contains("does not match")),
5000            other => panic!("expected cursor mismatch error, got {other:?}"),
5001        }
5002
5003        let wrong_format = SyncPullRequest {
5004            replica_id: "replica-a".into(),
5005            since_lsn: 0,
5006            max_units: 10,
5007            max_bytes: MAX_SYNC_PULL_BYTES,
5008            database_id: identity.database_id,
5009            primary_generation: identity.primary_generation,
5010            wal_format_version: identity.wal_format_version,
5011            catalog_version: identity.catalog_version,
5012            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION + 1,
5013        };
5014        match dispatch_sync_pull(&engine, wrong_format, true, Some(&principal)) {
5015            Message::Error { message } => assert!(message.contains("rebootstrap required")),
5016            other => panic!("expected format mismatch error, got {other:?}"),
5017        }
5018    }
5019}