use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use bytes::Bytes;
use delta_kernel::object_store::path::Path;
use delta_kernel::object_store::{
Attributes, CopyOptions, Error as ObjectStoreError, GetOptions, GetRange, GetResult,
GetResultPayload, ListResult, MultipartUpload, ObjectMeta, ObjectStore, PutMode,
PutMultipartOptions, PutOptions, PutPayload, PutResult, Result as ObjectStoreResult,
};
use futures::stream::BoxStream;
use futures::StreamExt as _;
use reqwest::header::HeaderMap;
use reqwest::Client;
use tracing::{info, warn};
use super::auth::AuthHeaderProvider;
use super::config::RestEndpointConfig;
use super::generic_error;
use super::response::{
classify_put_create_response, ensure_list_response, ensure_success_response,
is_retryable_http_status, is_transient, PutCreateDisposition, RetryFailure,
};
#[derive(Debug, Clone)]
pub struct RestObjectStore {
base_url: String,
client: Client,
auth: Arc<dyn AuthHeaderProvider>,
config: Arc<RestEndpointConfig>,
max_retries: u32,
verify_on_ambiguous: bool,
}
impl std::fmt::Display for RestObjectStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "RestObjectStore({})", self.base_url)
}
}
impl RestObjectStore {
pub fn new(
base_url: impl Into<String>,
client: Client,
auth: Arc<dyn AuthHeaderProvider>,
config: Arc<RestEndpointConfig>,
) -> Self {
Self {
base_url: base_url.into(),
client,
auth,
config,
max_retries: 0,
verify_on_ambiguous: false,
}
}
pub fn with_max_retries(mut self, n: u32) -> Self {
self.max_retries = n;
self
}
pub fn with_verify_on_ambiguous(mut self, verify: bool) -> Self {
self.verify_on_ambiguous = verify;
self
}
fn headers(&self) -> ObjectStoreResult<HeaderMap> {
self.auth.headers()
}
async fn send_idempotent(
&self,
target: &str,
make: impl Fn(&Client, HeaderMap) -> reqwest::RequestBuilder,
) -> ObjectStoreResult<reqwest::Response> {
let mut retries = 0u32;
let mut last_failure = None::<RetryFailure>;
loop {
let exhausted = retries >= self.max_retries;
let headers = self.headers()?;
match make(&self.client, headers).send().await {
Ok(resp) if !exhausted && is_retryable_http_status(resp.status()) => {
last_failure = Some(RetryFailure::ServerError(resp.status()));
}
Ok(resp) => {
if let Some(failure) = last_failure.as_ref().filter(|_| retries > 0) {
log_retry_outcome(target, retries, failure, true);
}
return Ok(resp);
}
Err(e) if !exhausted && is_transient(&e) => {
last_failure = Some(RetryFailure::Transport(e.to_string()));
}
Err(e) => {
if retries > 0 {
let transport_failure = RetryFailure::Transport(e.to_string());
let failure = last_failure.as_ref().unwrap_or(&transport_failure);
log_retry_outcome(target, retries, failure, false);
}
return Err(generic_error(e));
}
}
retries += 1;
backoff(retries).await;
}
}
async fn put_create_verified(
&self,
path: &str,
url: &str,
query: &[(String, String)],
body: Bytes,
) -> ObjectStoreResult<PutResult> {
let mut retries = 0u32;
let mut last_failure = None::<RetryFailure>;
loop {
let disposition = classify_put_create_response(
self.client
.put(url)
.query(query)
.headers(self.headers()?)
.body(body.clone())
.send()
.await,
path,
retries,
)?;
match disposition {
PutCreateDisposition::Success => {
if let Some(failure) = last_failure.as_ref().filter(|_| retries > 0) {
log_retry_outcome(path, retries, failure, true);
}
return Ok(put_result());
}
PutCreateDisposition::Terminal(err) => return Err(err),
PutCreateDisposition::Ambiguous { failure } => last_failure = Some(failure),
}
match self.read_back(path, &body).await {
Ok(WriteState::Matches) => {
if let Some(failure) = last_failure.as_ref().filter(|_| retries > 0) {
log_retry_outcome(path, retries, failure, true);
}
return Ok(put_result());
}
Ok(WriteState::Differs) => {
return Err(ObjectStoreError::AlreadyExists {
path: path.to_string(),
source: "verified conflicting write".into(),
})
}
Ok(WriteState::Absent) => {}
Err(e) if retries >= self.max_retries => {
let transport_failure = RetryFailure::Transport(e.to_string());
let failure = last_failure.as_ref().unwrap_or(&transport_failure);
log_retry_outcome(path, retries, failure, false);
return Err(e);
}
Err(e) => {
last_failure = Some(RetryFailure::Transport(e.to_string()));
}
}
if retries >= self.max_retries {
if let Some(failure) = &last_failure {
log_retry_outcome(path, retries, failure, false);
}
return Err(generic_error(format!(
"exceeded max retries ({}) during put for `{path}` without confirming the write landed",
self.max_retries
)));
}
retries += 1;
backoff(retries).await;
}
}
async fn read_back(&self, path: &str, expected: &Bytes) -> ObjectStoreResult<WriteState> {
match self.get_file(path, None).await {
Ok((bytes, _, _)) if bytes == *expected => Ok(WriteState::Matches),
Ok(_) => Ok(WriteState::Differs),
Err(ObjectStoreError::NotFound { .. }) => Ok(WriteState::Absent),
Err(e) => Err(e),
}
}
async fn delete_one(&self, location: &Path) -> ObjectStoreResult<()> {
let path = location.as_ref();
let url = self.config.file_url(&self.base_url, path);
let response = self
.send_idempotent(path, |c, h| c.delete(&url).headers(h))
.await?;
ensure_success_response(response, path)?;
Ok(())
}
async fn get_file(
&self,
path: &str,
range_header: Option<&str>,
) -> ObjectStoreResult<(Bytes, HeaderMap, reqwest::StatusCode)> {
let url = self.config.file_url(&self.base_url, path);
let range = range_header
.map(reqwest::header::HeaderValue::from_str)
.transpose()
.map_err(generic_error)?;
let response = self
.send_idempotent(path, |c, mut h| {
if let Some(v) = &range {
h.insert(reqwest::header::RANGE, v.clone());
}
c.get(&url).headers(h)
})
.await?;
let response = ensure_success_response(response, path)?;
let status = response.status();
let resp_headers = response.headers().clone();
let body = response.bytes().await.map_err(generic_error)?;
Ok((body, resp_headers, status))
}
async fn head_meta(&self, path: &str, location: &Path) -> ObjectStoreResult<ObjectMeta> {
let url = self.config.file_url(&self.base_url, path);
let response = self
.send_idempotent(path, |c, h| c.head(&url).headers(h))
.await?;
let response = ensure_success_response(response, path)?;
let headers = response.headers();
let size = match headers.get(reqwest::header::CONTENT_LENGTH) {
None => {
return Err(generic_error(format!(
"HEAD for `{path}` is missing a Content-Length header"
)));
}
Some(v) => {
let s = v
.to_str()
.map_err(|e| generic_error(format!("invalid Content-Length header: {e}")))?;
s.parse::<u64>().map_err(|e| {
generic_error(format!("invalid Content-Length header `{s}`: {e}"))
})?
}
};
let last_modified = parse_last_modified(headers);
let e_tag = parse_etag(headers);
Ok(ObjectMeta {
location: location.clone(),
last_modified,
size,
e_tag,
version: None,
})
}
fn list_paginated(
&self,
prefix: String,
start_from: Option<String>,
exclusive_offset: Option<Path>,
recursive: bool,
) -> BoxStream<'static, ObjectStoreResult<ObjectMeta>> {
let store = self.clone();
let stream = async_stream::stream! {
let mut page_token: Option<String> = None;
let mut start_from = start_from;
let mut last_path: Option<Path> = None;
'pages: loop {
let url = store.config.directory_url(&store.base_url, &prefix);
let query = store.config.list_query(
page_token.as_deref(),
start_from.as_deref(),
recursive,
);
start_from = None;
let list_target = if prefix.is_empty() {
"list".to_string()
} else {
format!("list `{prefix}`")
};
let response = match store
.send_idempotent(&list_target, |c, h| {
c.get(url.as_str()).query(&query).headers(h)
})
.await
{
Ok(r) => r,
Err(e) => { yield Err(e); break; }
};
let response = match ensure_list_response(
response,
&prefix,
page_token.as_deref(),
) {
Ok(Some(r)) => r,
Ok(None) => break,
Err(e) => { yield Err(e); break; }
};
let body = match response.bytes().await {
Ok(b) => b,
Err(e) => { yield Err(generic_error(e)); break; }
};
let page = match store.config.parse_list(&body) {
Ok(p) => p,
Err(e) => { yield Err(e); break; }
};
for meta in page.objects {
if let Some(off) = &exclusive_offset {
if meta.location <= *off {
continue;
}
}
if let Some(last) = &last_path {
if meta.location < *last {
yield Err(generic_error(format!(
"REST listing returned out-of-order entry `{}` after `{}`; \
RestEndpointConfig must return lexicographically sorted paths",
meta.location, last
)));
break 'pages;
}
}
last_path = Some(meta.location.clone());
yield Ok(meta);
}
match page.next_page_token {
Some(token) => page_token = Some(token),
None => break,
}
}
};
Box::pin(stream)
}
}
fn get_range_to_header(range: &GetRange) -> String {
match range {
GetRange::Bounded(r) => format!("bytes={}-{}", r.start, r.end.saturating_sub(1)),
GetRange::Offset(n) => format!("bytes={}-", n),
GetRange::Suffix(n) => format!("bytes=-{}", n),
}
}
fn parse_content_range(header: &str) -> ObjectStoreResult<(std::ops::Range<u64>, u64)> {
let invalid = || generic_error(format!("malformed Content-Range header: `{header}`"));
let (range_part, total_part) = header
.strip_prefix("bytes ")
.and_then(|inner| inner.split_once('/'))
.ok_or_else(invalid)?;
let total = total_part.parse::<u64>().map_err(|_| invalid())?;
let (start, end) = range_part.split_once('-').ok_or_else(invalid)?;
let start = start.parse::<u64>().map_err(|_| invalid())?;
let end = end.parse::<u64>().map_err(|_| invalid())?;
if end < start || start > total {
return Err(invalid());
}
Ok((start..end.saturating_add(1), total))
}
fn not_supported(op: &str) -> ObjectStoreError {
ObjectStoreError::NotSupported {
source: format!("RestObjectStore does not support {op}").into(),
}
}
fn put_result() -> PutResult {
PutResult {
e_tag: None,
version: None,
}
}
enum WriteState {
Matches,
Differs,
Absent,
}
fn log_retry_outcome(target: &str, retries: u32, last_failure: &RetryFailure, succeeded: bool) {
if succeeded {
info!(
target,
retries,
last_failure = %last_failure,
"REST request succeeded after retries"
);
} else {
warn!(
target,
retries,
last_failure = %last_failure,
"REST request failed after retries"
);
}
}
async fn backoff(n: u32) {
let ms = (50u64 << n.min(6)).min(2_000);
tokio::time::sleep(Duration::from_millis(ms)).await;
}
fn parse_last_modified(headers: &HeaderMap) -> chrono::DateTime<chrono::Utc> {
headers
.get(reqwest::header::LAST_MODIFIED)
.and_then(|v| v.to_str().ok())
.and_then(|s| chrono::DateTime::parse_from_rfc2822(s).ok())
.map(|dt| dt.with_timezone(&chrono::Utc))
.unwrap_or(chrono::DateTime::UNIX_EPOCH)
}
fn parse_etag(headers: &HeaderMap) -> Option<String> {
headers
.get(reqwest::header::ETAG)
.and_then(|v| v.to_str().ok())
.map(|s| s.to_string())
}
#[async_trait]
impl ObjectStore for RestObjectStore {
async fn get_opts(&self, location: &Path, options: GetOptions) -> ObjectStoreResult<GetResult> {
let path_str = location.as_ref();
if let Some(range) = &options.range {
range.is_valid().map_err(generic_error)?;
}
if options.head {
let meta = self.head_meta(path_str, location).await?;
options.check_preconditions(&meta)?;
let size = meta.size;
return Ok(GetResult {
payload: GetResultPayload::Stream(Box::pin(futures::stream::empty())),
range: 0..size,
meta,
attributes: Attributes::new(),
});
}
let range_header = options.range.as_ref().map(get_range_to_header);
let (content, headers, status) = self.get_file(path_str, range_header.as_deref()).await?;
let (range, total_size) = if let Some(requested) = &options.range {
if status != reqwest::StatusCode::PARTIAL_CONTENT {
return Err(generic_error(format!(
"ranged GET for `{path_str}` returned a non-partial response (HTTP {status}); \
expected 206 Partial Content"
)));
}
let content_range = headers
.get(reqwest::header::CONTENT_RANGE)
.and_then(|v| v.to_str().ok())
.ok_or_else(|| {
generic_error(format!(
"ranged GET for `{path_str}` is missing a Content-Range header"
))
})?;
let (actual, total) = parse_content_range(content_range)?;
let expected = requested
.as_range(total)
.map_err(|e| generic_error(format!("invalid range for `{path_str}`: {e}")))?;
if actual != expected {
return Err(generic_error(format!(
"ranged GET for `{path_str}` returned unexpected Content-Range \
`{content_range}`; expected bytes {expected:?}"
)));
}
(actual, total)
} else {
match headers
.get(reqwest::header::CONTENT_RANGE)
.and_then(|v| v.to_str().ok())
{
Some(cr) => parse_content_range(cr)?,
None => (0..content.len() as u64, content.len() as u64),
}
};
let last_modified = parse_last_modified(&headers);
let e_tag = parse_etag(&headers);
let meta = ObjectMeta {
location: location.clone(),
last_modified,
size: total_size,
e_tag,
version: None,
};
options.check_preconditions(&meta)?;
let stream = Box::pin(futures::stream::once(futures::future::ready(Ok(content))));
Ok(GetResult {
payload: GetResultPayload::Stream(stream),
meta,
range,
attributes: Attributes::new(),
})
}
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: PutOptions,
) -> ObjectStoreResult<PutResult> {
let overwrite = match opts.mode {
PutMode::Overwrite => true,
PutMode::Create => false,
PutMode::Update(_) => return Err(not_supported("PutMode::Update")),
};
let path_str = location.as_ref();
let url = self.config.file_url(&self.base_url, path_str);
let query = self.config.put_query(overwrite);
let body: Bytes = payload.into();
if !overwrite && self.verify_on_ambiguous {
return self.put_create_verified(path_str, &url, &query, body).await;
}
let response = self
.client
.put(&url)
.query(&query)
.headers(self.headers()?)
.body(body)
.send()
.await
.map_err(generic_error)?;
ensure_success_response(response, path_str)?;
Ok(put_result())
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, ObjectStoreResult<ObjectMeta>> {
let prefix = prefix.map(|p| p.as_ref().to_string()).unwrap_or_default();
self.list_paginated(prefix, None, None, true)
}
fn list_with_offset(
&self,
prefix: Option<&Path>,
offset: &Path,
) -> BoxStream<'static, ObjectStoreResult<ObjectMeta>> {
let prefix = prefix.map(|p| p.as_ref().to_string()).unwrap_or_default();
let offset_str = {
let raw = offset.as_ref();
if !prefix.is_empty() && raw.starts_with(&prefix) {
raw[prefix.len()..].trim_start_matches('/').to_string()
} else {
raw.to_string()
}
};
self.list_paginated(prefix, Some(offset_str), Some(offset.clone()), true)
}
fn delete_stream(
&self,
locations: BoxStream<'static, ObjectStoreResult<Path>>,
) -> BoxStream<'static, ObjectStoreResult<Path>> {
let store = self.clone();
Box::pin(locations.then(move |location| {
let store = store.clone();
async move {
let location = location?;
store.delete_one(&location).await?;
Ok(location)
}
}))
}
async fn list_with_delimiter(&self, _prefix: Option<&Path>) -> ObjectStoreResult<ListResult> {
Err(not_supported("list_with_delimiter"))
}
async fn put_multipart_opts(
&self,
_location: &Path,
_opts: PutMultipartOptions,
) -> ObjectStoreResult<Box<dyn MultipartUpload>> {
Err(not_supported("multipart upload"))
}
async fn copy_opts(
&self,
_from: &Path,
_to: &Path,
_options: CopyOptions,
) -> ObjectStoreResult<()> {
Err(not_supported("copy"))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn get_range_to_header_formats_bounded() {
assert_eq!(get_range_to_header(&GetRange::Bounded(2..6)), "bytes=2-5");
}
#[test]
fn get_range_to_header_formats_offset() {
assert_eq!(get_range_to_header(&GetRange::Offset(10)), "bytes=10-");
}
#[test]
fn get_range_to_header_formats_suffix() {
assert_eq!(get_range_to_header(&GetRange::Suffix(512)), "bytes=-512");
}
#[test]
fn parse_content_range_accepts_valid() {
let (range, total) = parse_content_range("bytes 2-5/10").unwrap();
assert_eq!(range, 2..6);
assert_eq!(total, 10);
}
#[test]
fn parse_content_range_rejects_reversed_range() {
assert!(parse_content_range("bytes 5-2/10").is_err());
}
#[test]
fn parse_content_range_rejects_start_past_total() {
assert!(parse_content_range("bytes 20-25/10").is_err());
}
}