Skip to content

Commit

Permalink
Allow the test utils of the plan to be used elsewhere
Browse files Browse the repository at this point in the history
  • Loading branch information
@mamcx
mamcx committed Jan 14, 2025
1 parent 928642f commit 33c8978
Show file tree
Hide file tree
Showing 5 changed files with 129 additions and 96 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 4 additions & 2 deletions crates/core/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@ spacetimedb-lib = { workspace = true, features = ["serde", "metrics_impls"] }
spacetimedb-client-api-messages.workspace = true
spacetimedb-commitlog.workspace = true
spacetimedb-durability.workspace = true
spacetimedb-expr.workspace = true
spacetimedb-execution.workspace = true
spacetimedb-metrics.workspace = true
spacetimedb-primitives.workspace = true
spacetimedb-paths.workspace = true
Expand All @@ -30,8 +32,6 @@ spacetimedb-schema.workspace = true
spacetimedb-table.workspace = true
spacetimedb-vm.workspace = true
spacetimedb-snapshot.workspace = true
spacetimedb-expr.workspace = true
spacetimedb-execution.workspace = true

anyhow = { workspace = true, features = ["backtrace"] }
arrayvec.workspace = true
Expand All @@ -50,6 +50,7 @@ derive_more.workspace = true
dirs.workspace = true
enum-as-inner.workspace = true
enum-map.workspace = true
expect-test.workspace = true
flate2.workspace = true
fs2.workspace = true
futures.workspace = true
Expand Down Expand Up @@ -119,6 +120,7 @@ default = ["unindexed_iter_by_col_range_warn"]
# Enable test helpers and utils
test = []


[dev-dependencies]
spacetimedb-lib = { path = "../lib", features = ["proptest"] }
spacetimedb-sats = { path = "../sats", features = ["proptest"] }
Expand Down
2 changes: 1 addition & 1 deletion crates/physical-plan/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ description = "The physical query plan for the SpacetimeDB query engine"
[dependencies]
anyhow.workspace = true
derive_more.workspace = true
expect-test.workspace = true
itertools.workspace = true
spacetimedb-lib.workspace = true
spacetimedb-primitives.workspace = true
Expand All @@ -19,4 +20,3 @@ spacetimedb-table.workspace = true

[dev-dependencies]
pretty_assertions.workspace = true
expect-test.workspace = true
122 changes: 57 additions & 65 deletions crates/physical-plan/src/plan.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -971,18 +971,54 @@ impl<'a> PhysicalCtx<'a> {
}
}

pub mod tests_utils {
use super::*;
use crate::compile::compile;
use crate::printer::{Explain, ExplainOptions};
use expect_test::Expect;
use spacetimedb_expr::check::{compile_sql_sub, SchemaView};
use spacetimedb_expr::statement::compile_sql_stmt;

fn sub<'a>(db: &'a impl SchemaView, sql: &'a str) -> PhysicalCtx<'a> {
let plan = compile_sql_sub(sql, db).unwrap();
compile(plan)
}

fn query<'a>(db: &'a impl SchemaView, sql: &'a str) -> PhysicalCtx<'a> {
let plan = compile_sql_stmt(sql, db).unwrap();
compile(plan)
}

fn check(plan: PhysicalCtx, options: ExplainOptions, expect: Expect) {
let plan = if options.optimize { plan.optimize() } else { plan };

let explain = Explain::new(&plan).with_options(options);

let explain = explain.build();
expect.assert_eq(&explain.to_string());
}

pub fn check_sub(db: &impl SchemaView, options: ExplainOptions, sql: &str, expect: Expect) {
let plan = sub(db, sql);
check(plan, options, expect);
}

pub fn check_query(db: &impl SchemaView, options: ExplainOptions, sql: &str, expect: Expect) {
let plan = query(db, sql);
check(plan, options, expect);
}
}

#[cfg(test)]
mod tests {
use super::*;

use crate::compile::compile;
use crate::compile::compile_project_plan;
use crate::plan::TupleField;
use crate::printer::Explain;
use crate::printer::ExplainOptions;
use expect_test::{expect, Expect};
use pretty_assertions::assert_eq;
use spacetimedb_expr::check::{compile_sql_sub, parse_and_type_sub, SchemaView};

use spacetimedb_expr::statement::compile_sql_stmt;
use spacetimedb_expr::check::{parse_and_type_sub, SchemaView};
use spacetimedb_lib::{
db::auth::{StAccess, StTableType},
AlgebraicType, AlgebraicValue,
Expand All @@ -994,45 +1030,28 @@ mod tests {
};
use spacetimedb_sql_parser::ast::BinOp;

use crate::compile::compile_project_plan;

use super::{PhysicalExpr, ProjectPlan};

struct SchemaViewer {
schemas: Vec<Arc<TableSchema>>,
optimize: bool,
show_source: bool,
show_schema: bool,
show_timings: bool,
options: ExplainOptions,
}

impl SchemaViewer {
fn new(schemas: Vec<Arc<TableSchema>>) -> Self {
Self {
schemas,
optimize: false,
show_source: false,
show_schema: false,
show_timings: false,
options: ExplainOptions::default(),
}
}

fn optimize(mut self) -> Self {
self.optimize = true;
self
}
fn with_source(mut self) -> Self {
self.show_source = true;
fn with_options(mut self, options: ExplainOptions) -> Self {
self.options = options;
self
}
fn with_schema(mut self) -> Self {
self.show_schema = true;
self
}
// TODO: Remove when we integrate it.
#[allow(dead_code)]
fn with_timings(mut self) -> Self {
self.show_timings = true;

fn optimize(mut self, optimize: bool) -> Self {
self.options = self.options.optimize(optimize);
self
}
}
Expand Down Expand Up @@ -1221,7 +1240,7 @@ Seq Scan on t
Some(0),
));

let db = SchemaViewer::new(vec![u.clone(), l.clone(), b.clone()]).optimize();
let db = SchemaViewer::new(vec![u.clone(), l.clone(), b.clone()]).optimize(true);

let sql = "
select b.*
Expand Down Expand Up @@ -1429,7 +1448,7 @@ Seq Scan on t
Some(0),
));

let db = SchemaViewer::new(vec![m.clone(), w.clone(), p.clone()]);
let db = SchemaViewer::new(vec![m.clone(), w.clone(), p.clone()]).optimize(false);

let sql = "
select p.*
Expand Down Expand Up @@ -1634,7 +1653,7 @@ Hash Join: All
None,
));

let db = SchemaViewer::new(vec![t.clone()]).optimize();
let db = SchemaViewer::new(vec![t.clone()]).optimize(true);

let sql = "select * from t where x = 3 and y = 4 and z = 5";
check_sub(
Expand Down Expand Up @@ -1772,47 +1791,20 @@ Index Scan using Index id 2: (x, y, z) on t
Some(0),
));

SchemaViewer::new(vec![m.clone(), w.clone(), p.clone()])
}

fn sub<'a>(db: &'a SchemaViewer, sql: &'a str) -> PhysicalCtx<'a> {
let plan = compile_sql_sub(sql, db).unwrap();
compile(plan)
}

fn query<'a>(db: &'a SchemaViewer, sql: &'a str) -> PhysicalCtx<'a> {
let plan = compile_sql_stmt(sql, db).unwrap();
compile(plan)
SchemaViewer::new(vec![m.clone(), w.clone(), p.clone()]).with_options(ExplainOptions::default().optimize(false))
}

fn check(db: &SchemaViewer, plan: PhysicalCtx, expect: Expect) {
let plan = if db.optimize { plan.optimize() } else { plan };

let explain = Explain::new(&plan);
let explain = if db.show_source { explain.with_source() } else { explain };
let explain = if db.show_schema { explain.with_schema() } else { explain };
let explain = if db.show_timings {
explain.with_timings()
} else {
explain
};

let explain = explain.build();
expect.assert_eq(&explain.to_string());
}
fn check_sub(db: &SchemaViewer, sql: &str, expect: Expect) {
let plan = sub(db, sql);
check(db, plan, expect);
tests_utils::check_sub(db, db.options, sql, expect);
}

fn check_query(db: &SchemaViewer, sql: &str, expect: Expect) {
let plan = query(db, sql);
check(db, plan, expect);
tests_utils::check_query(db, db.options, sql, expect);
}

#[test]
fn plan_metadata() {
let db = data().with_schema().with_source().optimize();
let db = data().with_options(ExplainOptions::new().with_schema().with_source().optimize(true));
check_query(
&db,
"SELECT m.* FROM m CROSS JOIN p WHERE m.employee = 1",
Expand Down Expand Up @@ -1920,7 +1912,7 @@ Index Scan using Index id 2: (x, y, z) on t

#[test]
fn index_scan_filter() {
let db = data().optimize();
let db = data().optimize(true);

check_sub(
&db,
Expand Down Expand Up @@ -1967,7 +1959,7 @@ Index Scan using Index id 2: (x, y, z) on t

#[test]
fn semi_join() {
let db = data().optimize();
let db = data().optimize(true);

check_sub(
&db,
Expand Down
Loading

0 comments on commit 33c8978

Please sign in to comment.