feat: support multiple connections (sessions) in a single test file

CHANGELOG.md

@@ -2,9 +2,16 @@
 
 All notable changes to this project will be documented in this file.
 
-The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.15.0] - 2023-07-06
+
+* Allow multiple connections to the database in a single test case, which is useful for testing the transaction behavior. This can be achieved by attaching a `connection foo` record before the query or statement.
+  - (parser) Add `Record::Connection`.
+  - (runner) **Breaking change**: Since the runner may establish multiple connections at runtime, `Runner::new` now takes a `impl MakeConnection`, which is usually a closure that returns a try-future of the `AsyncDB` instance.
+  - (bin) The connection to the database is now established lazily on the first query or statement.
+
 ## [0.14.0] - 2023-06-08
 
 * We enhanced how `skipif` and `onlyif` works. Previously it checks against `DB::engine_name()`, and `sqllogictest-bin` didn't implement it.

Cargo.toml

@@ -2,7 +2,7 @@
 members = ["examples/*", "sqllogictest", "sqllogictest-bin", "sqllogictest-engines", "tests"]
 
 [workspace.package]
-version = "0.14.0"
+version = "0.15.0"
 edition = "2021"
 homepage = "https://github.com/risinglightdb/sqllogictest-rs"
 keywords = ["sql", "database", "parser", "cli"]

examples/basic/examples/basic.rs

@@ -49,7 +49,7 @@ impl sqllogictest::DB for FakeDB {
 }
 
 fn main() {
-    let mut tester = sqllogictest::Runner::new(FakeDB);
+    let mut tester = sqllogictest::Runner::new(|| async { Ok(FakeDB) });
 
     let mut filename = PathBuf::from(file!());
     filename.pop();

examples/condition/examples/condition.rs

@@ -43,7 +43,7 @@ impl sqllogictest::DB for FakeDB {
 
 fn main() {
     for engine_name in ["risinglight", "otherdb"] {
-        let mut tester = sqllogictest::Runner::new(FakeDB { engine_name });
+        let mut tester = sqllogictest::Runner::new(|| async { Ok(FakeDB { engine_name }) });
 
         let mut filename = PathBuf::from(file!());
         filename.pop();

examples/connection/Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "connection"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[dependencies]
+sqllogictest = { path = "../../sqllogictest" }

examples/connection/connection.slt

@@ -0,0 +1,41 @@
+query I
+select counter()
+----
+1
+
+query I
+select counter()
+----
+2
+
+connection another
+query I
+select counter()
+----
+1
+
+# `default` is the name of the default connection if not specified
+connection default
+query I
+select counter()
+----
+3
+
+connection another
+query I
+select counter()
+----
+2
+
+# connection names are case sensitive
+connection AnOtHeR
+query I
+select counter()
+----
+1
+
+# connection only works for one record, the next one will use `default`
+query I
+select counter()
+----
+4

examples/connection/examples/connection.rs

@@ -0,0 +1,53 @@
+use std::path::PathBuf;
+
+use sqllogictest::{DBOutput, DefaultColumnType};
+
+pub struct FakeDB {
+    counter: u64,
+}
+
+impl FakeDB {
+    #[allow(clippy::unused_async)]
+    async fn connect() -> Result<Self, FakeDBError> {
+        Ok(Self { counter: 0 })
+    }
+}
+
+#[derive(Debug)]
+pub struct FakeDBError;
+
+impl std::fmt::Display for FakeDBError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{self:?}")
+    }
+}
+
+impl std::error::Error for FakeDBError {}
+
+impl sqllogictest::DB for FakeDB {
+    type Error = FakeDBError;
+    type ColumnType = DefaultColumnType;
+
+    fn run(&mut self, sql: &str) -> Result<DBOutput<Self::ColumnType>, FakeDBError> {
+        if sql == "select counter()" {
+            self.counter += 1;
+            Ok(DBOutput::Rows {
+                types: vec![DefaultColumnType::Integer],
+                rows: vec![vec![self.counter.to_string()]],
+            })
+        } else {
+            Err(FakeDBError)
+        }
+    }
+}
+
+fn main() {
+    let mut tester = sqllogictest::Runner::new(FakeDB::connect);
+
+    let mut filename = PathBuf::from(file!());
+    filename.pop();
+    filename.pop();
+    filename.push("connection.slt");
+
+    tester.run_file(filename).unwrap();
+}

examples/custom_type/examples/custom_type.rs

@@ -67,7 +67,7 @@ impl sqllogictest::DB for FakeDB {
 }
 
 fn main() {
-    let mut tester = sqllogictest::Runner::new(FakeDB);
+    let mut tester = sqllogictest::Runner::new(|| async { Ok(FakeDB) });
     tester.with_column_validator(strict_column_validator);
 
     let mut filename = PathBuf::from(file!());

examples/file_level_sort_mode/examples/file_level_sort_mode.rs

@@ -41,7 +41,7 @@ impl sqllogictest::DB for FakeDB {
 }
 
 fn main() {
-    let mut tester = sqllogictest::Runner::new(FakeDB);
+    let mut tester = sqllogictest::Runner::new(|| async { Ok(FakeDB) });
 
     let mut filename = PathBuf::from(file!());
     filename.pop();

examples/include/examples/include.rs

@@ -44,7 +44,7 @@ impl sqllogictest::DB for FakeDB {
 }
 
 fn main() {
-    let mut tester = sqllogictest::Runner::new(FakeDB);
+    let mut tester = sqllogictest::Runner::new(|| async { Ok(FakeDB) });
 
     let mut filename = PathBuf::from(file!());
     filename.pop();

examples/rowsort/examples/rowsort.rs

@@ -41,7 +41,7 @@ impl sqllogictest::DB for FakeDB {
 }
 
 fn main() {
-    let mut tester = sqllogictest::Runner::new(FakeDB);
+    let mut tester = sqllogictest::Runner::new(|| async { Ok(FakeDB) });
 
     let mut filename = PathBuf::from(file!());
     filename.pop();

examples/test_dir_escape/examples/test_dir_escape.rs

@@ -28,7 +28,7 @@ impl sqllogictest::DB for FakeDB {
 }
 
 fn main() {
-    let mut tester = sqllogictest::Runner::new(FakeDB);
+    let mut tester = sqllogictest::Runner::new(|| async { Ok(FakeDB) });
     // enable `__TEST_DIR__` override
     tester.enable_testdir();
 

examples/validator/examples/validator.rs

@@ -28,7 +28,7 @@ impl sqllogictest::DB for FakeDB {
 }
 
 fn main() {
-    let mut tester = sqllogictest::Runner::new(FakeDB);
+    let mut tester = sqllogictest::Runner::new(|| async { Ok(FakeDB) });
     // Validator will always return true.
     tester.with_validator(|_, _| true);
 

sqllogictest-bin/Cargo.toml

@@ -24,8 +24,8 @@ glob = "0.3"
 itertools = "0.10"
 quick-junit = { version = "0.2" }
 rand = "0.8"
-sqllogictest = { path = "../sqllogictest", version = "0.14" }
-sqllogictest-engines = { path = "../sqllogictest-engines", version = "0.14" }
+sqllogictest = { path = "../sqllogictest", version = "0.15" }
+sqllogictest-engines = { path = "../sqllogictest-engines", version = "0.15" }
 tokio = { version = "1", features = [
     "rt",
     "rt-multi-thread",

sqllogictest-bin/src/engines.rs

@@ -23,7 +23,7 @@ pub enum EngineConfig {
     External(String),
 }
 
-enum Engines {
+pub(crate) enum Engines {
     Postgres(PostgresSimple),
     PostgresExtended(PostgresExtended),
     External(ExternalDriver),
@@ -48,12 +48,21 @@ impl From<&DBConfig> for PostgresConfig {
     }
 }
 
-pub(super) async fn connect(engine: &EngineConfig, config: &DBConfig) -> Result<impl AsyncDB> {
+pub(crate) async fn connect(
+    engine: &EngineConfig,
+    config: &DBConfig,
+) -> Result<Engines, EnginesError> {
     Ok(match engine {
-        EngineConfig::Postgres => Engines::Postgres(PostgresSimple::connect(config.into()).await?),
-        EngineConfig::PostgresExtended => {
-            Engines::PostgresExtended(PostgresExtended::connect(config.into()).await?)
-        }
+        EngineConfig::Postgres => Engines::Postgres(
+            PostgresSimple::connect(config.into())
+                .await
+                .map_err(|e| EnginesError(e.into()))?,
+        ),
+        EngineConfig::PostgresExtended => Engines::PostgresExtended(
+            PostgresExtended::connect(config.into())
+                .await
+                .map_err(|e| EnginesError(e.into()))?,
+        ),
         EngineConfig::External(cmd_tmpl) => {
             let (host, port) = config.random_addr();
             let cmd_str = cmd_tmpl
@@ -64,21 +73,25 @@ pub(super) async fn connect(engine: &EngineConfig, config: &DBConfig) -> Result<
                 .replace("{pass}", &config.pass);
             let mut cmd = Command::new("bash");
             cmd.args(["-c", &cmd_str]);
-            Engines::External(ExternalDriver::connect(cmd).await?)
+            Engines::External(
+                ExternalDriver::connect(cmd)
+                    .await
+                    .map_err(|e| EnginesError(e.into()))?,
+            )
         }
     })
 }
 
 #[derive(Debug)]
-struct AnyhowError(anyhow::Error);
+pub(crate) struct EnginesError(anyhow::Error);
 
-impl Display for AnyhowError {
+impl Display for EnginesError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         self.0.fmt(f)
     }
 }
 
-impl std::error::Error for AnyhowError {
+impl std::error::Error for EnginesError {
     fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
         self.0.source()
     }
@@ -96,10 +109,10 @@ impl Engines {
 
 #[async_trait]
 impl AsyncDB for Engines {
-    type Error = AnyhowError;
+    type Error = EnginesError;
     type ColumnType = DefaultColumnType;
 
     async fn run(&mut self, sql: &str) -> Result<DBOutput<Self::ColumnType>, Self::Error> {
-        self.run(sql).await.map_err(AnyhowError)
+        self.run(sql).await.map_err(EnginesError)
     }
 }