fix(db): reject application models in migrations

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<TokenStream>,
@@ -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);

cot/src/db.rs

@@ -528,6 +528,15 @@ impl DatabaseError {
 /// An alias for [`Result`] that uses [`DatabaseError`] as the error type.
 pub type Result<T> = std::result::Result<T, DatabaseError>;
 
+#[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<DatabaseImpl>,
+    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,32 @@ 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<T: Model>(&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::<T>()
+                ))),
+            ),
+            _ => Ok(()),
+        }
+    }
+
     /// Closes the database connection.
     ///
     /// This method should be called when the database connection is no longer
@@ -1345,6 +1386,7 @@ impl Database {
     }
 
     async fn insert_or_update_impl<T: Model>(&self, data: &mut T, update: bool) -> Result<()> {
+        self.ensure_model_allowed::<T>()?;
         let column_identifiers = T::COLUMNS
             .iter()
             .map(|column| Identifier::from(column.name.as_str()));
@@ -1451,6 +1493,7 @@ impl Database {
     }
 
     async fn update_impl<T: Model>(&self, data: &mut T) -> Result<()> {
+        self.ensure_model_allowed::<T>()?;
         let column_identifiers = T::COLUMNS
             .iter()
             .map(|column| Identifier::from(column.name.as_str()));
@@ -1531,6 +1574,7 @@ impl Database {
     }
 
     async fn bulk_insert_impl<T: Model>(&self, data: &mut [T], update: bool) -> Result<()> {
+        self.ensure_model_allowed::<T>()?;
         // TODO: add transactions when implemented
 
         if data.is_empty() {
@@ -1741,6 +1785,7 @@ impl Database {
     ///
     /// Can return an error if the database connection is lost.
     pub async fn query<T: Model>(&self, query: &Query<T>) -> Result<Vec<T>> {
+        self.ensure_model_allowed::<T>()?;
         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 +1812,7 @@ impl Database {
     ///
     /// Can return an error if the database connection is lost.
     pub async fn get<T: Model>(&self, query: &Query<T>) -> Result<Option<T>> {
+        self.ensure_model_allowed::<T>()?;
         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 +1840,7 @@ impl Database {
     ///
     /// Can return an error if the database connection is lost.
     pub async fn exists<T: Model>(&self, query: &Query<T>) -> Result<bool> {
+        self.ensure_model_allowed::<T>()?;
         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 +1863,7 @@ impl Database {
     ///
     /// Can return an error if the database connection is lost.
     pub async fn delete<T: Model>(&self, query: &Query<T>) -> Result<StatementResult> {
+        self.ensure_model_allowed::<T>()?;
         let mut delete = sea_query::Query::delete();
         delete.from_table(T::TABLE_NAME);
         query.add_filter_to_statement(&mut delete);
@@ -2017,6 +2065,7 @@ impl Database {
         query: &str,
         values: &[&dyn ToDbValue],
     ) -> Result<Vec<T>> {
+        self.ensure_model_allowed::<T>()?;
         let values = values
             .iter()
             .map(ToDbValue::to_db_value)

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,