From c00b26d0786e8142d06eabb924267e4a29c0b545 Mon Sep 17 00:00:00 2001 From: Luther Monson Date: Thu, 9 Jul 2026 19:27:57 -0700 Subject: [PATCH] perf: LRU translate cache + prepare_cached + kill LIMIT-0 probe + utf8 fix Follows compat/session-statements-and-ci; that PR ships the CI + clippy cleanup and any conflicts should be resolved by merging PR A first. Translate LRU cache - New crate::cache module: bounded LRU keyed by (Dialect, SQL text) with default capacity 1024 -- enough to hold every distinct prepared statement WordPress + Woocommerce issue on a single page, small enough that worst-case memory is a few MB. - lib.rs exposes translate_cached(cache, sql, dialect) alongside translate(); errors are not stored so a repeatedly-bad statement always re-parses. - MysqlFrontend and PostgresFrontend each own an Arc shared across all accepted connections. - Micro-benchmark test (translate_cache_speedup) times cold vs cached on a WordPress-shaped SELECT with joins/ORDER BY/LIMIT/IN(?,?,?): debug : cold=191us/call warm=2.2us/call => 87x speedup release : cold=39us/call warm=277ns/call => 140x speedup prepare_cached in rusqlite backend - query()/execute() switched to conn.prepare_cached(), which interns the parsed statement in rusqlite's per-connection LRU. Second and later identical SQL avoid the sqlite3_prepare_v2 cost entirely. column_decltype + kill the LIMIT 0 probe - Enable the rusqlite `column_decltype` feature so Statement::columns() exposes decl_type(). - New Backend::describe_columns(sql) trait method with a default LIMIT-0 probe implementation; rusqlite overrides it to read metadata off the prepared statement without executing. - MySQL on_prepare: skip describe entirely for INSERT/UPDATE/DELETE (they have no result set); for SELECTs, use describe_columns instead of the old `SELECT ... LIMIT 0` probe. - PostgreSQL probe_columns: use describe_columns first; only fall back to a LIMIT 1 probe if any column is untyped (`SELECT 1 + 2` style). - Empty-result-set SELECTs now carry real column types on the wire, which fixes Laravel Query::exists() and PDO::getColumnMeta(). utf8 fallback in rusqlite backend - extract_value on ValueRef::Text used to lossy-decode via String::from_utf8_lossy, which silently corrupts non-UTF-8 bytes with U+FFFD. Invalid UTF-8 now surfaces as Value::Blob instead so the wire frontend can round-trip the raw bytes unchanged. Wire-format regression fix (uncovered by column_decltype flip) - MySQL row writer previously stringified integers/floats. That worked only because every column arrived at the wire as VAR_STRING. With real decltypes flowing through, integer columns become MYSQL_TYPE_LONGLONG and opensrv-mysql rejects a String payload for a LONGLONG column ("tried to use [49] as MYSQL_TYPE_LONGLONG"). Row writer now writes Value::Integer/Value::Float natively. Dialect gains Hash so it can be used as a HashMap/LruCache key. --- Cargo.lock | 2 + Cargo.toml | 6 +- crates/litewire-backend/src/lib.rs | 18 +++ .../litewire-backend/src/rusqlite_backend.rs | 134 +++++++++++++++-- crates/litewire-mysql/src/handler.rs | 41 ++++-- crates/litewire-mysql/src/lib.rs | 9 +- crates/litewire-postgres/src/handler.rs | 50 +++++-- crates/litewire-postgres/src/lib.rs | 9 +- crates/litewire-translate/Cargo.toml | 5 + crates/litewire-translate/src/cache.rs | 139 ++++++++++++++++++ crates/litewire-translate/src/lib.rs | 101 ++++++++++++- 11 files changed, 476 insertions(+), 38 deletions(-) create mode 100644 crates/litewire-translate/src/cache.rs diff --git a/Cargo.lock b/Cargo.lock index da16fc1..a7d6080 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1689,6 +1689,8 @@ dependencies = [ name = "litewire-translate" version = "0.1.0" dependencies = [ + "lru", + "parking_lot", "sqlparser", "thiserror 2.0.18", "tracing", diff --git a/Cargo.toml b/Cargo.toml index cec725c..929d126 100644 --- a/Cargo.toml +++ b/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" diff --git a/crates/litewire-backend/src/lib.rs b/crates/litewire-backend/src/lib.rs index c391fd3..03e3264 100644 --- a/crates/litewire-backend/src/lib.rs +++ b/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; + + /// 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, BackendError> { + let probe = format!("{sql} LIMIT 0"); + let rs = self.query(&probe, &[]).await?; + Ok(rs.columns) + } } /// Type alias for a shared backend reference. diff --git a/crates/litewire-backend/src/rusqlite_backend.rs b/crates/litewire-backend/src/rusqlite_backend.rs index ea9be27..448fa1b 100644 --- a/crates/litewire-backend/src/rusqlite_backend.rs +++ b/crates/litewire-backend/src/rusqlite_backend.rs @@ -66,17 +66,44 @@ fn bind_params(params: &[Value]) -> Vec> { } /// 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 { 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 { + 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 { @@ -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 = (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(¶ms); 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, 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 { 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(); diff --git a/crates/litewire-mysql/src/handler.rs b/crates/litewire-mysql/src/handler.rs index 06ed08c..1ca5553 100644 --- a/crates/litewire-mysql/src/handler.rs +++ b/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, /// Prepared statements keyed by the statement ID assigned during `on_prepare`. stmts: HashMap, /// 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) -> 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 AsyncMysqlShim 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 AsyncMysqlShim 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}"); diff --git a/crates/litewire-mysql/src/lib.rs b/crates/litewire-mysql/src/lib.rs index a224b79..c451944 100644 --- a/crates/litewire-mysql/src/lib.rs +++ b/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 diff --git a/crates/litewire-postgres/src/handler.rs b/crates/litewire-postgres/src/handler.rs index db9bc89..a8c1cc7 100644 --- a/crates/litewire-postgres/src/handler.rs +++ b/crates/litewire-postgres/src/handler.rs @@ -21,7 +21,7 @@ use pgwire::messages::data::DataRow; use tracing::{debug, warn}; use litewire_backend::{SharedBackend, Value}; -use litewire_translate::{self, Dialect, StatementKind, TranslateResult, classify}; +use litewire_translate::{self, Dialect, StatementKind, TranslateCache, TranslateResult, classify}; use crate::error_map; use crate::types::sqlite_to_pg_type; @@ -30,20 +30,23 @@ use crate::types::sqlite_to_pg_type; pub struct PostgresHandler { backend: SharedBackend, query_parser: Arc, + translate_cache: Arc, } impl PostgresHandler { - pub fn new(backend: SharedBackend) -> Self { + pub fn new(backend: SharedBackend, translate_cache: Arc) -> Self { Self { backend, query_parser: Arc::new(NoopQueryParser::new()), + translate_cache, } } /// Translate SQL from PostgreSQL dialect to SQLite and classify it. fn translate_sql(&self, query: &str) -> Result<(String, StatementKind), String> { let translated = - litewire_translate::translate(query, Dialect::PostgreSQL).map_err(|e| e.to_string())?; + litewire_translate::translate_cached(&self.translate_cache, query, Dialect::PostgreSQL) + .map_err(|e| e.to_string())?; let Some(result) = translated.into_iter().next() else { return Ok((String::new(), StatementKind::Other)); @@ -204,11 +207,37 @@ impl PostgresHandler { Ok(Response::Execution(tag)) } - /// Build column metadata for a query by probing with LIMIT 1. + /// Build column metadata for a query. /// - /// Uses LIMIT 1 (not LIMIT 0) so we can infer types from actual data - /// when columns lack declared types (e.g. `SELECT 1 + 2`). + /// Fast path: call the backend's `describe_columns`, which on rusqlite + /// reads column types off the prepared statement without executing. + /// If any column comes back without a declared type (typical for + /// expression columns like `SELECT 1 + 2`), fall back to a `LIMIT 1` + /// probe so we can infer from an actual value. async fn probe_columns(&self, sql: &str, format: &Format) -> PgWireResult> { + let cols = match self.backend.describe_columns(sql).await { + Ok(c) => c, + Err(_) => return Ok(vec![]), + }; + + let has_untyped = cols.iter().any(|c| c.decltype.is_none()); + if !has_untyped { + return Ok(cols + .iter() + .enumerate() + .map(|(idx, col)| { + FieldInfo::new( + col.name.clone(), + None, + None, + sqlite_to_pg_type(col.decltype.as_deref()), + format.format_for(idx), + ) + }) + .collect()); + } + + // At least one expression column; probe for a real row to infer. let probe = format!("{sql} LIMIT 1"); match self.backend.query(&probe, &[]).await { Ok(rs) => Ok(rs @@ -372,10 +401,11 @@ impl SimpleQueryHandler for PostgresHandler { debug!(sql = %query, "PG simple query"); let translated = - litewire_translate::translate(query, Dialect::PostgreSQL).map_err(|e| { - warn!("SQL translation error: {e}"); - pg_error(&e.to_string()) - })?; + litewire_translate::translate_cached(&self.translate_cache, query, Dialect::PostgreSQL) + .map_err(|e| { + warn!("SQL translation error: {e}"); + pg_error(&e.to_string()) + })?; let mut responses = Vec::new(); diff --git a/crates/litewire-postgres/src/lib.rs b/crates/litewire-postgres/src/lib.rs index 13640ea..0d6d158 100644 --- a/crates/litewire-postgres/src/lib.rs +++ b/crates/litewire-postgres/src/lib.rs @@ -12,6 +12,7 @@ use std::net::SocketAddr; use std::sync::Arc; use litewire_backend::SharedBackend; +use litewire_translate::TranslateCache; use pgwire::api::NoopErrorHandler; use pgwire::api::PgWireServerHandlers; use pgwire::api::auth::noop::NoopStartupHandler; @@ -53,8 +54,14 @@ impl PostgresFrontend { let listener = TcpListener::bind(self.config.listen).await?; info!(listen = %self.config.listen, "PostgreSQL frontend listening"); + // Shared parse+rewrite cache across every accepted connection -- + // same rationale as the MySQL frontend. + let translate_cache = Arc::new(TranslateCache::default()); let factory = Arc::new(LiteWireHandlerFactory { - handler: Arc::new(PostgresHandler::new(Arc::clone(&self.backend))), + handler: Arc::new(PostgresHandler::new( + Arc::clone(&self.backend), + translate_cache, + )), }); loop { diff --git a/crates/litewire-translate/Cargo.toml b/crates/litewire-translate/Cargo.toml index 04e49ef..257b617 100644 --- a/crates/litewire-translate/Cargo.toml +++ b/crates/litewire-translate/Cargo.toml @@ -10,3 +10,8 @@ license.workspace = true sqlparser = { workspace = true } thiserror = { workspace = true } tracing = { workspace = true } + +# LRU cache in front of the parse+rewrite pipeline. Small dep (no proc-macros, +# hashbrown-only), already transitively pulled in by other workspace members. +lru = "0.12" +parking_lot = { workspace = true } diff --git a/crates/litewire-translate/src/cache.rs b/crates/litewire-translate/src/cache.rs new file mode 100644 index 0000000..4a7d92f --- /dev/null +++ b/crates/litewire-translate/src/cache.rs @@ -0,0 +1,139 @@ +//! Bounded LRU cache for translated SQL. +//! +//! The parse-then-rewrite pipeline in [`crate::translate`] dominates +//! per-request latency for short queries; WordPress alone can issue ~60 +//! prepared-statement `SELECT`s to render a single page. Caching the exact +//! text -> `Vec` mapping shaves that cost off every repeat +//! statement. +//! +//! Design notes: +//! * The key is `(Dialect, String)` -- same text under different dialects +//! can produce different output (e.g. `SELECT $1` is a placeholder in PG +//! but a syntax error in MySQL), so dialect must be part of the key. +//! * Cache misses always call the underlying translator; there's no +//! negative caching, so a parse error is not stored. +//! * Values are cheap to clone (`String` + a small enum), so the LRU is +//! fine returning owned copies to the caller. + +use std::num::NonZeroUsize; + +use lru::LruCache; +use parking_lot::Mutex; + +use crate::{Dialect, TranslateResult}; + +/// Default per-process cache capacity. Sized to comfortably hold every +/// distinct prepared statement WordPress + Woocommerce issue in a request +/// (~200) plus headroom for admin dashboards; small enough that even with +/// worst-case 4 KB queries the total memory is a few MB. +pub const DEFAULT_CAPACITY: usize = 1024; + +/// A thread-safe bounded LRU cache keyed by dialect+SQL text. +pub struct TranslateCache { + inner: Mutex>>, +} + +impl TranslateCache { + /// Build a cache with the given capacity. Panics if `capacity` is zero. + #[must_use] + pub fn new(capacity: usize) -> Self { + let cap = NonZeroUsize::new(capacity).expect("cache capacity must be > 0"); + Self { + inner: Mutex::new(LruCache::new(cap)), + } + } + + /// Look up a previously translated result. Marks the entry as + /// recently-used on hit. + pub fn get(&self, dialect: Dialect, sql: &str) -> Option> { + // `lru::LruCache::get` requires `&mut`, hence the mutex. + let mut inner = self.inner.lock(); + inner.get(&(dialect, sql.to_string())).cloned() + } + + /// Insert a translated result. + pub fn put(&self, dialect: Dialect, sql: String, results: Vec) { + let mut inner = self.inner.lock(); + inner.put((dialect, sql), results); + } + + /// Return the number of entries currently held. + pub fn len(&self) -> usize { + self.inner.lock().len() + } + + /// Whether the cache is empty. + pub fn is_empty(&self) -> bool { + self.inner.lock().is_empty() + } +} + +impl Default for TranslateCache { + fn default() -> Self { + Self::new(DEFAULT_CAPACITY) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cache_hit_returns_cloned_value() { + let cache = TranslateCache::new(4); + cache.put( + Dialect::MySQL, + "SELECT 1".to_string(), + vec![TranslateResult::Sql("SELECT 1".to_string())], + ); + let hit = cache.get(Dialect::MySQL, "SELECT 1").expect("should hit"); + assert_eq!(hit.len(), 1); + assert!(matches!(hit[0], TranslateResult::Sql(ref s) if s == "SELECT 1")); + } + + #[test] + fn cache_miss_on_different_dialect() { + // Same SQL text under a different dialect must not collide. + let cache = TranslateCache::new(4); + cache.put( + Dialect::MySQL, + "SELECT 1".to_string(), + vec![TranslateResult::Sql("mysql".to_string())], + ); + assert!(cache.get(Dialect::PostgreSQL, "SELECT 1").is_none()); + } + + #[test] + fn cache_evicts_lru() { + let cache = TranslateCache::new(2); + cache.put( + Dialect::MySQL, + "A".to_string(), + vec![TranslateResult::Sql("a".into())], + ); + cache.put( + Dialect::MySQL, + "B".to_string(), + vec![TranslateResult::Sql("b".into())], + ); + // Touching A promotes it to most-recently-used. + let _ = cache.get(Dialect::MySQL, "A"); + // Adding C evicts B (LRU). + cache.put( + Dialect::MySQL, + "C".to_string(), + vec![TranslateResult::Sql("c".into())], + ); + assert!(cache.get(Dialect::MySQL, "B").is_none()); + assert!(cache.get(Dialect::MySQL, "A").is_some()); + assert!(cache.get(Dialect::MySQL, "C").is_some()); + } + + #[test] + fn len_tracks_inserts() { + let cache = TranslateCache::new(4); + assert!(cache.is_empty()); + cache.put(Dialect::MySQL, "A".into(), vec![TranslateResult::Noop]); + assert_eq!(cache.len(), 1); + } +} diff --git a/crates/litewire-translate/src/lib.rs b/crates/litewire-translate/src/lib.rs index 4e971bd..9e4cd44 100644 --- a/crates/litewire-translate/src/lib.rs +++ b/crates/litewire-translate/src/lib.rs @@ -3,6 +3,7 @@ //! Translates MySQL, PostgreSQL, and T-SQL dialects to SQLite-compatible SQL //! using `sqlparser-rs` for parsing and AST manipulation. +pub mod cache; pub mod common; pub mod emit; pub mod metadata; @@ -10,12 +11,14 @@ pub mod mysql; pub mod postgres; pub mod tds; +pub use cache::TranslateCache; + use sqlparser::ast::Statement; use sqlparser::dialect::{MsSqlDialect, MySqlDialect, PostgreSqlDialect}; use sqlparser::parser::Parser; /// Source SQL dialect for translation. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum Dialect { MySQL, PostgreSQL, @@ -33,7 +36,7 @@ pub enum TranslateError { } /// Result of translating a SQL statement. -#[derive(Debug)] +#[derive(Debug, Clone)] pub enum TranslateResult { /// Translated SQL string to execute against SQLite. Sql(String), @@ -89,6 +92,30 @@ pub fn translate(sql: &str, dialect: Dialect) -> Result, Tr Ok(results) } +/// Translate a SQL string, using a bounded LRU cache in front of the +/// parser + rewriter. See [`TranslateCache`]. +/// +/// Cache hits are O(1); misses fall through to [`translate`] and populate +/// the cache before returning. Errors are *not* cached -- a repeatedly-bad +/// statement will re-parse each time. +/// +/// # Errors +/// +/// Returns a [`TranslateError`] if the SQL cannot be parsed or contains +/// unsupported constructs. +pub fn translate_cached( + cache: &TranslateCache, + sql: &str, + dialect: Dialect, +) -> Result, TranslateError> { + if let Some(hit) = cache.get(dialect, sql) { + return Ok(hit); + } + let results = translate(sql, dialect)?; + cache.put(dialect, sql.to_string(), results.clone()); + Ok(results) +} + /// Rewrite MySQL/T-SQL-flavored transaction control statements to a form /// SQLite accepts. Returns the rewritten SQL if the input matched a known /// transaction shape, or `None` to fall through to the parser / no-op check. @@ -624,4 +651,74 @@ mod tests { expect_noop("LOCK TABLES users WRITE", Dialect::MySQL); expect_noop("UNLOCK TABLES", Dialect::MySQL); } + + // -- Cache ----------------------------------------------------------------- + + #[test] + fn translate_cached_hit_matches_uncached() { + let cache = TranslateCache::new(16); + let sql = "SELECT id, name FROM users WHERE id = ? ORDER BY id DESC LIMIT 10"; + let cold = translate(sql, Dialect::MySQL).unwrap(); + let warm = translate_cached(&cache, sql, Dialect::MySQL).unwrap(); + // Both should produce a single Sql result and be equal shape-wise. + assert_eq!(cold.len(), warm.len()); + // Second call must hit the cache. + assert_eq!(cache.len(), 1); + let _ = translate_cached(&cache, sql, Dialect::MySQL).unwrap(); + assert_eq!(cache.len(), 1); + } + + #[test] + fn translate_cached_errors_are_not_stored() { + let cache = TranslateCache::new(16); + let bad = "!!! NOT SQL @@@"; + assert!(translate_cached(&cache, bad, Dialect::MySQL).is_err()); + assert_eq!(cache.len(), 0); + } + + /// Micro-benchmark: translating the same WordPress-shaped query 5000 + /// times, cold vs cached. Not a criterion bench (we don't want a dev-dep + /// on criterion); numbers get printed via `cargo test -- --nocapture + /// translate_cache_speedup`. + #[test] + fn translate_cache_speedup() { + use std::time::Instant; + + // A representative WP query with joins, ORDER BY, LIMIT. + let sql = "SELECT wp_posts.* FROM wp_posts \ + INNER JOIN wp_term_relationships \ + ON wp_posts.ID = wp_term_relationships.object_id \ + WHERE wp_posts.post_status = 'publish' \ + AND wp_term_relationships.term_taxonomy_id IN (?, ?, ?) \ + ORDER BY wp_posts.post_date DESC LIMIT 10"; + let iters = 5_000; + + // Cold path: no cache, translate() every time. + let t0 = Instant::now(); + for _ in 0..iters { + let _ = translate(sql, Dialect::MySQL).unwrap(); + } + let cold = t0.elapsed(); + + // Warm path: shared cache, one miss then N hits. + let cache = TranslateCache::new(16); + let t0 = Instant::now(); + for _ in 0..iters { + let _ = translate_cached(&cache, sql, Dialect::MySQL).unwrap(); + } + let warm = t0.elapsed(); + + let cold_ns = cold.as_nanos() as f64 / f64::from(iters); + let warm_ns = warm.as_nanos() as f64 / f64::from(iters); + println!( + "translate_cache_speedup: {iters} iters | cold={cold_ns:.0}ns/call \ + warm={warm_ns:.0}ns/call speedup={:.1}x", + cold_ns / warm_ns + ); + // Sanity: cached path must be strictly faster. + assert!( + warm < cold, + "cached path is not faster: cold={cold:?} warm={warm:?}" + ); + } }