Skip to content

Commit a6fea95

Browse files
sitaowang1998sitaowang1998
authored
Merge branch 'main' into spider-compose-doc
2 parents 7466a2a + e6131a9 commit a6fea95

87 files changed

Lines changed: 2228 additions & 819 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎Cargo.lock‎

Lines changed: 284 additions & 534 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎components/api-server/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ path = "src/bin/api_server.rs"
1414
[dependencies]
1515
anyhow = "1.0.100"
1616
async-stream = "0.3.6"
17+
aws-sdk-s3 = "1.116.0"
1718
axum = { version = "0.8.6", features = ["json"] }
1819
clap = { version = "4.5.51", features = ["derive"] }
1920
clp-rust-utils = { path = "../clp-rust-utils" }

‎components/api-server/src/client.rs‎

Lines changed: 114 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ use std::pin::Pin;
22

33
use async_stream::stream;
44
use clp_rust_utils::{
5-
clp_config::package::{
6-
config::{Config, StorageEngine, StreamOutputStorage},
7-
credentials::Credentials,
5+
clp_config::{
6+
AwsAuthentication,
7+
package::{
8+
config::{Config, StorageEngine, StreamOutputStorage},
9+
credentials::Credentials,
10+
},
811
},
912
database::mysql::create_mysql_pool,
1013
job_config::{QUERY_JOBS_TABLE_NAME, QueryJobStatus, QueryJobType, SearchJobConfig},
@@ -165,6 +168,7 @@ impl Client {
165168
SearchResultStream<
166169
impl Stream<Item = Result<String, ClientError>> + use<>,
167170
impl Stream<Item = Result<String, ClientError>> + use<>,
171+
impl Stream<Item = Result<String, ClientError>> + use<>,
168172
>,
169173
ClientError,
170174
> {
@@ -186,11 +190,19 @@ impl Client {
186190
}
187191

188192
let job_config = self.get_job_config(search_job_id).await?;
193+
189194
if job_config.write_to_file {
190-
return Ok(SearchResultStream::File {
191-
inner: self.fetch_results_from_file(search_job_id),
192-
});
195+
let stream = match &self.config.stream_output.storage {
196+
StreamOutputStorage::Fs { .. } => SearchResultStream::File {
197+
inner: self.fetch_results_from_file(search_job_id),
198+
},
199+
StreamOutputStorage::S3 { .. } => SearchResultStream::S3 {
200+
inner: self.fetch_results_from_s3(search_job_id).await,
201+
},
202+
};
203+
return Ok(stream);
193204
}
205+
194206
self.fetch_results_from_mongo(search_job_id)
195207
.await
196208
.map(|s| SearchResultStream::Mongo { inner: s })
@@ -257,7 +269,7 @@ impl Client {
257269
search_job_id: u64,
258270
) -> impl Stream<Item = Result<String, ClientError>> + use<> {
259271
let StreamOutputStorage::Fs { directory } = &self.config.stream_output.storage else {
260-
todo!("Outputting query results to S3 is not supported for now.");
272+
unreachable!();
261273
};
262274
let search_job_output_dir = std::path::Path::new("/")
263275
.join(directory)
@@ -277,6 +289,92 @@ impl Client {
277289
stream
278290
}
279291

292+
/// Asynchronously fetches results of a completed search job from S3.
293+
///
294+
/// # Returns
295+
///
296+
/// A stream of the job's results on success. Each item in the stream is a [`Result`] that:
297+
///
298+
/// ## Returns
299+
///
300+
/// The log message in string representation on success.
301+
///
302+
/// ## Errors
303+
///
304+
/// Returns an error if:
305+
///
306+
/// * Forwards the pagination stream's error returned from S3 list-object-v2 operation's
307+
/// paginator (i.e., from
308+
/// [`aws_smithy_async::future::pagination_stream::PaginationStream::next`]).
309+
/// * Forwards S3 get-object operation's return values on failure (i.e., from
310+
/// [`aws_sdk_s3::operation::get_object::builders::GetObjectFluentBuilder::send`]).
311+
/// * Forwards [`aws_smithy_types::byte_stream::ByteStream::collect`]'s return values on
312+
/// failure.
313+
///
314+
/// # Panics
315+
///
316+
/// Panics if the stream output storage is not S3.
317+
async fn fetch_results_from_s3(
318+
&self,
319+
search_job_id: u64,
320+
) -> impl Stream<Item = Result<String, ClientError>> + use<> {
321+
tracing::info!("Streaming results from S3");
322+
let StreamOutputStorage::S3 { s3_config, .. } = &self.config.stream_output.storage else {
323+
unreachable!();
324+
};
325+
326+
let AwsAuthentication::Credentials { credentials } = &s3_config.aws_authentication;
327+
328+
let s3_config = s3_config.clone();
329+
let credentials = credentials.clone();
330+
331+
let s3_client = clp_rust_utils::s3::create_new_client(
332+
s3_config.region_code.as_str(),
333+
credentials.access_key_id.as_str(),
334+
credentials.secret_access_key.as_str(),
335+
None,
336+
)
337+
.await;
338+
339+
let key_prefix = format!("{}{}/", s3_config.key_prefix, search_job_id);
340+
tracing::info!("Streaming results from S3 prefix: {}", key_prefix);
341+
let mut object_pages = s3_client
342+
.list_objects_v2()
343+
.bucket(s3_config.bucket.as_str())
344+
.prefix(key_prefix)
345+
.into_paginator()
346+
.send();
347+
348+
stream! {
349+
while let Some(object_page) = object_pages.next().await {
350+
tracing::debug!("Received S3 object page: {:?}", object_page);
351+
for object in object_page?.contents() {
352+
let Some(key) = object.key() else {
353+
tracing::info!("S3 object {:?} doesn't have a key", object);
354+
continue;
355+
};
356+
if key.ends_with('/') {
357+
tracing::info!("Skipping S3 object with key {} as it is a directory", key);
358+
continue;
359+
}
360+
tracing::info!("Streaming results from S3 object with key: {}", key);
361+
let obj = s3_client
362+
.get_object()
363+
.bucket(s3_config.bucket.as_str())
364+
.key(key)
365+
.send()
366+
.await?;
367+
let bytes = obj.body.collect().await?.into_bytes();
368+
let mut deserializer = rmp_serde::Deserializer::from_read_ref(bytes.as_ref());
369+
while let Ok(event) = Deserialize::deserialize(&mut deserializer) {
370+
let event: (i64, String, String, String, i64) = event;
371+
yield Ok(event.1);
372+
}
373+
}
374+
}
375+
}
376+
}
377+
280378
/// Asynchronously fetches results of a completed search job from `MongoDB`.
281379
///
282380
/// # Returns
@@ -333,21 +431,26 @@ pin_project! {
333431
///
334432
/// * `FileStream`: Streaming from file system storage.
335433
/// * `MongoStream`: Streaming from MongoDB storage.
434+
/// * `S3Stream`: Streaming from S3 storage.
336435
#[project = SearchResultStreamProj]
337-
pub enum SearchResultStream<FileStream, MongoStream>
436+
pub enum SearchResultStream<FileStream, MongoStream, S3Stream>
338437
where
339438
FileStream: Stream<Item = Result<String, ClientError>>,
340-
MongoStream: Stream<Item = Result<String, ClientError>>
439+
MongoStream: Stream<Item = Result<String, ClientError>>,
440+
S3Stream: Stream<Item = Result<String, ClientError>>
341441
{
342442
File{#[pin] inner: FileStream},
343443
Mongo{#[pin] inner: MongoStream},
444+
S3{#[pin] inner: S3Stream},
344445
}
345446
}
346447

347-
impl<FileStream, MongoStream> Stream for SearchResultStream<FileStream, MongoStream>
448+
impl<FileStream, MongoStream, S3Stream> Stream
449+
for SearchResultStream<FileStream, MongoStream, S3Stream>
348450
where
349451
FileStream: Stream<Item = Result<String, ClientError>>,
350452
MongoStream: Stream<Item = Result<String, ClientError>>,
453+
S3Stream: Stream<Item = Result<String, ClientError>>,
351454
{
352455
type Item = Result<String, ClientError>;
353456

@@ -358,6 +461,7 @@ where
358461
let poll = match self.project() {
359462
SearchResultStreamProj::File { inner } => inner.poll_next(cx),
360463
SearchResultStreamProj::Mongo { inner } => inner.poll_next(cx),
464+
SearchResultStreamProj::S3 { inner } => inner.poll_next(cx),
361465
};
362466
if let std::task::Poll::Ready(Some(Err(err))) = &poll {
363467
tracing::error!("An error occurred when streaming results: {}", err);

‎components/api-server/src/error.rs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use aws_sdk_s3::{error::SdkError, primitives::ByteStreamError};
12
use num_enum::TryFromPrimitive;
23
use thiserror::Error;
34

@@ -17,6 +18,9 @@ pub enum ClientError {
1718

1819
#[error("Malformed data")]
1920
MalformedData,
21+
22+
#[error("`AwsError`: {description}")]
23+
Aws { description: String },
2024
}
2125

2226
/// Empty trait to mark errors that indicate malformed data.
@@ -36,6 +40,22 @@ impl<T: IsMalformedData> From<T> for ClientError {
3640
}
3741
}
3842

43+
impl<AwsSdkErrorType> From<SdkError<AwsSdkErrorType>> for ClientError {
44+
fn from(value: SdkError<AwsSdkErrorType>) -> Self {
45+
Self::Aws {
46+
description: value.to_string(),
47+
}
48+
}
49+
}
50+
51+
impl From<ByteStreamError> for ClientError {
52+
fn from(value: ByteStreamError) -> Self {
53+
Self::Aws {
54+
description: value.to_string(),
55+
}
56+
}
57+
}
58+
3959
impl From<clp_rust_utils::Error> for ClientError {
4060
fn from(value: clp_rust_utils::Error) -> Self {
4161
match value {

‎components/clp-mcp-server/pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ show_error_context = true
4747
show_error_end = true
4848

4949
[tool.ruff]
50-
extend = "../../tools/yscope-dev-utils/exports/lint-configs/python/ruff.toml"
50+
extend = "../../ruff.toml"
5151

5252
[tool.uv]
5353
default-groups = ["clp", "dev"]

‎components/clp-package-utils/pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ show_error_context = true
3838
show_error_end = true
3939

4040
[tool.ruff]
41-
extend = "../../tools/yscope-dev-utils/exports/lint-configs/python/ruff.toml"
41+
extend = "../../ruff.toml"
4242

4343
[tool.uv]
4444
default-groups = ["clp", "dev"]

‎components/clp-py-utils/pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ show_error_context = true
4141
show_error_end = true
4242

4343
[tool.ruff]
44-
extend = "../../tools/yscope-dev-utils/exports/lint-configs/python/ruff.toml"
44+
extend = "../../ruff.toml"
4545

4646
[tool.uv]
4747
default-groups = ["clp", "dev"]

‎components/clp-rust-utils/src/clp_config/package.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,3 +4,6 @@ pub mod credentials;
44
/// Mirror of configuration paths in `clp_py_utils.clp_config`.
55
pub const DEFAULT_CONFIG_FILE_PATH: &str = "etc/clp-config.yaml";
66
pub const DEFAULT_CREDENTIALS_FILE_PATH: &str = "etc/credentials.yaml";
7+
8+
/// Mirror of constants in `clp_py_utils.clp_config`.
9+
pub const DEFAULT_DATASET_NAME: &str = "default";

0 commit comments

Comments
 (0)