diff --git a/Cargo.lock b/Cargo.lock index 30e6cc6a..6a02a4a0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2270,7 +2270,7 @@ dependencies = [ [[package]] name = "opsqueue" -version = "0.35.3" +version = "0.36.3" dependencies = [ "anyhow", "arc-swap", @@ -2325,7 +2325,7 @@ dependencies = [ [[package]] name = "opsqueue_python" -version = "0.35.3" +version = "0.36.3" dependencies = [ "anyhow", "chrono", diff --git a/Cargo.toml b/Cargo.toml index 736d84ab..76d0baf6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,7 @@ members = [ ] [workspace.package] -version = "0.35.3" +version = "0.36.3" [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 f0c68508..aafd96e3 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 5598172a..59e13e90 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}, }; @@ -190,6 +190,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.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/conftest.py b/libs/opsqueue_python/tests/conftest.py index 4b0ebb41..32ec8c68 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 ad105414..1f6ac50c 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 @@ -508,3 +509,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 f30436eb..f48fea34 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 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 5355f646..b2cd4166 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -269,15 +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::{query, Sqlite}; - use axum_prometheus::metrics::{counter, histogram}; + use chunk::ChunkSize; + use sqlx::{query, query_scalar, QueryBuilder, Sqlite}; use super::*; @@ -544,6 +546,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, @@ -1035,6 +1090,7 @@ pub mod test { use chrono::Utc; use chunk::ChunkSize; use itertools::Itertools; + use sqlformat::{format, FormatOptions, QueryParams}; use sqlx::{Row, SqliteConnection}; use crate::common::StrategicMetadataMap; @@ -1069,6 +1125,41 @@ pub mod test { ); } + #[sqlx::test] + 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] 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 08acd428..c899d2b3 100644 --- a/opsqueue/src/config.rs +++ b/opsqueue/src/config.rs @@ -6,6 +6,13 @@ use std::num::NonZero; 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)] @@ -77,6 +84,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 +104,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 +114,7 @@ impl Default for Config { max_missable_heartbeats, max_chunk_retries, max_submission_age, + max_submissions_returned, } } } diff --git a/opsqueue/src/consumer/dispatcher/metastate.rs b/opsqueue/src/consumer/dispatcher/metastate.rs index 8e743a3f..a3a6668b 100644 --- a/opsqueue/src/consumer/dispatcher/metastate.rs +++ b/opsqueue/src/consumer/dispatcher/metastate.rs @@ -112,6 +112,7 @@ impl MetaStateField { pub fn too_high_counts(&self, max: usize) -> impl Iterator + '_ { tracing::debug!("metastate: {self:?}"); self.counts_to_vals + // Negatives ? .range((max, 0)..) .map(|entry| entry.value().1) } diff --git a/opsqueue/src/producer/client.rs b/opsqueue/src/producer/client.rs index b09e20b0..6bf85036 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -8,8 +8,9 @@ use http::StatusCode; use crate::{ common::{ errors::E::{L, R}, - errors::{SubmissionNotCancellable, SubmissionNotFound}, + errors::{SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions}, submission::{SubmissionId, SubmissionStatus}, + StrategicMetadataMap, }, tracing::CarrierMap, E, @@ -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 cfc9ec4b..58c2bec9 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)