dev: implement dump parsing logic using sqlite3_parser

Author: riazus

libsql-server/src/namespace/configurator/helpers.rs

@@ -7,10 +7,13 @@ use anyhow::Context as _;
 use bottomless::replicator::Options;
 use bytes::Bytes;
 use enclose::enclose;
+use fallible_iterator::FallibleIterator;
 use futures::Stream;
 use libsql_sys::EncryptionConfig;
 use rusqlite::hooks::{AuthAction, AuthContext, Authorization};
-use tokio::io::AsyncBufReadExt as _;
+use sqlite3_parser::ast::{Cmd, Stmt};
+use sqlite3_parser::lexer::sql::{Parser, ParserError};
+use tokio::io::AsyncReadExt;
 use tokio::task::JoinSet;
 use tokio_util::io::StreamReader;
 
@@ -33,9 +36,6 @@ use crate::{StatsSender, BLOCKING_RT, DB_CREATE_TIMEOUT, DEFAULT_AUTO_CHECKPOINT
 
 use super::{BaseNamespaceConfig, PrimaryConfig};
 
-const WASM_TABLE_CREATE: &str =
-    "CREATE TABLE libsql_wasm_func_table (name text PRIMARY KEY, body text) WITHOUT ROWID;";
-
 #[tracing::instrument(skip_all)]
 pub(super) async fn make_primary_connection_maker(
     primary_config: &PrimaryConfig,
@@ -290,184 +290,94 @@ async fn run_periodic_compactions(logger: Arc<ReplicationLogger>) -> anyhow::Res
     }
 }
 
-fn tokenize_sql_keywords(text: &str) -> Vec<String> {
-    let mut tokens = Vec::new();
-    let mut chars = text.chars().peekable();
-    let mut current_token = String::new();
-    let mut in_string_literal = false;
-    let mut string_delimiter = '\0';
-
-    while let Some(ch) = chars.next() {
-        match ch {
-            '\'' | '"' => {
-                if !in_string_literal {
-                    in_string_literal = true;
-                    string_delimiter = ch;
-                } else if ch == string_delimiter {
-                    in_string_literal = false;
-                }
-            }
-            c if c.is_whitespace() || "(){}[];,".contains(c) => {
-                if in_string_literal {
-                    continue;
-                }
-                if !current_token.is_empty() {
-                    tokens.push(current_token.to_uppercase());
-                    current_token.clear();
-                }
-            }
-            // Regular characters
-            _ => {
-                if !in_string_literal {
-                    current_token.push(ch);
-                }
-            }
-        }
-    }
-
-    if !current_token.is_empty() && !in_string_literal {
-        tokens.push(current_token.to_uppercase());
-    }
-
-    tokens
-}
-
-fn is_complete_sql_statement(sql: &str) -> bool {
-    let tokens = tokenize_sql_keywords(sql);
-    let mut begin_end_depth = 0;
-    let mut case_depth = 0;
-
-    for (i, token) in tokens.iter().enumerate() {
-        match token.as_str() {
-            "CASE" => {
-                case_depth += 1;
-            }
-            "BEGIN" => {
-                let next_token = tokens.get(i + 1).map(|s| s.as_str());
-                let is_transaction_keyword = matches!(
-                    next_token,
-                    Some("TRANSACTION") | Some("IMMEDIATE") | Some("EXCLUSIVE") | Some("DEFERRED")
-                );
-
-                if !is_transaction_keyword {
-                    begin_end_depth += 1;
-                }
-            }
-            "END" => {
-                if case_depth > 0 {
-                    case_depth -= 1;
-                } else {
-                    // This is a block-ending END (BEGIN/END, IF/END IF, etc.)
-                    let is_control_flow_end = tokens
-                        .get(i + 1)
-                        .map(|next| matches!(next.as_str(), "IF" | "LOOP" | "WHILE"))
-                        .unwrap_or(false);
-
-                    if !is_control_flow_end {
-                        begin_end_depth -= 1;
-                    }
-                }
-            }
-            _ => {}
-        }
-
-        if begin_end_depth < 0 {
-            return false;
-        }
-    }
-
-    begin_end_depth == 0 && case_depth == 0
-}
-
 async fn load_dump<S>(dump: S, conn: PrimaryConnection) -> crate::Result<(), LoadDumpError>
 where
     S: Stream<Item = std::io::Result<Bytes>> + Unpin,
 {
     let mut reader = tokio::io::BufReader::new(StreamReader::new(dump));
-    let mut curr = String::new();
-    let mut line = String::new();
+    let mut dump_content = String::new();
+    reader
+        .read_to_string(&mut dump_content)
+        .await
+        .map_err(|e| LoadDumpError::Internal(format!("Failed to read dump content: {}", e)))?;
+
+    if dump_content.to_lowercase().contains("attach") {
+        return Err(LoadDumpError::InvalidSqlInput(
+            "attach statements are not allowed in dumps".to_string(),
+        ));
+    }
+
+    let mut parser = Box::new(Parser::new(dump_content.as_bytes()));
     let mut skipped_wasm_table = false;
     let mut n_stmt = 0;
-    let mut line_id = 0;
 
-    while let Ok(n) = reader.read_line(&mut curr).await {
-        line_id += 1;
-        if n == 0 {
-            break;
-        }
-        let trimmed = curr.trim();
-        if trimmed.is_empty() || trimmed.starts_with("--") {
-            curr.clear();
-            continue;
-        }
+    loop {
+        match parser.next() {
+            Ok(Some(cmd)) => {
+                n_stmt += 1;
+
+                if !skipped_wasm_table {
+                    if let Cmd::Stmt(Stmt::CreateTable { tbl_name, .. }) = &cmd {
+                        if tbl_name.name.0 == "libsql_wasm_func_table" {
+                            skipped_wasm_table = true;
+                            tracing::debug!("Skipping WASM table creation");
+                            continue;
+                        }
+                    }
+                }
 
-        // we want to concat original(non-trimmed) lines as trimming will join all them in one
-        // single-line statement which is incorrect if comments in the end are present
-        line.push_str(&curr);
-        let statement_end = trimmed.ends_with(';') && is_complete_sql_statement(&line);
-        curr.clear();
-
-        // This is a hack to ignore the libsql_wasm_func_table table because it is already created
-        // by the system.
-        if !skipped_wasm_table && line.trim() == WASM_TABLE_CREATE {
-            skipped_wasm_table = true;
-            line.clear();
-            continue;
-        }
+                if n_stmt > 2 && conn.is_autocommit().await.unwrap() {
+                    return Err(LoadDumpError::NoTxn);
+                }
 
-        if statement_end {
-            n_stmt += 1;
-            // dump must be performd within a txn
-            if n_stmt > 2 && conn.is_autocommit().await.unwrap() {
-                return Err(LoadDumpError::NoTxn);
+                let stmt_sql = cmd.to_string();
+                tokio::task::spawn_blocking({
+                    let conn = conn.clone();
+                    move || -> crate::Result<(), LoadDumpError> {
+                        conn.with_raw(|conn| {
+                            conn.authorizer(Some(|auth: AuthContext<'_>| match auth.action {
+                                AuthAction::Attach { filename: _ } => Authorization::Deny,
+                                _ => Authorization::Allow,
+                            }));
+                            conn.execute(&stmt_sql, ())
+                        })
+                        .map_err(|e| match e {
+                            rusqlite::Error::SqlInputError {
+                                msg, sql, offset, ..
+                            } => LoadDumpError::InvalidSqlInput(format!(
+                                "msg: {}, sql: {}, offset: {}",
+                                msg, sql, offset
+                            )),
+                            e => LoadDumpError::Internal(format!(
+                                "statement: {}, error: {}",
+                                n_stmt, e
+                            )),
+                        })?;
+                        Ok(())
+                    }
+                })
+                .await??;
             }
+            Ok(None) => break,
+            Err(e) => {
+                let error_msg = match e {
+                    sqlite3_parser::lexer::sql::Error::ParserError(
+                        ParserError::SyntaxError { token_type, found },
+                        Some((line, col)),
+                    ) => {
+                        let near_token = found.as_deref().unwrap_or(&token_type);
+                        format!(
+                            "syntax error near '{}' at line {}, column {}",
+                            near_token, line, col
+                        )
+                    }
+                    _ => format!("parse error: {}", e),
+                };
 
-            line = tokio::task::spawn_blocking({
-                let conn = conn.clone();
-                move || -> crate::Result<String, LoadDumpError> {
-                    conn.with_raw(|conn| {
-                        conn.authorizer(Some(|auth: AuthContext<'_>| match auth.action {
-                            AuthAction::Attach { filename: _ } => Authorization::Deny,
-                            _ => Authorization::Allow,
-                        }));
-                        conn.execute(&line, ())
-                    })
-                    .map_err(|e| match e {
-                        rusqlite::Error::SqlInputError {
-                            msg, sql, offset, ..
-                        } => {
-                            let msg = if sql.to_lowercase().contains("attach") {
-                                format!(
-                                    "attach statements are not allowed in dumps, msg: {}, sql: {}, offset: {}",
-                                    msg,
-                                    sql,
-                                    offset
-                                )
-                            } else {
-                                format!("msg: {}, sql: {}, offset: {}", msg, sql, offset)
-                            };
-
-                            LoadDumpError::InvalidSqlInput(msg)
-                        }
-                        e => LoadDumpError::Internal(format!("line: {}, error: {}", line_id, e)),
-                    })?;
-                    Ok(line)
-                }
-            })
-            .await??;
-            line.clear();
-        } else {
-            line.push(' ');
+                return Err(LoadDumpError::InvalidSqlInput(error_msg));
+            }
         }
     }
-    tracing::debug!("loaded {} lines from dump", line_id);
-
-    if !line.trim().is_empty() {
-        return Err(LoadDumpError::InvalidSqlInput(format!(
-            "Incomplete SQL statement at end of dump: {}",
-            line.trim()
-        )));
-    }
 
     if !conn.is_autocommit().await.unwrap() {
         tokio::task::spawn_blocking({

libsql-server/tests/namespaces/snapshots/tests__namespaces__dumps__load_dump_with_attach_rejected.snap

@@ -1,6 +1,5 @@
 ---
 source: libsql-server/tests/namespaces/dumps.rs
 expression: resp.body_string().await?
-snapshot_kind: text
 ---
-{"error":"The passed dump sql is invalid: attach statements are not allowed in dumps, msg: near \"COMMIT\": syntax error, sql:     ATTACH foo/bar.sql\n     COMMIT;, offset: 28"}
+{"error":"The passed dump sql is invalid: attach statements are not allowed in dumps"}

libsql-server/tests/namespaces/snapshots/tests__namespaces__dumps__load_dump_with_invalid_sql.snap

@@ -3,4 +3,4 @@ source: libsql-server/tests/namespaces/dumps.rs
 expression: resp.body_string().await?
 snapshot_kind: text
 ---
-{"error":"The passed dump sql is invalid: msg: near \"COMMIT\": syntax error, sql:     SELECT abs(-9223372036854775808) \n     COMMIT;, offset: 43"}
+{"error":"The passed dump sql is invalid: syntax error near 'COMMIT' at line 7, column 11"}