mirror of
https://github.com/hcengineering/platform.git
synced 2026-09-22 01:25:00 +02:00
move upload code to blob.rs
Signed-off-by: Alexey Aristov <aav@acm.org>
This commit is contained in:
Alexey Aristov
+143
@@ -0,0 +1,143 @@
|
||||
use actix_web::web::Payload;
|
||||
use aws_sdk_s3::types::{CompletedMultipartUpload, CompletedPart};
|
||||
use blake3::Hasher;
|
||||
use bytes::BytesMut;
|
||||
use futures_util::StreamExt;
|
||||
use tracing::*;
|
||||
|
||||
use crate::s3::S3Client;
|
||||
use crate::{
|
||||
config::CONFIG,
|
||||
postgres::{self, Pool},
|
||||
};
|
||||
|
||||
use crate::handlers::{ApiError, HandlerResult};
|
||||
|
||||
pub struct Blob {
|
||||
pub s3_key: String,
|
||||
pub size: u64,
|
||||
}
|
||||
|
||||
// upload and deduplicate blob
|
||||
#[instrument(level = "debug", skip_all, fields(s3_bucket, s3_key, upload))]
|
||||
pub async fn upload(s3: &S3Client, pool: &Pool, mut payload: Payload) -> Result<Blob, ApiError> {
|
||||
let span = Span::current();
|
||||
|
||||
let s3_bucket = &CONFIG.s3_bucket;
|
||||
let s3_key = ksuid::Ksuid::generate().to_base62();
|
||||
|
||||
span.record("s3_bucket", &s3_bucket);
|
||||
span.record("s3_key", &s3_key);
|
||||
|
||||
let create_multipart = s3
|
||||
.create_multipart_upload()
|
||||
.bucket(s3_bucket)
|
||||
.key(&s3_key)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let upload_id = create_multipart.upload_id().unwrap();
|
||||
|
||||
span.record("upload", &upload_id[upload_id.len().saturating_sub(16)..]);
|
||||
|
||||
debug!("upload start");
|
||||
|
||||
let upload_part = async |number, buffer: BytesMut| -> HandlerResult<CompletedPart> {
|
||||
let upload = s3
|
||||
.upload_part()
|
||||
.bucket(s3_bucket)
|
||||
.key(&s3_key)
|
||||
.upload_id(upload_id)
|
||||
.body(buffer.freeze().into())
|
||||
.part_number(number)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let part = CompletedPart::builder()
|
||||
.e_tag(upload.e_tag.unwrap())
|
||||
.part_number(number)
|
||||
.build();
|
||||
|
||||
Ok(part)
|
||||
};
|
||||
|
||||
let mut buffer = BytesMut::with_capacity(1024 * 1024 * 6);
|
||||
let mut complete = CompletedMultipartUpload::builder();
|
||||
let mut part_number = 1;
|
||||
let mut hash = Hasher::new();
|
||||
let mut total_in = 0;
|
||||
let mut total_uploaded = 0;
|
||||
|
||||
while let Some(part) = payload.next().await {
|
||||
if let Ok(part) = part {
|
||||
hash.update(&part);
|
||||
|
||||
total_in += part.len();
|
||||
|
||||
buffer.extend_from_slice(&part);
|
||||
|
||||
// each part must be at least 5MB
|
||||
if buffer.len() > 1024 * 1024 * 5 {
|
||||
trace!(length = buffer.len(), part_number, "upload part");
|
||||
|
||||
total_uploaded += buffer.len();
|
||||
|
||||
let uploaded = upload_part(part_number, buffer).await?;
|
||||
|
||||
complete = complete.parts(uploaded);
|
||||
|
||||
buffer = BytesMut::new();
|
||||
part_number += 1;
|
||||
}
|
||||
} else {
|
||||
// TODO: cleanup incomplete upload
|
||||
panic!("read error")
|
||||
}
|
||||
}
|
||||
|
||||
// the last part
|
||||
if buffer.len() > 0 {
|
||||
total_uploaded += buffer.len();
|
||||
|
||||
trace!(length = buffer.len(), part_number, "upload part");
|
||||
let uploaded = upload_part(part_number, buffer).await?;
|
||||
complete = complete.parts(uploaded);
|
||||
}
|
||||
|
||||
assert_eq!(total_in, total_uploaded);
|
||||
|
||||
let _ = s3
|
||||
.complete_multipart_upload()
|
||||
.bucket(s3_bucket)
|
||||
.key(&s3_key)
|
||||
.multipart_upload(complete.build())
|
||||
.upload_id(upload_id)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let hash = hash.finalize().to_hex().to_string();
|
||||
|
||||
debug!(hash, "upload complete");
|
||||
|
||||
let s3_key = if let Some(s3_key_found) = postgres::find_blob_by_hash(&pool, &hash).await? {
|
||||
debug!(s3_key_found, "blob deduplicated");
|
||||
|
||||
// delete uploaded
|
||||
s3.delete_object()
|
||||
.bucket(s3_bucket)
|
||||
.key(s3_key)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
s3_key_found
|
||||
} else {
|
||||
debug!("blob created");
|
||||
postgres::insert_blob(&pool, &s3_key, &hash).await?;
|
||||
s3_key
|
||||
};
|
||||
|
||||
Ok(Blob {
|
||||
s3_key,
|
||||
size: total_uploaded as u64,
|
||||
})
|
||||
}
|
||||
+5
-137
@@ -10,22 +10,19 @@ use actix_web::{
|
||||
},
|
||||
web::{Data, Header, Path, Payload},
|
||||
};
|
||||
use blake3::Hasher;
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use futures_util::{Stream, StreamExt};
|
||||
use aws_sdk_s3::error::SdkError;
|
||||
use bytes::Bytes;
|
||||
use futures_util::Stream;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::*;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::blob::upload;
|
||||
use crate::s3::S3Client;
|
||||
use crate::{
|
||||
config::CONFIG,
|
||||
postgres::{self, Pool},
|
||||
};
|
||||
use aws_sdk_s3::{
|
||||
error::SdkError,
|
||||
types::{CompletedMultipartUpload, CompletedPart},
|
||||
};
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub struct ObjectPath {
|
||||
@@ -48,7 +45,7 @@ pub enum ApiError {
|
||||
Other(#[from] anyhow::Error),
|
||||
}
|
||||
|
||||
type HandlerResult<T> = Result<T, ApiError>;
|
||||
pub type HandlerResult<T> = Result<T, ApiError>;
|
||||
|
||||
impl actix_web::error::ResponseError for ApiError {
|
||||
fn error_response(&self) -> HttpResponse {
|
||||
@@ -73,11 +70,6 @@ impl<E: Display + std::error::Error + 'static, B: std::fmt::Debug> From<SdkError
|
||||
}
|
||||
}
|
||||
|
||||
struct Blob {
|
||||
s3_key: String,
|
||||
size: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize, serde::Deserialize, Debug)]
|
||||
struct PartData {
|
||||
workspace: Uuid,
|
||||
@@ -94,130 +86,6 @@ struct PartData {
|
||||
meta: Option<HashMap<String, String>>,
|
||||
}
|
||||
|
||||
// upload and deduplicate blob
|
||||
#[instrument(level = "debug", skip_all, fields(s3_bucket, s3_key, upload))]
|
||||
async fn upload(s3: &S3Client, pool: &Pool, mut payload: Payload) -> Result<Blob, ApiError> {
|
||||
let span = Span::current();
|
||||
|
||||
let s3_bucket = &CONFIG.s3_bucket;
|
||||
let s3_key = ksuid::Ksuid::generate().to_base62();
|
||||
|
||||
span.record("s3_bucket", &s3_bucket);
|
||||
span.record("s3_key", &s3_key);
|
||||
|
||||
let create_multipart = s3
|
||||
.create_multipart_upload()
|
||||
.bucket(s3_bucket)
|
||||
.key(&s3_key)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let upload_id = create_multipart.upload_id().unwrap();
|
||||
|
||||
span.record("upload", &upload_id[upload_id.len().saturating_sub(16)..]);
|
||||
|
||||
debug!("upload start");
|
||||
|
||||
let upload_part = async |number, buffer: BytesMut| -> HandlerResult<CompletedPart> {
|
||||
let upload = s3
|
||||
.upload_part()
|
||||
.bucket(s3_bucket)
|
||||
.key(&s3_key)
|
||||
.upload_id(upload_id)
|
||||
.body(buffer.freeze().into())
|
||||
.part_number(number)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let part = CompletedPart::builder()
|
||||
.e_tag(upload.e_tag.unwrap())
|
||||
.part_number(number)
|
||||
.build();
|
||||
|
||||
Ok(part)
|
||||
};
|
||||
|
||||
let mut buffer = BytesMut::with_capacity(1024 * 1024 * 6);
|
||||
let mut complete = CompletedMultipartUpload::builder();
|
||||
let mut part_number = 1;
|
||||
let mut hash = Hasher::new();
|
||||
let mut total_in = 0;
|
||||
let mut total_uploaded = 0;
|
||||
|
||||
while let Some(part) = payload.next().await {
|
||||
if let Ok(part) = part {
|
||||
hash.update(&part);
|
||||
|
||||
total_in += part.len();
|
||||
|
||||
buffer.extend_from_slice(&part);
|
||||
|
||||
// each part must be at least 5MB
|
||||
if buffer.len() > 1024 * 1024 * 5 {
|
||||
trace!(length = buffer.len(), part_number, "upload part");
|
||||
|
||||
total_uploaded += buffer.len();
|
||||
|
||||
let uploaded = upload_part(part_number, buffer).await?;
|
||||
|
||||
complete = complete.parts(uploaded);
|
||||
|
||||
buffer = BytesMut::new();
|
||||
part_number += 1;
|
||||
}
|
||||
} else {
|
||||
// TODO: cleanup incomplete upload
|
||||
panic!("read error")
|
||||
}
|
||||
}
|
||||
|
||||
// the last part
|
||||
if buffer.len() > 0 {
|
||||
total_uploaded += buffer.len();
|
||||
|
||||
trace!(length = buffer.len(), part_number, "upload part");
|
||||
let uploaded = upload_part(part_number, buffer).await?;
|
||||
complete = complete.parts(uploaded);
|
||||
}
|
||||
|
||||
assert_eq!(total_in, total_uploaded);
|
||||
|
||||
let _ = s3
|
||||
.complete_multipart_upload()
|
||||
.bucket(s3_bucket)
|
||||
.key(&s3_key)
|
||||
.multipart_upload(complete.build())
|
||||
.upload_id(upload_id)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let hash = hash.finalize().to_hex().to_string();
|
||||
|
||||
debug!(hash, "upload complete");
|
||||
|
||||
let s3_key = if let Some(s3_key_found) = postgres::find_blob_by_hash(&pool, &hash).await? {
|
||||
debug!(s3_key_found, "blob deduplicated");
|
||||
|
||||
// delete uploaded
|
||||
s3.delete_object()
|
||||
.bucket(s3_bucket)
|
||||
.key(s3_key)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
s3_key_found
|
||||
} else {
|
||||
debug!("blob created");
|
||||
postgres::insert_blob(&pool, &s3_key, &hash).await?;
|
||||
s3_key
|
||||
};
|
||||
|
||||
Ok(Blob {
|
||||
s3_key,
|
||||
size: total_uploaded as u64,
|
||||
})
|
||||
}
|
||||
|
||||
#[instrument(level = "debug", skip_all, fields(workspace, huly_key))]
|
||||
pub async fn put(request: HttpRequest, payload: Payload) -> HandlerResult<HttpResponse> {
|
||||
let span = Span::current();
|
||||
|
||||
@@ -14,6 +14,7 @@ use uuid::Uuid;
|
||||
|
||||
use hulyrs::services::jwt::actix::ServiceRequestExt;
|
||||
|
||||
mod blob;
|
||||
mod config;
|
||||
mod handlers;
|
||||
mod postgres;
|
||||
|
||||
Reference in New Issue
Block a user