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

@@ -26,6 +26,7 @@
     SubmissionFailed,
     ChunkFailed,
     SubmissionNotCancellable,
+    SubmissionPaused,
 )
 
 __all__ = [
@@ -38,6 +39,7 @@
     "SubmissionNotCancellable",
     "SubmissionNotCancellableError",
     "SubmissionNotFoundError",
+    "SubmissionPaused",
     "ChunkFailed",
 ]
 
@@ -324,7 +326,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.
 
@@ -335,6 +337,36 @@ def cancel_submission(self, submission_id: SubmissionId) -> None:
         """
         self.inner.cancel_submission(submission_id)
 
+    def pause_submission(self, submission_id: SubmissionId) -> None:
+        """
+        Pause a specific submission that is currently in progress.
+
+        Once paused, the submission's chunks will not be dispatched to consumers
+        until the submission is unpaused.
+
+        Returns None if the submission was successfully paused.
+
+        Raises:
+        - `SubmissionNotFoundError` if the submission is not currently in
+          the in-progress state (it may be already paused, completed, failed,
+          cancelled, or simply not found).
+        - `InternalProducerClientError` if there is a low-level internal error.
+        """
+        self.inner.pause_submission(submission_id)
+
+    def unpause_submission(self, submission_id: SubmissionId) -> None:
+        """
+        Unpause a specific submission that is currently paused,
+        making it available to consumers again.
+
+        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

@@ -348,6 +348,9 @@ pub enum SubmissionStatus {
     Cancelled {
         submission: SubmissionCancelled,
     },
+    Paused {
+        submission: SubmissionPaused,
+    },
 }
 
 impl From<opsqueue::common::submission::SubmissionStatus> for SubmissionStatus {
@@ -368,6 +371,9 @@ impl From<opsqueue::common::submission::SubmissionStatus> for SubmissionStatus {
             Cancelled(s) => SubmissionStatus::Cancelled {
                 submission: s.into(),
             },
+            Paused(s) => SubmissionStatus::Paused {
+                submission: s.into(),
+            },
         }
     }
 }
@@ -478,6 +484,42 @@ pub struct SubmissionCancelled {
     pub cancelled_at: DateTime<Utc>,
 }
 
+#[pyclass(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(frozen, module = "opsqueue")]

libs/opsqueue_python/src/errors.rs

@@ -2,24 +2,22 @@
 /// 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, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound,
+    IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound,
     UnexpectedOpsqueueConsumerServerResponse, E,
 };
 use pyo3::exceptions::PyBaseException;
 use pyo3::{import_exception, Bound, PyErr, Python};
 
 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, NewObjectStoreClientError);
@@ -166,22 +164,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())

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>()?;

libs/opsqueue_python/src/producer.rs

@@ -147,6 +147,59 @@ impl ProducerClient {
         })
     }
 
+    /// Pause an in-progress submission.
+    ///
+    /// Will return an error if the submission is not currently in progress
+    /// (e.g. already paused, completed, failed, or not found).
+    #[allow(clippy::result_large_err, clippy::type_complexity)]
+    pub fn pause_submission(
+        &self,
+        py: Python<'_>,
+        id: SubmissionId,
+    ) -> CPyResult<
+        (),
+        E![
+            FatalPythonException,
+            SubmissionNotFound,
+            InternalProducerClientError
+        ],
+    > {
+        py.allow_threads(|| {
+            self.block_unless_interrupted(async {
+                self.producer_client
+                    .pause_submission(id.into())
+                    .await
+                    .map_err(|e| CError(R(e)))
+            })
+        })
+    }
+
+    /// Unpause a paused submission, making it available to consumers again.
+    ///
+    /// Will return an error if the submission is not currently paused.
+    #[allow(clippy::result_large_err, clippy::type_complexity)]
+    pub fn unpause_submission(
+        &self,
+        py: Python<'_>,
+        id: SubmissionId,
+    ) -> CPyResult<
+        (),
+        E![
+            FatalPythonException,
+            SubmissionNotFound,
+            InternalProducerClientError
+        ],
+    > {
+        py.allow_threads(|| {
+            self.block_unless_interrupted(async {
+                self.producer_client
+                    .unpause_submission(id.into())
+                    .await
+                    .map_err(|e| CError(R(e)))
+            })
+        })
+    }
+
     /// Retrieve the status (in progress, completed or failed) of a specific submission.
     ///
     /// The returned SubmissionStatus object also includes the number of chunks finished so far,
@@ -212,6 +265,7 @@ impl ProducerClient {
                 },
                 metadata,
                 strategic_metadata,
+                paused: false,
             };
             self.block_unless_interrupted(async move {
                 self.producer_client
@@ -272,6 +326,7 @@ impl ProducerClient {
                     },
                     metadata,
                     strategic_metadata: strategic_metadata.unwrap_or_default(),
+                    paused: false,
                 };
                 self.producer_client
                     .insert_submission(&submission, &otel_trace_carrier)

libs/opsqueue_python/tests/test_roundtrip.py

@@ -14,6 +14,7 @@
     SubmissionNotFoundError,
     SubmissionNotCancellable,
     SubmissionNotCancellableError,
+    SubmissionPaused,
 )
 from opsqueue.consumer import ConsumerClient, Chunk
 from opsqueue.common import SerializationFormat
@@ -508,3 +509,94 @@ def consume(x: int) -> int | None:
         with pytest.raises(SubmissionFailedError) as exc_info:
             producer_client.blocking_stream_completed_submission(submission_id)
         assert exc_info.value.submission.chunks_done == len(chunks) - 1
+
+
+def test_pause_in_progress(opsqueue: OpsqueueProcess) -> None:
+    """Pausing an in-progress submission succeeds and its status becomes Paused."""
+    url = "file:///tmp/opsqueue/test_pause_in_progress"
+    producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+    submission_id = producer_client.insert_submission((1, 2, 3), chunk_size=1)
+    # Sanity-check: submission is in progress before pausing.
+    assert isinstance(
+        producer_client.get_submission_status(submission_id),
+        SubmissionStatus.InProgress,
+    )
+
+    producer_client.pause_submission(submission_id)
+
+    status = producer_client.get_submission_status(submission_id)
+    assert isinstance(status, SubmissionStatus.Paused)
+    assert isinstance(status.submission, SubmissionPaused)
+
+
+def test_pause_not_found(opsqueue: OpsqueueProcess) -> None:
+    """Pausing a non-existent submission raises SubmissionNotFoundError."""
+    url = "file:///tmp/opsqueue/test_pause_not_found"
+    producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+    with pytest.raises(SubmissionNotFoundError) as exc_info:
+        producer_client.pause_submission(SubmissionId(0))
+    assert exc_info.value.submission_id == 0
+
+
+def test_unpause_and_complete(
+    opsqueue: OpsqueueProcess, any_consumer_strategy: StrategyDescription
+) -> None:
+    """Unpausing a paused submission makes it available to consumers again,
+    and it can be completed normally afterwards."""
+    url = "file:///tmp/opsqueue/test_unpause_and_complete"
+    producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+    submission_id = producer_client.insert_submission((1, 2, 3), chunk_size=1)
+
+    producer_client.pause_submission(submission_id)
+    assert isinstance(
+        producer_client.get_submission_status(submission_id), SubmissionStatus.Paused
+    )
+
+    producer_client.unpause_submission(submission_id)
+    assert isinstance(
+        producer_client.get_submission_status(submission_id),
+        SubmissionStatus.InProgress,
+    )
+
+    def run_consumer() -> None:
+        consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url)
+        strategy = strategy_from_description(any_consumer_strategy)
+        consumer_client.run_each_op(lambda x: x, strategy=strategy)
+
+    with background_process(run_consumer):
+        producer_client.blocking_stream_completed_submission(submission_id)
+        assert isinstance(
+            producer_client.get_submission_status(submission_id),
+            SubmissionStatus.Completed,
+        )
+
+
+def test_unpause_not_found(opsqueue: OpsqueueProcess) -> None:
+    """Unpausing a submission that is not paused (e.g. in-progress) raises
+    SubmissionNotFoundError."""
+    url = "file:///tmp/opsqueue/test_unpause_not_found"
+    producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+    submission_id = producer_client.insert_submission((1, 2, 3), chunk_size=1)
+    assert isinstance(
+        producer_client.get_submission_status(submission_id),
+        SubmissionStatus.InProgress,
+    )
+    with pytest.raises(SubmissionNotFoundError):
+        producer_client.unpause_submission(submission_id)
+
+
+def test_cancel_paused(opsqueue: OpsqueueProcess) -> None:
+    """A paused submission can be cancelled; its status becomes Cancelled."""
+    url = "file:///tmp/opsqueue/test_cancel_paused"
+    producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+    submission_id = producer_client.insert_submission((1, 2, 3), chunk_size=1)
+
+    producer_client.pause_submission(submission_id)
+    assert isinstance(
+        producer_client.get_submission_status(submission_id), SubmissionStatus.Paused
+    )
+
+    producer_client.cancel_submission(submission_id)
+    assert isinstance(
+        producer_client.get_submission_status(submission_id), SubmissionStatus.Cancelled
+    )

opsqueue/migrations/20260715143000_pausing.down.sql

@@ -0,0 +1,2 @@
+DROP TABLE chunks_paused;
+DROP TABLE submissions_paused;