From 638806e67d0768d2c87400ec71905d2e0c8542ac Mon Sep 17 00:00:00 2001 From: "jeremy.barisch.rooney@channable.com" Date: Fri, 17 Jul 2026 09:54:53 +0200 Subject: [PATCH] Add ProducerClient.lookup_submission_ids_by_strategic_metadata --- Cargo.lock | 4 +- Cargo.toml | 2 +- .../python/opsqueue/exceptions.py | 28 +++++ .../python/opsqueue/producer.py | 21 ++++ libs/opsqueue_python/src/errors.rs | 9 +- libs/opsqueue_python/src/producer.rs | 30 +++++- libs/opsqueue_python/tests/benchmark.py | 1 + libs/opsqueue_python/tests/conftest.py | 2 + libs/opsqueue_python/tests/test_roundtrip.py | 92 +++++++++++++++- opsqueue/src/common/errors.rs | 4 + opsqueue/src/common/mod.rs | 48 +++++++++ opsqueue/src/common/submission.rs | 101 +++++++++++++++++- opsqueue/src/config.rs | 14 +++ opsqueue/src/producer/client.rs | 51 ++++++++- opsqueue/src/producer/server.rs | 41 ++++++- opsqueue/src/server.rs | 8 +- 16 files changed, 438 insertions(+), 18 deletions(-) create mode 100644 libs/opsqueue_python/tests/benchmark.py diff --git a/Cargo.lock b/Cargo.lock index 2b318677..236ed839 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2254,7 +2254,7 @@ dependencies = [ [[package]] name = "opsqueue" -version = "0.35.5" +version = "0.36.0" dependencies = [ "anyhow", "arc-swap", @@ -2309,7 +2309,7 @@ dependencies = [ [[package]] name = "opsqueue_python" -version = "0.35.5" +version = "0.36.0" dependencies = [ "anyhow", "chrono", diff --git a/Cargo.toml b/Cargo.toml index eb269fea..9be2ce9e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,7 @@ members = [ ] [workspace.package] -version = "0.35.5" +version = "0.36.0" [workspace.lints.clippy] cargo = { level = "warn", priority = -1 } diff --git a/libs/opsqueue_python/python/opsqueue/exceptions.py b/libs/opsqueue_python/python/opsqueue/exceptions.py index e1ae2bb3..615dc958 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 5c6a7f06..d72f0cd2 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", ] @@ -367,6 +369,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 diff --git a/libs/opsqueue_python/src/errors.rs b/libs/opsqueue_python/src/errors.rs index c3ddc2da..760c2d7e 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, E, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound, - UnexpectedOpsqueueConsumerServerResponse, + TooManyMatchingSubmissions, UnexpectedOpsqueueConsumerServerResponse, }; use pyo3::exceptions::PyBaseException; use pyo3::{Bound, PyErr, Python, import_exception}; @@ -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 c60ec56f..03dabc39 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -11,10 +11,9 @@ use futures::{StreamExt, TryStreamExt, stream::BoxStream}; use opsqueue::{ E, common::errors::E::{self, L, R}, - common::errors::{SubmissionNotCancellable, SubmissionNotFound}, + common::errors::{SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions}, common::{StrategicMetadataMap, chunk, submission}, - object_store::{ChunkRetrievalError, ChunkType}, - object_store::{ChunksStorageError, NewObjectStoreClientError}, + object_store::{ChunkRetrievalError, ChunkType, ChunksStorageError, NewObjectStoreClientError}, producer::ChunkContents, producer::client::{Client as ActualClient, InternalProducerClientError}, tracing::CarrierMap, @@ -188,6 +187,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, + E![ + FatalPythonException, + TooManyMatchingSubmissions, + InternalProducerClientError + ], + > { + py.detach(|| { + 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/benchmark.py b/libs/opsqueue_python/tests/benchmark.py new file mode 100644 index 00000000..e5a0d9b4 --- /dev/null +++ b/libs/opsqueue_python/tests/benchmark.py @@ -0,0 +1 @@ +#!/usr/bin/env python3 diff --git a/libs/opsqueue_python/tests/conftest.py b/libs/opsqueue_python/tests/conftest.py index 4b9a84b2..0052bca1 100644 --- a/libs/opsqueue_python/tests/conftest.py +++ b/libs/opsqueue_python/tests/conftest.py @@ -52,6 +52,7 @@ def opsqueue() -> Generator[OpsqueueProcess, None, None]: def opsqueue_service( *, port: int = 0, # The default of 0 means "pick any free port". + command_args: Iterable[str] = (), ) -> Generator[OpsqueueProcess, None, None]: # This will create a SQLite database in memory. # We need the `cache=shared` to allow sharing this DB between all threads within the same OS process. @@ -74,6 +75,7 @@ def opsqueue_service( str(write_fd), "--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 39505cfe..23d9a149 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 @@ -577,3 +578,92 @@ 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: list[SubmissionId] = [] + strategic_metadata = {"k": 1} + for _ in range(max_ + 1): + assert ( + inserted + == producer_client.lookup_submission_ids_by_strategic_metadata( + strategic_metadata + ) + ) + inserted.append( + producer_client.insert_submission( + [1], chunk_size=1, strategic_metadata=strategic_metadata + ) + ) + with pytest.raises(TooManyMatchingSubmissionsError) as exc: + assert len(inserted) == max_ + 1 + producer_client.lookup_submission_ids_by_strategic_metadata( + strategic_metadata + ) + assert exc.type is TooManyMatchingSubmissionsError + assert exc.value.max_submissions == max_ diff --git a/opsqueue/src/common/errors.rs b/opsqueue/src/common/errors.rs index 86d0086f..efc82dba 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:?}" diff --git a/opsqueue/src/common/mod.rs b/opsqueue/src/common/mod.rs index a6542749..4a305098 100644 --- a/opsqueue/src/common/mod.rs +++ b/opsqueue/src/common/mod.rs @@ -1,5 +1,6 @@ //! Common datatypes and errors shared across all parts of Opsqueue use rustc_hash::FxHashMap; +use std::num::NonZero; pub mod chunk; pub mod errors; @@ -12,3 +13,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. +/// Guarantees: 0 < MaxSubmissions 1 < i64::MAX; +#[derive(Debug, Clone, Copy)] +pub struct MaxSubmissions(NonZero); + +impl MaxSubmissions { + pub fn new(value: NonZero) -> Result { + if u64::from(value) < i64::MAX as u64 { + Ok(Self(value)) + } else { + Err(MaxSubmissionsTooLarge(value)) + } + } +} + +impl From for u64 { + fn from(max_submissions: MaxSubmissions) -> u64 { + u64::from(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 NonZero); + +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: NonZero = s.parse()?; + Ok(MaxSubmissions::new(value)?) + } +} diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index 3c2f9ee5..1ede03ec 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -270,15 +270,17 @@ impl Submission { pub mod db { use crate::{ common::{ - StrategicMetadataMap, - errors::{DatabaseError, E, SubmissionNotCancellable, SubmissionNotFound}, + MaxSubmissions, StrategicMetadataMap, + errors::{ + DatabaseError, E, SubmissionNotCancellable, SubmissionNotFound, + TooManyMatchingSubmissions, + }, }, db::{Connection, True, WriterConnection, WriterPool}, }; - use chunk::ChunkSize; - use sqlx::{Sqlite, query}; - use axum_prometheus::metrics::{counter, histogram}; + use chunk::ChunkSize; + use sqlx::{QueryBuilder, Sqlite, query, query_scalar}; use super::*; @@ -545,6 +547,59 @@ pub mod db { Ok(row.map(|row| row.id)) } + pub async fn lookup_ids_by_strategic_metadata( + strategic_metadata: StrategicMetadataMap, + max_submissions: MaxSubmissions, + mut conn: impl Connection, + ) -> Result, E> { + // MaxSubmissions provides us with the guarantee this won't overflow. + let limit = (u64::from(max_submissions) + 1) as i64; + // The main query to match on strategic_metadata will fail at run-time + // if strategic_metadata is empty, so we handle the empty case here. + let ids = if strategic_metadata.is_empty() { + query_scalar!( + r#"SELECT id AS "id: SubmissionId" FROM submissions ORDER BY id LIMIT ?"#, + limit + ) + .fetch_all(conn.get_inner()) + .await? + } else { + let mut query_builder: QueryBuilder = + lookup_ids_by_strategic_metadata_query(&strategic_metadata, limit); + query_builder + .build_query_scalar() + .fetch_all(conn.get_inner()) + .await? + }; + if ids.len() as u64 > u64::from(max_submissions) { + Err(E::R(TooManyMatchingSubmissions(u64::from(max_submissions)))) + } else { + Ok(ids) + } + } + + /// The query in 'lookup_ids_by_strategic_metadata', extracted for testing. + pub fn lookup_ids_by_strategic_metadata_query( + strategic_metadata: &StrategicMetadataMap, + limit: i64, + ) -> QueryBuilder<'_, Sqlite> { + let mut query_builder: QueryBuilder = + QueryBuilder::new("SELECT id FROM submissions"); + // Inner join for each piece of strategic metadata. + for (i, (key, value)) in strategic_metadata.iter().enumerate() { + query_builder.push(format!( + " INNER JOIN submissions_metadata AS s{i} ON s{i}.submission_id = submissions.id" + )); + query_builder.push(format!(" AND s{i}.metadata_key = ")); + query_builder.push_bind(key.clone()); + query_builder.push(format!(" AND s{i}.metadata_value = ")); + query_builder.push_bind(*value); + } + query_builder.push(" ORDER BY s0.submission_id LIMIT "); + query_builder.push_bind(limit); + query_builder + } + #[tracing::instrument(skip(conn))] pub async fn submission_status( id: SubmissionId, @@ -1036,6 +1091,7 @@ pub mod test { use chrono::Utc; use chunk::ChunkSize; use itertools::Itertools; + use sqlformat::{FormatOptions, QueryParams, format}; use sqlx::{Row, SqliteConnection}; use crate::common::StrategicMetadataMap; @@ -1070,6 +1126,41 @@ pub mod test { ); } + #[sqlx::test(migrator = "crate::MIGRATOR")] + pub async fn test_query_plan_lookup_by_strategic_metadata(db: sqlx::SqlitePool) { + let mut conn = db.acquire().await.unwrap(); + let strategic_metadata: StrategicMetadataMap = + [("company_id".to_string(), 1), ("project_id".to_string(), 2)] + .into_iter() + .collect(); + let qb = lookup_ids_by_strategic_metadata_query(&strategic_metadata, 100_000); + let options = FormatOptions::default(); + let formatted_query = format(qb.sql(), &QueryParams::None, &options); + insta::assert_snapshot!(formatted_query, @" + SELECT + id + FROM + submissions + INNER JOIN submissions_metadata AS s0 ON s0.submission_id = submissions.id + AND s0.metadata_key = ? + AND s0.metadata_value = ? + INNER JOIN submissions_metadata AS s1 ON s1.submission_id = submissions.id + AND s1.metadata_key = ? + AND s1.metadata_value = ? + ORDER BY + s0.submission_id + LIMIT + ? + "); + let explained = explain_query_plan(&formatted_query, &mut conn).await; + assert_non_regressing_query_plan(&formatted_query, &explained); + insta::assert_snapshot!(explained, @" + 8, 0, SEARCH s0 USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=?) + 16, 0, SEARCH submissions USING COVERING INDEX sqlite_autoindex_submissions_1 (id=?) + 21, 0, SEARCH s1 USING PRIMARY KEY (submission_id=? AND metadata_key=? AND metadata_value=?) + "); + } + #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_submission_status_in_progress(db: sqlx::SqlitePool) { let mut conn = db.acquire().await.unwrap(); diff --git a/opsqueue/src/config.rs b/opsqueue/src/config.rs index e4f171fa..92c2d693 100644 --- a/opsqueue/src/config.rs +++ b/opsqueue/src/config.rs @@ -12,6 +12,13 @@ use std::{ use clap::Parser; +use crate::common::MaxSubmissions; + +fn default_max_submissions_returned() -> MaxSubmissions { + MaxSubmissions::new(NonZero::new(100_000).expect("Non-zero u64")) + .expect("Valid MaxSubmissions default") +} + /// Making big work horizontally scalable. #[derive(Parser, Debug)] #[command(version, about, long_about = None)] @@ -93,6 +100,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 { @@ -109,6 +121,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, report_bound_port_pipe, @@ -119,6 +132,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 890459e3..b7c9983c 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -8,8 +8,9 @@ use http::StatusCode; use crate::{ E, common::{ + StrategicMetadataMap, errors::E::{L, R}, - errors::{SubmissionNotCancellable, SubmissionNotFound}, + errors::{SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions}, submission::{SubmissionId, SubmissionStatus}, }, tracing::CarrierMap, @@ -226,6 +227,54 @@ impl Client { .await } + pub async fn lookup_submission_ids_by_strategic_metadata( + &self, + strategic_metadata: &StrategicMetadataMap, + ) -> Result, E![TooManyMatchingSubmissions, InternalProducerClientError]> + { + (|| async { + let base_url = &self.base_url; + let response = self + .http_client + .post(format!( + "{base_url}/submissions/lookup_ids_by_strategic_metadata" + )) + .json(strategic_metadata) + .send() + .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(|e| match e { + L(_) => false, + R(client_err) => client_err.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 0298a2b0..2788d071 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::{MaxSubmissions, StrategicMetadataMap}; use crate::db::{self, DBPools}; +use axum::extract; use axum::extract::{Path, State}; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; @@ -13,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; } @@ -22,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) { @@ -60,6 +69,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 +146,30 @@ 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>, 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))] async fn insert_submission( State(state): State, diff --git a/opsqueue/src/server.rs b/opsqueue/src/server.rs index 2fbb4000..9f1224a4 100644 --- a/opsqueue/src/server.rs +++ b/opsqueue/src/server.rs @@ -94,8 +94,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)