Skip to content

Commit 2326db7

Browse files
mamcxmamcx
committed
count and limit
1 parent e2ab4cf commit 2326db7

File tree

6 files changed

+186
-37
lines changed

6 files changed

+186
-37
lines changed

‎crates/core/src/db/relational_db.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2051,7 +2051,7 @@ pub mod tests_utils {
20512051
pub fn expect_query_with_auth<T: TableSchemaView + StateView>(tx: &T, sql: &str, auth: AuthCtx, expect: Expect) {
20522052
let schema = SchemaViewer::new(tx, &auth);
20532053

2054-
check_query(&schema, ExplainOptions::default(), &auth, sql, expect)
2054+
check_query(&schema, ExplainOptions::default(), &auth, sql, expect).unwrap();
20552055
}
20562056

20572057
/// [`Expect`] the query `Explain` output, using `AuthCtx::for_testing`.
@@ -2064,7 +2064,7 @@ pub mod tests_utils {
20642064
let auth = AuthCtx::for_testing();
20652065
let schema = SchemaViewer::new(tx, &auth);
20662066

2067-
check_sub(&schema, ExplainOptions::default(), &auth, sql, expect)
2067+
check_sub(&schema, ExplainOptions::default(), &auth, sql, expect).unwrap()
20682068
}
20692069
}
20702070

‎crates/core/src/sql/compiler.rs

Lines changed: 82 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,9 @@ fn compile_statement(db: &RelationalDB, statement: SqlAst) -> Result<CrudExpr, P
223223
#[cfg(test)]
224224
mod tests {
225225
use super::*;
226-
use crate::db::relational_db::tests_utils::{begin_tx, expect_query, expect_sub, insert, with_auto_commit, TestDB};
226+
use crate::db::relational_db::tests_utils::{
227+
begin_tx, expect_query, expect_query_with_auth, expect_sub, insert, with_auto_commit, TestDB,
228+
};
227229
use crate::sql::execute::tests::run_for_testing;
228230
use expect_test::expect;
229231
use spacetimedb_datastore::execution_context::Workload;
@@ -232,6 +234,7 @@ mod tests {
232234
use spacetimedb_lib::{ConnectionId, Identity};
233235
use spacetimedb_primitives::col_list;
234236
use spacetimedb_sats::{product, AlgebraicType, GroundSpacetimeType as _};
237+
use spacetimedb_schema::schema::RowLevelSecuritySchema;
235238
use std::convert::From;
236239

237240
#[test]
@@ -447,8 +450,8 @@ Index Scan using Index test_b_idx_btree (test.b) on test
447450
"select * from test where b = 2 and a = 1",
448451
expect![
449452
r#"
450-
Index Scan using Index test_a_b_idx_btree (test.b, test.a) on test
451-
Index Cond: (test.b = U64(2), test.a = U64(1))
453+
Index Scan using Index test_a_b_idx_btree (test.a, test.b) on test
454+
Index Cond: (test.a = U64(1), test.b = U64(2))
452455
Output: test.a, test.b, test.c, test.d"#
453456
],
454457
);
@@ -785,8 +788,8 @@ Index Join: Rhs on lhs
785788
Inner Unique: false
786789
Join Cond: (rhs.b = lhs.b)
787790
Output: lhs.a, lhs.b
788-
-> Index Scan using Index rhs_b_c_idx_btree (rhs.c, rhs.b) on rhs
789-
Index Cond: (rhs.c = U64(2), rhs.b = U64(4))
791+
-> Index Scan using Index rhs_b_c_idx_btree (rhs.b, rhs.c) on rhs
792+
Index Cond: (rhs.b = U64(4), rhs.c = U64(2))
790793
Output: rhs.b, rhs.c, rhs.d
791794
-> Filter: (rhs.d = U64(3))"#
792795
],
@@ -868,4 +871,78 @@ Hash Join: Lhs
868871
);
869872
Ok(())
870873
}
874+
875+
#[test]
876+
fn compile_limit() -> ResultTest<()> {
877+
let db = TestDB::durable()?;
878+
db.create_table_for_test("test", &[("a", AlgebraicType::U64)], &[0.into()])?;
879+
880+
let tx = db.begin_tx(Workload::ForTests);
881+
expect_query(
882+
&tx,
883+
"select * from test limit 1",
884+
expect![
885+
r#"
886+
Limit: 1
887+
Output: test.a
888+
-> Seq Scan on test
889+
Output: test.a"#
890+
],
891+
);
892+
893+
Ok(())
894+
}
895+
896+
#[test]
897+
fn compile_count() -> ResultTest<()> {
898+
let db = TestDB::durable()?;
899+
db.create_table_for_test("test", &[("a", AlgebraicType::U64)], &[0.into()])?;
900+
901+
let tx = db.begin_tx(Workload::ForTests);
902+
expect_query(
903+
&tx,
904+
"select count(*) as n from test",
905+
expect![
906+
r#"
907+
Count
908+
Output: n
909+
-> Seq Scan on test
910+
Output: test.a"#
911+
],
912+
);
913+
914+
// TODO: Currently, we don't support `count` with a `limit` clause.
915+
Ok(())
916+
}
917+
918+
#[test]
919+
fn compile_rls() -> ResultTest<()> {
920+
let db = TestDB::in_memory()?;
921+
let table_id = db.create_table_for_test("test", &[("a", AlgebraicType::U64)], &[0.into()])?;
922+
923+
let mut tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::ForTests);
924+
925+
let rls = RowLevelSecuritySchema {
926+
sql: "SELECT * FROM test WHERE a = 1".into(),
927+
table_id,
928+
};
929+
tx.create_row_level_security(rls.clone())?;
930+
931+
let user = Identity::from_be_byte_array([3; 32]);
932+
expect_query_with_auth(
933+
&tx,
934+
"select count(*) as n from test",
935+
AuthCtx::new(Identity::ONE, user),
936+
expect![
937+
r#"
938+
Count
939+
Output: n
940+
-> Seq Scan on test
941+
Output: test.a"#
942+
],
943+
);
944+
945+
// TODO: Currently, we don't support `count` with a `limit` clause.
946+
Ok(())
947+
}
871948
}

‎crates/core/src/subscription/subscription.rs

Lines changed: 16 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -680,7 +680,7 @@ mod tests {
680680
// Create table [lhs] with index on [b]
681681
let schema = &[("a", AlgebraicType::U64), ("b", AlgebraicType::U64)];
682682
let indexes = &[1.into()];
683-
let _ = db.create_table_for_test("lhs", schema, indexes)?;
683+
db.create_table_for_test("lhs", schema, indexes)?;
684684

685685
// Create table [rhs] with index on [b, c]
686686
let schema = &[
@@ -716,11 +716,12 @@ mod tests {
716716
expect![
717717
r#"
718718
Index Join: Rhs on lhs
719-
-> Seq Scan on rhs
720-
Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))
721719
Inner Unique: false
722720
Join Cond: (rhs.b = lhs.b)
723-
Output: lhs.a, lhs.b"#
721+
Output: lhs.a, lhs.b
722+
-> Seq Scan on rhs
723+
Output: rhs.b, rhs.c, rhs.d
724+
-> Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))"#
724725
],
725726
);
726727

@@ -784,7 +785,7 @@ Index Join: Rhs on lhs
784785
("d", AlgebraicType::U64),
785786
];
786787
let indexes = &[0.into(), 1.into()];
787-
let _ = db.create_table_for_test("rhs", schema, indexes)?;
788+
db.create_table_for_test("rhs", schema, indexes)?;
788789

789790
let tx = begin_tx(&db);
790791
// Should generate an index join since there is an index on `lhs.b`.
@@ -812,11 +813,12 @@ Index Join: Rhs on lhs
812813
expect![
813814
r#"
814815
Index Join: Rhs on lhs
815-
-> Seq Scan on rhs
816-
Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))
817816
Inner Unique: false
818817
Join Cond: (rhs.b = lhs.b)
819-
Output: lhs.a, lhs.b"#
818+
Output: lhs.a, lhs.b
819+
-> Seq Scan on rhs
820+
Output: rhs.b, rhs.c, rhs.d
821+
-> Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))"#
820822
],
821823
);
822824

@@ -871,8 +873,7 @@ Index Join: Rhs on lhs
871873
// Create table [lhs] with index on [b]
872874
let schema = &[("a", AlgebraicType::U64), ("b", AlgebraicType::U64)];
873875
let indexes = &[1.into()];
874-
let _lhs_id = db
875-
.create_table_for_test("lhs", schema, indexes)
876+
db.create_table_for_test("lhs", schema, indexes)
876877
.expect("Failed to create_table_for_test lhs");
877878

878879
// Create table [rhs] with index on [b, c]
@@ -882,8 +883,7 @@ Index Join: Rhs on lhs
882883
("d", AlgebraicType::U64),
883884
];
884885
let indexes = &[0.into(), 1.into()];
885-
let _rhs_id = db
886-
.create_table_for_test("rhs", schema, indexes)
886+
db.create_table_for_test("rhs", schema, indexes)
887887
.expect("Failed to create_table_for_test rhs");
888888

889889
let tx = begin_tx(&db);
@@ -916,11 +916,12 @@ Index Join: Rhs on lhs
916916
expect![
917917
r#"
918918
Index Join: Rhs on lhs
919-
-> Seq Scan on rhs
920-
Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))
921919
Inner Unique: false
922920
Join Cond: (rhs.b = lhs.b)
923-
Output: lhs.a, lhs.b"#
921+
Output: lhs.a, lhs.b
922+
-> Seq Scan on rhs
923+
Output: rhs.b, rhs.c, rhs.d
924+
-> Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))"#
924925
],
925926
);
926927

‎crates/expr/src/errors.rs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,10 @@ pub enum TypingError {
149149
DuplicateName(#[from] DuplicateName),
150150
#[error(transparent)]
151151
FilterReturnType(#[from] FilterReturnType),
152-
#[error("SQL query exceeds maximum allowed length of {max_length}: \"{sql}...\"")]
153-
SqlMaxLengthExceeded { sql: String, max_length: usize },
152+
#[error("SQL query exceeds maximum allowed length: {length} > {max_length}: \"{sql}...\"")]
153+
SqlMaxLengthExceeded {
154+
sql: String,
155+
length: usize,
156+
max_length: usize,
157+
},
154158
}

‎crates/expr/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -387,6 +387,7 @@ pub fn check_sql_length(sql: &str) -> TypingResult<()> {
387387
if sql.len() > MAX_SQL_LENGTH {
388388
return Err(TypingError::SqlMaxLengthExceeded {
389389
sql: sql[..120].into(),
390+
length: sql.len(),
390391
max_length: MAX_SQL_LENGTH,
391392
});
392393
}

‎crates/physical-plan/src/plan.rs

Lines changed: 79 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1218,18 +1218,18 @@ pub mod tests_utils {
12181218
use crate::printer::{Explain, ExplainOptions};
12191219
use crate::PhysicalCtx;
12201220
use expect_test::Expect;
1221-
use spacetimedb_expr::check::{compile_sql_sub, SchemaView};
1221+
use spacetimedb_expr::check::{compile_sql_sub, SchemaView, TypingResult};
12221222
use spacetimedb_expr::statement::compile_sql_stmt_with_ctx;
12231223
use spacetimedb_lib::identity::AuthCtx;
12241224

1225-
fn sub<'a>(db: &'a impl SchemaView, auth: &AuthCtx, sql: &'a str) -> PhysicalCtx<'a> {
1226-
let plan = compile_sql_sub(sql, db, auth, true).unwrap();
1227-
compile(plan)
1225+
fn sub<'a>(db: &'a impl SchemaView, auth: &AuthCtx, sql: &'a str) -> TypingResult<PhysicalCtx<'a>> {
1226+
let plan = compile_sql_sub(sql, db, auth, true)?;
1227+
Ok(compile(plan))
12281228
}
12291229

1230-
fn query<'a>(db: &'a impl SchemaView, auth: &AuthCtx, sql: &'a str) -> PhysicalCtx<'a> {
1231-
let plan = compile_sql_stmt_with_ctx(sql, db, auth, true).unwrap();
1232-
compile(plan)
1230+
fn query<'a>(db: &'a impl SchemaView, auth: &AuthCtx, sql: &'a str) -> TypingResult<PhysicalCtx<'a>> {
1231+
let plan = compile_sql_stmt_with_ctx(sql, db, auth, true)?;
1232+
Ok(compile(plan))
12331233
}
12341234

12351235
fn check(plan: PhysicalCtx, options: ExplainOptions, expect: Expect) {
@@ -1242,14 +1242,28 @@ pub mod tests_utils {
12421242
expect.assert_eq(&explain.to_string());
12431243
}
12441244

1245-
pub fn check_sub(db: &impl SchemaView, options: ExplainOptions, auth: &AuthCtx, sql: &str, expect: Expect) {
1246-
let plan = sub(db, auth, sql);
1245+
pub fn check_sub(
1246+
db: &impl SchemaView,
1247+
options: ExplainOptions,
1248+
auth: &AuthCtx,
1249+
sql: &str,
1250+
expect: Expect,
1251+
) -> TypingResult<()> {
1252+
let plan = sub(db, auth, sql)?;
12471253
check(plan, options, expect);
1254+
Ok(())
12481255
}
12491256

1250-
pub fn check_query(db: &impl SchemaView, options: ExplainOptions, auth: &AuthCtx, sql: &str, expect: Expect) {
1251-
let plan = query(db, auth, sql);
1257+
pub fn check_query(
1258+
db: &impl SchemaView,
1259+
options: ExplainOptions,
1260+
auth: &AuthCtx,
1261+
sql: &str,
1262+
expect: Expect,
1263+
) -> TypingResult<()> {
1264+
let plan = query(db, auth, sql)?;
12521265
check(plan, options, expect);
1266+
Ok(())
12531267
}
12541268
}
12551269

@@ -1367,11 +1381,11 @@ mod tests {
13671381
}
13681382

13691383
fn check_sub(db: &SchemaViewer, sql: &str, expect: Expect) {
1370-
tests_utils::check_sub(db, db.options, &AuthCtx::for_testing(), sql, expect);
1384+
tests_utils::check_sub(db, db.options, &AuthCtx::for_testing(), sql, expect).unwrap();
13711385
}
13721386

13731387
fn check_query(db: &SchemaViewer, sql: &str, expect: Expect) {
1374-
tests_utils::check_query(db, db.options, &AuthCtx::for_testing(), sql, expect);
1388+
tests_utils::check_query(db, db.options, &AuthCtx::for_testing(), sql, expect).unwrap();
13751389
}
13761390

13771391
fn data() -> SchemaViewer {
@@ -2055,4 +2069,56 @@ Delete on p
20552069
-> Filter: (p.id = U64(1))"#]],
20562070
);
20572071
}
2072+
2073+
#[test]
2074+
fn count() {
2075+
let db = data().with_options(ExplainOptions::default().optimize(true));
2076+
2077+
check_query(
2078+
&db,
2079+
"SELECT count(*) as n FROM p",
2080+
expect![[r#"
2081+
Count
2082+
Output: n
2083+
-> Seq Scan on p
2084+
Output: p.id, p.name"#]],
2085+
);
2086+
2087+
check_query(
2088+
&db,
2089+
"SELECT count(*) as n FROM p WHERE id = 1",
2090+
expect![[r#"
2091+
Count
2092+
Output: n
2093+
-> Index Scan using Index id 0 Unique(p.id) on p
2094+
Index Cond: (p.id = U64(1))
2095+
Output: p.id, p.name"#]],
2096+
);
2097+
}
2098+
2099+
#[test]
2100+
fn limit() {
2101+
let db = data().with_options(ExplainOptions::default().optimize(true));
2102+
2103+
check_query(
2104+
&db,
2105+
"SELECT * FROM p LIMIT 10",
2106+
expect![[r#"
2107+
Limit: 10
2108+
Output: p.id, p.name
2109+
-> Seq Scan on p
2110+
Output: p.id, p.name"#]],
2111+
);
2112+
2113+
check_query(
2114+
&db,
2115+
"SELECT * FROM p WHERE id = 1 LIMIT 10",
2116+
expect![[r#"
2117+
Limit: 10
2118+
Output: p.id, p.name
2119+
-> Index Scan using Index id 0 Unique(p.id) on p
2120+
Index Cond: (p.id = U64(1))
2121+
Output: p.id, p.name"#]],
2122+
);
2123+
}
20582124
}

0 commit comments

Comments
 (0)