Lookup submision IDs by strategic metadata

libs/opsqueue_python/python/opsqueue/exceptions.py

@@ -153,6 +153,34 @@ class SubmissionNotCompletedYetError(IncorrectUsageError):
     pass
 
 
+class TooManyMatchingSubmissionsError(IncorrectUsageError):
+    """
+    Raised when a strategic-metadata lookup matches more submissions
+    than the server's configured maximum (``max_submissions_returned``).
+
+    Narrow the query with more specific strategic metadata, or raise the
+    server's configured maximum.
+    """
+
+    __slots__ = ["max_submissions"]
+
+    def __init__(
+        self,
+        max_submissions: int,
+    ):
+        super().__init__()
+        self.max_submissions = max_submissions
+
+    def __str__(self) -> str:
+        return (
+            f"The lookup matched more submissions than the configured "
+            f"maximum of {self.max_submissions}"
+        )
+
+    def __repr__(self) -> str:
+        return str(self)
+
+
 # Internal errors:
 
 

libs/opsqueue_python/python/opsqueue/producer.py

@@ -18,6 +18,7 @@
     SubmissionFailedError,
     SubmissionNotCancellableError,
     SubmissionNotFoundError,
+    TooManyMatchingSubmissionsError,
 )
 from .opsqueue_internal import (  # type: ignore[import-not-found]
     SubmissionId,
@@ -38,10 +39,15 @@
     "SubmissionNotCancellable",
     "SubmissionNotCancellableError",
     "SubmissionNotFoundError",
+    "TooManyMatchingSubmissionsError",
     "ChunkFailed",
 ]
 
 
+class LookupIdsWithEmptyStrategicMetadataError(Exception):
+    pass
+
+
 class ProducerClient:
     """
     Opsqueue producer client. Allows sending of large collections of operations ('submissions')
@@ -367,6 +373,25 @@ def lookup_submission_id_by_prefix(self, prefix: str) -> SubmissionId | None:
         """
         return self.inner.lookup_submission_id_by_prefix(prefix)
 
+    def lookup_submission_ids_by_strategic_metadata(
+        self, strategic_metadata: dict[str, int]
+    ) -> list[SubmissionId]:
+        """Attempts to find in-progress submissions where the strategic metadata
+        of that submission includes all of the key-value pairs of the given
+        'strategic_metadata'. A matching submission must include all of the
+        given key-value pairs, but it may also contain other key-value pairs.
+
+        Raises:
+        - `TooManyMatchingSubmissionsError` if the lookup matches more
+          submissions than the server's configured maximum. Narrow the query
+          with more specific strategic metadata.
+        - `InternalProducerClientError` if there is a low-level internal error.
+
+        """
+        return self.inner.lookup_submission_ids_by_strategic_metadata(  # type: ignore[no-any-return]
+            strategic_metadata
+        )
+
     def is_completed(self, submission_id: SubmissionId) -> bool:
         raise NotImplementedError
 

libs/opsqueue_python/src/errors.rs

@@ -5,7 +5,7 @@ use std::error::Error;
 use opsqueue::common::chunk::ChunkId;
 use opsqueue::common::errors::{
     ChunkNotFound, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound,
-    UnexpectedOpsqueueConsumerServerResponse, E,
+    TooManyMatchingSubmissions, UnexpectedOpsqueueConsumerServerResponse, E,
 };
 use pyo3::exceptions::PyBaseException;
 use pyo3::{import_exception, Bound, PyErr, Python};
@@ -22,6 +22,7 @@ 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);
 import_exception!(opsqueue.exceptions, NewObjectStoreClientError);
 import_exception!(opsqueue.exceptions, SubmissionNotCompletedYetError);
 
@@ -146,6 +147,12 @@ impl From<CError<SubmissionNotFound>> for PyErr {
     }
 }
 
+impl From<CError<TooManyMatchingSubmissions>> for PyErr {
+    fn from(value: CError<TooManyMatchingSubmissions>) -> Self {
+        TooManyMatchingSubmissionsError::new_err(value.0 .0)
+    }
+}
+
 pub struct SubmissionFailed(
     pub crate::common::SubmissionFailed,
     pub crate::common::ChunkFailed,

libs/opsqueue_python/src/producer.rs

@@ -10,7 +10,7 @@ use pyo3::{
 use futures::{stream::BoxStream, StreamExt, TryStreamExt};
 use opsqueue::{
     common::errors::E::{self, L, R},
-    common::errors::{SubmissionNotCancellable, SubmissionNotFound},
+    common::errors::{SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions},
     object_store::{ChunksStorageError, NewObjectStoreClientError},
     producer::client::{Client as ActualClient, InternalProducerClientError},
 };
@@ -189,6 +189,31 @@ impl ProducerClient {
         })
     }
 
+    /// Attempts to find the IDs of submission matching ALL key-values pairs of
+    /// the given strategic metadata.
+    pub fn lookup_submission_ids_by_strategic_metadata(
+        &self,
+        py: Python<'_>,
+        strategic_metadata: StrategicMetadataMap,
+    ) -> CPyResult<
+        Vec<SubmissionId>,
+        E![
+            FatalPythonException,
+            TooManyMatchingSubmissions,
+            InternalProducerClientError
+        ],
+    > {
+        py.allow_threads(|| {
+            self.block_unless_interrupted(async {
+                self.producer_client
+                    .lookup_submission_ids_by_strategic_metadata(&strategic_metadata)
+                    .await
+                    .map(|res| res.into_iter().map(Into::into).collect())
+                    .map_err(|e| CError(R(e)))
+            })
+        })
+    }
+
     /// Directly inserts a submission without sending the chunks to GCS
     /// (but immediately embedding them in the DB).
     /// NOTE: This does not support StrategicMetadata currently

libs/opsqueue_python/tests/conftest.py

@@ -53,7 +53,7 @@ def opsqueue() -> Generator[OpsqueueProcess, None, None]:
 
 @contextmanager
 def opsqueue_service(
-    *, port: int | None = None
+        *, port: int | None = None, command_args: Iterable[str] = ()
 ) -> Generator[OpsqueueProcess, None, None]:
     global test_opsqueue_port_offset
 
@@ -75,6 +75,7 @@ def opsqueue_service(
         str(port),
         "--database-filename",
         temp_dbname,
+        *command_args,
     ]
     env = os.environ.copy()  # We copy the env so e.g. RUST_LOG and other env vars are propagated from outside of the invocation of pytest
     if env.get("RUST_LOG") is None:

libs/opsqueue_python/tests/test_roundtrip.py

@@ -14,18 +14,19 @@
     SubmissionNotFoundError,
     SubmissionNotCancellable,
     SubmissionNotCancellableError,
+    TooManyMatchingSubmissionsError,
 )
 from opsqueue.consumer import ConsumerClient, Chunk
 from opsqueue.common import SerializationFormat
 from conftest import (
     background_process,
     multiple_background_processes,
     OpsqueueProcess,
+    opsqueue_service,
     StrategyDescription,
     strategy_from_description,
 )
 import logging
-
 import pytest
 
 
@@ -508,3 +509,77 @@ 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_lookup_submission_ids_by_strategic_metadata(opsqueue: OpsqueueProcess) -> None:
+    """Lookup of submission IDs should only match in progress submissions with
+    all pieces of strategic metadata.
+
+    """
+    url = "file:///tmp/opsqueue/test_lookup_submission_ids_by_strategic_metadata"
+    producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+    id_1 = producer_client.insert_submission(
+        [1], chunk_size=1, strategic_metadata={"foo": 1, "bar": 2, "wow": 3}
+    )
+    id_2 = producer_client.insert_submission(
+        [1], chunk_size=1, strategic_metadata={"foo": 1, "bar": 2, "moo": 3}
+    )
+    # Inserting some similar data to that above, which shouldn't get matched.
+    producer_client.insert_submission(
+        [1], chunk_size=1, strategic_metadata={"foo": 2, "bar": 1}
+    )
+
+    def test_lookup(
+        strategic_metadata: dict[str, int], expected_ids: list[int]
+    ) -> None:
+        found_ids = producer_client.lookup_submission_ids_by_strategic_metadata(
+            strategic_metadata
+        )
+        assert isinstance(found_ids, list)
+        assert all(map(lambda x: isinstance(x, SubmissionId), found_ids))
+        assert found_ids == expected_ids
+
+    test_lookup({"foo": 1}, [id_1, id_2])
+    test_lookup({"foo": 1, "bar": 2}, [id_1, id_2])
+    test_lookup({"foo": 1, "MISS": 2}, [])
+    test_lookup({"wow": 3}, [id_1])
+
+    # Should only match in-progress submission.
+    producer_client.cancel_submission(id_1)
+    test_lookup({"foo": 1}, [id_2])
+
+
+def test_lookup_submission_ids_by_empty_strategic_metadata(
+    opsqueue: OpsqueueProcess,
+) -> None:
+    """Lookup of submission IDs with empty strategic_metadata should NOT raise
+    an exception.
+
+    """
+    url = "file:///tmp/opsqueue/test_lookup_submission_ids_by_empty_strategic_metadata"
+    producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+    count = 6
+    for _ in range(count):
+        producer_client.insert_submission([1], chunk_size=1)
+    assert len(producer_client.lookup_submission_ids_by_strategic_metadata({})) == count
+
+def test_lookup_too_many_submission_ids_by_strategic_metadata() -> None:
+    """Lookup of too many submission IDs beyond the configured limit raises
+    TooManyMatchingSubmissionsError.
+
+    """
+    max_ = 2
+    # We didn't request the OpsQueueProcess as a parameter so an instance isn't
+    # started, instead we start one here with custom args.
+    with opsqueue_service(command_args=["--max-submissions-returned", str(max_)]) as opsqueue:
+        url = "file:///tmp/opsqueue/test_lookup_too_many_matching_submissions"
+        producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+        inserted = 0
+        for _ in range(max_ + 1):
+            producer_client.insert_submission(
+                [1], chunk_size=1, strategic_metadata={"k": 1}
+            )
+            inserted += 1
+        with pytest.raises(TooManyMatchingSubmissionsError):
+            assert inserted == max_ + 1 # Make sure Exception wasn't raised too early.
+            producer_client.lookup_submission_ids_by_strategic_metadata({"k": 1})

opsqueue/src/common/errors.rs

@@ -49,6 +49,10 @@ pub enum SubmissionNotCancellable {
     Cancelled(SubmissionCancelled),
 }
 
+#[derive(Error, Debug, Deserialize, Serialize)]
+#[error("Too many submissions matched the lookup, the maximum is {0:?}")]
+pub struct TooManyMatchingSubmissions(pub u64);
+
 #[derive(Error, Debug)]
 #[error("Unexpected opsqueue consumer server response. This indicates an error inside Opsqueue itself: {0:?}")]
 pub struct UnexpectedOpsqueueConsumerServerResponse(pub SyncServerToClientResponse);

opsqueue/src/common/mod.rs

@@ -12,3 +12,50 @@ pub mod submission;
 /// consider hashing them and using that hash as MetaStateVal.
 pub type MetaStateVal = i64;
 pub type StrategicMetadataMap = FxHashMap<String, MetaStateVal>;
+
+/// Maximum number of submissions a lookup may return.
+/// Constructor checks: MaxSubmissions + 1 <= i64::MAX;
+#[derive(Debug, Clone, Copy)]
+pub struct MaxSubmissions(u64);
+
+impl MaxSubmissions {
+    pub fn new(value: u64) -> Result<Self, MaxSubmissionsTooLarge> {
+        if value < i64::MAX as u64 {
+            Ok(Self(value))
+        } else {
+            Err(MaxSubmissionsTooLarge(value))
+        }
+    }
+}
+
+impl From<MaxSubmissions> for u64 {
+    fn from(max_submissions: MaxSubmissions) -> u64 {
+        max_submissions.0
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("max_submissions value {0} is too large; it must be at most i64::MAX - 1")]
+pub struct MaxSubmissionsTooLarge(pub u64);
+
+impl std::fmt::Display for MaxSubmissions {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum ParseMaxSubmissionsError {
+    #[error(transparent)]
+    NotANumber(#[from] std::num::ParseIntError),
+    #[error(transparent)]
+    TooLarge(#[from] MaxSubmissionsTooLarge),
+}
+
+impl std::str::FromStr for MaxSubmissions {
+    type Err = ParseMaxSubmissionsError;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let value: u64 = s.parse()?;
+        Ok(MaxSubmissions::new(value)?)
+    }
+}