tests: add hrana baton tests

libsql-server/tests/cluster/snapshots/tests__cluster__create_namespace.snap

@@ -1,5 +1,6 @@
 ---
-source: sqld/tests/cluster/mod.rs
+source: libsql-server/tests/cluster/mod.rs
 expression: e.to_string()
+snapshot_kind: text
 ---
-Hrana: `api error: `{"error":"Namespace `foo` doesn't exist"}``
+Hrana: `api error: `status=404 Not Found, body={"error":"Namespace `foo` doesn't exist"}``

libsql-server/tests/cluster/snapshots/tests__cluster__replication__replica_lazy_creation.snap

@@ -1,9 +1,10 @@
 ---
 source: libsql-server/tests/cluster/replication.rs
 expression: "conn.execute(\"create table test (x)\", ()).await.unwrap_err()"
+snapshot_kind: text
 ---
 Hrana(
     Api(
-        "{\"error\":\"Replicator error: Requested namespace doesn't exist\"}",
+        "status=404 Not Found, body={\"error\":\"Replicator error: Requested namespace doesn't exist\"}",
     ),
 )

libsql-server/tests/hrana/batch.rs

@@ -1,12 +1,15 @@
+use std::sync::Arc;
 use std::time::Duration;
 
 use futures::StreamExt;
-use insta::assert_json_snapshot;
+use insta::{assert_json_snapshot, assert_snapshot};
 use libsql::{params, Database};
+use libsql_server::config::UserApiConfig;
 use libsql_server::hrana_proto::{Batch, BatchStep, Stmt};
+use tokio::sync::Notify;
 
 use crate::common::http::Client;
-use crate::common::net::TurmoilConnector;
+use crate::common::net::{init_tracing, SimServer as _, TestServer, TurmoilConnector};
 
 #[test]
 fn sample_request() {
@@ -402,3 +405,179 @@ fn test_simulate_vector_index_load_from_dump() {
 
     sim.run().unwrap();
 }
+
+#[test]
+fn server_restart_query_execute_invalid_baton() {
+    server_restart(|notify, db| async move {
+        let conn = db.connect().unwrap();
+
+        conn.query("select 1;", ()).await.unwrap();
+
+        notify.notify_waiters();
+        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+
+        let err = conn.execute("select 1;", ()).await.unwrap_err();
+        assert_snapshot!(err);
+
+        Ok(())
+    });
+}
+
+#[test]
+fn server_restart_txn_execute_execute_invalid_baton() {
+    server_restart(|notify, db| async move {
+        let conn = db.connect().unwrap();
+
+        let txn = conn.transaction().await.unwrap();
+
+        txn.execute("select 1;", ()).await.unwrap();
+
+        notify.notify_waiters();
+
+        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+
+        let err = txn.execute("select 1;", ()).await.unwrap_err();
+
+        assert_snapshot!(err);
+
+        Ok(())
+    });
+}
+
+#[test]
+fn server_restart_txn_query_execute_invalid_baton() {
+    server_restart(|notify, db| async move {
+        let conn = db.connect().unwrap();
+
+        let txn = conn.transaction().await.unwrap();
+
+        txn.query("select 1;", ()).await.unwrap();
+
+        notify.notify_waiters();
+
+        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+
+        let err = txn.execute("select 1;", ()).await.unwrap_err();
+
+        assert_snapshot!(err);
+
+        Ok(())
+    });
+}
+
+#[test]
+fn server_restart_txn_execute_query_invalid_baton() {
+    server_restart(|notify, db| async move {
+        let conn = db.connect().unwrap();
+
+        let txn = conn.transaction().await.unwrap();
+
+        txn.execute("select 1;", ()).await.unwrap();
+
+        notify.notify_waiters();
+
+        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+
+        let err = txn.query("select 1;", ()).await.unwrap_err();
+
+        assert_snapshot!(err);
+
+        Ok(())
+    });
+}
+
+#[test]
+fn server_restart_execute_query() {
+    server_restart(|notify, db| async move {
+        let conn = db.connect().unwrap();
+        conn.execute("select 1;", ()).await.unwrap();
+
+        notify.notify_waiters();
+
+        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+
+        conn.query("select 1;", ()).await.unwrap();
+
+        Ok(())
+    });
+}
+
+#[test]
+fn server_timeout() {
+    server_restart(|_notify, db| async move {
+        let conn = db.connect().unwrap();
+        conn.query("select 1;", ()).await.unwrap();
+
+        tokio::time::sleep(std::time::Duration::from_secs(20)).await;
+
+        let err = conn.execute("select 1;", ()).await.unwrap_err();
+
+        assert_snapshot!(err);
+
+        let conn = db.connect().unwrap();
+        conn.execute("select 1;", ()).await.unwrap();
+
+        tokio::time::sleep(std::time::Duration::from_secs(20)).await;
+
+        conn.execute("select 1;", ()).await.unwrap();
+
+        Ok(())
+    });
+}
+
+#[track_caller]
+fn server_restart<F, Fut>(f: F)
+where
+    F: Fn(Arc<Notify>, Database) -> Fut + 'static,
+    Fut: std::future::Future<Output = Result<(), libsql::Error>> + 'static,
+{
+    let mut sim = turmoil::Builder::new()
+        .simulation_duration(Duration::from_secs(1000))
+        .build();
+
+    init_tracing();
+
+    let notify = Arc::new(Notify::new());
+
+    let notify_clone = notify.clone();
+
+    sim.host("primary", move || {
+        let notify = notify.clone();
+        async move {
+            let tmp = tempfile::tempdir()?;
+
+            let make_server = || TestServer {
+                path: tmp.path().to_owned().into(),
+                user_api_config: UserApiConfig {
+                    hrana_ws_acceptor: None,
+                    ..Default::default()
+                },
+                ..Default::default()
+            };
+
+            let server = make_server();
+
+            tokio::select! {
+                res = server.start_sim(8080) => {
+                    res.unwrap()
+                }
+                _ = notify.notified() => (),
+            }
+
+            let server = make_server();
+            server.start_sim(8080).await.unwrap();
+
+            Ok(())
+        }
+    });
+
+    sim.client("client", async move {
+        let db = Database::open_remote_with_connector("http://primary:8080", "", TurmoilConnector)?;
+
+        f(notify_clone, db).await.unwrap();
+
+        Ok(())
+    });
+
+    sim.run().unwrap();
+}

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

@@ -0,0 +1,6 @@
+---
+source: libsql-server/tests/hrana/batch.rs
+expression: err
+snapshot_kind: text
+---
+Hrana: `api error: `status=400 Bad Request, body=Received an invalid baton``

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

@@ -0,0 +1,6 @@
+---
+source: libsql-server/tests/hrana/batch.rs
+expression: err
+snapshot_kind: text
+---
+Hrana: `api error: `status=400 Bad Request, body={"message":"The stream has expired due to inactivity","code":"STREAM_EXPIRED"}``

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

@@ -0,0 +1,6 @@
+---
+source: libsql-server/tests/hrana/batch.rs
+expression: err
+snapshot_kind: text
+---
+Hrana: `api error: `status=400 Bad Request, body=Received an invalid baton``

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

@@ -0,0 +1,6 @@
+---
+source: libsql-server/tests/hrana/batch.rs
+expression: err
+snapshot_kind: text
+---
+Hrana: `api error: `status=400 Bad Request, body=Received an invalid baton``

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

@@ -0,0 +1,6 @@
+---
+source: libsql-server/tests/hrana/batch.rs
+expression: err
+snapshot_kind: text
+---
+Hrana: `api error: `status=400 Bad Request, body=Received an invalid baton``

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

@@ -1,9 +1,10 @@
 ---
 source: libsql-server/tests/namespaces/shared_schema.rs
 expression: "conn.execute(\"create table test (x)\", ()).await.unwrap_err()"
+snapshot_kind: text
 ---
 Hrana(
     Api(
-        "{\"error\":\"Authorization forbidden: Current session doesn't not have Write permission to namespace schema\"}",
+        "status=403 Forbidden, body={\"error\":\"Authorization forbidden: Current session doesn't not have Write permission to namespace schema\"}",
     ),
 )

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

@@ -1,9 +1,10 @@
 ---
 source: libsql-server/tests/namespaces/shared_schema.rs
 expression: "conn.execute(\"create table test (x)\", ()).await.unwrap_err()"
+snapshot_kind: text
 ---
 Hrana(
     Api(
-        "{\"error\":\"Authorization forbidden: DDL statements not permitted on namespace ns1\"}",
+        "status=403 Forbidden, body={\"error\":\"Authorization forbidden: DDL statements not permitted on namespace ns1\"}",
     ),
 )

libsql-server/tests/standalone/snapshots/tests__standalone__attach__attach_auth-2.snap

@@ -1,9 +1,10 @@
 ---
 source: libsql-server/tests/standalone/attach.rs
 expression: "txn.execute(\"ATTACH DATABASE bar as bar\", ()).await.unwrap_err()"
+snapshot_kind: text
 ---
 Hrana(
     Api(
-        "{\"error\":\"Authorization forbidden: Current session doesn't not have AttachRead permission to namespace bar\"}",
+        "status=403 Forbidden, body={\"error\":\"Authorization forbidden: Current session doesn't not have AttachRead permission to namespace bar\"}",
     ),
 )

libsql-server/tests/standalone/snapshots/tests__standalone__attach__attach_auth-3.snap.new

@@ -0,0 +1,10 @@
+---
+source: libsql-server/tests/standalone/attach.rs
+assertion_line: 182
+expression: "bar_conn.execute(\"ATTACH foo as foo\", ()).await.unwrap_err()"
+---
+Hrana(
+    Api(
+        "status=403 Forbidden, body={\"error\":\"Authorization forbidden: Namespace `foo` doesn't allow attach\"}",
+    ),
+)

libsql-server/tests/standalone/snapshots/tests__standalone__attach__attach_auth.snap

@@ -1,9 +1,10 @@
 ---
 source: libsql-server/tests/standalone/attach.rs
 expression: "bar_conn.execute(\"ATTACH foo as foo\", ()).await.unwrap_err()"
+snapshot_kind: text
 ---
 Hrana(
     Api(
-        "{\"error\":\"Authorization forbidden: Current session doesn't not have AttachRead permission to namespace foo\"}",
+        "status=403 Forbidden, body={\"error\":\"Authorization forbidden: Current session doesn't not have AttachRead permission to namespace foo\"}",
     ),
 )

libsql-server/tests/standalone/snapshots/tests__standalone__attach__attach_no_auth.snap

@@ -1,9 +1,10 @@
 ---
 source: libsql-server/tests/standalone/attach.rs
 expression: "bar_conn.execute(\"ATTACH foo as foo\", ()).await.unwrap_err()"
+snapshot_kind: text
 ---
 Hrana(
     Api(
-        "{\"error\":\"Authorization forbidden: Namespace `foo` doesn't allow attach\"}",
+        "status=403 Forbidden, body={\"error\":\"Authorization forbidden: Namespace `foo` doesn't allow attach\"}",
     ),
 )

libsql/src/rows.rs

@@ -89,6 +89,12 @@ impl Rows {
     }
 }
 
+impl fmt::Debug for Rows {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Rows").finish()
+    }
+}
+
 /// A libsql row.
 pub struct Row {
     pub(crate) inner: Box<dyn RowInner + Send + Sync>,