Introduce scheduler for sqllogictests

crates/sqllogictest/Cargo.toml

@@ -33,6 +33,8 @@ datafusion-sqllogictest = {workspace = true}
 enum-ordinalize = {workspace = true}
 sqllogictest = {workspace = true}
 indicatif = {workspace = true}
+serde = {workspace = true}
+tracing = {workspace = true}
 
 [package.metadata.cargo-machete]
 # These dependencies are added to ensure minimal dependency version

crates/sqllogictest/src/engine/datafusion.rs

@@ -26,25 +26,15 @@ use indicatif::ProgressBar;
 use sqllogictest::runner::AsyncDB;
 use toml::Table as TomlTable;
 
-use crate::engine::Engine;
+use crate::engine::EngineRunner;
 use crate::error::Result;
 
 pub struct DataFusionEngine {
     datafusion: DataFusion,
 }
 
 #[async_trait::async_trait]
-impl Engine for DataFusionEngine {
-    async fn new(config: TomlTable) -> Result<Self> {
-        let session_config = SessionConfig::new().with_target_partitions(4);
-        let ctx = SessionContext::new_with_config(session_config);
-        ctx.register_catalog("default", Self::create_catalog(&config).await?);
-
-        Ok(Self {
-            datafusion: DataFusion::new(ctx, PathBuf::from("testdata"), ProgressBar::new(100)),
-        })
-    }
-
+impl EngineRunner for DataFusionEngine {
     async fn run_slt_file(&mut self, path: &Path) -> Result<()> {
         let content = std::fs::read_to_string(path)
             .with_context(|| format!("Failed to read slt file {:?}", path))
@@ -61,6 +51,16 @@ impl Engine for DataFusionEngine {
 }
 
 impl DataFusionEngine {
+    pub async fn new(config: TomlTable) -> Result<Self> {
+        let session_config = SessionConfig::new().with_target_partitions(4);
+        let ctx = SessionContext::new_with_config(session_config);
+        ctx.register_catalog("default", Self::create_catalog(&config).await?);
+
+        Ok(Self {
+            datafusion: DataFusion::new(ctx, PathBuf::from("testdata"), ProgressBar::new(100)),
+        })
+    }
+
     async fn create_catalog(_: &TomlTable) -> anyhow::Result<Arc<dyn CatalogProvider>> {
         todo!()
     }

crates/sqllogictest/src/engine/mod.rs

@@ -21,10 +21,41 @@ use std::path::Path;
 
 use toml::Table as TomlTable;
 
+use crate::engine::datafusion::DataFusionEngine;
 use crate::error::Result;
 
+const KEY_TYPE: &str = "type";
+const TYPE_DATAFUSION: &str = "datafusion";
+
 #[async_trait::async_trait]
-pub trait Engine: Sized {
-    async fn new(config: TomlTable) -> Result<Self>;
+pub trait EngineRunner: Sized {
     async fn run_slt_file(&mut self, path: &Path) -> Result<()>;
 }
+
+pub enum Engine {
+    DataFusion(DataFusionEngine),
+}
+
+impl Engine {
+    pub async fn new(config: TomlTable) -> Result<Self> {
+        let engine_type = config
+            .get(KEY_TYPE)
+            .ok_or_else(|| anyhow::anyhow!("Missing required key: {KEY_TYPE}"))?
+            .as_str()
+            .ok_or_else(|| anyhow::anyhow!("Config value for {KEY_TYPE} must be a string"))?;
+
+        match engine_type {
+            TYPE_DATAFUSION => {
+                let engine = DataFusionEngine::new(config).await?;
+                Ok(Engine::DataFusion(engine))
+            }
+            _ => Err(anyhow::anyhow!("Unsupported engine type: {engine_type}").into()),
+        }
+    }
+
+    pub async fn run_slt_file(&mut self, path: &Path) -> Result<()> {
+        match self {
+            Engine::DataFusion(engine) => engine.run_slt_file(path).await,
+        }
+    }
+}

crates/sqllogictest/src/lib.rs

@@ -22,3 +22,5 @@
 mod engine;
 #[allow(dead_code)]
 mod error;
+#[allow(dead_code)]
+mod schedule;

crates/sqllogictest/src/schedule.rs

@@ -0,0 +1,103 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::HashMap;
+use std::path::PathBuf;
+
+use serde::{Deserialize, Serialize};
+
+use crate::engine::Engine;
+
+pub struct Schedule {
+    engines: HashMap<String, Engine>,
+    steps: Vec<Step>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Step {
+    /// Engine name
+    engine: String,
+    /// Slt file path
+    slt: String,
+}
+
+impl Schedule {
+    pub fn new(engines: HashMap<String, Engine>, steps: Vec<Step>) -> Self {
+        Self { engines, steps }
+    }
+
+    pub async fn run(mut self) -> anyhow::Result<()> {
+        for (idx, step) in self.steps.iter().enumerate() {
+            tracing::info!(
+                "Running step {}/{}, using engine {}, slt file path: {}",
+                idx + 1,
+                self.steps.len(),
+                &step.engine,
+                &step.slt
+            );
+
+            let engine = self
+                .engines
+                .get_mut(&step.engine)
+                .ok_or_else(|| anyhow::anyhow!("Engine {} not found", step.engine))?;
+
+            engine.run_slt_file(&PathBuf::from(&step.slt)).await?;
+            tracing::info!(
+                "Step {}/{}, engine {}, slt file path: {} finished",
+                idx + 1,
+                self.steps.len(),
+                &step.engine,
+                &step.slt
+            );
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use toml::Table as TomlTable;
+
+    use crate::schedule::Step;
+
+    #[test]
+    fn test_parse_steps() {
+        let steps = r#"
+            [[steps]]
+            engine = "datafusion"
+            slt = "test.slt"
+
+            [[steps]]
+            engine = "spark"
+            slt = "test2.slt"
+        "#;
+
+        let steps: Vec<Step> = toml::from_str::<TomlTable>(steps)
+            .unwrap()
+            .get("steps")
+            .unwrap()
+            .clone()
+            .try_into()
+            .unwrap();
+
+        assert_eq!(steps.len(), 2);
+        assert_eq!(steps[0].engine, "datafusion");
+        assert_eq!(steps[0].slt, "test.slt");
+        assert_eq!(steps[1].engine, "spark");
+        assert_eq!(steps[1].slt, "test2.slt");
+    }
+}