count and limit

Author: mamcx

crates/core/src/sql/compiler.rs

@@ -230,14 +230,17 @@ fn compile_statement(db: &RelationalDB, statement: SqlAst) -> Result<CrudExpr, P
 mod tests {
     use super::*;
     use crate::db::datastore::traits::IsolationLevel;
-    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::execution_context::Workload;
     use crate::sql::execute::tests::run_for_testing;
     use expect_test::expect;
     use spacetimedb_lib::error::{ResultTest, TestError};
     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]
@@ -396,10 +399,10 @@ Seq Scan on test
             "select * from test where b = 2 and a = 1",
             expect![
                 r#"
-Index Scan using Index test_b_idx_btree (test.b) on test)
-    Index Cond: (test.b = U64(2))
-    Output: test.a, test.b
-    -> Filter: (test.a = U64(1))"#
+Index Scan using Index test_b_idx_btree (test.b) on test
+  Index Cond: (test.b = U64(2))
+  Output: test.a, test.b
+  -> Filter: (test.a = U64(1))"#
             ],
         );
 
@@ -875,4 +878,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

@@ -679,7 +679,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 = &[
@@ -715,11 +715,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))"#
             ],
         );
 
@@ -783,7 +784,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`.
@@ -811,11 +812,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))"#
             ],
         );
 
@@ -870,8 +872,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]
@@ -881,8 +882,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/physical-plan/src/plan.rs

@@ -2055,4 +2055,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"#]],
+        );
+    }
 }