use axum::body::Bytes;
#[cfg(test)]
use shardline_index::FileRecord;
use shardline_protocol::{ByteRange, RepositoryScope};
use shardline_storage::ObjectStore;
use tokio::task;
use super::LocalBackend;
use crate::{
ServerError, ShardMetadataLimits,
model::UploadFileResponse,
object_store::{read_full_object, reconstruct_file_record_bytes},
upload_ingest::RequestBodyReader,
validation::validate_identifier,
xet_adapter::{FileReconstructionResponse, ShardUploadResponse, build_reconstruction_response},
};
impl LocalBackend {
pub async fn upload_file(
&self,
file_id: &str,
body: Bytes,
repository_scope: Option<&RepositoryScope>,
) -> Result<UploadFileResponse, ServerError> {
self.upload_file_stream(
file_id,
RequestBodyReader::from_bytes(body),
repository_scope,
None,
)
.await
}
pub(crate) async fn upload_file_stream(
&self,
file_id: &str,
mut body: RequestBodyReader,
repository_scope: Option<&RepositoryScope>,
expected_sha256: Option<&str>,
) -> Result<UploadFileResponse, ServerError> {
validate_identifier(file_id)?;
let object_store = self.object_store();
let mut ingestor = crate::upload_ingest::FileUploadIngestor::new_with_parallelism(
self.chunk_size,
expected_sha256.is_some(),
self.upload_max_in_flight_chunks,
);
while let Some(bytes) = body.next_bytes().await? {
ingestor.ingest_body_chunk(&object_store, &bytes).await?;
}
let (record, response) = ingestor
.finish(&object_store, file_id, repository_scope, expected_sha256)
.await?;
self.record_store
.commit_file_version_metadata(&record)
.await?;
Ok(response)
}
pub(crate) async fn upload_shard_stream(
&self,
mut body: RequestBodyReader,
repository_scope: Option<&RepositoryScope>,
shard_metadata_limits: ShardMetadataLimits,
) -> Result<ShardUploadResponse, ServerError> {
let uploaded_body = crate::upload_ingest::read_body_to_bytes(&mut body).await?;
let record_store = self.record_store.clone();
let object_store = self.object_store();
crate::xet_adapter::register_uploaded_shard_bytes(
&object_store,
&uploaded_body,
repository_scope,
shard_metadata_limits,
move |records, mappings| async move {
record_store
.commit_native_shard_metadata(&records, &mappings)
.await?;
Ok(())
},
)
.await
.map_err(ServerError::from)
}
pub async fn reconstruction(
&self,
file_id: &str,
content_hash: Option<&str>,
requested_range: Option<ByteRange>,
repository_scope: Option<&RepositoryScope>,
) -> Result<FileReconstructionResponse, ServerError> {
let record = self
.read_record(file_id, content_hash, repository_scope)
.await?;
Ok(build_reconstruction_response(
self.public_base_url(),
&record,
requested_range,
)?)
}
pub async fn file_total_bytes(
&self,
file_id: &str,
content_hash: Option<&str>,
repository_scope: Option<&RepositoryScope>,
) -> Result<u64, ServerError> {
let record = self
.read_record(file_id, content_hash, repository_scope)
.await?;
Ok(record.total_bytes)
}
#[cfg(test)]
pub(crate) async fn file_record(
&self,
file_id: &str,
content_hash: Option<&str>,
repository_scope: Option<&RepositoryScope>,
) -> Result<FileRecord, ServerError> {
self.read_record(file_id, content_hash, repository_scope)
.await
}
pub async fn download_file(
&self,
file_id: &str,
content_hash: Option<&str>,
repository_scope: Option<&RepositoryScope>,
) -> Result<Vec<u8>, ServerError> {
let record = self
.read_record(file_id, content_hash, repository_scope)
.await?;
let object_store = self.object_store();
let server_frontends = self.server_frontends.clone();
task::spawn_blocking(move || {
reconstruct_file_record_bytes(&object_store, &server_frontends, &record)
})
.await
.map_err(ServerError::BlockingTask)?
}
pub async fn read_chunk(&self, hash_hex: &str) -> Result<Vec<u8>, ServerError> {
let object_store = self.object_store();
let object_key = crate::chunk_store::chunk_object_key(hash_hex)?;
let metadata = object_store.metadata(&object_key)?;
let Some(metadata) = metadata else {
return Err(ServerError::NotFound);
};
task::spawn_blocking(move || {
read_full_object(&object_store, &object_key, metadata.length())
})
.await
.map_err(ServerError::BlockingTask)?
}
pub async fn read_chunk_for_file_version(
&self,
hash_hex: &str,
file_id: &str,
content_hash: &str,
repository_scope: Option<&RepositoryScope>,
) -> Result<Vec<u8>, ServerError> {
let record = self
.read_record(file_id, Some(content_hash), repository_scope)
.await?;
if !record.chunks.iter().any(|chunk| chunk.hash == hash_hex) {
return Err(ServerError::NotFound);
}
self.read_chunk(hash_hex).await
}
pub async fn chunk_length(&self, hash_hex: &str) -> Result<u64, ServerError> {
let object_store = self.object_store();
let object_key = crate::chunk_store::chunk_object_key(hash_hex)?;
let metadata = object_store.metadata(&object_key)?;
let Some(metadata) = metadata else {
return Err(ServerError::NotFound);
};
Ok(metadata.length())
}
}