Have the scheduled_id/at columns be specified in the schema, not by t…

…he column name (#1861)

crates/bindings-macro/src/lib.rs

@@ -1030,7 +1030,19 @@ fn table_impl(mut args: TableArgs, mut item: MutItem<syn::DeriveInput>) -> syn::
     let unique_col_ids = unique_columns.iter().map(|col| col.index);
     let primary_col_id = primary_key_column.iter().map(|col| col.index);
     let sequence_col_ids = sequenced_columns.iter().map(|col| col.index);
-    let scheduled_reducer_ident = args.scheduled.iter();
+
+    let schedule = args
+        .scheduled
+        .as_ref()
+        .map(|reducer| {
+            // scheduled_at was inserted as the last field
+            let scheduled_at_id = (fields.len() - 1) as u16;
+            quote!(spacetimedb::table::ScheduleDesc {
+                reducer_name: <#reducer as spacetimedb::rt::ReducerInfo>::NAME,
+                scheduled_at_column: #scheduled_at_id,
+            })
+        })
+        .into_iter();
 
     let unique_err = if !unique_columns.is_empty() {
         quote!(spacetimedb::UniqueConstraintViolation)
@@ -1063,7 +1075,7 @@ fn table_impl(mut args: TableArgs, mut item: MutItem<syn::DeriveInput>) -> syn::
             const INDEXES: &'static [spacetimedb::table::IndexDesc<'static>] = &[#(#index_descs),*];
             #(const PRIMARY_KEY: Option<u16> = Some(#primary_col_id);)*
             const SEQUENCES: &'static [u16] = &[#(#sequence_col_ids),*];
-            #(const SCHEDULED_REDUCER_NAME: Option<&'static str> = Some(<#scheduled_reducer_ident as spacetimedb::rt::ReducerInfo>::NAME);)*
+            #(const SCHEDULE: Option<spacetimedb::table::ScheduleDesc<'static>> = Some(#schedule);)*
 
             #table_id_from_name_func
         }

crates/bindings/src/lib.rs

@@ -35,6 +35,7 @@ pub use spacetimedb_lib::ser::Serialize;
 pub use spacetimedb_lib::Address;
 pub use spacetimedb_lib::AlgebraicValue;
 pub use spacetimedb_lib::Identity;
+pub use spacetimedb_lib::ScheduleAt;
 pub use spacetimedb_primitives::TableId;
 pub use sys::Errno;
 pub use table::{AutoIncOverflow, BTreeIndex, Table, TryInsertError, UniqueColumn, UniqueConstraintViolation};

crates/bindings/src/rt.rs

@@ -165,6 +165,20 @@ pub trait TableColumn {
 }
 impl<T: SpacetimeType> TableColumn for T {}
 
+/// Assert that the primary_key column of a scheduled table is a u64.
+pub const fn assert_scheduled_table_primary_key<T: ScheduledTablePrimaryKey>() {}
+
+mod sealed {
+    pub trait Sealed {}
+}
+#[diagnostic::on_unimplemented(
+    message = "scheduled table primary key must be a `u64`",
+    label = "should be `u64`, not `{Self}`"
+)]
+pub trait ScheduledTablePrimaryKey: sealed::Sealed {}
+impl sealed::Sealed for u64 {}
+impl ScheduledTablePrimaryKey for u64 {}
+
 /// Used in the last type parameter of `Reducer` to indicate that the
 /// context argument *should* be passed to the reducer logic.
 pub struct ContextArg;
@@ -331,8 +345,8 @@ pub fn register_table<T: Table>() {
         for &col in T::SEQUENCES {
             table = table.with_column_sequence(col, None);
         }
-        if let Some(scheduled_reducer) = T::SCHEDULED_REDUCER_NAME {
-            table = table.with_schedule(scheduled_reducer, None);
+        if let Some(schedule) = T::SCHEDULE {
+            table = table.with_schedule(schedule.reducer_name, schedule.scheduled_at_column, None);
         }
 
         table.finish();

crates/bindings/src/table.rs

@@ -124,7 +124,7 @@ pub trait TableInternal: Sized {
     const INDEXES: &'static [IndexDesc<'static>];
     const PRIMARY_KEY: Option<u16> = None;
     const SEQUENCES: &'static [u16];
-    const SCHEDULED_REDUCER_NAME: Option<&'static str> = None;
+    const SCHEDULE: Option<ScheduleDesc<'static>> = None;
 
     /// Returns the ID of this table.
     fn table_id() -> TableId;
@@ -143,6 +143,11 @@ pub enum IndexAlgo<'a> {
     BTree { columns: &'a [u16] },
 }
 
+pub struct ScheduleDesc<'a> {
+    pub reducer_name: &'a str,
+    pub scheduled_at_column: u16,
+}
+
 #[doc(hidden)]
 pub trait __MapRowTypeToTable {
     type Table: Table;

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

@@ -1391,6 +1391,7 @@ 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_SCHEDULED_ID.into(), pos: 4, name: "at_column", ty: AlgebraicType::U16 },
 
             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 },

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

@@ -156,6 +156,7 @@ impl MutTxId {
                 schedule_id: ScheduleId::SENTINEL,
                 schedule_name: schedule.schedule_name,
                 reducer_name: schedule.reducer_name,
+                at_column: schedule.at_column,
             };
             let (generated, ..) = self.insert(ST_SCHEDULED_ID, &mut row.into(), database_address)?;
             let id = generated.as_u32();

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

@@ -266,6 +266,7 @@ st_fields_enum!(enum StScheduledFields {
     "table_id", TableId = 1,
     "reducer_name", ReducerName = 2,
     "schedule_name", ScheduleName = 3,
+    "at_column", AtColumn = 4,
 });
 
 /// Helper method to check that a system table has the correct fields.
@@ -1285,6 +1286,7 @@ pub struct StScheduledRow {
     pub(crate) table_id: TableId,
     pub(crate) reducer_name: Box<str>,
     pub(crate) schedule_name: Box<str>,
+    pub(crate) at_column: ColId,
 }
 
 impl TryFrom<RowRef<'_>> for StScheduledRow {
@@ -1307,6 +1309,7 @@ impl From<StScheduledRow> for ScheduleSchema {
             reducer_name: row.reducer_name,
             schedule_id: row.schedule_id,
             schedule_name: row.schedule_name,
+            at_column: row.at_column,
         }
     }
 }

crates/core/src/db/relational_db.rs

@@ -18,7 +18,7 @@ use crate::error::{DBError, DatabaseError, TableError};
 use crate::execution_context::ExecutionContext;
 use crate::messages::control_db::HostType;
 use crate::util::spawn_rayon;
-use anyhow::anyhow;
+use anyhow::{anyhow, Context};
 use fs2::FileExt;
 use futures::channel::mpsc;
 use futures::StreamExt;
@@ -527,14 +527,18 @@ impl RelationalDB {
             .get_all_tables_tx(&ExecutionContext::internal(self.address), tx)
     }
 
-    pub fn is_scheduled_table(
+    pub fn table_scheduled_id_and_at(
         &self,
         ctx: &ExecutionContext,
-        tx: &mut MutTx,
+        tx: &impl StateView,
         table_id: TableId,
-    ) -> Result<bool, DBError> {
-        tx.schema_for_table(ctx, table_id)
-            .map(|schema| schema.schedule.is_some())
+    ) -> Result<Option<(ColId, ColId)>, DBError> {
+        let schema = tx.schema_for_table(ctx, table_id)?;
+        let Some(sched) = &schema.schedule else { return Ok(None) };
+        let primary_key = schema
+            .primary_key
+            .context("scheduled table doesn't have a primary key?")?;
+        Ok(Some((primary_key, sched.at_column)))
     }
 
     pub fn decode_column(

crates/core/src/host/instance_env.rs

@@ -128,14 +128,14 @@ impl InstanceEnv {
                 }
             })?;
 
-        if stdb.is_scheduled_table(ctx, tx, table_id)? {
+        if let Some((id_column, at_column)) = stdb.table_scheduled_id_and_at(ctx, tx, table_id)? {
             let row_ref = tx.get(table_id, row_ptr)?.unwrap();
-            let (schedule_id, schedule_at) = get_schedule_from_row(tx, stdb, table_id, &row_ref)
+            let (schedule_id, schedule_at) = get_schedule_from_row(&row_ref, id_column, at_column)
                 // NOTE(centril): Should never happen,
                 // as we successfully inserted and thus `ret` is verified against the table schema.
                 .map_err(|e| NodesError::ScheduleError(ScheduleError::DecodingError(e)))?;
             self.scheduler
-                .schedule(table_id, schedule_id, schedule_at)
+                .schedule(table_id, schedule_id, schedule_at, id_column, at_column)
                 .map_err(NodesError::ScheduleError)?;
         }