count and limit

Author: mamcx

crates/core/src/db/relational_db.rs

@@ -2051,7 +2051,7 @@ pub mod tests_utils {
     pub fn expect_query_with_auth<T: TableSchemaView + StateView>(tx: &T, sql: &str, auth: AuthCtx, expect: Expect) {
         let schema = SchemaViewer::new(tx, &auth);
 
-        check_query(&schema, ExplainOptions::default(), &auth, sql, expect)
+        check_query(&schema, ExplainOptions::default(), &auth, sql, expect).unwrap();
     }
 
     /// [`Expect`] the query `Explain` output, using `AuthCtx::for_testing`.
@@ -2064,7 +2064,7 @@ pub mod tests_utils {
         let auth = AuthCtx::for_testing();
         let schema = SchemaViewer::new(tx, &auth);
 
-        check_sub(&schema, ExplainOptions::default(), &auth, sql, expect)
+        check_sub(&schema, ExplainOptions::default(), &auth, sql, expect).unwrap()
     }
 }
 

crates/core/src/sql/compiler.rs

@@ -223,7 +223,9 @@ fn compile_statement(db: &RelationalDB, statement: SqlAst) -> Result<CrudExpr, P
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::db::relational_db::tests_utils::{begin_tx, expect_query, expect_sub, insert, with_auto_commit, TestDB};
+    use crate::db::relational_db::tests_utils::{
+        begin_tx, expect_query, expect_query_with_auth, expect_sub, insert, with_auto_commit, TestDB,
+    };
     use crate::sql::execute::tests::run_for_testing;
     use expect_test::expect;
     use spacetimedb_datastore::execution_context::Workload;
@@ -232,6 +234,7 @@ mod tests {
     use spacetimedb_lib::{ConnectionId, Identity};
     use spacetimedb_primitives::col_list;
     use spacetimedb_sats::{product, AlgebraicType, GroundSpacetimeType as _};
+    use spacetimedb_schema::schema::RowLevelSecuritySchema;
     use std::convert::From;
 
     #[test]
@@ -447,8 +450,8 @@ Index Scan using Index test_b_idx_btree (test.b) on test
             "select * from test where b = 2 and a = 1",
             expect![
                 r#"
-Index Scan using Index test_a_b_idx_btree (test.b, test.a) on test
-  Index Cond: (test.b = U64(2), test.a = U64(1))
+Index Scan using Index test_a_b_idx_btree (test.a, test.b) on test
+  Index Cond: (test.a = U64(1), test.b = U64(2))
   Output: test.a, test.b, test.c, test.d"#
             ],
         );
@@ -785,8 +788,8 @@ Index Join: Rhs on lhs
   Inner Unique: false
   Join Cond: (rhs.b = lhs.b)
   Output: lhs.a, lhs.b
-  -> Index Scan using Index rhs_b_c_idx_btree (rhs.c, rhs.b) on rhs
-     Index Cond: (rhs.c = U64(2), rhs.b = U64(4))
+  -> Index Scan using Index rhs_b_c_idx_btree (rhs.b, rhs.c) on rhs
+     Index Cond: (rhs.b = U64(4), rhs.c = U64(2))
      Output: rhs.b, rhs.c, rhs.d
      -> Filter: (rhs.d = U64(3))"#
             ],
@@ -868,4 +871,78 @@ Hash Join: Lhs
         );
         Ok(())
     }
+
+    #[test]
+    fn compile_limit() -> ResultTest<()> {
+        let db = TestDB::durable()?;
+        db.create_table_for_test("test", &[("a", AlgebraicType::U64)], &[0.into()])?;
+
+        let tx = db.begin_tx(Workload::ForTests);
+        expect_query(
+            &tx,
+            "select * from test limit 1",
+            expect![
+                r#"
+Limit: 1
+  Output: test.a
+  -> Seq Scan on test
+     Output: test.a"#
+            ],
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn compile_count() -> ResultTest<()> {
+        let db = TestDB::durable()?;
+        db.create_table_for_test("test", &[("a", AlgebraicType::U64)], &[0.into()])?;
+
+        let tx = db.begin_tx(Workload::ForTests);
+        expect_query(
+            &tx,
+            "select count(*) as n from test",
+            expect![
+                r#"
+Count
+  Output: n
+  -> Seq Scan on test
+     Output: test.a"#
+            ],
+        );
+
+        // TODO: Currently, we don't support `count` with a `limit` clause.
+        Ok(())
+    }
+
+    #[test]
+    fn compile_rls() -> ResultTest<()> {
+        let db = TestDB::in_memory()?;
+        let table_id = db.create_table_for_test("test", &[("a", AlgebraicType::U64)], &[0.into()])?;
+
+        let mut tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::ForTests);
+
+        let rls = RowLevelSecuritySchema {
+            sql: "SELECT * FROM test WHERE a = 1".into(),
+            table_id,
+        };
+        tx.create_row_level_security(rls.clone())?;
+
+        let user = Identity::from_be_byte_array([3; 32]);
+        expect_query_with_auth(
+            &tx,
+            "select count(*) as n from test",
+            AuthCtx::new(Identity::ONE, user),
+            expect![
+                r#"
+Count
+  Output: n
+  -> Seq Scan on test
+     Output: test.a"#
+            ],
+        );
+
+        // TODO: Currently, we don't support `count` with a `limit` clause.
+        Ok(())
+    }
 }

crates/core/src/subscription/subscription.rs

@@ -680,7 +680,7 @@ mod tests {
         // Create table [lhs] with index on [b]
         let schema = &[("a", AlgebraicType::U64), ("b", AlgebraicType::U64)];
         let indexes = &[1.into()];
-        let _ = db.create_table_for_test("lhs", schema, indexes)?;
+        db.create_table_for_test("lhs", schema, indexes)?;
 
         // Create table [rhs] with index on [b, c]
         let schema = &[
@@ -716,11 +716,12 @@ mod tests {
             expect![
                 r#"
 Index Join: Rhs on lhs
-  -> Seq Scan on rhs
-    Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))
   Inner Unique: false
   Join Cond: (rhs.b = lhs.b)
-  Output: lhs.a, lhs.b"#
+  Output: lhs.a, lhs.b
+  -> Seq Scan on rhs
+     Output: rhs.b, rhs.c, rhs.d
+     -> Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))"#
             ],
         );
 
@@ -784,7 +785,7 @@ Index Join: Rhs on lhs
             ("d", AlgebraicType::U64),
         ];
         let indexes = &[0.into(), 1.into()];
-        let _ = db.create_table_for_test("rhs", schema, indexes)?;
+        db.create_table_for_test("rhs", schema, indexes)?;
 
         let tx = begin_tx(&db);
         // Should generate an index join since there is an index on `lhs.b`.
@@ -812,11 +813,12 @@ Index Join: Rhs on lhs
             expect![
                 r#"
 Index Join: Rhs on lhs
-  -> Seq Scan on rhs
-    Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))
   Inner Unique: false
   Join Cond: (rhs.b = lhs.b)
-  Output: lhs.a, lhs.b"#
+  Output: lhs.a, lhs.b
+  -> Seq Scan on rhs
+     Output: rhs.b, rhs.c, rhs.d
+     -> Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))"#
             ],
         );
 
@@ -871,8 +873,7 @@ Index Join: Rhs on lhs
         // Create table [lhs] with index on [b]
         let schema = &[("a", AlgebraicType::U64), ("b", AlgebraicType::U64)];
         let indexes = &[1.into()];
-        let _lhs_id = db
-            .create_table_for_test("lhs", schema, indexes)
+        db.create_table_for_test("lhs", schema, indexes)
             .expect("Failed to create_table_for_test lhs");
 
         // Create table [rhs] with index on [b, c]
@@ -882,8 +883,7 @@ Index Join: Rhs on lhs
             ("d", AlgebraicType::U64),
         ];
         let indexes = &[0.into(), 1.into()];
-        let _rhs_id = db
-            .create_table_for_test("rhs", schema, indexes)
+        db.create_table_for_test("rhs", schema, indexes)
             .expect("Failed to create_table_for_test rhs");
 
         let tx = begin_tx(&db);
@@ -916,11 +916,12 @@ Index Join: Rhs on lhs
             expect![
                 r#"
 Index Join: Rhs on lhs
-  -> Seq Scan on rhs
-    Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))
   Inner Unique: false
   Join Cond: (rhs.b = lhs.b)
-  Output: lhs.a, lhs.b"#
+  Output: lhs.a, lhs.b
+  -> Seq Scan on rhs
+     Output: rhs.b, rhs.c, rhs.d
+     -> Filter: (rhs.c > U64(2) AND rhs.c < U64(4) AND rhs.d = U64(3))"#
             ],
         );
 

crates/expr/src/errors.rs

@@ -149,6 +149,10 @@ pub enum TypingError {
     DuplicateName(#[from] DuplicateName),
     #[error(transparent)]
     FilterReturnType(#[from] FilterReturnType),
-    #[error("SQL query exceeds maximum allowed length of {max_length}: \"{sql}...\"")]
-    SqlMaxLengthExceeded { sql: String, max_length: usize },
+    #[error("SQL query exceeds maximum allowed length: {length} > {max_length}: \"{sql}...\"")]
+    SqlMaxLengthExceeded {
+        sql: String,
+        length: usize,
+        max_length: usize,
+    },
 }

crates/expr/src/lib.rs

@@ -387,6 +387,7 @@ pub fn check_sql_length(sql: &str) -> TypingResult<()> {
     if sql.len() > MAX_SQL_LENGTH {
         return Err(TypingError::SqlMaxLengthExceeded {
             sql: sql[..120].into(),
+            length: sql.len(),
             max_length: MAX_SQL_LENGTH,
         });
     }

crates/physical-plan/src/plan.rs

@@ -1218,18 +1218,18 @@ pub mod tests_utils {
     use crate::printer::{Explain, ExplainOptions};
     use crate::PhysicalCtx;
     use expect_test::Expect;
-    use spacetimedb_expr::check::{compile_sql_sub, SchemaView};
+    use spacetimedb_expr::check::{compile_sql_sub, SchemaView, TypingResult};
     use spacetimedb_expr::statement::compile_sql_stmt_with_ctx;
     use spacetimedb_lib::identity::AuthCtx;
 
-    fn sub<'a>(db: &'a impl SchemaView, auth: &AuthCtx, sql: &'a str) -> PhysicalCtx<'a> {
-        let plan = compile_sql_sub(sql, db, auth, true).unwrap();
-        compile(plan)
+    fn sub<'a>(db: &'a impl SchemaView, auth: &AuthCtx, sql: &'a str) -> TypingResult<PhysicalCtx<'a>> {
+        let plan = compile_sql_sub(sql, db, auth, true)?;
+        Ok(compile(plan))
     }
 
-    fn query<'a>(db: &'a impl SchemaView, auth: &AuthCtx, sql: &'a str) -> PhysicalCtx<'a> {
-        let plan = compile_sql_stmt_with_ctx(sql, db, auth, true).unwrap();
-        compile(plan)
+    fn query<'a>(db: &'a impl SchemaView, auth: &AuthCtx, sql: &'a str) -> TypingResult<PhysicalCtx<'a>> {
+        let plan = compile_sql_stmt_with_ctx(sql, db, auth, true)?;
+        Ok(compile(plan))
     }
 
     fn check(plan: PhysicalCtx, options: ExplainOptions, expect: Expect) {
@@ -1242,14 +1242,28 @@ pub mod tests_utils {
         expect.assert_eq(&explain.to_string());
     }
 
-    pub fn check_sub(db: &impl SchemaView, options: ExplainOptions, auth: &AuthCtx, sql: &str, expect: Expect) {
-        let plan = sub(db, auth, sql);
+    pub fn check_sub(
+        db: &impl SchemaView,
+        options: ExplainOptions,
+        auth: &AuthCtx,
+        sql: &str,
+        expect: Expect,
+    ) -> TypingResult<()> {
+        let plan = sub(db, auth, sql)?;
         check(plan, options, expect);
+        Ok(())
     }
 
-    pub fn check_query(db: &impl SchemaView, options: ExplainOptions, auth: &AuthCtx, sql: &str, expect: Expect) {
-        let plan = query(db, auth, sql);
+    pub fn check_query(
+        db: &impl SchemaView,
+        options: ExplainOptions,
+        auth: &AuthCtx,
+        sql: &str,
+        expect: Expect,
+    ) -> TypingResult<()> {
+        let plan = query(db, auth, sql)?;
         check(plan, options, expect);
+        Ok(())
     }
 }
 
@@ -1367,11 +1381,11 @@ mod tests {
     }
 
     fn check_sub(db: &SchemaViewer, sql: &str, expect: Expect) {
-        tests_utils::check_sub(db, db.options, &AuthCtx::for_testing(), sql, expect);
+        tests_utils::check_sub(db, db.options, &AuthCtx::for_testing(), sql, expect).unwrap();
     }
 
     fn check_query(db: &SchemaViewer, sql: &str, expect: Expect) {
-        tests_utils::check_query(db, db.options, &AuthCtx::for_testing(), sql, expect);
+        tests_utils::check_query(db, db.options, &AuthCtx::for_testing(), sql, expect).unwrap();
     }
 
     fn data() -> SchemaViewer {
@@ -2055,4 +2069,56 @@ Delete on p
      -> Filter: (p.id = U64(1))"#]],
         );
     }
+
+    #[test]
+    fn count() {
+        let db = data().with_options(ExplainOptions::default().optimize(true));
+
+        check_query(
+            &db,
+            "SELECT count(*) as n FROM p",
+            expect![[r#"
+Count
+  Output: n
+  -> Seq Scan on p
+     Output: p.id, p.name"#]],
+        );
+
+        check_query(
+            &db,
+            "SELECT count(*) as n FROM p WHERE id = 1",
+            expect![[r#"
+Count
+  Output: n
+  -> Index Scan using Index id 0 Unique(p.id) on p
+     Index Cond: (p.id = U64(1))
+     Output: p.id, p.name"#]],
+        );
+    }
+
+    #[test]
+    fn limit() {
+        let db = data().with_options(ExplainOptions::default().optimize(true));
+
+        check_query(
+            &db,
+            "SELECT * FROM p LIMIT 10",
+            expect![[r#"
+Limit: 10
+  Output: p.id, p.name
+  -> Seq Scan on p
+     Output: p.id, p.name"#]],
+        );
+
+        check_query(
+            &db,
+            "SELECT * FROM p WHERE id = 1 LIMIT 10",
+            expect![[r#"
+Limit: 10
+  Output: p.id, p.name
+  -> Index Scan using Index id 0 Unique(p.id) on p
+     Index Cond: (p.id = U64(1))
+     Output: p.id, p.name"#]],
+        );
+    }
 }