From 1a57ae3141876c329d6bd0c5c2a6a2f469d259ca Mon Sep 17 00:00:00 2001 From: Guflly <145608489+Guflly@users.noreply.github.com> Date: Mon, 27 Jul 2026 11:39:09 -0700 Subject: [PATCH 1/3] fix(db): reject application models in migrations --- cot-macros/src/model.rs | 5 ++++ cot/src/db.rs | 34 +++++++++++++++++++++ cot/src/db/migrations.rs | 64 ++++++++++++++++++++++++++++++++++++++-- 3 files changed, 100 insertions(+), 3 deletions(-) diff --git a/cot-macros/src/model.rs b/cot-macros/src/model.rs index 91006c40c..8fe86a77a 100644 --- a/cot-macros/src/model.rs +++ b/cot-macros/src/model.rs @@ -75,6 +75,7 @@ struct ModelBuilder { name: Ident, vis: syn::Visibility, table_name: String, + is_application_model: bool, pk_field: Field, fields_struct_name: Ident, fields_as_columns: Vec, @@ -109,6 +110,7 @@ impl ModelBuilder { name: model.name.clone(), vis: model.vis, table_name, + is_application_model: model.model_type == ModelType::Application, pk_field: model.pk_field.clone(), fields_struct_name: format_ident!("{}Fields", model.name), fields_as_columns: Vec::with_capacity(field_count), @@ -166,6 +168,7 @@ impl ModelBuilder { let name = &self.name; let app_name = &self.app_name; let table_name = &self.table_name; + let is_application_model = self.is_application_model; let fields_struct_name = &self.fields_struct_name; let fields_as_columns = &self.fields_as_columns; let pk_field_name = &self.pk_field.name; @@ -185,6 +188,8 @@ impl ModelBuilder { const COLUMNS: &'static [#orm_ident::Column] = &[ #(#fields_as_columns,)* ]; + #[doc(hidden)] + const IS_APPLICATION_MODEL: bool = #is_application_model; const APP_NAME: &'static str = #app_name; const TABLE_NAME: #orm_ident::Identifier = #orm_ident::Identifier::new(#table_name); const PRIMARY_KEY_NAME: #orm_ident::Identifier = #orm_ident::Identifier::new(#pk_column_name); diff --git a/cot/src/db.rs b/cot/src/db.rs index 6e0ec4156..70cfef932 100644 --- a/cot/src/db.rs +++ b/cot/src/db.rs @@ -585,6 +585,9 @@ pub trait Model: Sized + Send + 'static { /// The columns of the model. const COLUMNS: &'static [Column]; + #[doc(hidden)] + const IS_APPLICATION_MODEL: bool = true; + /// Creates a model instance from a database row. /// /// # Errors @@ -1192,6 +1195,7 @@ pub trait SqlxValueRef<'r>: Sized { #[derive(Debug, Clone)] pub struct Database { inner: Arc, + migration_context: bool, } #[derive(Debug)] @@ -1254,6 +1258,7 @@ impl Database { let inner = DatabaseSqlite::new(&url).await?; return Ok(Self { inner: Arc::new(DatabaseImpl::Sqlite(inner)), + migration_context: false, }); } @@ -1262,6 +1267,7 @@ impl Database { let inner = DatabasePostgres::new(&url).await?; return Ok(Self { inner: Arc::new(DatabaseImpl::Postgres(inner)), + migration_context: false, }); } @@ -1270,12 +1276,32 @@ impl Database { let inner = DatabaseMySql::new(&url).await?; return Ok(Self { inner: Arc::new(DatabaseImpl::MySql(inner)), + migration_context: false, }); } panic!("Unsupported database URL: {url}"); } + fn for_migration(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + migration_context: true, + } + } + + fn ensure_model_allowed(&self) -> Result<()> { + if self.migration_context && T::IS_APPLICATION_MODEL { + return Err(DatabaseError::MigrationError( + migrations::MigrationEngineError::Custom(format!( + "application model `{}` cannot be used in migrations; use a migration model", + std::any::type_name::() + )), + )); + } + Ok(()) + } + /// Closes the database connection. /// /// This method should be called when the database connection is no longer @@ -1345,6 +1371,7 @@ impl Database { } async fn insert_or_update_impl(&self, data: &mut T, update: bool) -> Result<()> { + self.ensure_model_allowed::()?; let column_identifiers = T::COLUMNS .iter() .map(|column| Identifier::from(column.name.as_str())); @@ -1451,6 +1478,7 @@ impl Database { } async fn update_impl(&self, data: &mut T) -> Result<()> { + self.ensure_model_allowed::()?; let column_identifiers = T::COLUMNS .iter() .map(|column| Identifier::from(column.name.as_str())); @@ -1531,6 +1559,7 @@ impl Database { } async fn bulk_insert_impl(&self, data: &mut [T], update: bool) -> Result<()> { + self.ensure_model_allowed::()?; // TODO: add transactions when implemented if data.is_empty() { @@ -1741,6 +1770,7 @@ impl Database { /// /// Can return an error if the database connection is lost. pub async fn query(&self, query: &Query) -> Result> { + self.ensure_model_allowed::()?; let columns_to_get: Vec<_> = T::COLUMNS.iter().map(|column| column.name).collect(); let mut select = sea_query::Query::select(); select.columns(columns_to_get).from(T::TABLE_NAME); @@ -1767,6 +1797,7 @@ impl Database { /// /// Can return an error if the database connection is lost. pub async fn get(&self, query: &Query) -> Result> { + self.ensure_model_allowed::()?; let columns_to_get: Vec<_> = T::COLUMNS.iter().map(|column| column.name).collect(); let mut select = sea_query::Query::select(); select.columns(columns_to_get).from(T::TABLE_NAME); @@ -1794,6 +1825,7 @@ impl Database { /// /// Can return an error if the database connection is lost. pub async fn exists(&self, query: &Query) -> Result { + self.ensure_model_allowed::()?; let mut select = sea_query::Query::select(); select.expr(sea_query::Expr::value(1)).from(T::TABLE_NAME); query.add_filter_to_statement(&mut select); @@ -1816,6 +1848,7 @@ impl Database { /// /// Can return an error if the database connection is lost. pub async fn delete(&self, query: &Query) -> Result { + self.ensure_model_allowed::()?; let mut delete = sea_query::Query::delete(); delete.from_table(T::TABLE_NAME); query.add_filter_to_statement(&mut delete); @@ -2017,6 +2050,7 @@ impl Database { query: &str, values: &[&dyn ToDbValue], ) -> Result> { + self.ensure_model_allowed::()?; let values = values .iter() .map(ToDbValue::to_db_value) diff --git a/cot/src/db/migrations.rs b/cot/src/db/migrations.rs index 1ac4590d1..7ce64799e 100644 --- a/cot/src/db/migrations.rs +++ b/cot/src/db/migrations.rs @@ -532,7 +532,8 @@ impl Operation { forwards, backwards: _, } => { - let context = MigrationContext::new(database); + let database = database.for_migration(); + let context = MigrationContext::new(&database); forwards(context).await?; } } @@ -617,7 +618,8 @@ impl Operation { backwards, } => { if let Some(backwards) = backwards { - let context = MigrationContext::new(database); + let database = database.for_migration(); + let context = MigrationContext::new(&database); backwards(context).await?; } else { return Err(crate::db::DatabaseError::MigrationError( @@ -2025,7 +2027,19 @@ mod tests { use cot::test::TestDatabase; use super::*; - use crate::db::{ColumnType, DatabaseField, Identifier}; + use crate::db::{ColumnType, DatabaseError, DatabaseField, Identifier, Model}; + + #[model] + struct ApplicationModel { + #[model(primary_key)] + id: i32, + } + + #[model(model_type = "migration")] + struct _MigrationModel { + #[model(primary_key)] + id: i32, + } struct TestMigration; @@ -2168,6 +2182,50 @@ mod tests { assert!(result.is_ok()); } + #[cot::test] + #[cfg_attr( + miri, + ignore = "unsupported operation: can't call foreign function `sqlite3_open_v2`" + )] + async fn operation_custom_model_access() { + let test_db = TestDatabase::new_sqlite().await.unwrap(); + + #[migration_op] + async fn application_model(ctx: MigrationContext<'_>) -> Result<()> { + ApplicationModel::objects().all(ctx.db).await?; + Ok(()) + } + + let result = Operation::custom(application_model) + .build() + .forwards(&test_db.database()) + .await; + assert!(matches!( + result, + Err(DatabaseError::MigrationError(MigrationEngineError::Custom( + _ + ))) + )); + + test_db + .database() + .raw("CREATE TABLE cot__migration_model (id INTEGER PRIMARY KEY)") + .await + .unwrap(); + + #[migration_op] + async fn migration_model(ctx: MigrationContext<'_>) -> Result<()> { + _MigrationModel::objects().all(ctx.db).await?; + Ok(()) + } + + Operation::custom(migration_model) + .build() + .forwards(&test_db.database()) + .await + .unwrap(); + } + #[cot::test] #[cfg_attr( miri, From 313a12a466d8cc5cc74fcba1e6e9c82ff4c7acb6 Mon Sep 17 00:00:00 2001 From: Guflly <145608489+Guflly@users.noreply.github.com> Date: Mon, 27 Jul 2026 13:59:47 -0700 Subject: [PATCH 2/3] refactor(db): use enums for model and database context --- cot-macros/src/model.rs | 13 ++++++++----- cot/src/db.rs | 41 ++++++++++++++++++++++++++++------------- 2 files changed, 36 insertions(+), 18 deletions(-) diff --git a/cot-macros/src/model.rs b/cot-macros/src/model.rs index 8fe86a77a..ef60f116f 100644 --- a/cot-macros/src/model.rs +++ b/cot-macros/src/model.rs @@ -75,7 +75,7 @@ struct ModelBuilder { name: Ident, vis: syn::Visibility, table_name: String, - is_application_model: bool, + model_type: ModelType, pk_field: Field, fields_struct_name: Ident, fields_as_columns: Vec, @@ -110,7 +110,7 @@ impl ModelBuilder { name: model.name.clone(), vis: model.vis, table_name, - is_application_model: model.model_type == ModelType::Application, + model_type: model.model_type, pk_field: model.pk_field.clone(), fields_struct_name: format_ident!("{}Fields", model.name), fields_as_columns: Vec::with_capacity(field_count), @@ -168,7 +168,11 @@ impl ModelBuilder { let name = &self.name; let app_name = &self.app_name; let table_name = &self.table_name; - let is_application_model = self.is_application_model; + let model_type = match self.model_type { + ModelType::Application => quote!(#orm_ident::ModelType::Application), + ModelType::Migration => quote!(#orm_ident::ModelType::Migration), + ModelType::Internal => quote!(#orm_ident::ModelType::Internal), + }; let fields_struct_name = &self.fields_struct_name; let fields_as_columns = &self.fields_as_columns; let pk_field_name = &self.pk_field.name; @@ -188,8 +192,7 @@ impl ModelBuilder { const COLUMNS: &'static [#orm_ident::Column] = &[ #(#fields_as_columns,)* ]; - #[doc(hidden)] - const IS_APPLICATION_MODEL: bool = #is_application_model; + const MODEL_TYPE: #orm_ident::ModelType = #model_type; const APP_NAME: &'static str = #app_name; const TABLE_NAME: #orm_ident::Identifier = #orm_ident::Identifier::new(#table_name); const PRIMARY_KEY_NAME: #orm_ident::Identifier = #orm_ident::Identifier::new(#pk_column_name); diff --git a/cot/src/db.rs b/cot/src/db.rs index 70cfef932..4974c5ebd 100644 --- a/cot/src/db.rs +++ b/cot/src/db.rs @@ -528,6 +528,15 @@ impl DatabaseError { /// An alias for [`Result`] that uses [`DatabaseError`] as the error type. pub type Result = std::result::Result; +#[expect(missing_docs)] +#[non_exhaustive] +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +pub enum ModelType { + Application, + Migration, + Internal, +} + /// A model trait for database models. /// /// This trait is used to define a model that can be stored in a database. @@ -585,8 +594,8 @@ pub trait Model: Sized + Send + 'static { /// The columns of the model. const COLUMNS: &'static [Column]; - #[doc(hidden)] - const IS_APPLICATION_MODEL: bool = true; + #[expect(missing_docs)] + const MODEL_TYPE: ModelType = ModelType::Application; /// Creates a model instance from a database row. /// @@ -1187,6 +1196,12 @@ pub trait SqlxValueRef<'r>: Sized { } } +#[derive(Debug, Clone, Copy)] +enum DatabaseContext { + Default, + InMigration, +} + /// A database connection structure that holds the connection to the database. /// /// It is used to execute queries and interact with the database. The connection @@ -1195,7 +1210,7 @@ pub trait SqlxValueRef<'r>: Sized { #[derive(Debug, Clone)] pub struct Database { inner: Arc, - migration_context: bool, + context: DatabaseContext, } #[derive(Debug)] @@ -1258,7 +1273,7 @@ impl Database { let inner = DatabaseSqlite::new(&url).await?; return Ok(Self { inner: Arc::new(DatabaseImpl::Sqlite(inner)), - migration_context: false, + context: DatabaseContext::Default, }); } @@ -1267,7 +1282,7 @@ impl Database { let inner = DatabasePostgres::new(&url).await?; return Ok(Self { inner: Arc::new(DatabaseImpl::Postgres(inner)), - migration_context: false, + context: DatabaseContext::Default, }); } @@ -1276,7 +1291,7 @@ impl Database { let inner = DatabaseMySql::new(&url).await?; return Ok(Self { inner: Arc::new(DatabaseImpl::MySql(inner)), - migration_context: false, + context: DatabaseContext::Default, }); } @@ -1286,20 +1301,20 @@ impl Database { fn for_migration(&self) -> Self { Self { inner: Arc::clone(&self.inner), - migration_context: true, + context: DatabaseContext::InMigration, } } fn ensure_model_allowed(&self) -> Result<()> { - if self.migration_context && T::IS_APPLICATION_MODEL { - return Err(DatabaseError::MigrationError( - migrations::MigrationEngineError::Custom(format!( + match (self.context, T::MODEL_TYPE) { + (DatabaseContext::InMigration, ModelType::Application) => Err( + DatabaseError::MigrationError(migrations::MigrationEngineError::Custom(format!( "application model `{}` cannot be used in migrations; use a migration model", std::any::type_name::() - )), - )); + ))), + ), + _ => Ok(()), } - Ok(()) } /// Closes the database connection. From ce735a6a173c2f8cb9c70542d3b570a9f7bfc053 Mon Sep 17 00:00:00 2001 From: Guflly <145608489+Guflly@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:09:48 -0700 Subject: [PATCH 3/3] fix(db): reject migration models outside migrations --- cot/src/db.rs | 6 ++++++ cot/src/db/migrations.rs | 8 ++++++++ 2 files changed, 14 insertions(+) diff --git a/cot/src/db.rs b/cot/src/db.rs index 4974c5ebd..fa3968452 100644 --- a/cot/src/db.rs +++ b/cot/src/db.rs @@ -1313,6 +1313,12 @@ impl Database { std::any::type_name::() ))), ), + (DatabaseContext::Default, ModelType::Migration) => Err(DatabaseError::MigrationError( + migrations::MigrationEngineError::Custom(format!( + "migration model `{}` cannot be used outside migrations; use an application model", + std::any::type_name::() + )), + )), _ => Ok(()), } } diff --git a/cot/src/db/migrations.rs b/cot/src/db/migrations.rs index 7ce64799e..a99171874 100644 --- a/cot/src/db/migrations.rs +++ b/cot/src/db/migrations.rs @@ -2190,6 +2190,14 @@ mod tests { async fn operation_custom_model_access() { let test_db = TestDatabase::new_sqlite().await.unwrap(); + let result = _MigrationModel::objects().all(&test_db.database()).await; + assert!(matches!( + result, + Err(DatabaseError::MigrationError(MigrationEngineError::Custom( + _ + ))) + )); + #[migration_op] async fn application_model(ctx: MigrationContext<'_>) -> Result<()> { ApplicationModel::objects().all(ctx.db).await?;