Implement pausing for submissions

libs/opsqueue_python/python/opsqueue/exceptions.py

@@ -92,15 +92,6 @@ class TryFromIntError(IncorrectUsageError):
     pass
 
 
-class ChunkNotFoundError(IncorrectUsageError):
-    """
-    Raised when a method is used to look up information about a chunk
-    but the chunk doesn't exist within the Opsqueue.
-    """
-
-    pass
-
-
 class SubmissionNotFoundError(IncorrectUsageError):
     """
     Raised when a method is used to look up information about a submission

libs/opsqueue_python/python/opsqueue/producer.py

@@ -27,6 +27,7 @@
     SubmissionFailed,
     ChunkFailed,
     SubmissionNotCancellable,
+    SubmissionPaused,
 )
 
 __all__ = [
@@ -39,6 +40,7 @@
     "SubmissionNotCancellable",
     "SubmissionNotCancellableError",
     "SubmissionNotFoundError",
+    "SubmissionPaused",
     "TooManyMatchingSubmissionsError",
     "ChunkFailed",
 ]
@@ -96,6 +98,7 @@ def run_submission(
         serialization_format: SerializationFormat = DEFAULT_SERIALIZATION_FORMAT,
         metadata: None | bytes = None,
         strategic_metadata: None | dict[str, int] = None,
+        timeout: float | None = None,
     ) -> Iterator[Any]:
         """
         Inserts a submission into the queue, and blocks until it is completed.
@@ -116,6 +119,7 @@ def run_submission(
                 metadata=metadata,
                 strategic_metadata=strategic_metadata,
                 chunk_size=chunk_size,
+                timeout=timeout,
             )
             return _unchunk_iterator(results_iter, serialization_format)
 
@@ -146,6 +150,7 @@ def insert_submission(
         serialization_format: SerializationFormat = DEFAULT_SERIALIZATION_FORMAT,
         metadata: None | bytes = None,
         strategic_metadata: None | dict[str, int] = None,
+        paused: bool = False,
     ) -> SubmissionId:
         """
         Inserts a submission into the queue,
@@ -162,13 +167,15 @@ def insert_submission(
             metadata=metadata,
             strategic_metadata=strategic_metadata,
             chunk_size=chunk_size,
+            paused=paused,
         )
 
     def blocking_stream_completed_submission(
         self,
         submission_id: SubmissionId,
         *,
         serialization_format: SerializationFormat = DEFAULT_SERIALIZATION_FORMAT,
+        timeout: float | None = None,
     ) -> Iterator[Any]:
         """
         Blocks until the submission is completed.
@@ -181,7 +188,7 @@ def blocking_stream_completed_submission(
           (after retrying a consumer kept failing on one of the chunks)
         """
         return _unchunk_iterator(
-            self.blocking_stream_completed_submission_chunks(submission_id),
+            self.blocking_stream_completed_submission_chunks(submission_id, timeout),
             serialization_format,
         )
 
@@ -211,6 +218,7 @@ def run_submission_chunks(
         metadata: None | bytes = None,
         strategic_metadata: None | dict[str, int] = None,
         chunk_size: None | int = None,
+        timeout: float | None = None,
     ) -> Iterator[bytes]:
         """
         Inserts an already-chunked submission into the queue, and blocks until it is completed.
@@ -229,7 +237,7 @@ def run_submission_chunks(
             strategic_metadata=strategic_metadata,
             chunk_size=chunk_size,
         )
-        return self.blocking_stream_completed_submission_chunks(submission_id)
+        return self.blocking_stream_completed_submission_chunks(submission_id, timeout)
 
     async def async_run_submission_chunks(
         self,
@@ -259,6 +267,7 @@ def insert_submission_chunks(
         metadata: None | bytes = None,
         strategic_metadata: None | dict[str, int] = None,
         chunk_size: None | int = None,
+        paused: bool = False,
     ) -> SubmissionId:
         """
         Inserts an already-chunked submission into the queue,
@@ -275,10 +284,13 @@ def insert_submission_chunks(
             strategic_metadata=strategic_metadata,
             chunk_size=chunk_size,
             otel_trace_carrier=otel_trace_carrier,
+            paused=paused,
         )
 
     def blocking_stream_completed_submission_chunks(
-        self, submission_id: SubmissionId
+        self,
+        submission_id: SubmissionId,
+        timeout: float | None = None,
     ) -> Iterator[bytes]:
         """
         Blocks until the submission is completed, and returns an iterator that lazily
@@ -289,7 +301,9 @@ def blocking_stream_completed_submission_chunks(
         - `SubmissionFailedError` if the submission failed permanently
           (after retrying a consumer kept failing on one of the chunks)
         """
-        return self.inner.blocking_stream_completed_submission_chunks(submission_id)  # type: ignore[no-any-return]
+        return self.inner.blocking_stream_completed_submission_chunks(  # type: ignore[no-any-return]
+            submission_id, timeout
+        )
 
     async def async_stream_completed_submission_chunks(
         self, submission_id: SubmissionId
@@ -326,7 +340,7 @@ def count_submissions(self) -> int:
 
     def cancel_submission(self, submission_id: SubmissionId) -> None:
         """
-        Cancel a specific submission that is in progress.
+        Cancel a specific submission that is in progress or paused.
 
         Returns None if the submission was successfully cancelled.
 
@@ -337,6 +351,19 @@ def cancel_submission(self, submission_id: SubmissionId) -> None:
         """
         self.inner.cancel_submission(submission_id)
 
+    def unpause_submission(self, submission_id: SubmissionId) -> None:
+        """
+        Unpause a specific submission that is currently paused,
+        making it available to consumers.
+
+        Returns None if the submission was successfully unpaused.
+
+        Raises:
+        - `SubmissionNotFoundError` if the submission is not currently paused.
+        - `InternalProducerClientError` if there is a low-level internal error.
+        """
+        self.inner.unpause_submission(submission_id)
+
     def get_submission_status(
         self, submission_id: SubmissionId
     ) -> SubmissionStatus | None:

libs/opsqueue_python/src/common.rs

@@ -364,12 +364,15 @@ pub enum SubmissionStatus {
     Cancelled {
         submission: SubmissionCancelled,
     },
+    Paused {
+        submission: SubmissionPaused,
+    },
 }
 
 impl From<opsqueue::common::submission::SubmissionStatus> for SubmissionStatus {
     fn from(value: opsqueue::common::submission::SubmissionStatus) -> Self {
         use opsqueue::common::submission::SubmissionStatus::{
-            Cancelled, Completed, Failed, InProgress,
+            Cancelled, Completed, Failed, InProgress, Paused,
         };
         match value {
             InProgress(s) => SubmissionStatus::InProgress {
@@ -386,6 +389,9 @@ impl From<opsqueue::common::submission::SubmissionStatus> for SubmissionStatus {
             Cancelled(s) => SubmissionStatus::Cancelled {
                 submission: s.into(),
             },
+            Paused(s) => SubmissionStatus::Paused {
+                submission: s.into(),
+            },
         }
     }
 }
@@ -510,6 +516,42 @@ pub struct SubmissionCancelled {
     pub cancelled_at: DateTime<Utc>,
 }
 
+#[pyclass(from_py_object, frozen, get_all, module = "opsqueue")]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct SubmissionPaused {
+    pub id: SubmissionId,
+    pub chunks_total: u64,
+    pub chunks_done: u64,
+    pub metadata: Option<submission::Metadata>,
+    pub strategic_metadata: StrategicMetadataMap,
+}
+
+impl From<opsqueue::common::submission::SubmissionPaused> for SubmissionPaused {
+    fn from(value: opsqueue::common::submission::SubmissionPaused) -> Self {
+        Self {
+            id: value.id.into(),
+            chunks_total: value.chunks_total.into(),
+            chunks_done: value.chunks_done.into(),
+            metadata: value.metadata,
+            strategic_metadata: value.strategic_metadata,
+        }
+    }
+}
+
+#[pymethods]
+impl SubmissionPaused {
+    fn __repr__(&self) -> String {
+        format!(
+            "SubmissionPaused(id={0}, chunks_total={1}, chunks_done={2}, metadata={3:?}, strategic_metadata={4:?})",
+            self.id.__repr__(),
+            self.chunks_total,
+            self.chunks_done,
+            self.metadata,
+            self.strategic_metadata
+        )
+    }
+}
+
 /// Submission could not be cancelled because it was already completed, failed
 /// or cancelled.
 #[pyclass(from_py_object, frozen, module = "opsqueue")]

libs/opsqueue_python/src/errors.rs

@@ -2,24 +2,21 @@
 /// so we have nice IDE support for docs-on-hover and for 'go to definition'.
 use std::error::Error;
 
-use opsqueue::common::chunk::ChunkId;
 use opsqueue::common::errors::{
-    ChunkNotFound, E, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound,
-    TooManyMatchingSubmissions, UnexpectedOpsqueueConsumerServerResponse,
+    E, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions,
+    UnexpectedOpsqueueConsumerServerResponse,
 };
-use pyo3::exceptions::PyBaseException;
+use pyo3::exceptions::{PyBaseException, PyTimeoutError};
 use pyo3::{Bound, PyErr, Python, import_exception};
 
 use crate::common;
-use crate::common::{ChunkIndex, SubmissionId};
 
 // Expected errors:
 import_exception!(opsqueue.exceptions, SubmissionFailedError);
 
 // Incorrect usage errors:
 import_exception!(opsqueue.exceptions, IncorrectUsageError);
 import_exception!(opsqueue.exceptions, TryFromIntError);
-import_exception!(opsqueue.exceptions, ChunkNotFoundError);
 import_exception!(opsqueue.exceptions, SubmissionNotFoundError);
 import_exception!(opsqueue.exceptions, SubmissionNotCancellableError);
 import_exception!(opsqueue.exceptions, TooManyMatchingSubmissionsError);
@@ -173,22 +170,6 @@ impl From<CError<crate::producer::SubmissionNotCompletedYetError>> for PyErr {
     }
 }
 
-impl From<CError<ChunkNotFound>> for PyErr {
-    fn from(value: CError<ChunkNotFound>) -> Self {
-        let ChunkId {
-            submission_id,
-            chunk_index,
-        } = value.0.0;
-        ChunkNotFoundError::new_err((
-            value.0.to_string(),
-            (
-                SubmissionId::from(submission_id),
-                ChunkIndex::from(chunk_index),
-            ),
-        ))
-    }
-}
-
 impl From<CError<opsqueue::object_store::NewObjectStoreClientError>> for PyErr {
     fn from(value: CError<opsqueue::object_store::NewObjectStoreClientError>) -> Self {
         NewObjectStoreClientError::new_err(value.0.to_string())
@@ -201,6 +182,12 @@ impl From<CError<UnexpectedOpsqueueConsumerServerResponse>> for PyErr {
     }
 }
 
+impl From<CError<tokio::time::error::Elapsed>> for PyErr {
+    fn from(_value: CError<tokio::time::error::Elapsed>) -> Self {
+        PyTimeoutError::new_err("timeout was reached")
+    }
+}
+
 impl<T> From<PyErr> for CError<E<FatalPythonException, T>> {
     fn from(value: PyErr) -> Self {
         CError(E::L(FatalPythonException(value)))

libs/opsqueue_python/src/lib.rs

@@ -24,6 +24,7 @@ fn opsqueue_internal(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<common::SubmissionCancelled>()?;
     m.add_class::<common::SubmissionCompleted>()?;
     m.add_class::<common::SubmissionFailed>()?;
+    m.add_class::<common::SubmissionPaused>()?;
     m.add_class::<common::SubmissionNotCancellable>()?;
     m.add_class::<producer::PyChunksIter>()?;
     m.add_class::<consumer::ConsumerClient>()?;