WIP feat(sqlite): add position information to DB errors

abonander · abonander · commit a73b332ec2d5 · 2024-07-20T14:44:08.000-07:00
diff --git a/sqlx-sqlite/src/connection/execute.rs b/sqlx-sqlite/src/connection/execute.rs
@@ -1,3 +1,4 @@
+use std::sync::Arc;
 use crate::connection::{ConnectionHandle, ConnectionState};
 use crate::error::Error;
 use crate::logger::QueryLogger;
@@ -20,14 +21,14 @@ pub struct ExecuteIter<'a> {
 
 pub(crate) fn iter<'a>(
     conn: &'a mut ConnectionState,
-    query: &'a str,
+    query: &'a Arc<str>,
     args: Option<SqliteArguments<'a>>,
     persistent: bool,
 ) -> Result<ExecuteIter<'a>, Error> {
     // fetch the cached statement or allocate a new one
-    let statement = conn.statements.get(query, persistent)?;
+    let statement = conn.statements.get(query.clone(), persistent)?;
 
-    let logger = QueryLogger::new(query, conn.log_settings.clone());
+    let logger = QueryLogger::new(&query, conn.log_settings.clone());
 
     Ok(ExecuteIter {
         handle: &mut conn.handle,
diff --git a/sqlx-sqlite/src/connection/mod.rs b/sqlx-sqlite/src/connection/mod.rs
@@ -6,7 +6,7 @@ use std::os::raw::{c_int, c_void};
 use std::panic::catch_unwind;
 use std::ptr;
 use std::ptr::NonNull;
-
+use std::sync::Arc;
 use futures_core::future::BoxFuture;
 use futures_intrusive::sync::MutexGuard;
 use futures_util::future;
@@ -395,12 +395,12 @@ impl Statements {
         }
     }
 
-    fn get(&mut self, query: &str, persistent: bool) -> Result<&mut VirtualStatement, Error> {
+    fn get(&mut self, query: Arc<str>, persistent: bool) -> Result<&mut VirtualStatement, Error> {
         if !persistent || !self.cached.is_enabled() {
             return Ok(self.temp.insert(VirtualStatement::new(query, false)?));
         }
 
-        let exists = self.cached.contains_key(query);
+        let exists = self.cached.contains_key(&query);
 
         if !exists {
             let statement = VirtualStatement::new(query, true)?;
diff --git a/sqlx-sqlite/src/connection/worker.rs b/sqlx-sqlite/src/connection/worker.rs
@@ -40,15 +40,15 @@ pub(crate) struct WorkerSharedState {
 
 enum Command {
     Prepare {
-        query: Box<str>,
+        query: Arc<str>,
         tx: oneshot::Sender<Result<SqliteStatement<'static>, Error>>,
     },
     Describe {
-        query: Box<str>,
+        query: Arc<str>,
         tx: oneshot::Sender<Result<Describe<Sqlite>, Error>>,
     },
     Execute {
-        query: Box<str>,
+        query: Arc<str>,
         arguments: Option<SqliteArguments<'static>>,
         persistent: bool,
         tx: flume::Sender<Result<Either<SqliteQueryResult, SqliteRow>, Error>>,
@@ -119,7 +119,7 @@ impl ConnectionWorker {
                     let _guard = span.enter();
                     match cmd {
                         Command::Prepare { query, tx } => {
-                            tx.send(prepare(&mut conn, &query).map(|prepared| {
+                            tx.send(prepare(&mut conn, query).map(|prepared| {
                                 update_cached_statements_size(
                                     &conn,
                                     &shared.cached_statements_size,
@@ -394,7 +394,7 @@ impl ConnectionWorker {
     }
 }
 
-fn prepare(conn: &mut ConnectionState, query: &str) -> Result<SqliteStatement<'static>, Error> {
+fn prepare(conn: &mut ConnectionState, query: Arc<str>) -> Result<SqliteStatement<'static>, Error> {
     // prepare statement object (or checkout from cache)
     let statement = conn.statements.get(query, true)?;
 
diff --git a/sqlx-sqlite/src/error.rs b/sqlx-sqlite/src/error.rs
@@ -5,9 +5,9 @@ use std::os::raw::c_int;
 use std::{borrow::Cow, str::from_utf8_unchecked};
 
 use libsqlite3_sys::{
-    sqlite3, sqlite3_errmsg, sqlite3_extended_errcode, SQLITE_CONSTRAINT_CHECK,
-    SQLITE_CONSTRAINT_FOREIGNKEY, SQLITE_CONSTRAINT_NOTNULL, SQLITE_CONSTRAINT_PRIMARYKEY,
-    SQLITE_CONSTRAINT_UNIQUE,
+    sqlite3, sqlite3_errmsg, sqlite3_error_offset, sqlite3_extended_errcode,
+    SQLITE_CONSTRAINT_CHECK, SQLITE_CONSTRAINT_FOREIGNKEY, SQLITE_CONSTRAINT_NOTNULL,
+    SQLITE_CONSTRAINT_PRIMARYKEY, SQLITE_CONSTRAINT_UNIQUE,
 };
 
 pub(crate) use sqlx_core::error::*;
@@ -19,6 +19,8 @@ pub(crate) use sqlx_core::error::*;
 pub struct SqliteError {
     code: c_int,
     message: String,
+    offset: Option<usize>,
+    error_pos: Option<ErrorPosition>,
 }
 
 impl SqliteError {
@@ -34,9 +36,26 @@ impl SqliteError {
             from_utf8_unchecked(CStr::from_ptr(msg).to_bytes())
         };
 
+        // returns `-1` if not applicable
+        let offset = unsafe { sqlite3_error_offset(handle) }.try_into().ok();
+
         Self {
             code,
             message: message.to_owned(),
+            offset,
+            error_pos,
+        }
+    }
+
+    pub(crate) fn add_offset(&mut self, offset: usize) {
+        if let Some(prev_offset) = self.offset {
+            self.offset = prev_offset.checked_add(offset);
+        }
+    }
+
+    pub(crate) fn find_error_pos(&mut self, query: &str) {
+        if let Some(offset) = self.offset {
+            self.error_pos = ErrorPosition::find(query, PositionBasis::ByteOffset(offset));
         }
     }
 
@@ -72,6 +91,10 @@ impl DatabaseError for SqliteError {
         Some(format!("{}", self.code).into())
     }
 
+    fn position(&self) -> Option<ErrorPosition> {
+        self.error_pos
+    }
+
     #[doc(hidden)]
     fn as_error(&self) -> &(dyn StdError + Send + Sync + 'static) {
         self
diff --git a/sqlx-sqlite/src/statement/virtual.rs b/sqlx-sqlite/src/statement/virtual.rs
@@ -31,8 +31,11 @@ pub struct VirtualStatement {
     /// there are no more statements to execute and `reset()` must be called
     index: Option<usize>,
 
-    /// tail of the most recently prepared SQL statement within this container
-    tail: Bytes,
+    /// The full input SQL.
+    sql: Arc<str>,
+
+    /// The byte offset of the next statement to prepare in `sql`.
+    tail_offset: usize,
 
     /// underlying sqlite handles for each inner statement
     /// a SQL query string in SQLite is broken up into N statements
@@ -44,6 +47,9 @@ pub struct VirtualStatement {
 
     // each set of column names
     pub(crate) column_names: SmallVec<[Arc<HashMap<UStr, usize>>; 1]>,
+
+    /// Offsets into `sql` for each statement.
+    pub(crate) sql_offsets: SmallVec<[usize; 1]>,
 }
 
 pub struct PreparedStatement<'a> {
@@ -53,9 +59,7 @@ pub struct PreparedStatement<'a> {
 }
 
 impl VirtualStatement {
-    pub(crate) fn new(mut query: &str, persistent: bool) -> Result<Self, Error> {
-        query = query.trim();
-
+    pub(crate) fn new(query: Arc<str>, persistent: bool) -> Result<Self, Error> {
         if query.len() > i32::max_value() as usize {
             return Err(err_protocol!(
                 "query string must be smaller than {} bytes",
@@ -65,11 +69,13 @@ impl VirtualStatement {
 
         Ok(Self {
             persistent,
-            tail: Bytes::from(String::from(query)),
+            sql: query,
+            tail_offset: 0,
             handles: SmallVec::with_capacity(1),
             index: None,
             columns: SmallVec::with_capacity(1),
             column_names: SmallVec::with_capacity(1),
+            sql_offsets: SmallVec::with_capacity(1),
         })
     }
 
@@ -84,11 +90,33 @@ impl VirtualStatement {
             .or(Some(0));
 
         while self.handles.len() <= self.index.unwrap_or(0) {
-            if self.tail.is_empty() {
+            let sql_offset = self.tail_offset;
+
+            let query = self.sql.get(sql_offset..).unwrap_or("");
+
+            if query.is_empty() {
                 return Ok(None);
             }
 
-            if let Some(statement) = prepare(conn.as_ptr(), &mut self.tail, self.persistent)? {
+            let (consumed, maybe_statement) = try_prepare(
+                conn.as_ptr(),
+                query,
+                self.persistent,
+            ).map_err(|mut e| {
+                // `sqlite3_offset()` returns the offset into the passed string,
+                // but we want the offset into the original SQL string.
+                e.add_offset(sql_offset);
+                e.find_error_pos(&self.sql);
+                e
+            })?;
+
+            self.tail_offset = self.tail_offset
+                .checked_add(consumed)
+                // Highly unlikely, but since we're dealing with `unsafe` here
+                // it's best not to fool around.
+                .ok_or_else(|| Error::Protocol(format!("overflow adding {n:?} bytes to tail_offset {tail_offset:?}")))?;
+
+            if let Some(statement) = maybe_statement {
                 let num = statement.column_count();
 
                 let mut columns = Vec::with_capacity(num);
@@ -112,6 +140,7 @@ impl VirtualStatement {
                 self.handles.push(statement);
                 self.columns.push(Arc::new(columns));
                 self.column_names.push(Arc::new(column_names));
+                self.sql_offsets.push(sql_offset);
             }
         }
 
@@ -140,11 +169,13 @@ impl VirtualStatement {
     }
 }
 
-fn prepare(
+/// Attempt to prepare one statement, returning the number of bytes consumed from `sql`,
+/// and the statement handle if successful.
+fn try_prepare(
     conn: *mut sqlite3,
-    query: &mut Bytes,
+    query: &str,
     persistent: bool,
-) -> Result<Option<StatementHandle>, Error> {
+) -> Result<(usize, Option<StatementHandle>), SqliteError> {
     let mut flags = 0;
 
     // For some reason, when building with the `sqlcipher` feature enabled
@@ -158,40 +189,37 @@ fn prepare(
         flags |= SQLITE_PREPARE_PERSISTENT as u32;
     }
 
-    while !query.is_empty() {
-        let mut statement_handle: *mut sqlite3_stmt = null_mut();
-        let mut tail: *const c_char = null();
-
-        let query_ptr = query.as_ptr() as *const c_char;
-        let query_len = query.len() as i32;
-
-        // <https://www.sqlite.org/c3ref/prepare.html>
-        let status = unsafe {
-            sqlite3_prepare_v3(
-                conn,
-                query_ptr,
-                query_len,
-                flags,
-                &mut statement_handle,
-                &mut tail,
-            )
-        };
-
-        if status != SQLITE_OK {
-            return Err(SqliteError::new(conn).into());
-        }
-
-        // tail should point to the first byte past the end of the first SQL
-        // statement in zSql. these routines only compile the first statement,
-        // so tail is left pointing to what remains un-compiled.
+    let mut statement_handle: *mut sqlite3_stmt = null_mut();
+    let mut tail_ptr: *const c_char = null();
+
+    let query_ptr = query.as_ptr() as *const c_char;
+    let query_len = query.len() as i32;
+
+    // <https://www.sqlite.org/c3ref/prepare.html>
+    let status = unsafe {
+        sqlite3_prepare_v3(
+            conn,
+            query_ptr,
+            query_len,
+            flags,
+            &mut statement_handle,
+            &mut tail_ptr,
+        )
+    };
+
+    if status != SQLITE_OK {
+        // Note: `offset` and `error_pos` will be updated in `VirtualStatement::prepare_next()`.
+        return Err(SqliteError::new(conn));
+    }
 
-        let n = (tail as usize) - (query_ptr as usize);
-        query.advance(n);
+    // tail should point to the first byte past the end of the first SQL
+    // statement in zSql. these routines only compile the first statement,
+    // so tail is left pointing to what remains un-compiled.
 
-        if let Some(handle) = NonNull::new(statement_handle) {
-            return Ok(Some(StatementHandle::new(handle)));
-        }
-    }
+    let consumed = (tail_ptr as usize) - (query_ptr as usize);
 
-    Ok(None)
+    Ok((
+        consumed,
+        NonNull::new(statement_handle).map(StatementHandle::new),
+    ))
 }
diff --git a/tests/sqlite/error.rs b/tests/sqlite/error.rs
@@ -1,4 +1,5 @@
 use sqlx::{error::ErrorKind, sqlite::Sqlite, Connection, Executor};
+use sqlx_sqlite::{SqliteConnection, SqliteError};
 use sqlx_test::new;
 
 #[sqlx_macros::test]
@@ -70,3 +71,24 @@ async fn it_fails_with_check_violation() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[sqlx_macros::test]
+async fn it_fails_with_useful_information() -> anyhow::Result<()> {
+    let mut conn = SqliteConnection::connect(":memory:").await?;
+
+    let err: sqlx::Error = sqlx::query("SELECT foo FORM bar")
+        .execute(&mut conn)
+        .await
+        .unwrap_err();
+
+    let sqlx::Error::Database(dbe) = err else {
+        panic!("unexpected error kind: {err:?}")
+    };
+
+    let dbe= dbe.downcast::<SqliteError>();
+
+    eprintln!("{dbe}");
+    eprintln!("{dbe:?}");
+
+    Ok(())
+}
