perf: LRU translate cache + prepare_cached + kill LIMIT-0 probe + utf8 fix

Cargo.toml

@@ -35,8 +35,10 @@ pgwire = "0.28"
 # SQL parsing
 sqlparser = { version = "0.57", features = ["serde"] }
 
-# Backends
-rusqlite = { version = "0.32", features = ["bundled"] }
+# Backends. `column_decltype` exposes `Statement::columns()`, which we need
+# to fill in `Column.decltype` on empty result sets (an unexecuted prepared
+# SELECT). Without it, wire-frontend type inference has to LIMIT-0 probe.
+rusqlite = { version = "0.32", features = ["bundled", "column_decltype"] }
 
 # HTTP (Hrana)
 axum = "0.8"

crates/litewire-backend/src/lib.rs

@@ -77,6 +77,24 @@ pub trait Backend: Send + Sync + 'static {
 
     /// Execute a statement that modifies data.
     async fn execute(&self, sql: &str, params: &[Value]) -> Result<ExecuteResult, BackendError>;
+
+    /// Describe the columns a prepared SELECT would produce, *without*
+    /// executing it. Used by wire frontends to answer `COM_STMT_PREPARE`
+    /// (MySQL) / `Describe` (Postgres) without a `LIMIT 0` round trip
+    /// through the query planner.
+    ///
+    /// Default implementation falls back to running the statement with a
+    /// `LIMIT 0` wrapper (the previous behaviour). Backends that can
+    /// inspect the prepared statement directly (rusqlite) override this.
+    ///
+    /// # Errors
+    ///
+    /// Returns the underlying [`BackendError`] on parse / prepare failure.
+    async fn describe_columns(&self, sql: &str) -> Result<Vec<Column>, BackendError> {
+        let probe = format!("{sql} LIMIT 0");
+        let rs = self.query(&probe, &[]).await?;
+        Ok(rs.columns)
+    }
 }
 
 /// Type alias for a shared backend reference.

crates/litewire-backend/src/rusqlite_backend.rs

@@ -66,17 +66,44 @@ fn bind_params(params: &[Value]) -> Vec<Box<dyn rusqlite::types::ToSql>> {
 }
 
 /// Extract a [`Value`] from a rusqlite row at the given column index.
+///
+/// TEXT handling: SQLite `TEXT` cells hold raw bytes that are *supposed* to
+/// be UTF-8. If they are not, this used to lossy-decode via
+/// `String::from_utf8_lossy`, which silently replaces bytes with U+FFFD and
+/// corrupts round-trips (e.g. arbitrary latin-1 or WTF-8 stored by a
+/// previous client). Instead, we surface invalid UTF-8 as a `Blob` so the
+/// caller / wire frontend can decide how to send it back to the client.
 fn extract_value(row: &rusqlite::Row<'_>, idx: usize) -> Result<Value, rusqlite::Error> {
     use rusqlite::types::ValueRef;
     match row.get_ref(idx)? {
         ValueRef::Null => Ok(Value::Null),
         ValueRef::Integer(i) => Ok(Value::Integer(i)),
         ValueRef::Real(f) => Ok(Value::Float(f)),
-        ValueRef::Text(s) => Ok(Value::Text(String::from_utf8_lossy(s).into_owned())),
+        ValueRef::Text(s) => match std::str::from_utf8(s) {
+            Ok(v) => Ok(Value::Text(v.to_string())),
+            Err(_) => Ok(Value::Blob(s.to_vec())),
+        },
         ValueRef::Blob(b) => Ok(Value::Blob(b.to_vec())),
     }
 }
 
+/// Read the column metadata off a prepared `Statement`, using rusqlite's
+/// `column_decltype` feature to populate `Column.decltype`.
+///
+/// This is called *without* stepping the statement -- SELECTs on empty
+/// tables now get real column types on the wire, which fixes clients
+/// (Laravel `Query::exists()`, PDO `getColumnMeta()`) that would otherwise
+/// see everything as `NULL` / `TEXT`.
+fn describe_columns(stmt: &rusqlite::Statement<'_>) -> Vec<Column> {
+    stmt.columns()
+        .iter()
+        .map(|c| Column {
+            name: c.name().to_string(),
+            decltype: c.decl_type().map(str::to_string),
+        })
+        .collect()
+}
+
 #[async_trait::async_trait]
 impl Backend for Rusqlite {
     async fn query(&self, sql: &str, params: &[Value]) -> Result<ResultSet, BackendError> {
@@ -86,17 +113,16 @@ impl Backend for Rusqlite {
 
         task::spawn_blocking(move || {
             let conn = conn.lock();
+            // `prepare_cached` interns the parsed statement in rusqlite's
+            // per-connection LRU. Repeated identical SQL (which is the norm
+            // for prepared-statement heavy workloads and for the KV/session
+            // handler path) then avoids the sqlite3_prepare_v2 cost.
             let mut stmt = conn
-                .prepare(&sql)
+                .prepare_cached(&sql)
                 .map_err(|e| BackendError::Sqlite(e.to_string()))?;
 
-            let col_count = stmt.column_count();
-            let columns: Vec<Column> = (0..col_count)
-                .map(|i| Column {
-                    name: stmt.column_name(i).unwrap_or("?").to_string(),
-                    decltype: None, // rusqlite 0.32 does not expose column_decltype on Statement
-                })
-                .collect();
+            let columns = describe_columns(&stmt);
+            let col_count = columns.len();
 
             let bound = bind_params(&params);
             let param_refs: Vec<&dyn rusqlite::types::ToSql> =
@@ -129,6 +155,21 @@ impl Backend for Rusqlite {
         .map_err(|e| BackendError::Other(format!("spawn_blocking join error: {e}")))?
     }
 
+    async fn describe_columns(&self, sql: &str) -> Result<Vec<Column>, BackendError> {
+        let conn = Arc::clone(&self.conn);
+        let sql = sql.to_string();
+
+        task::spawn_blocking(move || {
+            let conn = conn.lock();
+            let stmt = conn
+                .prepare_cached(&sql)
+                .map_err(|e| BackendError::Sqlite(e.to_string()))?;
+            Ok(describe_columns(&stmt))
+        })
+        .await
+        .map_err(|e| BackendError::Other(format!("spawn_blocking join error: {e}")))?
+    }
+
     async fn execute(&self, sql: &str, params: &[Value]) -> Result<ExecuteResult, BackendError> {
         let conn = Arc::clone(&self.conn);
         let sql = sql.to_string();
@@ -141,8 +182,11 @@ impl Backend for Rusqlite {
             let param_refs: Vec<&dyn rusqlite::types::ToSql> =
                 bound.iter().map(|b| b.as_ref()).collect();
 
-            let affected = conn
-                .execute(&sql, param_refs.as_slice())
+            let mut stmt = conn
+                .prepare_cached(&sql)
+                .map_err(|e| BackendError::Sqlite(e.to_string()))?;
+            let affected = stmt
+                .execute(param_refs.as_slice())
                 .map_err(|e| BackendError::Sqlite(e.to_string()))?;
 
             let last_id = conn.last_insert_rowid();
@@ -407,6 +451,74 @@ mod tests {
         assert_eq!(rs.columns[2].name, "email");
     }
 
+    #[tokio::test]
+    async fn describe_columns_returns_decltypes() {
+        // Verify the fix for the previous "decltype is always None" behaviour:
+        // describe_columns must fill in the declared type without executing
+        // the statement.
+        let backend = Rusqlite::memory().unwrap();
+        backend
+            .execute(
+                "CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL, score REAL)",
+                &[],
+            )
+            .await
+            .unwrap();
+
+        let cols = backend
+            .describe_columns("SELECT id, name, score FROM users")
+            .await
+            .unwrap();
+        assert_eq!(cols.len(), 3);
+        assert_eq!(cols[0].name, "id");
+        assert_eq!(cols[0].decltype.as_deref(), Some("INTEGER"));
+        assert_eq!(cols[1].name, "name");
+        assert_eq!(cols[1].decltype.as_deref(), Some("TEXT"));
+        assert_eq!(cols[2].name, "score");
+        assert_eq!(cols[2].decltype.as_deref(), Some("REAL"));
+    }
+
+    #[tokio::test]
+    async fn describe_columns_null_decltype_for_expressions() {
+        // Bare expressions like `SELECT 1 + 2` have no declared type.
+        let backend = Rusqlite::memory().unwrap();
+        let cols = backend.describe_columns("SELECT 1 + 2 AS x").await.unwrap();
+        assert_eq!(cols.len(), 1);
+        assert!(cols[0].decltype.is_none(), "got: {:?}", cols[0].decltype);
+    }
+
+    #[tokio::test]
+    async fn invalid_utf8_in_text_column_surfaces_as_blob() {
+        // Store a byte sequence that is *not* valid UTF-8 into a TEXT column.
+        // The old code used `String::from_utf8_lossy`, which would silently
+        // corrupt the bytes with U+FFFD replacements; the new code must
+        // surface the raw bytes intact as a Blob so callers can round-trip.
+        let backend = Rusqlite::memory().unwrap();
+        backend
+            .execute("CREATE TABLE t (v TEXT)", &[])
+            .await
+            .unwrap();
+
+        // Direct-write invalid UTF-8 via CAST -- x'ffff' is 0xFF 0xFF which
+        // is not a valid UTF-8 sequence.
+        backend
+            .execute("INSERT INTO t VALUES (CAST(x'ffff' AS TEXT))", &[])
+            .await
+            .unwrap();
+
+        let rs = backend.query("SELECT v FROM t", &[]).await.unwrap();
+        assert_eq!(rs.rows.len(), 1);
+        match &rs.rows[0][0] {
+            Value::Blob(b) => assert_eq!(b.as_slice(), &[0xFF, 0xFF][..]),
+            Value::Text(s) => panic!(
+                "expected Blob fallback for invalid UTF-8, got Text({:?}) with bytes {:?}",
+                s,
+                s.as_bytes()
+            ),
+            other => panic!("expected Blob for invalid UTF-8, got: {other:?}"),
+        }
+    }
+
     #[tokio::test]
     async fn query_with_alias() {
         let backend = Rusqlite::memory().unwrap();

crates/litewire-mysql/src/handler.rs

@@ -4,9 +4,10 @@
 //! prepared statement prepare/execute/close.
 
 use std::collections::HashMap;
+use std::sync::Arc;
 
 use litewire_backend::{SharedBackend, Value};
-use litewire_translate::{self, Dialect, StatementKind, TranslateResult, classify};
+use litewire_translate::{self, Dialect, StatementKind, TranslateCache, TranslateResult, classify};
 use opensrv_mysql::*;
 use tokio::io::AsyncWrite;
 use tracing::{debug, warn};
@@ -49,6 +50,8 @@ struct PreparedStmt {
 /// Handler for a single MySQL client connection.
 pub struct LiteWireHandler {
     backend: SharedBackend,
+    /// Shared translation cache across all connections on this frontend.
+    translate_cache: Arc<TranslateCache>,
     /// Prepared statements keyed by the statement ID assigned during `on_prepare`.
     stmts: HashMap<u32, PreparedStmt>,
     /// Next statement ID to assign.
@@ -58,9 +61,10 @@ pub struct LiteWireHandler {
 }
 
 impl LiteWireHandler {
-    pub fn new(backend: SharedBackend) -> Self {
+    pub fn new(backend: SharedBackend, translate_cache: Arc<TranslateCache>) -> Self {
         Self {
             backend,
+            translate_cache,
             stmts: HashMap::new(),
             next_stmt_id: 1,
             in_transaction: false,
@@ -91,10 +95,18 @@ impl LiteWireHandler {
 
                 for row in &rs.rows {
                     for val in row {
+                        // Write each value in its native form so the binary
+                        // (prepared-statement) protocol accepts it against
+                        // the declared column type. Previously integers were
+                        // stringified, which worked only because every
+                        // column was declared as VAR_STRING. Now that
+                        // decltype flows through from `column_decltype`,
+                        // integer columns arrive at the wire as LONGLONG
+                        // and opensrv rejects a `String` payload.
                         match val {
                             Value::Null => rw.write_col(None::<&str>)?,
-                            Value::Integer(i) => rw.write_col(i.to_string())?,
-                            Value::Float(f) => rw.write_col(f.to_string())?,
+                            Value::Integer(i) => rw.write_col(*i)?,
+                            Value::Float(f) => rw.write_col(*f)?,
                             Value::Text(s) => rw.write_col(s.as_str())?,
                             Value::Blob(b) => rw.write_col(b.as_slice())?,
                         }
@@ -155,7 +167,8 @@ impl LiteWireHandler {
     /// Translate SQL and return the first translated result, or an error string.
     fn translate_sql(&self, query: &str) -> Result<(String, StatementKind), String> {
         let translated =
-            litewire_translate::translate(query, Dialect::MySQL).map_err(|e| e.to_string())?;
+            litewire_translate::translate_cached(&self.translate_cache, query, Dialect::MySQL)
+                .map_err(|e| e.to_string())?;
 
         let Some(result) = translated.into_iter().next() else {
             return Ok((String::new(), StatementKind::Other));
@@ -233,12 +246,14 @@ impl<W: AsyncWrite + Send + Unpin> AsyncMysqlShim<W> for LiteWireHandler {
             })
             .collect();
 
-        // Try to determine output columns for queries.
+        // Determine output columns. INSERT/UPDATE/DELETE have no result
+        // set -- skip describing them entirely. For SELECTs, use the
+        // backend's `describe_columns`, which on the rusqlite backend
+        // reads column metadata off the prepared statement without
+        // executing it (was: `SELECT ... LIMIT 0` round trip).
         let columns = if kind == StatementKind::Query && !sqlite_sql.is_empty() {
-            let probe_sql = format!("{sqlite_sql} LIMIT 0");
-            match self.backend.query(&probe_sql, &[]).await {
-                Ok(rs) => rs
-                    .columns
+            match self.backend.describe_columns(&sqlite_sql).await {
+                Ok(cols) => cols
                     .iter()
                     .map(|c| Column {
                         table: String::new(),
@@ -332,7 +347,11 @@ impl<W: AsyncWrite + Send + Unpin> AsyncMysqlShim<W> for LiteWireHandler {
     ) -> Result<(), Self::Error> {
         debug!(sql = %query, "COM_QUERY");
 
-        let translated = match litewire_translate::translate(query, Dialect::MySQL) {
+        let translated = match litewire_translate::translate_cached(
+            &self.translate_cache,
+            query,
+            Dialect::MySQL,
+        ) {
             Ok(r) => r,
             Err(e) => {
                 warn!("SQL translation error: {e}");

crates/litewire-mysql/src/lib.rs

@@ -13,6 +13,7 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 
 use litewire_backend::SharedBackend;
+use litewire_translate::TranslateCache;
 use tokio::net::TcpListener;
 use tracing::{debug, info, warn};
 
@@ -50,6 +51,11 @@ impl MysqlFrontend {
         info!(listen = %self.config.listen, "MySQL frontend listening");
 
         let backend = Arc::clone(&self.backend);
+        // Shared parse+rewrite cache across every accepted connection.
+        // Hot workloads (WordPress, Laravel) re-issue the same handful of
+        // prepared statements repeatedly; caching drops sqlparser off the
+        // hot path entirely.
+        let translate_cache = Arc::new(TranslateCache::default());
 
         loop {
             let (stream, peer) = match listener.accept().await {
@@ -66,8 +72,9 @@ impl MysqlFrontend {
             debug!(%peer, "MySQL client connected");
 
             let be = Arc::clone(&backend);
+            let cache = Arc::clone(&translate_cache);
             tokio::spawn(async move {
-                let handler = LiteWireHandler::new(be);
+                let handler = LiteWireHandler::new(be, cache);
                 let (reader, writer) = stream.into_split();
                 if let Err(e) =
                     opensrv_mysql::AsyncMysqlIntermediary::run_on(handler, reader, writer).await