Adding the system table for row level security (#1746)

crates/core/src/db/datastore/locking_tx_datastore/committed_state.rs

@@ -11,9 +11,9 @@ use crate::{
                 system_tables, StColumnRow, StConstraintData, StConstraintRow, StIndexAlgorithm, StIndexRow,
                 StSequenceRow, StTableFields, StTableRow, SystemTable, ST_CLIENT_ID, ST_CLIENT_IDX, ST_COLUMN_ID,
                 ST_COLUMN_IDX, ST_COLUMN_NAME, ST_CONSTRAINT_ID, ST_CONSTRAINT_IDX, ST_CONSTRAINT_NAME, ST_INDEX_ID,
-                ST_INDEX_IDX, ST_INDEX_NAME, ST_MODULE_ID, ST_MODULE_IDX, ST_RESERVED_SEQUENCE_RANGE, ST_SCHEDULED_ID,
-                ST_SCHEDULED_IDX, ST_SEQUENCE_ID, ST_SEQUENCE_IDX, ST_SEQUENCE_NAME, ST_TABLE_ID, ST_TABLE_IDX,
-                ST_VAR_ID, ST_VAR_IDX,
+                ST_INDEX_IDX, ST_INDEX_NAME, ST_MODULE_ID, ST_MODULE_IDX, ST_RESERVED_SEQUENCE_RANGE,
+                ST_ROW_LEVEL_SECURITY_ID, ST_ROW_LEVEL_SECURITY_IDX, ST_SCHEDULED_ID, ST_SCHEDULED_IDX, ST_SEQUENCE_ID,
+                ST_SEQUENCE_IDX, ST_SEQUENCE_NAME, ST_TABLE_ID, ST_TABLE_IDX, ST_VAR_ID, ST_VAR_IDX,
             },
             traits::TxData,
         },
@@ -225,6 +225,10 @@ impl CommittedState {
 
         self.create_table(ST_SCHEDULED_ID, schemas[ST_SCHEDULED_IDX].clone());
 
+        self.create_table(ST_ROW_LEVEL_SECURITY_ID, schemas[ST_ROW_LEVEL_SECURITY_IDX].clone());
+
+        // IMPORTANT: It is crucial that the `st_sequences` table is created last
+
         // Insert the sequences into `st_sequences`
         let (st_sequences, blob_store) =
             self.get_table_and_blob_store_or_create(ST_SEQUENCE_ID, &schemas[ST_SEQUENCE_IDX]);

crates/core/src/db/datastore/locking_tx_datastore/datastore.rs

@@ -61,7 +61,7 @@ pub struct Locking {
     /// The state of sequence generation in this database.
     sequence_state: Arc<Mutex<SequencesState>>,
     /// The address of this database.
-    database_address: Address,
+    pub(crate) database_address: Address,
 }
 
 impl Locking {
@@ -379,6 +379,10 @@ impl MutTxDatastore for Locking {
         tx.table_id_from_name(table_name, self.database_address)
     }
 
+    fn table_id_exists_mut_tx(&self, tx: &Self::MutTx, table_id: &TableId) -> bool {
+        tx.table_name(*table_id).is_some()
+    }
+
     fn table_name_from_id_mut_tx<'a>(
         &'a self,
         ctx: &'a ExecutionContext,
@@ -511,10 +515,6 @@ impl MutTxDatastore for Locking {
         Ok((gens, row_ref.collapse()))
     }
 
-    fn table_id_exists_mut_tx(&self, tx: &Self::MutTx, table_id: &TableId) -> bool {
-        tx.table_name(*table_id).is_some()
-    }
-
     fn metadata_mut_tx(&self, tx: &Self::MutTx) -> Result<Option<Metadata>> {
         let ctx = ExecutionContext::internal(self.database_address);
         tx.iter(&ctx, ST_MODULE_ID)?.next().map(metadata_from_row).transpose()
@@ -927,9 +927,10 @@ mod tests {
     use super::*;
     use crate::db::datastore::system_tables::{
         system_tables, StColumnRow, StConstraintData, StConstraintFields, StConstraintRow, StIndexAlgorithm,
-        StIndexFields, StIndexRow, StScheduledFields, StSequenceFields, StSequenceRow, StTableRow, StVarFields,
-        StVarValue, ST_CLIENT_NAME, ST_COLUMN_ID, ST_COLUMN_NAME, ST_CONSTRAINT_ID, ST_CONSTRAINT_NAME, ST_INDEX_ID,
-        ST_INDEX_NAME, ST_MODULE_NAME, ST_RESERVED_SEQUENCE_RANGE, ST_SCHEDULED_ID, ST_SCHEDULED_NAME, ST_SEQUENCE_ID,
+        StIndexFields, StIndexRow, StRowLevelSecurityFields, StScheduledFields, StSequenceFields, StSequenceRow,
+        StTableRow, StVarFields, StVarValue, ST_CLIENT_NAME, ST_COLUMN_ID, ST_COLUMN_NAME, ST_CONSTRAINT_ID,
+        ST_CONSTRAINT_NAME, ST_INDEX_ID, ST_INDEX_NAME, ST_MODULE_NAME, ST_RESERVED_SEQUENCE_RANGE,
+        ST_ROW_LEVEL_SECURITY_ID, ST_ROW_LEVEL_SECURITY_NAME, ST_SCHEDULED_ID, ST_SCHEDULED_NAME, ST_SEQUENCE_ID,
         ST_SEQUENCE_NAME, ST_TABLE_NAME, ST_VAR_ID, ST_VAR_NAME,
     };
     use crate::db::datastore::traits::{IsolationLevel, MutTx};
@@ -944,7 +945,9 @@ mod tests {
     use spacetimedb_primitives::{col_list, ColId, ScheduleId};
     use spacetimedb_sats::{product, AlgebraicType, GroundSpacetimeType};
     use spacetimedb_schema::def::{BTreeAlgorithm, ConstraintData, IndexAlgorithm, UniqueConstraintData};
-    use spacetimedb_schema::schema::{ColumnSchema, ConstraintSchema, IndexSchema, SequenceSchema};
+    use spacetimedb_schema::schema::{
+        ColumnSchema, ConstraintSchema, IndexSchema, RowLevelSecuritySchema, SequenceSchema,
+    };
     use spacetimedb_table::table::UniqueConstraintViolation;
 
     /// For the first user-created table, sequences in the system tables start
@@ -1320,6 +1323,7 @@ mod tests {
             TableRow { id: ST_CLIENT_ID.into(), name: ST_CLIENT_NAME, ty: StTableType::System, access: StAccess::Public, primary_key: None },
             TableRow { id: ST_VAR_ID.into(), name: ST_VAR_NAME, ty: StTableType::System, access: StAccess::Public, primary_key: Some(StVarFields::Name.into()) },
             TableRow { id: ST_SCHEDULED_ID.into(), name: ST_SCHEDULED_NAME, ty: StTableType::System, access: StAccess::Public, primary_key: Some(StScheduledFields::ScheduleId.into()) },
+            TableRow { id: ST_ROW_LEVEL_SECURITY_ID.into(), name: ST_ROW_LEVEL_SECURITY_NAME, ty: StTableType::System, access: StAccess::Public, primary_key: Some(StRowLevelSecurityFields::Sql.into()) },
         ]));
         #[rustfmt::skip]
         assert_eq!(query.scan_st_columns()?, map_array([
@@ -1371,6 +1375,9 @@ mod tests {
             ColRow { table: ST_SCHEDULED_ID.into(), pos: 1, name: "table_id", ty: TableId::get_type() },
             ColRow { table: ST_SCHEDULED_ID.into(), pos: 2, name: "reducer_name", ty: AlgebraicType::String },
             ColRow { table: ST_SCHEDULED_ID.into(), pos: 3, name: "schedule_name", ty: AlgebraicType::String },
+
+            ColRow { table: ST_ROW_LEVEL_SECURITY_ID.into(), pos: 0, name: "table_id", ty: TableId::get_type() },
+            ColRow { table: ST_ROW_LEVEL_SECURITY_ID.into(), pos: 1, name: "sql", ty: AlgebraicType::String },
         ]));
         #[rustfmt::skip]
         assert_eq!(query.scan_st_indexes()?, map_array([
@@ -1384,6 +1391,8 @@ mod tests {
             IndexRow { id: 8, table: ST_VAR_ID.into(), col: col(0), name: "idx_st_var_name_unique", },
             IndexRow { id: 9, table: ST_SCHEDULED_ID.into(), col: col(0), name: "idx_st_scheduled_schedule_id_unique", },
             IndexRow { id: 10, table: ST_SCHEDULED_ID.into(), col: col(1), name: "idx_st_scheduled_table_id_unique", },
+            IndexRow { id: 11, table: ST_ROW_LEVEL_SECURITY_ID.into(), col: col(0), name: "idx_st_row_level_security_btree_table_id"},
+            IndexRow { id: 12, table: ST_ROW_LEVEL_SECURITY_ID.into(), col: col(1), name: "idx_st_row_level_security_sql_unique"},
         ]));
         let start = FIRST_NON_SYSTEM_ID as i128;
         #[rustfmt::skip]
@@ -1412,6 +1421,7 @@ mod tests {
             ConstraintRow { constraint_id: 8, table_id: ST_VAR_ID.into(), unique_columns: col(0), constraint_name: "ct_st_var_name_unique" },
             ConstraintRow { constraint_id: 9, table_id: ST_SCHEDULED_ID.into(), unique_columns: col(0), constraint_name: "ct_st_scheduled_schedule_id_unique" },
             ConstraintRow { constraint_id: 10, table_id: ST_SCHEDULED_ID.into(), unique_columns: col(1), constraint_name: "ct_st_scheduled_table_id_unique" },
+            ConstraintRow { constraint_id: 11, table_id: ST_ROW_LEVEL_SECURITY_ID.into(), unique_columns: col(1), constraint_name: "ct_st_row_level_security_sql_unique" },
         ]));
 
         // Verify we get back the tables correctly with the proper ids...
@@ -1823,6 +1833,8 @@ mod tests {
             IndexRow { id: 8, table: ST_VAR_ID.into(), col: col(0), name: "idx_st_var_name_unique", },
             IndexRow { id: 9, table: ST_SCHEDULED_ID.into(), col: col(0), name: "idx_st_scheduled_schedule_id_unique", },
             IndexRow { id: 10, table: ST_SCHEDULED_ID.into(), col: col(1), name: "idx_st_scheduled_table_id_unique", },
+            IndexRow { id: 11, table: ST_ROW_LEVEL_SECURITY_ID.into(), col: col(0), name: "idx_st_row_level_security_btree_table_id"},
+            IndexRow { id: 12, table: ST_ROW_LEVEL_SECURITY_ID.into(), col: col(1), name: "idx_st_row_level_security_sql_unique"},
             IndexRow { id: seq_start,     table: FIRST_NON_SYSTEM_ID, col: col(0), name: "id_idx",  },
             IndexRow { id: seq_start + 1, table: FIRST_NON_SYSTEM_ID, col: col(1), name: "name_idx",  },
             IndexRow { id: seq_start + 2, table: FIRST_NON_SYSTEM_ID, col: col(2), name: "age_idx",  },
@@ -1876,6 +1888,8 @@ mod tests {
             IndexRow { id: 8, table: ST_VAR_ID.into(), col: col(0), name: "idx_st_var_name_unique", },
             IndexRow { id: 9, table: ST_SCHEDULED_ID.into(), col: col(0), name: "idx_st_scheduled_schedule_id_unique", },
             IndexRow { id: 10, table: ST_SCHEDULED_ID.into(), col: col(1), name: "idx_st_scheduled_table_id_unique", },
+            IndexRow { id: 11, table: ST_ROW_LEVEL_SECURITY_ID.into(), col: col(0), name: "idx_st_row_level_security_btree_table_id"},
+            IndexRow { id: 12, table: ST_ROW_LEVEL_SECURITY_ID.into(), col: col(1), name: "idx_st_row_level_security_sql_unique"},
             IndexRow { id: seq_start    , table: FIRST_NON_SYSTEM_ID, col: col(0), name: "id_idx" },
             IndexRow { id: seq_start + 1, table: FIRST_NON_SYSTEM_ID, col: col(1), name: "name_idx" },
             IndexRow { id: seq_start + 2, table: FIRST_NON_SYSTEM_ID, col: col(2), name: "age_idx" },
@@ -1930,6 +1944,8 @@ mod tests {
             IndexRow { id: 8, table: ST_VAR_ID.into(), col: col(0), name: "idx_st_var_name_unique", },
             IndexRow { id: 9, table: ST_SCHEDULED_ID.into(), col: col(0), name: "idx_st_scheduled_schedule_id_unique", },
             IndexRow { id: 10, table: ST_SCHEDULED_ID.into(), col: col(1), name: "idx_st_scheduled_table_id_unique", },
+            IndexRow { id: 11, table: ST_ROW_LEVEL_SECURITY_ID.into(), col: col(0), name: "idx_st_row_level_security_btree_table_id"},
+            IndexRow { id: 12, table: ST_ROW_LEVEL_SECURITY_ID.into(), col: col(1), name: "idx_st_row_level_security_sql_unique"},
             IndexRow { id: seq_start,     table: FIRST_NON_SYSTEM_ID, col: col(0), name: "id_idx" },
             IndexRow { id: seq_start + 1, table: FIRST_NON_SYSTEM_ID, col: col(1), name: "name_idx" },
         ].map(Into::into));
@@ -2019,6 +2035,32 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_row_level_security() -> ResultTest<()> {
+        let (_, mut tx, table_id) = setup_table()?;
+
+        let rls = RowLevelSecuritySchema {
+            sql: "SELECT * FROM bar".into(),
+            table_id,
+        };
+        let ctx = ExecutionContext::default();
+        tx.create_row_level_security(&ctx, rls.clone())?;
+
+        let result = tx.row_level_security_for_table_id(&ctx, table_id)?;
+        assert_eq!(
+            result,
+            vec![RowLevelSecuritySchema {
+                sql: "SELECT * FROM bar".into(),
+                table_id,
+            }]
+        );
+
+        tx.drop_row_level_security(&ctx, rls.sql)?;
+        assert_eq!(tx.row_level_security_for_table_id(&ctx, table_id)?, []);
+
+        Ok(())
+    }
+
     // TODO: Add the following tests
     // - Create index with unique constraint and immediately insert a row that violates the constraint before committing.
     // - Create a tx that inserts 2000 rows with an auto_inc column

crates/core/src/db/datastore/locking_tx_datastore/mut_tx.rs

@@ -7,6 +7,7 @@ use super::{
     tx_state::{DeleteTable, IndexIdMap, TxState},
     SharedMutexGuard, SharedWriteGuard,
 };
+use crate::db::datastore::system_tables::{StRowLevelSecurityFields, StRowLevelSecurityRow, ST_ROW_LEVEL_SECURITY_ID};
 use crate::db::datastore::{
     system_tables::{
         StColumnFields, StColumnRow, StConstraintFields, StConstraintRow, StFields as _, StIndexFields, StIndexRow,
@@ -22,6 +23,7 @@ use crate::{
 use core::ops::RangeBounds;
 use core::{iter, ops::Bound};
 use smallvec::SmallVec;
+use spacetimedb_lib::db::raw_def::v9::RawSql;
 use spacetimedb_lib::{
     address::Address,
     bsatn::Deserializer,
@@ -36,7 +38,7 @@ use spacetimedb_sats::{
 };
 use spacetimedb_schema::{
     def::{BTreeAlgorithm, IndexAlgorithm},
-    schema::{ConstraintSchema, IndexSchema, SequenceSchema, TableSchema},
+    schema::{ConstraintSchema, IndexSchema, RowLevelSecuritySchema, SequenceSchema, TableSchema},
 };
 use spacetimedb_table::{
     blob_store::{BlobStore, HashMapBlobStore},
@@ -946,6 +948,92 @@ impl MutTxId {
         })
     }
 
+    /// Create a row level security policy.
+    ///
+    /// Requires:
+    /// - `row_level_security_schema.table_id != TableId::SENTINEL`
+    /// - `row_level_security_schema.sql` must be unique.
+    ///
+    /// Ensures:
+    ///
+    /// - The row level security policy metadata is inserted into the system tables (and other data structures reflecting them).
+    /// - The returned `sql` is unique.
+    pub fn create_row_level_security(
+        &mut self,
+        ctx: &ExecutionContext,
+        row_level_security_schema: RowLevelSecuritySchema,
+    ) -> Result<RawSql> {
+        if row_level_security_schema.table_id == TableId::SENTINEL {
+            return Err(anyhow::anyhow!(
+                "`table_id` must not be `TableId::SENTINEL` in `{:#?}`",
+                row_level_security_schema
+            )
+            .into());
+        }
+
+        log::trace!(
+            "ROW LEVEL SECURITY CREATING for table: {}",
+            row_level_security_schema.table_id
+        );
+
+        // Insert the row into st_row_level_security
+        // NOTE: Because st_row_level_security has a unique index on sql, this will
+        // fail if already exists.
+        let row = StRowLevelSecurityRow {
+            table_id: row_level_security_schema.table_id,
+            sql: row_level_security_schema.sql,
+        };
+
+        let row = self.insert(ST_ROW_LEVEL_SECURITY_ID, &mut ProductValue::from(row), ctx.database())?;
+        let row_level_security_sql = row.1.collapse().read_col(StRowLevelSecurityFields::Sql)?;
+        let existed = matches!(row.1, RowRefInsertion::Existed(_));
+
+        // Add the row level security to the transaction's insert table.
+        self.get_or_create_insert_table_mut(row_level_security_schema.table_id)?;
+
+        if existed {
+            log::trace!("ROW LEVEL SECURITY ALREADY EXISTS: {row_level_security_sql}");
+        } else {
+            log::trace!("ROW LEVEL SECURITY CREATED: {row_level_security_sql}");
+        }
+
+        Ok(row_level_security_sql)
+    }
+
+    pub fn row_level_security_for_table_id(
+        &self,
+        ctx: &ExecutionContext,
+        table_id: TableId,
+    ) -> Result<Vec<RowLevelSecuritySchema>> {
+        Ok(self
+            .iter_by_col_eq(
+                ctx,
+                ST_ROW_LEVEL_SECURITY_ID,
+                StRowLevelSecurityFields::TableId,
+                &table_id.into(),
+            )?
+            .map(|row| {
+                let row = StRowLevelSecurityRow::try_from(row).unwrap();
+                row.into()
+            })
+            .collect())
+    }
+
+    pub fn drop_row_level_security(&mut self, ctx: &ExecutionContext, sql: RawSql) -> Result<()> {
+        let st_rls_ref = self
+            .iter_by_col_eq(
+                ctx,
+                ST_ROW_LEVEL_SECURITY_ID,
+                StRowLevelSecurityFields::Sql,
+                &sql.clone().into(),
+            )?
+            .next()
+            .ok_or_else(|| TableError::RawSqlNotFound(SystemTable::st_row_level_security, sql))?;
+        self.delete(ST_ROW_LEVEL_SECURITY_ID, st_rls_ref.pointer())?;
+
+        Ok(())
+    }
+
     // TODO(perf, deep-integration):
     //   When all of [`Table::read_row`], [`RowRef::new`], [`CommittedState::get`]
     //   and [`TxState::get`] become unsafe,