From ddee24389a93959c3613195009bb10cb7a2dc380 Mon Sep 17 00:00:00 2001 From: "jeremy.barisch.rooney@channable.com" Date: Mon, 8 Jun 2026 13:57:48 +0200 Subject: [PATCH 1/5] Add ProducerClient.lookup_submission_ids_by_strategic_metadata --- .../python/opsqueue/producer.py | 23 ++++++++ libs/opsqueue_python/src/producer.rs | 18 +++++++ libs/opsqueue_python/tests/test_roundtrip.py | 52 +++++++++++++++++++ opsqueue/src/common/submission.rs | 32 +++++++++++- opsqueue/src/producer/client.rs | 28 ++++++++++ opsqueue/src/producer/server.rs | 16 ++++++ 6 files changed, 167 insertions(+), 2 deletions(-) diff --git a/libs/opsqueue_python/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index 5c6a7f0..307a264 100644 --- a/libs/opsqueue_python/python/opsqueue/producer.py +++ b/libs/opsqueue_python/python/opsqueue/producer.py @@ -42,6 +42,10 @@ ] +class LookupIdsWithEmptyStrategicMetadataError(Exception): + pass + + class ProducerClient: """ Opsqueue producer client. Allows sending of large collections of operations ('submissions') @@ -367,6 +371,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: + - `LookupIdsWithEmptyStrategicMetadataError` if the provided + 'strategic_metadata' contained no key-value pairs to look for. + + """ + if len(strategic_metadata) == 0: + raise LookupIdsWithEmptyStrategicMetadataError() + 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 diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index 7d4266c..a0a70ed 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -189,6 +189,24 @@ 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, E> { + 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 diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index ad10541..ddb97d3 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -4,6 +4,7 @@ from collections.abc import Iterator, Sequence from opsqueue.producer import ( + LookupIdsWithEmptyStrategicMetadataError, SubmissionId, ProducerClient, SubmissionCompleted, @@ -508,3 +509,54 @@ 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 throw a + LookupIdsWithEmptyStrategicMetadataError. + + """ + url = "file:///tmp/opsqueue/test_lookup_submission_ids_by_empty_strategic_metadata" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + with pytest.raises(LookupIdsWithEmptyStrategicMetadataError): + producer_client.lookup_submission_ids_by_strategic_metadata({}) diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index e5998f6..ff31938 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -275,8 +275,7 @@ pub mod db { db::{Connection, True, WriterConnection, WriterPool}, }; use chunk::ChunkSize; - use sqlx::{query, Sqlite}; - + use sqlx::{query, query_as, QueryBuilder, Row, Sqlite}; use axum_prometheus::metrics::{counter, histogram}; use super::*; @@ -550,6 +549,35 @@ pub mod db { Ok(row.map(|row| row.id)) } + pub async fn lookup_ids_by_strategic_metadata( + strategic_metadata: StrategicMetadataMap, + mut conn: impl Connection, + ) -> Result, DatabaseError> { + // SQLx currently only supports "WHERE X IN (a, ...)" queries for postgres: + // https://github.com/transact-rs/sqlx/blob/main/FAQ.md#how-can-i-do-a-select--where-foo-in--query + // So we workaround this by manually building the query, foregoing + // sqlx's nice type-checking. + let mut query_builder: QueryBuilder = QueryBuilder::new( + " + SELECT submission_id + FROM submissions_metadata + INNER JOIN submissions on submissions.id = submission_id + WHERE (metadata_key, metadata_value) IN ( + ", + ); + query_builder.push_values(strategic_metadata.iter(), |mut b, sm| { + b.push_bind(sm.0).push_bind(sm.1); + }); + query_builder.push(") GROUP BY submission_id HAVING count(*) = "); + query_builder.push_bind(strategic_metadata.len() as i64); + query_builder.push(" ORDER BY submission_id"); + let rows = query_builder.build().fetch_all(conn.get_inner()).await?; + Ok(rows + .into_iter() + .map(|row| row.get("submission_id")) + .collect()) + } + #[tracing::instrument(skip(conn))] pub async fn submission_status( id: SubmissionId, diff --git a/opsqueue/src/producer/client.rs b/opsqueue/src/producer/client.rs index b09e20b..2bfa905 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -10,6 +10,7 @@ use crate::{ errors::E::{L, R}, errors::{SubmissionNotCancellable, SubmissionNotFound}, submission::{SubmissionId, SubmissionStatus}, + StrategicMetadataMap, }, tracing::CarrierMap, E, @@ -226,6 +227,33 @@ impl Client { .await } + pub async fn lookup_submission_ids_by_strategic_metadata( + &self, + strategic_metadata: &StrategicMetadataMap, + ) -> Result, InternalProducerClientError> { + (|| async { + let base_url = &self.base_url; + let resp = self + .http_client + .post(format!( + "{base_url}/submissions/lookup_ids_by_strategic_metadata" + )) + .json(strategic_metadata) + .send() + .await? + .error_for_status()?; + let bytes = resp.bytes().await?; + let body = serde_json::from_slice(&bytes)?; + Ok(body) + }) + .retry(retry_policy()) + .when(InternalProducerClientError::is_ephemeral) + .notify(|err, dur| { + tracing::debug!("retrying error {err:?} with sleeping {dur:?}"); + }) + .await + } + /// Get the server's version from the `/version` endpoint. /// /// A successful result will be the value of [`VERSION_CARGO_SEMVER`][crate::VERSION_CARGO_SEMVER] diff --git a/opsqueue/src/producer/server.rs b/opsqueue/src/producer/server.rs index 0298a2b..4986505 100644 --- a/opsqueue/src/producer/server.rs +++ b/opsqueue/src/producer/server.rs @@ -2,7 +2,9 @@ use std::sync::Arc; use crate::common::errors::E::{L, R}; use crate::common::submission::{self, SubmissionId}; +use crate::common::StrategicMetadataMap; use crate::db::{self, DBPools}; +use axum::extract; use axum::extract::{Path, State}; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; @@ -60,6 +62,10 @@ impl ServerState { "/submissions/lookup_id_by_prefix/{prefix}", get(lookup_submission_id_by_prefix), ) + .route( + "/submissions/lookup_ids_by_strategic_metadata", + post(lookup_submission_ids_by_strategic_metadata), + ) .route("/submissions/{submission_id}", get(submission_status)) .route("/version", get(crate::server::version_endpoint)) // We're also exposing it here so the producer client can view it .with_state(self) @@ -133,6 +139,16 @@ async fn lookup_submission_id_by_prefix( Ok(Json(submission_id)) } +async fn lookup_submission_ids_by_strategic_metadata( + State(state): State, + extract::Json(strategic_metadata): extract::Json, +) -> Result>, ServerError> { + let mut conn = state.pool.reader_conn().await?; + let submission_ids = + submission::db::lookup_ids_by_strategic_metadata(strategic_metadata, &mut conn).await?; + Ok(Json(submission_ids)) +} + #[tracing::instrument(level = "debug", skip(state))] async fn insert_submission( State(state): State, From 36606c16a2299b3b7b5572d34d0f8933973db707 Mon Sep 17 00:00:00 2001 From: "jeremy.barisch.rooney@channable.com" Date: Tue, 7 Jul 2026 13:55:45 +0200 Subject: [PATCH 2/5] fixup! Add ProducerClient.lookup_submission_ids_by_strategic_metadata Document InternalProducerClientError raised --- libs/opsqueue_python/python/opsqueue/producer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/libs/opsqueue_python/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index 307a264..340c1af 100644 --- a/libs/opsqueue_python/python/opsqueue/producer.py +++ b/libs/opsqueue_python/python/opsqueue/producer.py @@ -382,6 +382,7 @@ def lookup_submission_ids_by_strategic_metadata( Raises: - `LookupIdsWithEmptyStrategicMetadataError` if the provided 'strategic_metadata' contained no key-value pairs to look for. + - `InternalProducerClientError` if there is a low-level internal error. """ if len(strategic_metadata) == 0: From f329863fe9ccb194ef0ea9ffa0dd2825399ee182 Mon Sep 17 00:00:00 2001 From: "jeremy.barisch.rooney@channable.com" Date: Tue, 7 Jul 2026 14:42:57 +0200 Subject: [PATCH 3/5] fixup! Add ProducerClient.lookup_submission_ids_by_strategic_metadata Support empty strategic metadata --- libs/opsqueue_python/python/opsqueue/producer.py | 4 ---- libs/opsqueue_python/tests/test_roundtrip.py | 10 ++++++---- opsqueue/src/common/submission.rs | 13 ++++++++++++- 3 files changed, 18 insertions(+), 9 deletions(-) diff --git a/libs/opsqueue_python/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index 340c1af..11adfe7 100644 --- a/libs/opsqueue_python/python/opsqueue/producer.py +++ b/libs/opsqueue_python/python/opsqueue/producer.py @@ -380,13 +380,9 @@ def lookup_submission_ids_by_strategic_metadata( given key-value pairs, but it may also contain other key-value pairs. Raises: - - `LookupIdsWithEmptyStrategicMetadataError` if the provided - 'strategic_metadata' contained no key-value pairs to look for. - `InternalProducerClientError` if there is a low-level internal error. """ - if len(strategic_metadata) == 0: - raise LookupIdsWithEmptyStrategicMetadataError() return self.inner.lookup_submission_ids_by_strategic_metadata( # type: ignore[no-any-return] strategic_metadata ) diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index ddb97d3..9ea4f2f 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -552,11 +552,13 @@ def test_lookup( def test_lookup_submission_ids_by_empty_strategic_metadata( opsqueue: OpsqueueProcess, ) -> None: - """Lookup of submission IDs with empty strategic_metadata should throw a - LookupIdsWithEmptyStrategicMetadataError. + """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) - with pytest.raises(LookupIdsWithEmptyStrategicMetadataError): - producer_client.lookup_submission_ids_by_strategic_metadata({}) + 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 diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index ff31938..b16dccb 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -275,7 +275,7 @@ pub mod db { db::{Connection, True, WriterConnection, WriterPool}, }; use chunk::ChunkSize; - use sqlx::{query, query_as, QueryBuilder, Row, Sqlite}; + use sqlx::{QueryBuilder, Row, Sqlite, query, query_scalar}; use axum_prometheus::metrics::{counter, histogram}; use super::*; @@ -553,6 +553,17 @@ pub mod db { strategic_metadata: StrategicMetadataMap, mut conn: impl Connection, ) -> Result, DatabaseError> { + // The main query below to match on strategic_metadata will fail at + // run-time if strategic_metadata is empty (because the query will + // produce `IN ()`). So we handle the empty case here separately. + if strategic_metadata.is_empty() { + let ids = query_scalar!( + r#"SELECT id AS "id: SubmissionId" FROM submissions ORDER by id"#, + ) + .fetch_all(conn.get_inner()) + .await?; + return Ok(ids); + } // SQLx currently only supports "WHERE X IN (a, ...)" queries for postgres: // https://github.com/transact-rs/sqlx/blob/main/FAQ.md#how-can-i-do-a-select--where-foo-in--query // So we workaround this by manually building the query, foregoing From 54bc50bb974199124ffd4f53e136b1e4d4ad81ee Mon Sep 17 00:00:00 2001 From: "jeremy.barisch.rooney@channable.com" Date: Tue, 7 Jul 2026 17:47:56 +0200 Subject: [PATCH 4/5] fixup! Add ProducerClient.lookup_submission_ids_by_strategic_metadata Add configurable limit to submissions lookup --- justfile | 4 +- .../python/opsqueue/exceptions.py | 28 +++++++++++ .../python/opsqueue/producer.py | 5 ++ libs/opsqueue_python/src/errors.rs | 9 +++- libs/opsqueue_python/src/producer.rs | 11 ++++- libs/opsqueue_python/tests/test_roundtrip.py | 1 - opsqueue/src/common/errors.rs | 4 ++ opsqueue/src/common/mod.rs | 47 +++++++++++++++++++ opsqueue/src/common/submission.rs | 47 ++++++++++++++----- opsqueue/src/config.rs | 13 +++++ opsqueue/src/producer/client.rs | 39 +++++++++++---- opsqueue/src/producer/server.rs | 37 +++++++++++---- opsqueue/src/server.rs | 8 +++- 13 files changed, 216 insertions(+), 37 deletions(-) diff --git a/justfile b/justfile index ee1baa4..981d1d0 100644 --- a/justfile +++ b/justfile @@ -47,7 +47,7 @@ test-integration *TEST_ARGS: build-bin build-python cd libs/opsqueue_python source "./.setup_local_venv.sh" - timeout 600 pytest --color=yes {{TEST_ARGS}} + timeout 60 pytest --color=yes {{TEST_ARGS}} # Python integration test suite, using artefacts built through Nix. Args are forwarded to pytest [group('nix')] @@ -61,7 +61,7 @@ nix-test-integration *TEST_ARGS: nix-build-bin export OPSQUEUE_VIA_NIX=true export RUST_LOG="opsqueue=debug" - timeout 600 pytest --color=yes {{TEST_ARGS}} + timeout 60 pytest --color=yes {{TEST_ARGS}} # Run all linters, fast and slow [group('lint')] diff --git a/libs/opsqueue_python/python/opsqueue/exceptions.py b/libs/opsqueue_python/python/opsqueue/exceptions.py index e1ae2bb..615dc95 100644 --- a/libs/opsqueue_python/python/opsqueue/exceptions.py +++ b/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: diff --git a/libs/opsqueue_python/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index 11adfe7..7a93da5 100644 --- a/libs/opsqueue_python/python/opsqueue/producer.py +++ b/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,6 +39,7 @@ "SubmissionNotCancellable", "SubmissionNotCancellableError", "SubmissionNotFoundError", + "TooManyMatchingSubmissionsError", "ChunkFailed", ] @@ -380,6 +382,9 @@ def lookup_submission_ids_by_strategic_metadata( 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. """ diff --git a/libs/opsqueue_python/src/errors.rs b/libs/opsqueue_python/src/errors.rs index f0c6850..aafd96e 100644 --- a/libs/opsqueue_python/src/errors.rs +++ b/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> for PyErr { } } +impl From> for PyErr { + fn from(value: CError) -> Self { + TooManyMatchingSubmissionsError::new_err(value.0 .0) + } +} + pub struct SubmissionFailed( pub crate::common::SubmissionFailed, pub crate::common::ChunkFailed, diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index a0a70ed..1a45750 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/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}, }; @@ -195,7 +195,14 @@ impl ProducerClient { &self, py: Python<'_>, strategic_metadata: StrategicMetadataMap, - ) -> CPyResult, E> { + ) -> CPyResult< + Vec, + E![ + FatalPythonException, + TooManyMatchingSubmissions, + InternalProducerClientError + ], + > { py.allow_threads(|| { self.block_unless_interrupted(async { self.producer_client diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index 9ea4f2f..a5d0eb6 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -4,7 +4,6 @@ from collections.abc import Iterator, Sequence from opsqueue.producer import ( - LookupIdsWithEmptyStrategicMetadataError, SubmissionId, ProducerClient, SubmissionCompleted, diff --git a/opsqueue/src/common/errors.rs b/opsqueue/src/common/errors.rs index f30436e..f48fea3 100644 --- a/opsqueue/src/common/errors.rs +++ b/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); diff --git a/opsqueue/src/common/mod.rs b/opsqueue/src/common/mod.rs index a654274..65aadba 100644 --- a/opsqueue/src/common/mod.rs +++ b/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; + +/// 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 { + if value < i64::MAX as u64 { + Ok(Self(value)) + } else { + Err(MaxSubmissionsTooLarge(value)) + } + } +} + +impl From 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 { + let value: u64 = s.parse()?; + Ok(MaxSubmissions::new(value)?) + } +} diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index b16dccb..d1706e9 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -269,14 +269,17 @@ impl Submission { pub mod db { use crate::{ common::{ - errors::{DatabaseError, SubmissionNotCancellable, SubmissionNotFound, E}, - StrategicMetadataMap, + errors::{ + DatabaseError, SubmissionNotCancellable, SubmissionNotFound, + TooManyMatchingSubmissions, E, + }, + MaxSubmissions, StrategicMetadataMap, }, db::{Connection, True, WriterConnection, WriterPool}, }; - use chunk::ChunkSize; - use sqlx::{QueryBuilder, Row, Sqlite, query, query_scalar}; use axum_prometheus::metrics::{counter, histogram}; + use chunk::ChunkSize; + use sqlx::{query, query_scalar, QueryBuilder, Row, Sqlite}; use super::*; @@ -551,18 +554,35 @@ pub mod db { pub async fn lookup_ids_by_strategic_metadata( strategic_metadata: StrategicMetadataMap, + max_submissions: MaxSubmissions, mut conn: impl Connection, - ) -> Result, DatabaseError> { + ) -> Result, E> { + // MaxSubmissions provides us with the guarantee this won't overflow. + let limit = (u64::from(max_submissions) + 1) as i64; + + // If 'max_submissions' results is exceeded return an explicit error. + fn error_if_too_many( + max_submissions: MaxSubmissions, + ids: Vec, + ) -> Result, E> { + if ids.len() as u64 > u64::from(max_submissions) { + Err(E::R(TooManyMatchingSubmissions(u64::from(max_submissions)))) + } else { + Ok(ids) + } + } + // The main query below to match on strategic_metadata will fail at // run-time if strategic_metadata is empty (because the query will // produce `IN ()`). So we handle the empty case here separately. if strategic_metadata.is_empty() { let ids = query_scalar!( - r#"SELECT id AS "id: SubmissionId" FROM submissions ORDER by id"#, + r#"SELECT id AS "id: SubmissionId" FROM submissions ORDER by id LIMIT ?"#, + limit ) .fetch_all(conn.get_inner()) .await?; - return Ok(ids); + return error_if_too_many(max_submissions, ids); } // SQLx currently only supports "WHERE X IN (a, ...)" queries for postgres: // https://github.com/transact-rs/sqlx/blob/main/FAQ.md#how-can-i-do-a-select--where-foo-in--query @@ -581,12 +601,15 @@ pub mod db { }); query_builder.push(") GROUP BY submission_id HAVING count(*) = "); query_builder.push_bind(strategic_metadata.len() as i64); - query_builder.push(" ORDER BY submission_id"); + query_builder.push(" ORDER BY submission_id limit "); + query_builder.push_bind(limit); let rows = query_builder.build().fetch_all(conn.get_inner()).await?; - Ok(rows - .into_iter() - .map(|row| row.get("submission_id")) - .collect()) + error_if_too_many( + max_submissions, + rows.into_iter() + .map(|row| row.get("submission_id")) + .collect(), + ) } #[tracing::instrument(skip(conn))] diff --git a/opsqueue/src/config.rs b/opsqueue/src/config.rs index 08acd42..4af004c 100644 --- a/opsqueue/src/config.rs +++ b/opsqueue/src/config.rs @@ -6,6 +6,12 @@ use std::num::NonZero; use clap::Parser; +use crate::common::MaxSubmissions; + +fn default_max_submissions_returned() -> MaxSubmissions { + MaxSubmissions::new(100_000).expect("Valid MaxSubmissions default") +} + /// Making big work horizontally scalable. #[derive(Parser, Debug)] #[command(version, about, long_about = None)] @@ -77,6 +83,11 @@ pub struct Config { #[arg(long, default_value = "1 hour")] pub max_submission_age: humantime::Duration, + + /// Maximum number of submission IDs that a single + /// `lookup_submission_ids_by_strategic_metadata` request may return. + #[arg(long, default_value_t = default_max_submissions_returned())] + pub max_submissions_returned: MaxSubmissions, } impl Default for Config { @@ -92,6 +103,7 @@ impl Default for Config { let max_missable_heartbeats = 3; let max_chunk_retries = 10; let max_submission_age = humantime::Duration::from_str("1 hour").expect("valid humantime"); + let max_submissions_returned = default_max_submissions_returned(); Config { port, database_filename, @@ -101,6 +113,7 @@ impl Default for Config { max_missable_heartbeats, max_chunk_retries, max_submission_age, + max_submissions_returned, } } } diff --git a/opsqueue/src/producer/client.rs b/opsqueue/src/producer/client.rs index 2bfa905..6bf8503 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -8,7 +8,7 @@ use http::StatusCode; use crate::{ common::{ errors::E::{L, R}, - errors::{SubmissionNotCancellable, SubmissionNotFound}, + errors::{SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions}, submission::{SubmissionId, SubmissionStatus}, StrategicMetadataMap, }, @@ -230,24 +230,45 @@ impl Client { pub async fn lookup_submission_ids_by_strategic_metadata( &self, strategic_metadata: &StrategicMetadataMap, - ) -> Result, InternalProducerClientError> { + ) -> Result, E![TooManyMatchingSubmissions, InternalProducerClientError]> + { (|| async { let base_url = &self.base_url; - let resp = self + let response = self .http_client .post(format!( "{base_url}/submissions/lookup_ids_by_strategic_metadata" )) .json(strategic_metadata) .send() - .await? - .error_for_status()?; - let bytes = resp.bytes().await?; - let body = serde_json::from_slice(&bytes)?; - Ok(body) + .await + .map_err(|e| R(e.into()))?; + let status = response.status(); + match status { + // 200, the lookup succeeded. + StatusCode::OK => { + let submission_ids = response + .json::>() + .await + .map_err(|e| R(e.into()))?; + Ok(submission_ids) + } + // 400, matched more submissions than the configured maximum. + StatusCode::BAD_REQUEST => { + let too_many_err = response + .json::() + .await + .map_err(|e| R(e.into()))?; + Err(L(too_many_err)) + } + _ => Err(R(InternalProducerClientError::UnexpectedStatus(status))), + } }) .retry(retry_policy()) - .when(InternalProducerClientError::is_ephemeral) + .when(|e| match e { + L(_) => false, + R(client_err) => client_err.is_ephemeral(), + }) .notify(|err, dur| { tracing::debug!("retrying error {err:?} with sleeping {dur:?}"); }) diff --git a/opsqueue/src/producer/server.rs b/opsqueue/src/producer/server.rs index 4986505..2788d07 100644 --- a/opsqueue/src/producer/server.rs +++ b/opsqueue/src/producer/server.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use crate::common::errors::E::{L, R}; use crate::common::submission::{self, SubmissionId}; -use crate::common::StrategicMetadataMap; +use crate::common::{MaxSubmissions, StrategicMetadataMap}; use crate::db::{self, DBPools}; use axum::extract; use axum::extract::{Path, State}; @@ -15,7 +15,8 @@ use tokio::sync::Notify; use super::common::{ChunkContents, InsertSubmission}; pub async fn serve_for_tests(database_pool: DBPools, server_addr: Box) { - ServerState::new(database_pool, Arc::new(Notify::new())) + let max_submissions = crate::config::Config::default().max_submissions_returned; + ServerState::new(database_pool, Arc::new(Notify::new()), max_submissions) .serve_for_tests(server_addr) .await; } @@ -24,13 +25,19 @@ pub async fn serve_for_tests(database_pool: DBPools, server_addr: Box) { pub struct ServerState { pool: DBPools, notify_on_insert: Arc, + max_submissions: MaxSubmissions, } impl ServerState { - pub fn new(pool: DBPools, notify_on_insert: Arc) -> Self { + pub fn new( + pool: DBPools, + notify_on_insert: Arc, + max_submissions: MaxSubmissions, + ) -> Self { ServerState { pool, notify_on_insert, + max_submissions, } } pub async fn serve_for_tests(self, server_addr: Box) { @@ -139,14 +146,28 @@ async fn lookup_submission_id_by_prefix( Ok(Json(submission_id)) } +/// 200 if the query was successful. +/// 400 if the query exceeded the maximum lookup amount. async fn lookup_submission_ids_by_strategic_metadata( State(state): State, extract::Json(strategic_metadata): extract::Json, -) -> Result>, ServerError> { - let mut conn = state.pool.reader_conn().await?; - let submission_ids = - submission::db::lookup_ids_by_strategic_metadata(strategic_metadata, &mut conn).await?; - Ok(Json(submission_ids)) +) -> Result>, Response> { + let mut conn = state + .pool + .reader_conn() + .await + .map_err(|e| ServerError(e.into()).into_response())?; + match submission::db::lookup_ids_by_strategic_metadata( + strategic_metadata, + state.max_submissions, + &mut conn, + ) + .await + { + Ok(submission_ids) => Ok(Json(submission_ids)), + Err(L(db_err)) => Err(ServerError(db_err.into()).into_response()), + Err(R(too_many_err)) => Err((StatusCode::BAD_REQUEST, Json(too_many_err)).into_response()), + } } #[tracing::instrument(level = "debug", skip(state))] diff --git a/opsqueue/src/server.rs b/opsqueue/src/server.rs index cfc9ec4..58c2bec 100644 --- a/opsqueue/src/server.rs +++ b/opsqueue/src/server.rs @@ -80,8 +80,12 @@ pub fn build_router( ) .run_background() .build_router(); - let producer_routes = - crate::producer::server::ServerState::new(pool, notify_on_insert).build_router(); + let producer_routes = crate::producer::server::ServerState::new( + pool, + notify_on_insert, + config.max_submissions_returned, + ) + .build_router(); let routes = Router::new() .nest("/producer", producer_routes) From be5f6d2db02efddabab29228f63ef75d64c4be5a Mon Sep 17 00:00:00 2001 From: "jeremy.barisch.rooney@channable.com" Date: Tue, 7 Jul 2026 18:07:26 +0200 Subject: [PATCH 5/5] fixup! fixup! Add ProducerClient.lookup_submission_ids_by_strategic_metadata Test it --- justfile | 4 ++-- libs/opsqueue_python/tests/conftest.py | 3 ++- libs/opsqueue_python/tests/test_roundtrip.py | 24 +++++++++++++++++++- 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/justfile b/justfile index 981d1d0..ee1baa4 100644 --- a/justfile +++ b/justfile @@ -47,7 +47,7 @@ test-integration *TEST_ARGS: build-bin build-python cd libs/opsqueue_python source "./.setup_local_venv.sh" - timeout 60 pytest --color=yes {{TEST_ARGS}} + timeout 600 pytest --color=yes {{TEST_ARGS}} # Python integration test suite, using artefacts built through Nix. Args are forwarded to pytest [group('nix')] @@ -61,7 +61,7 @@ nix-test-integration *TEST_ARGS: nix-build-bin export OPSQUEUE_VIA_NIX=true export RUST_LOG="opsqueue=debug" - timeout 60 pytest --color=yes {{TEST_ARGS}} + timeout 600 pytest --color=yes {{TEST_ARGS}} # Run all linters, fast and slow [group('lint')] diff --git a/libs/opsqueue_python/tests/conftest.py b/libs/opsqueue_python/tests/conftest.py index 4b0ebb4..e6945ea 100644 --- a/libs/opsqueue_python/tests/conftest.py +++ b/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: diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index a5d0eb6..4db7636 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -14,6 +14,7 @@ SubmissionNotFoundError, SubmissionNotCancellable, SubmissionNotCancellableError, + TooManyMatchingSubmissionsError, ) from opsqueue.consumer import ConsumerClient, Chunk from opsqueue.common import SerializationFormat @@ -21,11 +22,11 @@ background_process, multiple_background_processes, OpsqueueProcess, + opsqueue_service, StrategyDescription, strategy_from_description, ) import logging - import pytest @@ -561,3 +562,24 @@ def test_lookup_submission_ids_by_empty_strategic_metadata( 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})