feat: support cluster and singleton

cnosdb · May 24, 2024 · 9514e40 · 9514e40
1 parent ec6147e
commit 9514e40
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 1 deletion.
diff --git a/query_server/sqllogicaltests/src/instance.rs b/query_server/sqllogicaltests/src/instance.rs
@@ -21,18 +21,21 @@ use crate::error::{Result, SqlError};
 use crate::utils::normalize;
 
 pub struct CnosDBClient {
+    engine_name: String,
     relative_path: PathBuf,
     options: SqlClientOptions,
     create_options: CreateOptions,
 }
 
 impl CnosDBClient {
     pub fn new(
+        engine_name: impl Into<String>,
         relative_path: impl Into<PathBuf>,
         options: SqlClientOptions,
         create_options: CreateOptions,
     ) -> Result<Self, SqlError> {
         Ok(Self {
+            engine_name: engine_name.into(),
             relative_path: relative_path.into(),
             options,
             create_options,
@@ -66,7 +69,7 @@ impl sqllogictest::AsyncDB for CnosDBClient {
     }
 
     fn engine_name(&self) -> &str {
-        "CnosDB"
+        &self.engine_name
     }
 
     async fn sleep(dur: Duration) {

diff --git a/query_server/sqllogicaltests/src/main.rs b/query_server/sqllogicaltests/src/main.rs
@@ -1,5 +1,6 @@
 use std::error::Error;
 use std::path::{Path, PathBuf};
+use std::process::exit;
 
 use clap::Parser;
 
@@ -33,6 +34,11 @@ const CNOSDB_TARGET_PARTITIONS_DEFAULT: usize = 8;
 pub async fn main() -> Result<(), Box<dyn Error>> {
     let options = Options::new();
 
+    if options.mode != "singleton" && options.mode != "cluster" {
+        eprintln!("Unsupported mode: {}", options.mode);
+        exit(1);
+    }
+
     let db_options = SqlClientOptions {
         flight_host: options.flight_host.clone(),
         flight_port: options.flight_port,
@@ -71,6 +77,7 @@ pub async fn main() -> Result<(), Box<dyn Error>> {
         }
         if options.complete_mode {
             os::run_complete_file(
+                options.mode.clone(),
                 &path,
                 relative_path,
                 db_options.clone(),
@@ -79,6 +86,7 @@ pub async fn main() -> Result<(), Box<dyn Error>> {
             .await?;
         } else {
             os::run_test_file(
+                options.mode.clone(),
                 &path,
                 relative_path,
                 db_options.clone(),
@@ -138,6 +146,8 @@ struct Options {
     /// Number of replication sets
     replica_count: u32,
 
+    mode: String,
+
     flight_host: String,
     flight_port: u16,
     http_host: String,
@@ -157,6 +167,9 @@ struct CliOptions {
 
     #[arg(long = "complete", default_value = "false")]
     complete_mode: bool,
+
+    #[arg(short = 'M', long = "mode", default_value = "singleton")]
+    mode: String,
 }
 
 impl Options {
@@ -183,6 +196,7 @@ impl Options {
             complete_mode: args.complete_mode,
             shard_count: args.shard_count,
             replica_count: args.replica_count,
+            mode: args.mode,
             flight_host,
             flight_port,
             http_host,

diff --git a/query_server/sqllogicaltests/src/os/unix.rs b/query_server/sqllogicaltests/src/os/unix.rs
@@ -7,6 +7,7 @@ use trace::info;
 use crate::instance::{CnosDBClient, CreateOptions, SqlClientOptions};
 
 pub async fn run_test_file(
+    engine_name: String,
     path: &Path,
     relative_path: PathBuf,
     options: SqlClientOptions,
@@ -15,6 +16,7 @@ pub async fn run_test_file(
     info!("Running with DataFusion runner: {}", path.display());
     let mut runner = sqllogictest::Runner::new(|| async {
         CnosDBClient::new(
+            engine_name.clone(),
             relative_path.clone(),
             options.clone(),
             createoptions.clone(),
@@ -25,6 +27,7 @@ pub async fn run_test_file(
 }
 
 pub async fn run_complete_file(
+    engine_name: String,
     path: &Path,
     relative_path: PathBuf,
     options: SqlClientOptions,
@@ -35,6 +38,7 @@ pub async fn run_complete_file(
     // let mut data = 3;
     let mut runner = sqllogictest::Runner::new(|| async {
         CnosDBClient::new(
+            engine_name.clone(),
             relative_path.clone(),
             options.clone(),
             createoptions.clone(),