Skip to main content

aven_core/operations/tasks/
attachment_creation.rs

1use super::creation::{insert_task, record_task_creation_undo, validate_task_draft};
2use super::{TaskCreationOptions, TaskCreationUndo, TaskDraft, TaskOutcome};
3use crate::db::{Database, begin_immediate};
4use crate::ids::TaskId;
5use crate::refs::get_task_in_workspace;
6use crate::undo::task_snapshot;
7use crate::workspaces::Workspace;
8use anyhow::{Result, bail};
9use sqlx::SqliteConnection;
10use std::collections::BTreeMap;
11use std::path::Path;
12use tracing::{info, warn};
13
14impl Database {
15    pub async fn create_task_with_attachments(
16        &self,
17        workspace: &Workspace,
18        blob_dir: &Path,
19        lifecycle_policy: crate::attachments::lifecycle::LifecyclePolicy,
20        draft: TaskDraft,
21        attachments: Vec<crate::operations::TaskAttachmentAddInput>,
22    ) -> Result<TaskOutcome> {
23        self.create_task_with_attachments_and_undo(
24            workspace,
25            blob_dir,
26            lifecycle_policy,
27            draft,
28            attachments,
29            TaskCreationUndo::None,
30        )
31        .await
32    }
33
34    pub async fn create_task_with_attachments_and_undo(
35        &self,
36        workspace: &Workspace,
37        blob_dir: &Path,
38        lifecycle_policy: crate::attachments::lifecycle::LifecyclePolicy,
39        draft: TaskDraft,
40        attachments: Vec<crate::operations::TaskAttachmentAddInput>,
41        undo: TaskCreationUndo,
42    ) -> Result<TaskOutcome> {
43        self.create_task_with_attachments_and_options(
44            workspace,
45            blob_dir,
46            lifecycle_policy,
47            draft,
48            attachments,
49            TaskCreationOptions::standalone(undo),
50        )
51        .await
52    }
53
54    pub async fn create_task_with_attachments_and_options(
55        &self,
56        workspace: &Workspace,
57        blob_dir: &Path,
58        lifecycle_policy: crate::attachments::lifecycle::LifecyclePolicy,
59        draft: TaskDraft,
60        attachments: Vec<crate::operations::TaskAttachmentAddInput>,
61        options: TaskCreationOptions,
62    ) -> Result<TaskOutcome> {
63        let mut conn = self.acquire_writer().await?;
64        create_task_with_attachments_and_epic(
65            &mut conn,
66            workspace,
67            blob_dir,
68            lifecycle_policy,
69            draft,
70            attachments,
71            options,
72        )
73        .await
74    }
75
76    pub async fn create_task_with_attachments_for_epic(
77        &self,
78        workspace: &Workspace,
79        blob_dir: &Path,
80        lifecycle_policy: crate::attachments::lifecycle::LifecyclePolicy,
81        draft: TaskDraft,
82        attachments: Vec<crate::operations::TaskAttachmentAddInput>,
83        epic_id: &TaskId,
84    ) -> Result<TaskOutcome> {
85        self.create_task_with_attachments_for_epic_and_undo(
86            workspace,
87            blob_dir,
88            lifecycle_policy,
89            draft,
90            attachments,
91            TaskCreationOptions::for_epic(epic_id.clone(), TaskCreationUndo::None),
92        )
93        .await
94    }
95
96    pub async fn create_task_with_attachments_for_epic_and_undo(
97        &self,
98        workspace: &Workspace,
99        blob_dir: &Path,
100        lifecycle_policy: crate::attachments::lifecycle::LifecyclePolicy,
101        draft: TaskDraft,
102        attachments: Vec<crate::operations::TaskAttachmentAddInput>,
103        options: TaskCreationOptions,
104    ) -> Result<TaskOutcome> {
105        let mut conn = self.acquire_writer().await?;
106        create_task_with_attachments_and_epic(
107            &mut conn,
108            workspace,
109            blob_dir,
110            lifecycle_policy,
111            draft,
112            attachments,
113            options,
114        )
115        .await
116    }
117}
118
119async fn create_task_with_attachments_and_epic(
120    conn: &mut SqliteConnection,
121    workspace: &Workspace,
122    blob_dir: &Path,
123    lifecycle_policy: crate::attachments::lifecycle::LifecyclePolicy,
124    draft: TaskDraft,
125    attachments: Vec<crate::operations::attachments::TaskAttachmentAddInput>,
126    options: TaskCreationOptions,
127) -> Result<TaskOutcome> {
128    validate_task_draft(&draft)?;
129    let TaskCreationOptions {
130        epic_id,
131        undo,
132        create_missing_labels,
133        require_existing_project,
134        capture_undo_snapshot,
135        require_existing_epic,
136    } = options;
137    let mut prepared = Vec::with_capacity(attachments.len());
138    for attachment in attachments {
139        prepared.push(crate::operations::attachments::prepare_task_attachment(attachment).await?);
140    }
141
142    let mut unique = BTreeMap::new();
143    for attachment in &prepared {
144        unique
145            .entry(attachment.sha256.clone())
146            .or_insert_with(|| attachment.clone());
147    }
148    let attachment_hashes = unique.keys().cloned().collect::<Vec<_>>();
149
150    let mut capacity_reservations = Vec::new();
151    for attachment in unique.values() {
152        let available: bool = sqlx::query_scalar(
153            "SELECT EXISTS(SELECT 1 FROM blob_inventory WHERE sha256 = ? AND available = 1)",
154        )
155        .bind(&attachment.sha256)
156        .fetch_one(&mut *conn)
157        .await?;
158        if available {
159            continue;
160        }
161        match crate::attachments::lifecycle::ensure_local_capacity(
162            conn,
163            blob_dir,
164            &attachment.sha256,
165            attachment.byte_size,
166            lifecycle_policy,
167            &crate::attachments::lifecycle::SystemClock,
168        )
169        .await
170        {
171            Ok(Some(reservation_id)) => capacity_reservations.push(reservation_id),
172            Ok(None) => {}
173            Err(error) => {
174                for reservation_id in capacity_reservations {
175                    let _ =
176                        crate::attachments::lifecycle::release_reservation(conn, &reservation_id)
177                            .await;
178                }
179                return Err(error);
180            }
181        }
182    }
183
184    let mut staging_leases = Vec::with_capacity(unique.len());
185    for attachment in unique.values() {
186        match crate::attachments::lifecycle::acquire_lease(
187            conn,
188            &attachment.sha256,
189            "staging",
190            &crate::attachments::lifecycle::SystemClock,
191        )
192        .await
193        {
194            Ok(lease_id) => staging_leases.push(lease_id),
195            Err(error) => {
196                for lease_id in staging_leases {
197                    let _ = crate::attachments::lifecycle::release_lease(conn, &lease_id).await;
198                }
199                for reservation_id in capacity_reservations {
200                    let _ =
201                        crate::attachments::lifecycle::release_reservation(conn, &reservation_id)
202                            .await;
203                }
204                return Err(error);
205            }
206        }
207    }
208
209    let mut created_hashes = Vec::new();
210    for attachment in unique.values() {
211        match crate::attachments::storage::stage_blob(
212            blob_dir,
213            &attachment.sha256,
214            &attachment.bytes,
215        )
216        .await
217        {
218            Ok(staged) if staged.byte_size == attachment.byte_size => {
219                if staged.created {
220                    created_hashes.push(staged.sha256);
221                }
222            }
223            Ok(_) => {
224                cleanup_attachment_guards(conn, &staging_leases, &capacity_reservations).await;
225                cleanup_created_objects(conn, blob_dir, &created_hashes).await;
226                bail!("error attachment-staged-size-mismatch");
227            }
228            Err(error) => {
229                cleanup_attachment_guards(conn, &staging_leases, &capacity_reservations).await;
230                cleanup_created_objects(conn, blob_dir, &created_hashes).await;
231                return Err(error);
232            }
233        }
234    }
235
236    let database_result = async {
237        let mut tx = begin_immediate(conn).await?;
238        if require_existing_epic && let Some(epic_id) = &epic_id {
239            crate::operations::epics::require_consumer_epic(&mut tx, workspace, epic_id).await?;
240        }
241        for attachment in unique.values() {
242            crate::attachments::storage::upsert_inventory_available(
243                &mut tx,
244                &attachment.sha256,
245                attachment.byte_size,
246                &attachment.facts.media_type,
247            )
248            .await?;
249        }
250        let inserted = insert_task(
251            &mut tx,
252            workspace,
253            draft,
254            create_missing_labels,
255            require_existing_project,
256        )
257        .await?;
258        if let Some(epic_id) = epic_id.as_ref() {
259            crate::operations::add_task_to_epic_in_transaction(
260                &mut tx,
261                workspace,
262                &inserted.id,
263                epic_id,
264            )
265            .await?;
266        }
267        let mut attachment_change_ids = Vec::with_capacity(prepared.len());
268        let attachment_base = crate::ids::now_utc();
269        for (index, attachment) in prepared.iter().enumerate() {
270            let created_at = (attachment_base
271                + chrono::TimeDelta::microseconds(i64::try_from(index)?))
272            .to_rfc3339_opts(chrono::SecondsFormat::Micros, true);
273            attachment_change_ids.push(
274                crate::operations::attachments::insert_prepared_attachment(
275                    &mut tx,
276                    workspace,
277                    &inserted.id,
278                    attachment,
279                    &created_at,
280                )
281                .await?,
282            );
283        }
284        crate::attachments::lifecycle::reconcile_liveness_for_hashes_in_transaction(
285            &mut tx,
286            &attachment_hashes,
287            &crate::attachments::lifecycle::SystemClock,
288        )
289        .await?;
290        let task = get_task_in_workspace(&mut tx, workspace, &inserted.id).await?;
291        let attachment_ids = prepared
292            .iter()
293            .map(|attachment| attachment.attachment_id.clone())
294            .collect();
295        record_task_creation_undo(
296            &mut tx,
297            workspace,
298            &task,
299            &inserted,
300            attachment_ids,
301            attachment_change_ids.clone(),
302            undo,
303        )
304        .await?;
305        let undo_snapshot = if capture_undo_snapshot {
306            Some(task_snapshot(&mut tx, &workspace.id, &task.id).await?)
307        } else {
308            None
309        };
310        tx.commit().await?;
311        Ok::<_, anyhow::Error>((inserted, attachment_change_ids, task, undo_snapshot))
312    }
313    .await;
314
315    let (inserted, attachment_change_ids, task, undo_snapshot) = match database_result {
316        Ok(value) => value,
317        Err(error) => {
318            cleanup_attachment_guards(conn, &staging_leases, &capacity_reservations).await;
319            cleanup_created_objects(conn, blob_dir, &created_hashes).await;
320            return Err(error);
321        }
322    };
323    cleanup_attachment_guards(conn, &staging_leases, &capacity_reservations).await;
324    info!(
325        task_id = %inserted.id,
326        project_key = %inserted.project_key,
327        label_count = inserted.label_count,
328        attachment_count = prepared.len(),
329        "task created"
330    );
331    Ok(TaskOutcome {
332        task,
333        create_change_id: Some(inserted.change_id),
334        attachment_change_ids,
335        undo_snapshot,
336    })
337}
338
339async fn cleanup_attachment_guards(
340    conn: &mut SqliteConnection,
341    leases: &[String],
342    reservations: &[String],
343) {
344    for lease_id in leases {
345        if let Err(error) = crate::attachments::lifecycle::release_lease(conn, lease_id).await {
346            warn!(%error, "failed to release attachment staging lease");
347        }
348    }
349    for reservation_id in reservations {
350        if let Err(error) =
351            crate::attachments::lifecycle::release_reservation(conn, reservation_id).await
352        {
353            warn!(%error, "failed to release attachment capacity reservation");
354        }
355    }
356}
357async fn cleanup_created_objects(conn: &mut SqliteConnection, blob_dir: &Path, hashes: &[String]) {
358    for sha256 in hashes {
359        crate::attachments::storage::remove_staged_blob_if_unreferenced(conn, blob_dir, sha256)
360            .await;
361    }
362}