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 5c6a7f0..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,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 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 7d4266c..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}, }; @@ -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, + 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 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 ad10541..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 @@ -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}) 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 e5998f6..d1706e9 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, Row, Sqlite}; use super::*; @@ -550,6 +552,66 @@ 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; + + // 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 LIMIT ?"#, + limit + ) + .fetch_all(conn.get_inner()) + .await?; + 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 + // 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 limit "); + query_builder.push_bind(limit); + let rows = query_builder.build().fetch_all(conn.get_inner()).await?; + error_if_too_many( + max_submissions, + 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/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 b09e20b..6bf8503 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 0298a2b..2788d07 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 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)