diff --git a/cot-macros/src/model.rs b/cot-macros/src/model.rs index 91006c40..ef60f116 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, + model_type: ModelType, 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, + 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), @@ -166,6 +168,11 @@ impl ModelBuilder { let name = &self.name; let app_name = &self.app_name; let table_name = &self.table_name; + 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; @@ -185,6 +192,7 @@ impl ModelBuilder { const COLUMNS: &'static [#orm_ident::Column] = &[ #(#fields_as_columns,)* ]; + 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 6e0ec415..fa396845 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,6 +594,9 @@ pub trait Model: Sized + Send + 'static { /// The columns of the model. const COLUMNS: &'static [Column]; + #[expect(missing_docs)] + const MODEL_TYPE: ModelType = ModelType::Application; + /// Creates a model instance from a database row. /// /// # Errors @@ -1184,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 @@ -1192,6 +1210,7 @@ pub trait SqlxValueRef<'r>: Sized { #[derive(Debug, Clone)] pub struct Database { inner: Arc, + context: DatabaseContext, } #[derive(Debug)] @@ -1254,6 +1273,7 @@ impl Database { let inner = DatabaseSqlite::new(&url).await?; return Ok(Self { inner: Arc::new(DatabaseImpl::Sqlite(inner)), + context: DatabaseContext::Default, }); } @@ -1262,6 +1282,7 @@ impl Database { let inner = DatabasePostgres::new(&url).await?; return Ok(Self { inner: Arc::new(DatabaseImpl::Postgres(inner)), + context: DatabaseContext::Default, }); } @@ -1270,12 +1291,38 @@ impl Database { let inner = DatabaseMySql::new(&url).await?; return Ok(Self { inner: Arc::new(DatabaseImpl::MySql(inner)), + context: DatabaseContext::Default, }); } panic!("Unsupported database URL: {url}"); } + fn for_migration(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + context: DatabaseContext::InMigration, + } + } + + fn ensure_model_allowed(&self) -> Result<()> { + 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::() + ))), + ), + (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(()), + } + } + /// Closes the database connection. /// /// This method should be called when the database connection is no longer @@ -1345,6 +1392,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 +1499,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 +1580,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 +1791,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 +1818,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 +1846,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 +1869,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 +2071,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 1ac4590d..a9917187 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,58 @@ 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(); + + 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?; + 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,