Skip to content

Commit

Permalink
update Rust toolchain to 1.79, add basic healthcheck API
Browse files Browse the repository at this point in the history
  • Loading branch information
auguwu committed Jul 23, 2024
1 parent c045fbe commit 74fe9f7
Show file tree
Hide file tree
Showing 5 changed files with 51 additions and 19 deletions.
21 changes: 19 additions & 2 deletions crates/s3/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,8 @@
use aws_sdk_s3::{
operation::{
create_bucket::CreateBucketError, delete_object::DeleteObjectError, get_object::GetObjectError,
head_object::HeadObjectError, list_buckets::ListBucketsError, list_objects_v2::ListObjectsV2Error,
put_object::PutObjectError,
head_bucket::HeadBucketError, head_object::HeadObjectError, list_buckets::ListBucketsError,
list_objects_v2::ListObjectsV2Error, put_object::PutObjectError,
},
primitives::SdkBody,
};
Expand Down Expand Up @@ -120,6 +120,10 @@ pub enum Error {
/// Occurs when an error occurred when transforming AWS S3's responses.
ByteStream(aws_sdk_s3::primitives::ByteStreamError),

/// Occurs when `remi-s3` cannot perform a HEAD request to the current bucket. This is mainly
/// used in healthchecks to determine if the storage service is ok.
HeadBucket(HeadBucketError),

/// Something that `remi-s3` has emitted on its own.
Library(Cow<'static, str>),
}
Expand Down Expand Up @@ -150,6 +154,7 @@ impl Display for Error {
E::ListBuckets(err) => Display::fmt(err, f),
E::ListObjectsV2(err) => Display::fmt(err, f),
E::PutObject(err) => Display::fmt(err, f),
E::HeadBucket(err) => Display::fmt(err, f),
E::Library(msg) => f.write_str(msg),
}
}
Expand Down Expand Up @@ -253,6 +258,18 @@ impl From<SdkError<PutObjectError, Response<SdkBody>>> for Error {
}
}

impl From<SdkError<HeadBucketError, Response<SdkBody>>> for Error {
fn from(value: SdkError<HeadBucketError, Response<SdkBody>>) -> Self {
match value {
SdkError::ConstructionFailure(err) => Self::ConstructionFailure(err),
SdkError::DispatchFailure(err) => Self::DispatchFailure(err),
SdkError::TimeoutError(err) => Self::TimeoutError(err),
SdkError::ResponseError(err) => Self::Response(err),
err => Error::HeadBucket(err.into_service_error()),
}
}
}

impl From<aws_sdk_s3::primitives::ByteStreamError> for Error {
fn from(value: aws_sdk_s3::primitives::ByteStreamError) -> Self {
Self::ByteStream(value)
Expand Down
26 changes: 20 additions & 6 deletions crates/s3/src/service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ use bytes::Bytes;
use remi::{Blob, Directory, File, ListBlobsRequest, UploadRequest};
use std::{borrow::Cow, path::Path};

const DEFAULT_CONTENT_TYPE: &str = "application/octet; charset=utf-8";
const DEFAULT_CONTENT_TYPE: &str = "application/octet-stream";

/// Represents an implementation of [`StorageService`] for Amazon Simple Storage Service.
#[derive(Debug, Clone)]
Expand Down Expand Up @@ -99,10 +99,6 @@ impl StorageService {

#[async_trait]
impl remi::StorageService for StorageService {
// this has to stay `io::Error` since `SdkError` requires too much information
// and this can narrow down.
//
// TODO(@auguwu): this can be a flat error if we could do?
type Error = crate::Error;

fn name(&self) -> Cow<'static, str> {
Expand Down Expand Up @@ -483,6 +479,7 @@ impl remi::StorageService for StorageService {

Ok(true)
}

Err(e) => {
let inner = e.into_service_error();
if inner.is_not_found() {
Expand Down Expand Up @@ -525,7 +522,7 @@ impl remi::StorageService for StorageService {

self.client
.put_object()
.bucket(self.config.bucket.clone())
.bucket(&self.config.bucket)
.key(normalized)
.acl(
self.config
Expand All @@ -545,6 +542,23 @@ impl remi::StorageService for StorageService {
.map(|_| ())
.map_err(From::from)
}

#[cfg_attr(feature = "tracing", tracing::instrument(name = "remi.s3.healthcheck", skip_all))]
async fn healthcheck(&self) -> Result<(), Self::Error> {
#[cfg(feature = "log")]
log::trace!("performing healthcheck...");

#[cfg(feature = "tracing")]
tracing::trace!("performing healthcheck...");

self.client
.head_bucket()
.bucket(&self.config.bucket)
.send()
.await
.map(|_| ())
.map_err(From::from)
}
}

#[cfg(test)]
Expand Down
15 changes: 6 additions & 9 deletions flake.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 5 additions & 1 deletion remi/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,6 @@ pub trait StorageService: Send + Sync {
/// Returns the name of the storage service.
///
/// * since 0.1.0
#[allow(deprecated)]
fn name(&self) -> Cow<'static, str>
where
Self: Sized;
Expand Down Expand Up @@ -115,6 +114,11 @@ pub trait StorageService: Send + Sync {
async fn upload<P: AsRef<Path> + Send>(&self, path: P, options: UploadRequest) -> Result<(), Self::Error>
where
Self: Sized;

/// Performs any healthchecks to determine the storage service's health.
async fn healthcheck(&self) -> Result<(), Self::Error> {
Ok(())
}
}

#[cfg(test)]
Expand Down
2 changes: 1 addition & 1 deletion rust-toolchain.toml
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
# SOFTWARE.

[toolchain]
channel = "1.78"
channel = "1.79"
profile = "minimal"
components = [
"rustc",
Expand Down

0 comments on commit 74fe9f7

Please sign in to comment.