Merge #259

bors[bot] · MarinPostma · web-flow · commit 950cedecfa8d · 2023-03-01T13:57:25.000Z
259: fix dump bug r=MarinPostma a=MarinPostma

Fix bug that prevented large dumps to be loaded


Co-authored-by: ad hoc &lt;postma.marin@protonmail.com&gt;
diff --git a/sqld/src/http/mod.rs b/sqld/src/http/mod.rs
@@ -229,18 +229,27 @@ fn handle_version() -> Response<Body> {
     Response::new(Body::from(version.as_bytes()))
 }
 
+async fn load_body(mut req: Request<Body>) -> anyhow::Result<Bytes> {
+    let mut bytes = BytesMut::new();
+    while let Some(data) = req.data().await {
+        let data = data?;
+        bytes.extend(data);
+    }
+
+    Ok(bytes.freeze())
+}
+
 async fn load_dump(
-    mut req: Request<Body>,
+    req: Request<Body>,
     sender: mpsc::Sender<Message>,
 ) -> anyhow::Result<Response<Body>> {
-    match req.data().await {
-        Some(Ok(data)) => {
+    match load_body(req).await {
+        Ok(data) => {
             // FIXME: Dumps may not fit in memory. A better way would be to stream the payload, and
             // have a dedicated path to load the dump from it.
             let mut queries = Vec::new();
-            let s = String::from_utf8(data.to_vec())?;
-            for line in s.lines() {
-                let stmt = Statement::new_unchecked(line);
+            for stmt in Statement::parse_unchecked(&data) {
+                let stmt = stmt?;
                 queries.push(Query {
                     stmt,
                     params: Params::empty(),
@@ -258,8 +267,7 @@ async fn load_dump(
                 Err(e) => Ok(error(&e.to_string(), StatusCode::INTERNAL_SERVER_ERROR)),
             }
         }
-        Some(Err(e)) => Ok(error(&e.to_string(), StatusCode::INTERNAL_SERVER_ERROR)),
-        None => Ok(Response::new(Body::empty())),
+        Err(e) => Ok(error(&e.to_string(), StatusCode::INTERNAL_SERVER_ERROR)),
     }
 }
 
diff --git a/sqld/src/query_analysis.rs b/sqld/src/query_analysis.rs
@@ -94,14 +94,33 @@ impl Statement {
     /// Returns a statement instance without performing any validation to the input
     /// It is always assumed that such a statement will be a write, and will always be handled by
     /// the primary
-    pub fn new_unchecked(s: &str) -> Self {
+    pub fn new_unchecked(stmt: String) -> Self {
         Self {
-            stmt: s.to_string(),
+            stmt,
             kind: StmtKind::Write,
             is_iud: false,
         }
     }
 
+    /// parses a series of statements into unchecked statemments
+    pub fn parse_unchecked(s: &[u8]) -> impl Iterator<Item = Result<Self>> + '_ {
+        let mut parser = Box::new(Parser::new(s));
+        std::iter::from_fn(move || match parser.next() {
+            Ok(Some(cmd)) => Some(Ok(Self::new_unchecked(cmd.to_string()))),
+            Ok(None) => None,
+            Err(sqlite3_parser::lexer::sql::Error::ParserError(
+                ParserError::SyntaxError {
+                    token_type: _,
+                    found: Some(found),
+                },
+                Some((line, col)),
+            )) => Some(Err(anyhow::anyhow!(
+                "syntax error around L{line}:{col}: `{found}`"
+            ))),
+            Err(e) => Some(Err(e.into())),
+        })
+    }
+
     pub fn parse(s: &str) -> impl Iterator<Item = Result<Self>> + '_ {
         fn parse_inner(c: Cmd) -> Result<Statement> {
             let kind =
diff --git a/testing/end-to-end/integration/basic_cluster/dump.rs b/testing/end-to-end/integration/basic_cluster/dump.rs
@@ -65,3 +65,45 @@ async fn load_dump_from_replica(app: App) {
 
     insta::assert_json_snapshot!(resp.json::<serde_json::Value>().await.unwrap());
 }
+
+#[octopod::test(app = "simple-cluster")]
+async fn load_large_dump(app: App) {
+    let primary_ip = app.service("primary").unwrap().ip().await.unwrap();
+    let url = format!("http://{primary_ip}:8080/load-dump");
+    let client = reqwest::Client::new();
+
+    let dump = std::iter::once("CREATE TABLE test (who);\n")
+        .chain(std::iter::repeat("INSERT INTO test VALUES ('hello world');\n").take(10000))
+        .fold(String::new(), |acc, s| acc + s);
+
+    let resp = client.post(url).body(dump).send().await.unwrap();
+    assert!(resp.status().is_success(), "{}", resp.text().await.unwrap());
+
+    // wait for dump to be loaded and replicated
+    tokio::time::sleep(Duration::from_secs(3)).await;
+
+    let url = format!("http://{primary_ip}:8080");
+    let resp = client
+        .post(url)
+        .json(&json!({
+            "statements": ["select count(*) from test"]
+        }))
+        .send()
+        .await
+        .unwrap();
+
+    insta::assert_json_snapshot!(resp.json::<serde_json::Value>().await.unwrap());
+
+    // ensure replica is up to date
+    let replica_id = app.service("replica").unwrap().ip().await.unwrap();
+    let url = format!("http://{replica_id}:8080");
+    let resp = client
+        .post(url)
+        .json(&json!({
+            "statements": ["select count(*) from test"]
+        }))
+        .send()
+        .await
+        .unwrap();
+    insta::assert_json_snapshot!(resp.json::<serde_json::Value>().await.unwrap());
+}
diff --git a/testing/end-to-end/run-macos.sh b/testing/end-to-end/run-macos.sh
@@ -3,5 +3,6 @@
 podman build ../.. -t sqld
 podman save -o sqld.tar sqld
 podman build . -t end-to-end --isolation=chroot
+echo "here"
 podman run --privileged -v $PWD:/end-to-end -it --rm end-to-end
 rm sqld.tar
