Merge branch 'main' into cancel-none-fix

Author: sitaowang1998

Cargo.lock

@@ -25,6 +25,7 @@ clp-rust-utils = { path = "../clp-rust-utils" }
 futures = "0.3.31"
 mongodb = "3.3.0"
 num_enum = "0.7.5"
+non-empty-string = { version = "0.2.6", features = ["serde"] }
 pin-project-lite = "0.2.16"
 rmp-serde = "1.3.0"
 secrecy = "0.10.3"

components/api-server/Cargo.toml

@@ -2,6 +2,7 @@ use std::pin::Pin;
 
 use async_stream::stream;
 use clp_rust_utils::{
+    aws::AWS_DEFAULT_REGION,
     clp_config::{
         AwsAuthentication,
         package::{
@@ -161,6 +162,7 @@ impl Client {
     /// * Forwards [`Client::get_status`]'s return values on failure.
     /// * Forwards [`Client::get_job_config`]'s return values on failure.
     /// * Forwards [`Client::fetch_results_from_mongo`]'s return values on failure.
+    /// * Forwards [`Client::fetch_results_from_s3`]'s return values on failure.
     pub async fn fetch_results(
         &self,
         search_job_id: u64,
@@ -197,7 +199,7 @@ impl Client {
                     inner: self.fetch_results_from_file(search_job_id),
                 },
                 StreamOutputStorage::S3 { .. } => SearchResultStream::S3 {
-                    inner: self.fetch_results_from_s3(search_job_id).await,
+                    inner: self.fetch_results_from_s3(search_job_id).await?,
                 },
             };
             return Ok(stream);
@@ -311,13 +313,20 @@ impl Client {
     /// * Forwards [`aws_smithy_types::byte_stream::ByteStream::collect`]'s return values on
     ///   failure.
     ///
+    /// # Errors
+    ///
+    /// Return an error if:
+    ///
+    /// * [`ClientError::Aws`] if a region code is not provided when using the default AWS S3
+    ///   endpoint.
+    ///
     /// # Panics
     ///
     /// Panics if the stream output storage is not S3.
     async fn fetch_results_from_s3(
         &self,
         search_job_id: u64,
-    ) -> impl Stream<Item = Result<String, ClientError>> + use<> {
+    ) -> Result<impl Stream<Item = Result<String, ClientError>> + use<>, ClientError> {
         tracing::info!("Streaming results from S3");
         let StreamOutputStorage::S3 { s3_config, .. } = &self.config.stream_output.storage else {
             unreachable!();
@@ -326,13 +335,22 @@ impl Client {
         let AwsAuthentication::Credentials { credentials } = &s3_config.aws_authentication;
 
         let s3_config = s3_config.clone();
-        let credentials = credentials.clone();
+        if s3_config.region_code.is_none() && s3_config.endpoint_url.is_none() {
+            return Err(ClientError::Aws {
+                description: "a region code must be given when using the default AWS S3 endpoint"
+                    .to_owned(),
+            });
+        }
 
+        let credentials = credentials.clone();
         let s3_client = clp_rust_utils::s3::create_new_client(
-            s3_config.region_code.as_str(),
             credentials.access_key_id.as_str(),
             credentials.secret_access_key.as_str(),
-            None,
+            s3_config
+                .region_code
+                .as_ref()
+                .map_or(AWS_DEFAULT_REGION, non_empty_string::NonEmptyString::as_str),
+            s3_config.endpoint_url.as_ref(),
         )
         .await;
 
@@ -345,7 +363,7 @@ impl Client {
             .into_paginator()
             .send();
 
-        stream! {
+        Ok(stream! {
             while let Some(object_page) = object_pages.next().await {
                 tracing::debug!("Received S3 object page: {:?}", object_page);
                 for object in object_page?.contents() {
@@ -372,7 +390,7 @@ impl Client {
                     }
                 }
             }
-        }
+        })
     }
 
     /// Asynchronously fetches results of a completed search job from `MongoDB`.

components/api-server/src/client.rs

@@ -0,0 +1 @@
+pub const AWS_DEFAULT_REGION: &str = "us-east-1";

components/clp-rust-utils/src/aws.rs

@@ -5,8 +5,9 @@ use serde::{Deserialize, Serialize};
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct S3Config {
     pub bucket: NonEmptyString,
-    pub region_code: String,
+    pub region_code: Option<NonEmptyString>,
     pub key_prefix: NonEmptyString,
+    pub endpoint_url: Option<NonEmptyString>,
     pub aws_authentication: AwsAuthentication,
 }
 

components/clp-rust-utils/src/clp_config/s3_config.rs

@@ -6,9 +6,6 @@ pub mod s3 {
     /// Base configuration for ingesting logs from S3.
     #[derive(Clone, Debug, Serialize, Deserialize, ToSchema)]
     pub struct BaseConfig {
-        /// AWS service region.
-        pub region: String,
-
         /// The S3 bucket to ingest from.
         #[schema(value_type = String, min_length = 1)]
         pub bucket_name: NonEmptyString,
@@ -17,6 +14,17 @@ pub mod s3 {
         #[schema(value_type = String, min_length = 1)]
         pub key_prefix: NonEmptyString,
 
+        /// AWS service region. Must be provided if using the default AWS S3 endpoint.
+        #[serde(default)]
+        #[schema(value_type = String, min_length = 1)]
+        pub region: Option<NonEmptyString>,
+
+        /// The endpoint URL for custom S3-compatible object stores (e.g., `MinIO`, `LocalStack`).
+        /// Use the default AWS S3 endpoint if not provided.
+        #[serde(default)]
+        #[schema(value_type = String, min_length = 1)]
+        pub endpoint_url: Option<NonEmptyString>,
+
         /// The dataset to ingest into. Defaults to `None` (which uses the default dataset).
         #[serde(default)]
         #[schema(value_type = String, min_length = 1)]

components/clp-rust-utils/src/job_config/ingestion.rs

@@ -1,3 +1,4 @@
+pub mod aws;
 pub mod clp_config;
 pub mod database;
 mod error;

components/clp-rust-utils/src/lib.rs

@@ -3,6 +3,7 @@ use aws_sdk_s3::{
     Client,
     config::{Builder, Credentials, Region},
 };
+use non_empty_string::NonEmptyString;
 
 /// Creates a new S3 client.
 ///
@@ -16,10 +17,10 @@ use aws_sdk_s3::{
 /// A newly created S3 client.
 #[must_use]
 pub async fn create_new_client(
-    region_id: &str,
     access_key_id: &str,
     secret_access_key: &str,
-    endpoint: Option<&str>,
+    region_id: &str,
+    endpoint: Option<&NonEmptyString>,
 ) -> Client {
     let credential = Credentials::new(
         access_key_id,
@@ -28,13 +29,12 @@ pub async fn create_new_client(
         None,
         "clp-credential-provider",
     );
-    let region = Region::new(region_id.to_owned());
     let base_config = aws_config::defaults(BehaviorVersion::latest()).load().await;
     let mut config_builder = Builder::from(&base_config)
-        .region(region)
         .credentials_provider(credential)
+        .region(Some(Region::new(region_id.to_string())))
         .force_path_style(true);
-    config_builder.set_endpoint_url(endpoint.map(std::borrow::ToOwned::to_owned));
+    config_builder.set_endpoint_url(endpoint.map(std::string::ToString::to_string));
     let config = config_builder.build();
     Client::from_conf(config)
 }

components/clp-rust-utils/src/s3/client.rs

@@ -3,6 +3,7 @@ use aws_sdk_sqs::{
     Client,
     config::{Builder, Credentials, Region},
 };
+use non_empty_string::NonEmptyString;
 
 /// Creates a new SQS client.
 /// The client is configured using the latest AWS SDK behavior version.
@@ -12,10 +13,10 @@ use aws_sdk_sqs::{
 /// A newly created SQS client.
 #[must_use]
 pub async fn create_new_client(
-    region_id: &str,
     access_key_id: &str,
     secret_access_key: &str,
-    endpoint: Option<&str>,
+    region_id: &str,
+    endpoint: Option<&NonEmptyString>,
 ) -> Client {
     let credential = Credentials::new(
         access_key_id,
@@ -24,11 +25,10 @@ pub async fn create_new_client(
         None,
         "clp-credential-provider",
     );
-    let region = Region::new(region_id.to_owned());
     let base_config = aws_config::defaults(BehaviorVersion::latest()).load().await;
     let mut config_builder = Builder::from(&base_config)
         .credentials_provider(credential)
-        .region(region);
-    config_builder.set_endpoint_url(endpoint.map(std::borrow::ToOwned::to_owned));
+        .region(Some(Region::new(region_id.to_string())));
+    config_builder.set_endpoint_url(endpoint.map(std::string::ToString::to_string));
     Client::from_conf(config_builder.build())
 }

components/clp-rust-utils/src/sqs/client.rs

@@ -11,8 +11,9 @@ use serde_json::Value;
 fn test_clp_io_config_serialization() {
     let s3_config = S3Config {
         bucket: NonEmptyString::from_static_str("yscope"),
-        region_code: "us-east-2".into(),
+        region_code: Some(NonEmptyString::from_static_str("us-east-2")),
         key_prefix: NonEmptyString::from_static_str("sample-logs/cockroachdb.clp.zst"),
+        endpoint_url: None,
         aws_authentication: AwsAuthentication::Credentials {
             credentials: AwsCredentials {
                 access_key_id: "ACCESS_KEY_ID".into(),
@@ -42,13 +43,12 @@ fn test_clp_io_config_serialization() {
 
     let brotli_compressed_msgpack = BrotliMsgpack::serialize(&config)
         .expect("Brotli-compressed MessagePack serialized config.");
-    let expected = "1ba00100e4ffdf9f43284b650e496850ba5f1eeefb53844a05d074faa66eb23ebef2dc45638\
-                275e9c24cb3bccba29c9bfc9d95db42175d52eecc81793cb3bc3c4ed0bf604c56e5c9a24581d9e65080\
-                1fd7263a8fb774fa362adf02eecc5b9d99532b8be8be173f6b659a9538c6c56a15571bc9856e20d0267\
-                b1591599975a75cdeb2aea30b83c8b486f3a2b3a74b419d6f99db0742a1482603a9480912e1336f2780\
-                dd9c3391503a9205a89a755bfe2c0d3a6be4c98ef0489c0b7e7f2d50b85f8f6e671a54d5dc6fa16d1ac\
-                cbaaffc5c3f1fb140f21ba0dce6ff0e8bc5f2da3c58426a9947046ca3cf9a06c7c8219e25a6ad0c4c67\
-                b6aceb8c88c782293b";
+    let expected = "1bae0100e4ffdf9f43284b650e496850ba5f98ef7b53044b04d074faa66eb23ebef25c0d1a13ac\
+        4b17669ac5cbe254cffc4e7423edc455d2f61c988e679697db09ba171ce575912c9a14986d0e05f87e4dd1babda\
+        9d5b5d1c5267067deea8cacf9928cf673f1a75a627945e06c7cbe6c6b21d98945e16df94ba0d5fce7c2d61158ca\
+        6541545e74dd5e3cddf21ea4832251d93cf6b3f69e3c05ab7eabc2dee3f305120988057210f19f59dc01ece49c0\
+        8f8e2812440d5acf6d2cfd2ae334791ec8888c43ae0f7870485fb756b76b4485c29fc164a27fd592de69feb5724\
+        16287f03545be2df6230185e931b8348a98a88804df492d3e01fd9c4b3c4b4b081e1cc970aaa2319b07cca0c";
     assert_eq!(expected, hex::encode(brotli_compressed_msgpack));
 
     let json_serialized_result = serde_json::to_string_pretty(&config);
@@ -60,6 +60,7 @@ fn test_clp_io_config_serialization() {
         "bucket": "yscope",
         "region_code": "us-east-2",
         "key_prefix": "sample-logs/cockroachdb.clp.zst",
+        "endpoint_url": null,
         "aws_authentication": {
           "type": "credentials",
           "credentials": {