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,
+    is_application_model: bool,
     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,
+            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);

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<DatabaseImpl>,
+    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<T: Model>(&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::<T>()
+                )),
+            ));
+        }
+        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<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 +1478,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 +1559,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 +1770,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 +1797,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 +1825,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 +1848,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 +2050,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,