Skip to main content

sqlness_runner/cmd/
compat.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::path::PathBuf;
16use std::sync::Arc;
17
18use clap::{Parser, ValueEnum};
19use sqlness::QueryContext;
20use sqlness::interceptor::template::DELIMITER as TEMPLATE_DELIMITER;
21use sqlness::interceptor::{InterceptorRef, Registry};
22
23use crate::cmd::bare::ServerAddr;
24use crate::cmd::compat_case::{self, CompatCase, try_infer_version, version_matches_range};
25use crate::env::bare::{Env, StoreConfig, WalConfig};
26use crate::protocol_interceptor::{self, POSTGRES, PROTOCOL_KEY};
27use crate::util;
28
29const COMMENT_PREFIX: &str = "--";
30const INTERCEPTOR_PREFIX: &str = "-- SQLNESS";
31const QUERY_DELIMITER: char = ';';
32
33#[derive(Debug, Clone, Copy, ValueEnum, PartialEq, Eq)]
34enum CompatTopology {
35    Distributed,
36    Standalone,
37}
38
39impl CompatTopology {
40    fn as_str(self) -> &'static str {
41        match self {
42            Self::Distributed => "distributed",
43            Self::Standalone => "standalone",
44        }
45    }
46}
47
48/// Run compatibility tests in bare mode.
49///
50/// Starts a "from" distributed cluster, runs setup SQLs,
51/// then restarts the cluster with a "to" version on preserved state,
52/// and runs verify SQLs comparing results against `verify.result` files.
53///
54/// PR1 notes:
55/// - Sqlness interceptor comments are supported for each statement.
56/// - The runner starts the full distributed topology, including flownode.
57#[derive(Debug, Parser)]
58pub struct CompatCommand {
59    /// Version of the "from" GreptimeDB binary (e.g. "v0.9.5") or "current".
60    /// If neither --from-version nor --from-bins-dir is specified, the
61    /// current debug build is used for both from and to.
62    #[clap(long)]
63    from_version: Option<String>,
64
65    /// Path to the directory containing the "from" GreptimeDB binary.
66    #[clap(long)]
67    from_bins_dir: Option<PathBuf>,
68
69    /// Path to the directory containing the "to" GreptimeDB binary.
70    /// Defaults to the current debug build.
71    #[clap(long)]
72    to_bins_dir: Option<PathBuf>,
73
74    /// Version of the "to" GreptimeDB binary (e.g. "v1.1.4") or "current".
75    /// Downloads the release binary when needed. Cannot be used with
76    /// `--to-bins-dir`.
77    #[clap(long)]
78    to_version: Option<String>,
79
80    /// Directory of compatibility test cases.
81    /// Defaults to `tests/compatibility/cases` relative to workspace root.
82    #[clap(long)]
83    case_dir: Option<PathBuf>,
84
85    /// Name of test cases to run. Accepts a regexp.
86    #[clap(long, default_value = ".*")]
87    test_filter: String,
88
89    /// Require exactly this many cases after all filters are applied.
90    #[clap(long)]
91    expect_cases: Option<usize>,
92
93    /// Topology to start for the compatibility test.
94    #[clap(long, value_enum, default_value_t = CompatTopology::Distributed)]
95    topology: CompatTopology,
96
97    /// Fail this run as soon as one case fails.
98    #[clap(long, default_value = "false")]
99    fail_fast: bool,
100
101    /// Preserve persistent state in the temporary directory after run.
102    /// Etcd is always cleaned up regardless of this flag.
103    #[clap(long, default_value = "false")]
104    preserve_state: bool,
105
106    /// Pull different versions of GreptimeDB on need.
107    #[clap(long, default_value = "true")]
108    pull_version_on_need: bool,
109
110    /// Whether to set up etcd via Docker. Required for PR1 distributed compat.
111    /// External metadata stores are not supported by the compat MVP yet.
112    #[clap(long, default_value = "true")]
113    setup_etcd: bool,
114
115    /// Perform discovery and filtering only; print what would run without
116    /// starting any services, mutating files, or running setup/verify.
117    #[clap(long, default_value = "false")]
118    dry_run: bool,
119}
120
121impl CompatCommand {
122    pub async fn run(self) {
123        let dry_run = self.dry_run;
124        let topology = self.topology;
125
126        if self.to_bins_dir.is_some() && self.to_version.is_some() {
127            panic!("--to-version cannot be used with --to-bins-dir");
128        }
129
130        // ---- 1. Validate MVP runtime constraints ----
131        if !dry_run && topology == CompatTopology::Distributed && !self.setup_etcd {
132            panic!(
133                "compat MVP requires Docker etcd (--setup-etcd=true); external metadata stores are not supported yet"
134            );
135        }
136
137        // ---- 2. Resolve case directory ----
138        let case_dir = self.case_dir.unwrap_or_else(default_compat_case_dir);
139
140        if !case_dir.is_dir() {
141            panic!("Case directory not found: {}", case_dir.display());
142        }
143
144        // ---- 3. Discover cases ----
145        let mut cases = compat_case::discover_cases(&case_dir).unwrap_or_else(|e| panic!("{e}"));
146
147        // Filter by test_filter
148        let filter_re = regex::Regex::new(&self.test_filter)
149            .unwrap_or_else(|e| panic!("Invalid test filter regex '{}': {e}", self.test_filter));
150        cases.retain(|c| filter_re.is_match(&c.metadata.name));
151
152        // Filter by topology
153        cases.retain(|c| {
154            c.metadata
155                .topologies
156                .iter()
157                .any(|candidate| candidate == topology.as_str())
158        });
159
160        if cases.is_empty() {
161            if let Some(expected_cases) = self.expect_cases {
162                panic!(
163                    "Expected {expected_cases} compatibility cases after name and topology filtering, found 0"
164                );
165            }
166            if dry_run {
167                println!(
168                    "DRY-RUN: no compat cases found matching filter '{}' and topology '{}'",
169                    self.test_filter,
170                    topology.as_str()
171                );
172            } else {
173                println!(
174                    "No compat cases found matching filter '{}' and topology '{}'",
175                    self.test_filter,
176                    topology.as_str()
177                );
178            }
179            return;
180        }
181
182        // ---- 3b. Validate metadata (incl. version constraints) before filtering ----
183        // Must run before version-range filtering so invalid constraints like
184        // `>=not-a-version` cause a hard error instead of silent skip.
185        compat_case::validate_cases_metadata(&cases).unwrap_or_else(|e| panic!("{e}"));
186
187        // ---- 3c. Validate namespace dedup before version filtering ----
188        // Validate globally for all selected topology/name cases so duplicated
189        // namespaces cannot hide behind version filters.
190        compat_case::validate_case_namespaces(&cases).unwrap_or_else(|e| panic!("{e}"));
191
192        // ---- 4. Resolve "from" and "to" versions ----
193        let (from_bins_dir, from_version, from_ver_parsed, to_bins_dir, to_version, to_ver_parsed) =
194            if dry_run {
195                // Dry-run: resolve versions without panicking on missing binaries.
196                // try_infer_version returns None gracefully when the binary is absent.
197                let dry_run_from_bins_dir = self
198                    .from_bins_dir
199                    .clone()
200                    .unwrap_or_else(|| util::get_binary_dir("debug"));
201                let from_ver_str = self
202                    .from_version
203                    .as_deref()
204                    .and_then(|v| {
205                        if v == "current" {
206                            None
207                        } else {
208                            Some(v.to_string())
209                        }
210                    })
211                    .or_else(|| try_infer_version(&dry_run_from_bins_dir).map(|v| v.to_string()));
212                let from_ver_parsed = from_ver_str
213                    .as_deref()
214                    .and_then(|s| compat_case::Version::parse(s).ok());
215
216                let dry_run_to_bins_dir = self.to_bins_dir.clone().unwrap_or_else(|| {
217                    self.to_version
218                        .as_deref()
219                        .filter(|version| *version != "current")
220                        .map(|version| PathBuf::from(util::get_workspace_root()).join(version))
221                        .unwrap_or_else(|| util::get_binary_dir("debug"))
222                });
223                let to_ver_str = self
224                    .to_version
225                    .as_deref()
226                    .filter(|version| *version != "current")
227                    .map(str::to_string)
228                    .or_else(|| try_infer_version(&dry_run_to_bins_dir).map(|v| v.to_string()));
229                let to_ver_parsed = to_ver_str
230                    .as_deref()
231                    .and_then(|s| compat_case::Version::parse(s).ok());
232
233                (
234                    Some(dry_run_from_bins_dir),
235                    from_ver_str,
236                    from_ver_parsed,
237                    Some(dry_run_to_bins_dir),
238                    to_ver_str,
239                    to_ver_parsed,
240                )
241            } else {
242                // Normal path: resolve bins (may panic if binary not found).
243                let from_bins_dir = resolve_bins(
244                    self.from_bins_dir.as_ref(),
245                    self.from_version.as_deref(),
246                    self.pull_version_on_need,
247                )
248                .await;
249
250                let from_version = if let Some(ref ver) = self.from_version {
251                    if ver != "current" {
252                        Some(ver.clone())
253                    } else {
254                        try_infer_version(&from_bins_dir).map(|v| v.to_string())
255                    }
256                } else {
257                    try_infer_version(&from_bins_dir).map(|v| v.to_string())
258                };
259
260                let from_ver_parsed = from_version
261                    .as_deref()
262                    .and_then(|s| compat_case::Version::parse(s).ok());
263
264                let to_bins_dir = resolve_bins(
265                    self.to_bins_dir.as_ref(),
266                    self.to_version.as_deref(),
267                    self.pull_version_on_need,
268                )
269                .await;
270                let to_version = self
271                    .to_version
272                    .as_deref()
273                    .filter(|version| *version != "current")
274                    .map(str::to_string)
275                    .or_else(|| try_infer_version(&to_bins_dir).map(|v| v.to_string()));
276                let to_ver_parsed = to_version
277                    .as_deref()
278                    .and_then(|s| compat_case::Version::parse(s).ok());
279
280                (
281                    Some(from_bins_dir),
282                    from_version,
283                    from_ver_parsed,
284                    Some(to_bins_dir),
285                    to_version,
286                    to_ver_parsed,
287                )
288            };
289
290        // ---- 5b. Filter by version range ----
291        let pre_filter_count = cases.len();
292        cases.retain(|c| {
293            let from_ok = version_matches_range(from_ver_parsed.as_ref(), &c.metadata.from_range);
294            if !from_ok {
295                let from_label = from_ver_parsed
296                    .as_ref()
297                    .map(|v| v.to_string())
298                    .unwrap_or_else(|| "unknown".to_string());
299                println!(
300                    "Skipping case '{}': from_range {:?} does not match version '{}'",
301                    c.metadata.name, c.metadata.from_range, from_label
302                );
303            }
304            from_ok
305        });
306        cases.retain(|c| {
307            let to_ok = version_matches_range(to_ver_parsed.as_ref(), &c.metadata.to_range);
308            if !to_ok {
309                let to_label = to_ver_parsed
310                    .as_ref()
311                    .map(|v| v.to_string())
312                    .unwrap_or_else(|| "unknown".to_string());
313                println!(
314                    "Skipping case '{}': to_range {:?} does not match version '{}'",
315                    c.metadata.name, c.metadata.to_range, to_label
316                );
317            }
318            to_ok
319        });
320
321        if pre_filter_count != cases.len() {
322            println!(
323                "Version-range filtering: {} → {} cases",
324                pre_filter_count,
325                cases.len()
326            );
327        }
328
329        if let Some(expected_cases) = self.expect_cases {
330            assert_eq!(
331                cases.len(),
332                expected_cases,
333                "Expected {expected_cases} compatibility cases after all filtering, found {}",
334                cases.len()
335            );
336        }
337
338        if cases.is_empty() {
339            if dry_run {
340                println!("DRY-RUN: no compat cases would run after version-range filtering");
341            } else {
342                println!("No compat cases remaining after version-range filtering");
343            }
344            return;
345        }
346
347        if dry_run {
348            println!("DRY-RUN: would run {} compat case(s)", cases.len());
349            println!("  topology:     {}", topology.as_str());
350            println!(
351                "  from version: {}",
352                from_version.as_deref().unwrap_or(
353                    "unknown (use --from-version, --from-bins-dir, or build debug binary)"
354                )
355            );
356            println!(
357                "  to version:   {}",
358                to_version
359                    .as_deref()
360                    .unwrap_or("unknown (use --to-bins-dir or build debug binary)")
361            );
362            if pre_filter_count != cases.len() {
363                println!();
364                println!(
365                    "Version-range filtering reduced {} → {} cases (see 'Skipping case' messages above)",
366                    pre_filter_count,
367                    cases.len()
368                );
369            }
370            println!();
371            for c in &cases {
372                println!("  case:        {}", c.metadata.name);
373                println!("    namespace:   {}", c.namespace);
374                println!("    from_range:  {:?}", c.metadata.from_range);
375                println!("    to_range:    {:?}", c.metadata.to_range);
376                println!("    features:    {:?}", c.metadata.features);
377            }
378            println!();
379            println!("Dry run complete. Remove --dry-run to execute.");
380            return;
381        }
382
383        println!(
384            "Running {} compat case(s) with topology {}:",
385            cases.len(),
386            topology.as_str()
387        );
388        for c in &cases {
389            println!(
390                "  - {} (namespace: {}, topologies: {:?})",
391                c.metadata.name, c.namespace, c.metadata.topologies
392            );
393        }
394
395        // ---- 6. Create temp directory (after filtering so early exits don't leave empty dirs) ----
396        let temp_dir = tempfile::Builder::new()
397            .prefix("sqlness-compat")
398            .tempdir()
399            .unwrap();
400        let sqlness_home = temp_dir.keep();
401        unsafe {
402            std::env::set_var(
403                "SQLNESS_HOME",
404                sqlness_home.join("copy").display().to_string(),
405            );
406        }
407
408        // ---- 7. Build interceptor registry ----
409        let interceptor_registry = create_interceptor_registry();
410
411        // ---- 7b. Create Env for the selected topology ----
412        let setup_etcd = topology == CompatTopology::Distributed && self.setup_etcd;
413        let store_config = StoreConfig {
414            store_addrs: if setup_etcd {
415                vec!["127.0.0.1:2379".to_string()]
416            } else {
417                vec![]
418            },
419            setup_etcd,
420            setup_pg: None,
421            setup_mysql: None,
422            enable_flat_format: false,
423            enable_gc: false,
424        };
425
426        let env = Env::new(
427            sqlness_home.clone(),
428            ServerAddr::default(),
429            WalConfig::RaftEngine,
430            self.pull_version_on_need,
431            from_bins_dir,
432            store_config,
433            vec![],
434        );
435
436        // ---- 7c. Etcd cleanup guard ----
437        // Arm this only immediately before starting the cluster. Earlier validation
438        // failures should not stop an unrelated local container named `etcd`.
439        let mut etcd_guard = if setup_etcd {
440            Some(EtcdGuard::new())
441        } else {
442            None
443        };
444
445        // ---- 8. Run setup phase on the from-version cluster ----
446        println!("Starting from-version {} cluster...", topology.as_str());
447        let mut db = match topology {
448            CompatTopology::Distributed => env.compat_start_distributed(0).await,
449            CompatTopology::Standalone => env.compat_start_standalone(0).await,
450        };
451
452        println!("Running setup phase...");
453        for case in &cases {
454            run_compat_phase(&db, case, &interceptor_registry, CompatPhase::Setup)
455                .await
456                .unwrap_or_else(|e| panic!("Setup failed for case '{}': {e}", case.metadata.name));
457            println!("  Setup: {} - OK", case.metadata.name);
458        }
459
460        // ---- 9. Switch to the to-version binary and restart cluster ----
461        // to_bins_dir was already resolved during version-range filtering
462        println!("Restarting cluster with to-version binary on preserved state...");
463        let to_bins_dir = to_bins_dir.expect("to_bins_dir must be resolved in non-dry-run mode");
464        env.compat_restart(&db, to_bins_dir).await;
465
466        // ---- 10. Run verify phase on the to-version cluster ----
467        println!("Running verify phase...");
468        let mut failed = Vec::new();
469        for case in &cases {
470            match run_compat_phase(&db, case, &interceptor_registry, CompatPhase::Verify).await {
471                Ok(()) => println!("  Verify: {} - PASSED", case.metadata.name),
472                Err(e) => {
473                    println!("  Verify: {} - FAILED: {e}", case.metadata.name);
474                    failed.push(case.metadata.name.clone());
475                    if self.fail_fast {
476                        break;
477                    }
478                }
479            }
480        }
481
482        // ---- 11. Stop cluster ----
483        db.compat_stop();
484
485        // ---- 12. Cleanup ----
486        // Etcd is always cleaned up; --preserve-state only preserves sqlness_home.
487        if setup_etcd {
488            println!("Stopping etcd");
489            util::stop_rm_etcd();
490        }
491
492        if !self.preserve_state {
493            println!("Removing state in {:?}", sqlness_home);
494            tokio::fs::remove_dir_all(sqlness_home)
495                .await
496                .unwrap_or_else(|e| println!("Warning: failed to clean up temp dir: {e}"));
497        }
498
499        // Disarm the etcd guard now that we've done normal cleanup.
500        if let Some(mut guard) = etcd_guard.take() {
501            guard.disarm();
502        }
503
504        if failed.is_empty() {
505            println!("\n\x1b[32mAll compat tests passed!\x1b[0m");
506        } else {
507            println!("\n\x1b[31mFailed cases: {}\x1b[0m", failed.join(", "));
508            // Explicitly drop the guard before exit so it doesn't double-cleanup.
509            std::process::exit(1);
510        }
511    }
512}
513
514/// Guard that stops/removes Docker etcd on drop (panic or early exit).
515/// Disarm before normal cleanup to avoid double-cleanup.
516///
517/// The guard refuses to arm if a container named `etcd` already exists, so a
518/// failed compat run never deletes a developer-owned container with that name.
519struct EtcdGuard {
520    active: bool,
521}
522
523impl EtcdGuard {
524    fn new() -> Self {
525        let inspect_status = std::process::Command::new("docker")
526            .args(["container", "inspect", "etcd"])
527            .stdout(std::process::Stdio::null())
528            .stderr(std::process::Stdio::null())
529            .status();
530        if inspect_status.is_ok_and(|status| status.success()) {
531            panic!(
532                "A Docker container named `etcd` already exists. \
533                 Remove it before running compat tests so the cleanup guard \
534                 cannot delete a container it did not create."
535            );
536        }
537        Self { active: true }
538    }
539
540    fn disarm(&mut self) {
541        self.active = false;
542    }
543}
544
545impl Drop for EtcdGuard {
546    fn drop(&mut self) {
547        if self.active {
548            println!("EtcdGuard: emergency etcd cleanup (panic or early exit)");
549            // Best-effort: don't panic in Drop
550            let _ = std::process::Command::new("docker")
551                .args(["container", "stop", "etcd"])
552                .status();
553            let _ = std::process::Command::new("docker")
554                .args(["container", "rm", "etcd"])
555                .status();
556        }
557    }
558}
559
560/// Phase of compat execution.
561#[derive(Clone, Copy, PartialEq, Eq)]
562enum CompatPhase {
563    Setup,
564    Verify,
565}
566
567/// Create an interceptor registry matching the ordinary sqlness runner.
568fn create_interceptor_registry() -> Registry {
569    let mut interceptor_registry: Registry = Default::default();
570    interceptor_registry.register(
571        protocol_interceptor::PREFIX,
572        Arc::new(protocol_interceptor::ProtocolInterceptorFactory),
573    );
574    interceptor_registry
575}
576
577/// Resolve binary directory: explicit path takes priority, then version (pulls if needed),
578/// otherwise default to current debug build.
579///
580/// Validates that `<dir>/greptime` exists after resolution and canonicalizes the path.
581async fn resolve_bins(
582    bins_dir: Option<&PathBuf>,
583    version: Option<&str>,
584    pull_version_on_need: bool,
585) -> PathBuf {
586    let dir = if let Some(dir) = bins_dir {
587        dir.clone()
588    } else if let Some(ver) = version {
589        if ver == "current" {
590            util::get_binary_dir("debug")
591        } else {
592            util::maybe_pull_binary(ver, pull_version_on_need).await;
593            let root = std::path::PathBuf::from(util::get_workspace_root());
594            std::path::PathBuf::from_iter([root, std::path::PathBuf::from(ver)])
595        }
596    } else {
597        // Default: current debug build
598        util::get_binary_dir("debug")
599    };
600
601    // Canonicalize when possible (may fail if dir doesn't exist)
602    let dir = match dir.canonicalize() {
603        Ok(canon) => canon,
604        Err(e) => panic!(
605            "Cannot resolve binary directory '{}': {e}. \
606             Use --from-bins-dir / --to-bins-dir to specify the correct path, \
607             or --from-version to pull a release.",
608            dir.display()
609        ),
610    };
611
612    if !dir.join(util::PROGRAM).is_file() {
613        panic!(
614            "greptime binary not found in '{}'. \
615             Use --from-bins-dir / --to-bins-dir to specify the correct directory, \
616             or build greptime first (e.g. `cargo build -p greptime`). \
617             Note: if you use a custom target-dir, the binary may be elsewhere; \
618             pass the actual directory with --from-bins-dir or --to-bins-dir.",
619            dir.display()
620        );
621    }
622
623    dir
624}
625
626/// Default case directory: `tests/compatibility/cases` relative to workspace root.
627fn default_compat_case_dir() -> PathBuf {
628    let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
629    // CARGO_MANIFEST_DIR is tests/runner
630    // Pop to tests/
631    path.pop();
632    path.push("compatibility");
633    path.push("cases");
634    path
635}
636
637/// Run a single compat phase (setup or verify) for one case.
638async fn run_compat_phase(
639    db: &crate::env::bare::GreptimeDB,
640    case: &CompatCase,
641    registry: &Registry,
642    phase: CompatPhase,
643) -> Result<(), String> {
644    let sql_file = match phase {
645        CompatPhase::Setup => case.dir.join("setup.sql"),
646        CompatPhase::Verify => case.dir.join("verify.sql"),
647    };
648
649    let sql_content = std::fs::read_to_string(&sql_file)
650        .map_err(|e| format!("Failed to read {}: {e}", sql_file.display()))?;
651
652    let mut statements = parse_sql_file(&sql_content, registry)?;
653
654    // Execute statements
655    let mut verify_output = String::new();
656
657    for statement in &mut statements {
658        let (display, results) = statement.execute(db, &case.namespace).await?;
659
660        match phase {
661            CompatPhase::Setup => {
662                // Setup: just check for success (already returned Ok)
663            }
664            CompatPhase::Verify => {
665                verify_output.push_str(&display);
666                for result in results {
667                    verify_output.push_str(&result);
668                    verify_output.push('\n');
669                    verify_output.push('\n');
670                }
671            }
672        }
673    }
674
675    if phase == CompatPhase::Verify {
676        trim_trailing_blank_lines(&mut verify_output);
677
678        let result_path = case.dir.join("verify.result");
679
680        // If verify.result doesn't exist, generate it from actual output but
681        // return an error so the author must review, commit, and rerun.
682        if !result_path.is_file() {
683            std::fs::write(&result_path, &verify_output)
684                .map_err(|e| format!("Failed to create {}: {e}", result_path.display()))?;
685            return Err(format!(
686                "Created missing verify.result for case '{}'; review the generated file, commit it, and rerun",
687                case.metadata.name
688            ));
689        }
690
691        let expected = std::fs::read_to_string(&result_path)
692            .map_err(|e| format!("Failed to read {}: {e}", result_path.display()))?;
693
694        if verify_output != expected {
695            // Update the result file with actual output to aid local update.
696            std::fs::write(&result_path, &verify_output)
697                .map_err(|e| format!("Failed to update {}: {e}", result_path.display()))?;
698
699            // Generate a simple diff
700            let diff = simple_diff(&expected, &verify_output);
701            return Err(format!(
702                "Result mismatch for case '{}'.\nDiff:\n{diff}",
703                case.metadata.name
704            ));
705        }
706    }
707
708    Ok(())
709}
710
711/// Keep generated snapshots compatible with `git diff --check` by avoiding a
712/// trailing blank line at EOF while preserving the final newline.
713fn trim_trailing_blank_lines(output: &mut String) {
714    while output.ends_with("\n\n") {
715        output.pop();
716    }
717}
718
719/// Execute the namespace prelude (CREATE DATABASE IF NOT EXISTS + USE) for a case.
720/// This is NOT written into verify.result.
721///
722/// The prelude is protocol-aware:
723/// - `CREATE DATABASE` is always sent via gRPC so it works regardless of
724///   statement-level protocol directives.
725/// - For Postgres-protocol statements, `SET search_path` selects the case
726///   namespace instead of running `USE` (which is not valid PG SQL).
727/// - For MySQL and default/gRPC statements, `USE <ns>` runs through the
728///   statement's effective context.
729async fn run_namespace_prelude(
730    db: &crate::env::bare::GreptimeDB,
731    namespace: &str,
732    query_ctx: &QueryContext,
733) -> Result<(), String> {
734    // CREATE DATABASE always via gRPC — no protocol override
735    let create_db = format!("CREATE DATABASE IF NOT EXISTS {namespace}");
736    let default_ctx = QueryContext::default();
737    db.compat_query(&create_db, &default_ctx).await?;
738
739    // Postgres: select the namespace via search_path instead of USE.
740    if query_ctx
741        .context
742        .get(PROTOCOL_KEY)
743        .is_some_and(|p| p == POSTGRES)
744    {
745        let set_search_path = format!("SET search_path TO '{namespace}'");
746        db.compat_query(&set_search_path, query_ctx).await?;
747        return Ok(());
748    }
749
750    // MySQL / default (gRPC): execute USE
751    let use_db = format!("USE {namespace}");
752    db.compat_query(&use_db, query_ctx).await?;
753
754    Ok(())
755}
756
757/// A parsed SQL statement with sqlness comments and interceptors.
758struct ParsedStatement {
759    comment_lines: Vec<String>,
760    display_query: Vec<String>,
761    execute_query: Vec<String>,
762    interceptors: Vec<InterceptorRef>,
763}
764
765impl ParsedStatement {
766    fn new() -> Self {
767        Self {
768            comment_lines: Vec::new(),
769            display_query: Vec::new(),
770            execute_query: Vec::new(),
771            interceptors: Vec::new(),
772        }
773    }
774
775    fn push_comment(&mut self, line: String) {
776        self.comment_lines.push(line);
777    }
778
779    fn push_interceptor(&mut self, line: &str, registry: &Registry) -> Result<(), String> {
780        let Some((_, remaining)) = line.split_once(INTERCEPTOR_PREFIX) else {
781            return Err(format!(
782                "Missing sqlness interceptor prefix in line: {line}"
783            ));
784        };
785        let interceptor = registry.create(remaining).map_err(|e| e.to_string())?;
786        self.interceptors.push(interceptor);
787        Ok(())
788    }
789
790    fn append_query_line(&mut self, line: &str) {
791        self.display_query.push(line.to_string());
792        self.execute_query.push(line.to_string());
793    }
794
795    fn is_empty(&self) -> bool {
796        self.comment_lines.is_empty()
797            && self.display_query.is_empty()
798            && self.execute_query.is_empty()
799            && self.interceptors.is_empty()
800    }
801
802    fn has_query(&self) -> bool {
803        !self.execute_query.is_empty()
804    }
805
806    fn display_text(&self) -> String {
807        let mut output = String::new();
808        for comment in &self.comment_lines {
809            output.push_str(comment);
810            output.push('\n');
811        }
812        for line in &self.display_query {
813            output.push_str(line);
814        }
815        output.push('\n');
816        output.push('\n');
817        output
818    }
819
820    fn concat_query_lines(&self) -> String {
821        self.execute_query
822            .iter()
823            .fold(String::new(), |query, line| query + line)
824            .trim_start()
825            .to_string()
826    }
827
828    async fn before_execute_intercept(&mut self) -> QueryContext {
829        let mut context = QueryContext::default();
830        for interceptor in &self.interceptors {
831            interceptor
832                .before_execute_async(&mut self.execute_query, &mut context)
833                .await;
834        }
835        context
836    }
837
838    async fn after_execute_intercept(&self, result: &mut String) {
839        for interceptor in &self.interceptors {
840            interceptor.after_execute_async(result).await;
841        }
842    }
843
844    async fn execute(
845        &mut self,
846        db: &crate::env::bare::GreptimeDB,
847        namespace: &str,
848    ) -> Result<(String, Vec<String>), String> {
849        let display = self.display_text();
850        let context = self.before_execute_intercept().await;
851        db.compat_prepare_query_context(&context).await;
852        run_namespace_prelude(db, namespace, &context).await?;
853        let sql = self.concat_query_lines();
854        let mut results = Vec::new();
855
856        for sql in sql.split(TEMPLATE_DELIMITER) {
857            if sql.trim().is_empty() {
858                continue;
859            }
860            let sql = if sql.ends_with(QUERY_DELIMITER) {
861                sql.to_string()
862            } else {
863                format!("{sql};")
864            };
865            let mut result = db.compat_query(&sql, &context).await?;
866            self.after_execute_intercept(&mut result).await;
867            results.push(result);
868        }
869
870        Ok((display, results))
871    }
872}
873
874/// Parse a SQL file into statements using the same sqlness comment/interceptor
875/// conventions as the ordinary runner.
876fn parse_sql_file(content: &str, registry: &Registry) -> Result<Vec<ParsedStatement>, String> {
877    let mut statements = Vec::new();
878    let mut current_stmt = ParsedStatement::new();
879
880    for line in content.lines() {
881        if line.starts_with(COMMENT_PREFIX) {
882            current_stmt.push_comment(line.to_string());
883
884            if line.starts_with(INTERCEPTOR_PREFIX) {
885                current_stmt.push_interceptor(line, registry)?;
886            }
887            continue;
888        }
889
890        if line.is_empty() {
891            continue;
892        }
893
894        current_stmt.append_query_line(line);
895
896        // Check for statement terminator
897        if line.ends_with(QUERY_DELIMITER) {
898            if current_stmt.has_query() {
899                statements.push(current_stmt);
900            }
901            current_stmt = ParsedStatement::new();
902        } else {
903            current_stmt.append_query_line("\n");
904        }
905    }
906
907    // Flush any remaining statement
908    if !current_stmt.is_empty() && current_stmt.has_query() {
909        statements.push(current_stmt);
910    }
911
912    if statements.is_empty() {
913        return Err("No SQL statements found in file".to_string());
914    }
915
916    Ok(statements)
917}
918
919/// Generate a simple line-based diff between expected and actual.
920fn simple_diff(expected: &str, actual: &str) -> String {
921    let mut diff = String::new();
922    let expected_lines: Vec<&str> = expected.lines().collect();
923    let actual_lines: Vec<&str> = actual.lines().collect();
924    let max_len = expected_lines.len().max(actual_lines.len());
925
926    for i in 0..max_len {
927        let exp = expected_lines.get(i).unwrap_or(&"(missing)");
928        let act = actual_lines.get(i).unwrap_or(&"(missing)");
929        if exp != act {
930            diff.push_str(&format!("  Line {}:\n", i + 1));
931            diff.push_str(&format!("    expected: {exp}\n"));
932            diff.push_str(&format!("    actual:   {act}\n"));
933        }
934    }
935
936    if diff.is_empty() {
937        diff.push_str("  (files differ but no line-level diff found — may be whitespace)\n");
938    }
939
940    diff
941}
942
943#[cfg(test)]
944mod tests {
945    use super::trim_trailing_blank_lines;
946
947    #[test]
948    fn test_trim_trailing_blank_lines_preserves_single_final_newline() {
949        let mut output = "SELECT 1;\n\n+---+\n\n".to_string();
950        trim_trailing_blank_lines(&mut output);
951
952        assert_eq!(output, "SELECT 1;\n\n+---+\n");
953    }
954}