Add rust cancel_submission function

Author: jerbaroo

opsqueue/src/common/submission.rs

@@ -634,6 +634,58 @@ pub mod db {
         .await
     }
 
+    #[tracing::instrument(skip(conn))]
+    pub async fn cancel_submission(
+        id: SubmissionId,
+        mut conn: impl WriterConnection,
+    ) -> sqlx::Result<()> {
+        conn.transaction(move |mut tx| {
+            Box::pin(
+                async move { cancel_submission_notx(id, &mut tx).await },
+            )
+        })
+        .await
+    }
+
+    /// Do not call directly! Must be called inside a transaction.
+    pub async fn cancel_submission_notx(
+        id: SubmissionId,
+        mut conn: impl WriterConnection<Transaction = True>,
+    ) -> sqlx::Result<()> {
+        cancel_submission_raw(id, &mut conn).await?;
+        super::chunk::db::skip_remaining_chunks(id, conn).await?;
+        Ok(())
+    }
+
+    #[tracing::instrument(skip(conn))]
+    pub(super) async fn cancel_submission_raw(
+        id: SubmissionId,
+        mut conn: impl WriterConnection,
+    ) -> sqlx::Result<()> {
+        let now = chrono::prelude::Utc::now();
+
+        query!(
+            "
+    INSERT INTO submissions_cancelled
+    (id, chunks_total, prefix, metadata, cancelled_at, chunks_done)
+    SELECT id, chunks_total, prefix, metadata, julianday($1), chunks_done FROM submissions WHERE id = $3;
+
+    DELETE FROM submissions WHERE id = $4 RETURNING *;
+    ",
+            now,
+            id,
+            id,
+        )
+        .fetch_one(conn.get_inner())
+        .await?;
+        counter!(crate::prometheus::SUBMISSIONS_CANCELLED_COUNTER).increment(1);
+        histogram!(crate::prometheus::SUBMISSIONS_DURATION_CANCEL_HISTOGRAM).record(
+            crate::prometheus::time_delta_as_f64(Utc::now() - id.timestamp()),
+        );
+
+        Ok(())
+    }
+
     #[tracing::instrument(skip(conn))]
     /// Do not call directly! MUST be called inside a transaction.
     pub(super) async fn complete_submission_raw(

opsqueue/src/prometheus.rs

@@ -17,8 +17,10 @@ use crate::db::DBPools;
 pub const SUBMISSIONS_TOTAL_COUNTER: &str = "submissions_total_count";
 pub const SUBMISSIONS_COMPLETED_COUNTER: &str = "submissions_completed_count";
 pub const SUBMISSIONS_FAILED_COUNTER: &str = "submissions_failed_count";
+pub const SUBMISSIONS_CANCELLED_COUNTER: &str = "submissions_failed_count";
 pub const SUBMISSIONS_DURATION_COMPLETE_HISTOGRAM: &str = "submissions_complete_duration_seconds";
 pub const SUBMISSIONS_DURATION_FAIL_HISTOGRAM: &str = "submissions_fail_duration_seconds";
+pub const SUBMISSIONS_DURATION_CANCEL_HISTOGRAM: &str = "submissions_cancel_duration_seconds";
 
 pub const CHUNKS_COMPLETED_COUNTER: &str = "chunks_completed_count";
 pub const CHUNKS_FAILED_COUNTER: &str = "chunks_failed_count";