db/models/
assignment_file.rs

1// models/assignment_file.rs
2
3use chrono::{DateTime, Utc};
4// use code_runner::ExecutionConfig;
5use sea_orm::{ActiveValue::Set, DbErr, DatabaseConnection};
6use sea_orm::entity::prelude::*;
7use std::env;
8use std::fs;
9use std::path::PathBuf;
10use strum::{Display, EnumIter, EnumString};
11use util::execution_config::ExecutionConfig;
12
13/// Represents a file associated with an assignment, such as a spec, main file, memo, or submission.
14#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
15#[sea_orm(table_name = "assignment_files")]
16pub struct Model {
17    #[sea_orm(primary_key)]
18    pub id: i64,
19
20    /// Foreign key reference to an assignment.
21    pub assignment_id: i64,
22
23    /// Original file name.
24    pub filename: String,
25
26    /// Relative path to the stored file from the storage root.
27    pub path: String,
28
29    /// Type of the file (spec, main, memo, submission).
30    pub file_type: FileType,
31
32    pub created_at: DateTime<Utc>,
33    pub updated_at: DateTime<Utc>,
34}
35
36/// Enum representing the type/category of an assignment file.
37#[derive(Debug, Clone, PartialEq, EnumIter, EnumString, Display, DeriveActiveEnum)]
38#[strum(ascii_case_insensitive)]
39#[sea_orm(
40    rs_type = "String",
41    db_type = "Enum",
42    enum_name = "assignment_file_type"
43)]
44pub enum FileType {
45    #[strum(serialize = "spec")]
46    #[sea_orm(string_value = "spec")]
47    Spec,
48    #[strum(serialize = "main")]
49    #[sea_orm(string_value = "main")]
50    Main,
51    #[strum(serialize = "memo")]
52    #[sea_orm(string_value = "memo")]
53    Memo,
54    #[strum(serialize = "makefile")]
55    #[sea_orm(string_value = "makefile")]
56    Makefile,
57    #[strum(serialize = "mark_allocator")]
58    #[sea_orm(string_value = "mark_allocator")]
59    MarkAllocator,
60    #[strum(serialize = "config")]
61    #[sea_orm(string_value = "config")]
62    Config,
63}
64
65#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
66pub enum Relation {
67    #[sea_orm(
68        belongs_to = "super::assignment::Entity",
69        from = "Column::AssignmentId",
70        to = "super::assignment::Column::Id"
71    )]
72    Assignment,
73}
74
75impl ActiveModelBehavior for ActiveModel {}
76
77impl Model {
78    /// Loads and returns the `ExecutionConfig` if the file type is `Config`.
79    /// Requires `module_id` because it's not stored in the DB.
80    pub fn load_execution_config(&self, module_id: i64) -> Result<ExecutionConfig, String> {
81        if self.file_type != FileType::Config {
82            return Err("File is not of type 'config'".to_string());
83        }
84
85        ExecutionConfig::get_execution_config(module_id, self.assignment_id)
86    }
87
88    /// Returns the base directory for assignment file storage from the environment.
89    pub fn storage_root() -> PathBuf {
90        env::var("ASSIGNMENT_STORAGE_ROOT")
91            .map(PathBuf::from)
92            .unwrap_or_else(|_| PathBuf::from("data/assignment_files"))
93    }
94
95    /// Computes the full directory path based on module ID, assignment ID, and file type.
96    pub fn full_directory_path(
97        module_id: i64,
98        assignment_id: i64,
99        file_type: &FileType,
100    ) -> PathBuf {
101        Self::storage_root()
102            .join(format!("module_{module_id}"))
103            .join(format!("assignment_{assignment_id}"))
104            .join(file_type.to_string())
105    }
106
107    pub fn full_path(&self) -> PathBuf {
108        Self::storage_root().join(&self.path)
109    }
110
111    pub async fn save_file(
112        db: &DatabaseConnection,
113        assignment_id: i64,
114        module_id: i64,
115        file_type: FileType,
116        filename: &str,
117        bytes: &[u8],
118    ) -> Result<Self, DbErr> {
119        let now = Utc::now();
120
121        use sea_orm::ColumnTrait;
122        use sea_orm::EntityTrait;
123        use sea_orm::QueryFilter;
124
125        use crate::models::assignment_file::{Column, Entity as AssignmentFile};
126
127        if let Some(existing) = AssignmentFile::find()
128            .filter(Column::AssignmentId.eq(assignment_id))
129            .filter(Column::FileType.eq(file_type.clone()))
130            .one(db)
131            .await?
132        {
133            let existing_path = Self::storage_root().join(&existing.path);
134            let _ = fs::remove_file(existing_path); // Silently ignore failure
135
136            existing.delete(db).await?;
137        }
138
139        let partial = ActiveModel {
140            assignment_id: Set(assignment_id),
141            filename: Set(filename.to_string()),
142            path: Set("".to_string()), // will be updated after write
143            file_type: Set(file_type.clone()),
144            created_at: Set(now),
145            updated_at: Set(now),
146            ..Default::default()
147        };
148
149        let inserted: Model = partial.insert(db).await?;
150
151        let ext = PathBuf::from(filename)
152            .extension()
153            .map(|e| e.to_string_lossy().to_string());
154
155        let stored_filename = match ext {
156            Some(ext) => format!("{}.{}", inserted.id, ext),
157            None => inserted.id.to_string(),
158        };
159
160        let dir_path = Self::full_directory_path(module_id, assignment_id, &file_type);
161        fs::create_dir_all(&dir_path)
162            .map_err(|e| DbErr::Custom(format!("Failed to create directory: {e}")))?;
163
164        let file_path = dir_path.join(&stored_filename);
165        let relative_path = file_path
166            .strip_prefix(Self::storage_root())
167            .unwrap()
168            .to_string_lossy()
169            .to_string();
170
171        fs::write(&file_path, bytes)
172            .map_err(|e| DbErr::Custom(format!("Failed to write file: {e}")))?;
173
174        let mut model: ActiveModel = inserted.into();
175        model.path = Set(relative_path);
176        model.updated_at = Set(Utc::now());
177
178        model.update(db).await
179    }
180
181    /// Loads the file contents from disk based on the path stored in the model.
182    pub fn load_file(&self) -> Result<Vec<u8>, std::io::Error> {
183        let full_path = Self::storage_root().join(&self.path);
184        fs::read(full_path)
185    }
186
187    /// Deletes the file from disk (but not the DB record).
188    pub fn delete_file_only(&self) -> Result<(), std::io::Error> {
189        let full_path = Self::storage_root().join(&self.path);
190        fs::remove_file(full_path)
191    }
192    
193    // TODO: Change this to get the skeleton files instead of the memo files
194    pub async fn get_base_files(
195        db: &DatabaseConnection,
196        assignment_id: i64,
197    ) -> Result<Vec<Self>, DbErr> {
198        Entity::find()
199            .filter(Column::AssignmentId.eq(assignment_id))
200            .filter(Column::FileType.eq(FileType::Main))
201            .all(db)
202            .await
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209    use crate::models::assignment::AssignmentType;
210    use crate::test_utils::setup_test_db;
211    use chrono::Utc;
212    use sea_orm::Set;
213    use std::env;
214    use tempfile::TempDir;
215
216    fn fake_bytes() -> Vec<u8> {
217        vec![0x50, 0x4B, 0x03, 0x04] // ZIP file signature
218    }
219
220    fn override_storage_dir(temp: &TempDir) {
221        unsafe {
222            env::set_var("ASSIGNMENT_STORAGE_ROOT", temp.path());
223        }
224    }
225
226    #[tokio::test]
227    #[ignore]
228    async fn test_save_and_load_file() {
229        let temp_dir = TempDir::new().unwrap();
230        override_storage_dir(&temp_dir);
231        let db = setup_test_db().await;
232
233        // Insert dummy module so assignment FK passes
234        let _module = crate::models::module::ActiveModel {
235            code: Set("COS301".to_string()),
236            year: Set(2025),
237            description: Set(Some("Capstone".to_string())),
238            created_at: Set(Utc::now()),
239            updated_at: Set(Utc::now()),
240            ..Default::default()
241        }
242        .insert(&db)
243        .await
244        .expect("Insert module failed");
245
246        // Insert dummy assignment using enum value for assignment_type
247        let _assignment = crate::models::assignment::Model::create(
248            &db,
249            1,
250            "Test Assignment",
251            Some("Desc"),
252            AssignmentType::Practical,
253            Utc::now(),
254            Utc::now(),
255        )
256        .await
257        .expect("Insert assignment failed");
258
259        let content = fake_bytes();
260        let filename = "test_file.zip";
261        let saved = Model::save_file(
262            &db,
263            1, // assignment_id
264            1, // module_id
265            FileType::Spec,
266            filename,
267            &content,
268        )
269        .await
270        .unwrap();
271
272        assert_eq!(saved.assignment_id, 1);
273        assert_eq!(saved.filename, filename);
274        assert_eq!(saved.file_type, FileType::Spec);
275
276        // Confirm file on disk
277        let full_path = Model::storage_root().join(&saved.path);
278        //The error was this line
279        assert!(full_path.exists());
280
281        // Load contents
282        let bytes = saved.load_file().unwrap();
283        assert_eq!(bytes, content);
284
285        // Delete file only
286        saved.delete_file_only().unwrap();
287        assert!(!full_path.exists());
288    }
289}
db/models/assignment_file.rs

db/models/
assignment_file.rs
