feat: support include sql (#8)

* feat: support include sql

Signed-off-by: Alex Chi <[email protected]>

* bump sqlplannertest version

Signed-off-by: Alex Chi Z <[email protected]>

---------

Signed-off-by: Alex Chi Z <[email protected]>

Author: skyzh

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "sqlplannertest"
-version = "0.2.0"
+version = "0.3.0"
 edition = "2021"
 description = "A yaml-based SQL planner test framework."
 license = "MIT OR Apache-2.0"

naivedb/tests/include.sql

@@ -0,0 +1 @@
+CREATE TABLE t3(v3 int);

naivedb/tests/ok.planner.sql

@@ -18,6 +18,8 @@ SELECT * FROM t1, t2 WHERE t1.v1 = t2.v2;
 /*
 === logical
 I'm a naive db, so I don't now how to process
+CREATE TABLE t3(v3 int);
+
 CREATE TABLE t1(v1 int);
 
 CREATE TABLE t2(v2 int);
@@ -28,6 +30,8 @@ SELECT * FROM t1, t2 WHERE t1.v1 = t2.v2;
 
 === physical
 I'm a naive db, so I don't now how to process
+CREATE TABLE t3(v3 int);
+
 CREATE TABLE t1(v1 int);
 
 CREATE TABLE t2(v2 int);

naivedb/tests/ok.yml

@@ -7,7 +7,15 @@
 - sql: |
     SELECT * FROM t1, t2 WHERE t1.v1 = t2.v2;
   desc: Test whether join works correctly.
-  before: ["*sql1", "*sql2", "CREATE TABLE t3(v3 int);"]
+  before: ["include_sql:include.sql", "*sql1", "*sql2", "CREATE TABLE t3(v3 int);"]
+  tasks:
+  - logical
+  - physical
+- sql: |
+    SELECT * FROM t1, t2 WHERE t1.v1 = t2.v2;
+  desc: Test whether join works correctly.
+  before: ["include_sql:include.sql", "*sql1", "*sql2", "CREATE TABLE t3(v3 int);"]
+  no_capture: true
   tasks:
   - logical
   - physical

src/apply.rs

@@ -62,11 +62,20 @@ where
             tokio::spawn(async move {
                 let testcases = tokio::fs::read(&path).await?;
                 let testcases: Vec<TestCase> = serde_yaml::from_slice(&testcases)?;
-                let testcases = parse_test_cases(testcases)?;
+                let testcases = parse_test_cases(
+                    {
+                        let mut path = path.clone();
+                        path.pop();
+                        path
+                    },
+                    testcases,
+                )?;
                 let mut generated_result = String::new();
                 for testcase in testcases {
                     let runner_result = runner.run(&testcase).await;
-                    generate_result(&testcase, &runner_result, &mut generated_result)?;
+                    if !testcase.no_capture {
+                        generate_result(&testcase, &runner_result, &mut generated_result)?;
+                    }
                 }
                 let path = {
                     let mut path = path;

src/lib.rs

@@ -19,6 +19,7 @@ pub struct TestCase {
     pub desc: Option<String>,
     pub sql: String,
     pub before: Option<Vec<String>>,
+    pub no_capture: Option<bool>,
     pub tasks: Option<Vec<String>>,
 }
 
@@ -29,6 +30,7 @@ pub struct ParsedTestCase {
     pub desc: Option<String>,
     pub sql: String,
     pub before_sql: Vec<String>,
+    pub no_capture: bool,
     pub tasks: Vec<String>,
 }
 
@@ -39,8 +41,11 @@ pub trait PlannerTestRunner: Send {
     async fn run(&mut self, test_case: &ParsedTestCase) -> Result<String>;
 }
 
-pub fn parse_test_cases(tests: Vec<TestCase>) -> Result<Vec<ParsedTestCase>> {
-    resolve_testcase_id(tests)
+pub fn parse_test_cases(
+    base_path: impl AsRef<Path>,
+    tests: Vec<TestCase>,
+) -> Result<Vec<ParsedTestCase>> {
+    resolve_testcase_id(base_path, tests)
 }
 
 const TEST_SUFFIX: &str = ".yml";
@@ -77,7 +82,7 @@ mod tests {
         "#
         .trim();
         let test_case: Vec<TestCase> = serde_yaml::from_str(x).unwrap();
-        let test_case = parse_test_cases(test_case).unwrap();
+        let test_case = parse_test_cases("", test_case).unwrap();
         assert_eq!(test_case.len(), 3);
         assert_eq!(test_case[2].before_sql.len(), 3);
         assert_eq!(

src/resolve_id.rs

@@ -12,13 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::HashMap;
+use std::{collections::HashMap, path::Path};
 
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, Context, Result};
 
 use crate::{ParsedTestCase, TestCase};
 
-pub fn resolve_testcase_id(testcases: Vec<TestCase>) -> Result<Vec<ParsedTestCase>> {
+pub fn resolve_testcase_id(
+    base_path: impl AsRef<Path>,
+    testcases: Vec<TestCase>,
+) -> Result<Vec<ParsedTestCase>> {
     let mut testcases_with_ids = HashMap::new();
     for testcase in &testcases {
         if let Some(id) = &testcase.id {
@@ -33,14 +36,18 @@ pub fn resolve_testcase_id(testcases: Vec<TestCase>) -> Result<Vec<ParsedTestCas
                 Some(
                     before
                         .iter()
-                        .map(|id| {
-                            if let Some(id) = id.strip_prefix('*') {
+                        .map(|maybe_sql| {
+                            if let Some(id) = maybe_sql.strip_prefix('*') {
                                 testcases_with_ids
                                     .get(id)
                                     .map(|case| case.sql.clone())
                                     .ok_or_else(|| anyhow!("failed to resolve {}: not found", id))
+                            } else if let Some(include) = maybe_sql.strip_prefix("include_sql:") {
+                                let path = base_path.as_ref().join(include);
+                                std::fs::read_to_string(&path)
+                                    .with_context(|| format!("failed to read: {}", include))
                             } else {
-                                Ok(id.to_string())
+                                Ok(maybe_sql.to_string())
                             }
                         })
                         .collect::<Result<Vec<_>>>()?,
@@ -54,6 +61,7 @@ pub fn resolve_testcase_id(testcases: Vec<TestCase>) -> Result<Vec<ParsedTestCas
                 id: testcase.id,
                 desc: testcase.desc,
                 sql: testcase.sql,
+                no_capture: testcase.no_capture == Some(true),
                 tasks: testcase.tasks.unwrap_or_default(),
             })
         })

src/test_runner.rs

@@ -79,11 +79,20 @@ where
         let mut runner = runner_fn().await?;
         let testcases = tokio::fs::read(&path).await?;
         let testcases: Vec<TestCase> = serde_yaml::from_slice(&testcases)?;
-        let testcases = parse_test_cases(testcases)?;
+        let testcases = parse_test_cases(
+            {
+                let mut path = path.clone();
+                path.pop();
+                path
+            },
+            testcases,
+        )?;
         let mut generated_result = String::new();
         for testcase in testcases {
             let runner_result = runner.run(&testcase).await;
-            generate_result(&testcase, &runner_result, &mut generated_result)?;
+            if !testcase.no_capture {
+                generate_result(&testcase, &runner_result, &mut generated_result)?;
+            }
         }
         let path = {
             let mut path = path;