Merge pull request #579 from tursodatabase/fix-batch-replication-index

fix batch replication index

libsql-server/src/hrana/proto.rs

@@ -12,7 +12,7 @@ pub struct Error {
     pub code: String,
 }
 
-#[derive(Deserialize, prost::Message)]
+#[derive(Deserialize, Serialize, prost::Message)]
 pub struct Stmt {
     #[serde(default)]
     #[prost(string, optional, tag = "1")]
@@ -34,7 +34,7 @@ pub struct Stmt {
     pub replication_index: Option<u64>,
 }
 
-#[derive(Deserialize, prost::Message)]
+#[derive(Deserialize, Serialize, prost::Message)]
 pub struct NamedArg {
     #[prost(string, tag = "1")]
     pub name: String,
@@ -53,8 +53,6 @@ pub struct StmtResult {
     #[serde(with = "option_i64_as_str")]
     #[prost(sint64, optional, tag = "4")]
     pub last_insert_rowid: Option<i64>,
-    #[prost(uint64, optional, tag = "5")]
-    pub replication_index: Option<u64>,
 }
 
 #[derive(Serialize, prost::Message)]
@@ -72,7 +70,7 @@ pub struct Row {
     pub values: Vec<Value>,
 }
 
-#[derive(Deserialize, prost::Message)]
+#[derive(Deserialize, Serialize, prost::Message)]
 pub struct Batch {
     #[prost(message, repeated, tag = "1")]
     pub steps: Vec<BatchStep>,
@@ -81,7 +79,7 @@ pub struct Batch {
     pub replication_index: Option<u64>,
 }
 
-#[derive(Deserialize, prost::Message)]
+#[derive(Deserialize, Serialize, prost::Message)]
 pub struct BatchStep {
     #[serde(default)]
     #[prost(message, optional, tag = "1")]
@@ -94,9 +92,10 @@ pub struct BatchStep {
 pub struct BatchResult {
     pub step_results: Vec<Option<StmtResult>>,
     pub step_errors: Vec<Option<Error>>,
+    pub replication_index: Option<u64>,
 }
 
-#[derive(Deserialize, Debug, Default)]
+#[derive(Deserialize, Serialize, Debug, Default)]
 #[serde(tag = "type", rename_all = "snake_case")]
 pub enum BatchCond {
     #[serde(skip_deserializing)]
@@ -116,7 +115,7 @@ pub enum BatchCond {
     IsAutocommit {},
 }
 
-#[derive(Deserialize, prost::Message)]
+#[derive(Deserialize, Serialize, prost::Message)]
 pub struct BatchCondList {
     #[prost(message, repeated, tag = "1")]
     pub conds: Vec<BatchCond>,

libsql-server/src/hrana/result_builder.rs

@@ -24,7 +24,6 @@ pub struct SingleStatementBuilder {
     current_size: u64,
     max_response_size: u64,
     max_total_response_size: u64,
-    last_frame_no: Option<FrameNo>,
 }
 
 struct SizeFormatter {
@@ -227,10 +226,9 @@ impl QueryResultBuilder for SingleStatementBuilder {
 
     fn finish(
         &mut self,
-        last_frame_no: Option<FrameNo>,
+        _last_frame_no: Option<FrameNo>,
         _is_autocommit: bool,
     ) -> Result<(), QueryResultBuilderError> {
-        self.last_frame_no = last_frame_no;
         Ok(())
     }
 
@@ -242,7 +240,6 @@ impl QueryResultBuilder for SingleStatementBuilder {
                 rows: std::mem::take(&mut self.rows),
                 affected_row_count: std::mem::take(&mut self.affected_row_count),
                 last_insert_rowid: std::mem::take(&mut self.last_insert_rowid),
-                replication_index: self.last_frame_no,
             }),
         }
     }
@@ -268,6 +265,7 @@ pub struct HranaBatchProtoBuilder {
     current_size: u64,
     max_response_size: u64,
     step_empty: bool,
+    last_frame_no: Option<FrameNo>,
 }
 
 impl QueryResultBuilder for HranaBatchProtoBuilder {
@@ -350,16 +348,18 @@ impl QueryResultBuilder for HranaBatchProtoBuilder {
 
     fn finish(
         &mut self,
-        _last_frame_no: Option<FrameNo>,
+        last_frame_no: Option<FrameNo>,
         _is_autocommit: bool,
     ) -> Result<(), QueryResultBuilderError> {
+        self.last_frame_no = last_frame_no;
         Ok(())
     }
 
     fn into_ret(self) -> Self::Ret {
         proto::BatchResult {
             step_results: self.step_results,
             step_errors: self.step_errors,
+            replication_index: self.last_frame_no,
         }
     }
 }

libsql-server/src/lib.rs

@@ -48,6 +48,8 @@ pub mod net;
 pub mod rpc;
 pub mod version;
 
+pub use hrana::proto as hrana_proto;
+
 mod auth;
 mod database;
 mod error;

libsql-server/tests/hrana/batch.rs

@@ -0,0 +1,36 @@
+use insta::assert_json_snapshot;
+use sqld::hrana_proto::{Batch, BatchStep, Stmt};
+
+use crate::common::http::Client;
+
+#[test]
+fn sample_request() {
+    let mut sim = turmoil::Builder::new().build();
+    sim.host("primary", super::make_standalone_server);
+    sim.client("client", async {
+        let batch = Batch {
+            steps: vec![BatchStep {
+                condition: None,
+                stmt: Stmt {
+                    sql: Some("create table test (x)".to_string()),
+                    ..Default::default()
+                },
+            }],
+            replication_index: None,
+        };
+        let client = Client::new();
+
+        let resp = client
+            .post(
+                "http://primary:8080/v1/batch",
+                serde_json::json!({ "batch": batch }),
+            )
+            .await
+            .unwrap();
+        assert_json_snapshot!(resp.json_value().await.unwrap());
+
+        Ok(())
+    });
+
+    sim.run().unwrap();
+}

libsql-server/tests/hrana/mod.rs

@@ -0,0 +1,24 @@
+//! Test hrana related functionalities
+
+use sqld::config::UserApiConfig;
+use tempfile::tempdir;
+
+use crate::common::net::{init_tracing, SimServer, TestServer};
+mod batch;
+
+async fn make_standalone_server() -> Result<(), Box<dyn std::error::Error>> {
+    init_tracing();
+    let tmp = tempdir()?;
+    let server = TestServer {
+        path: tmp.path().to_owned().into(),
+        user_api_config: UserApiConfig {
+            hrana_ws_acceptor: None,
+            ..Default::default()
+        },
+        ..Default::default()
+    };
+
+    server.start_sim(8080).await?;
+
+    Ok(())
+}

libsql-server/tests/hrana/snapshots/tests__hrana__batch__sample_request.snap

@@ -0,0 +1,20 @@
+---
+source: libsql-server/tests/hrana/batch.rs
+expression: resp.json_value().await.unwrap()
+---
+{
+  "result": {
+    "step_results": [
+      {
+        "cols": [],
+        "rows": [],
+        "affected_row_count": 0,
+        "last_insert_rowid": null
+      }
+    ],
+    "step_errors": [
+      null
+    ],
+    "replication_index": 1
+  }
+}

libsql-server/tests/tests.rs

@@ -5,5 +5,6 @@ mod common;
 
 mod cluster;
 mod embedded_replica;
+mod hrana;
 mod namespaces;
 mod standalone;