Merge pull request #23 from RelationalAI/ag-list-offset

Add offset to list and prevent checksum generation when using S3 unsigned payload

Author: andrebsguedes

Cargo.lock

@@ -20,6 +20,10 @@ bench = false
 [profile.release]
 debug = 1
 
+[features]
+default = ["julia"]
+julia = []
+
 [dependencies]
 tokio = { version = "1.32.0", features = ["rt-multi-thread", "macros", "signal", "time"] }
 tokio-util = { version = "0.7", default-features = false, features = ["io"] }
@@ -28,9 +32,9 @@ tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 futures-util = "0.3"
 reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "hickory-dns"] }
-# object_store = { version = "0.8", features = ["azure", "aws"] }
-# Pinned to a specific commit while waiting for the next release
-object_store = { version = "0.10.1", features = ["azure", "aws"] }
+# object_store = { version = "0.10.1", features = ["azure", "aws"] }
+# Pinned to a specific commit while waiting for upstream
+object_store = { git = "https://github.com/andrebsguedes/arrow-rs.git", branch = "unsigned-payload-and-azure-list-offset", features = ["azure", "aws", "experimental-azure-list-offset"] }
 thiserror = "1"
 anyhow = { version = "1", features = ["backtrace"] }
 once_cell = "1.18"

Cargo.toml

@@ -1,12 +1,11 @@
-use crate::{CResult, Config, NotifyGuard, SQ, clients, dyn_connect, static_config, Request, util::cstr_to_path, Context, RawResponse, ResponseGuard};
+use crate::{CResult, Client, RawConfig, NotifyGuard, SQ, static_config, Request, util::cstr_to_path, Context, RawResponse, ResponseGuard};
 
 use object_store::{path::Path, ObjectStore};
 
 use anyhow::anyhow;
 use std::ffi::{c_char, c_void};
 use futures_util::StreamExt;
 use tokio::io::AsyncWriteExt;
-use std::sync::Arc;
 
 // The type used to give Julia the result of an async request. It will be allocated
 // by Julia as part of the request and filled in by Rust.
@@ -43,15 +42,15 @@ impl RawResponse for Response {
     }
 }
 
-async fn multipart_get(slice: &mut [u8], path: &Path, client: &dyn ObjectStore) -> anyhow::Result<usize> {
-    let result = client.head(&path).await?;
+async fn multipart_get(slice: &mut [u8], path: &Path, client: &Client) -> anyhow::Result<usize> {
+    let result = client.store.head(&path).await?;
     if result.size > slice.len() {
         return Err(anyhow!("Supplied buffer was too small"));
     }
 
     let part_ranges = crate::util::size_to_ranges(result.size);
 
-    let result_vec = client.get_ranges(&path, &part_ranges).await?;
+    let result_vec = client.store.get_ranges(&path, &part_ranges).await?;
     let mut accum: usize = 0;
     for i in 0..result_vec.len() {
         slice[accum..accum + result_vec[i].len()].copy_from_slice(&result_vec[i]);
@@ -61,9 +60,9 @@ async fn multipart_get(slice: &mut [u8], path: &Path, client: &dyn ObjectStore)
     return Ok(accum);
 }
 
-async fn multipart_put(slice: &[u8], path: &Path, client: Arc<dyn ObjectStore>) -> anyhow::Result<()> {
+async fn multipart_put(slice: &[u8], path: &Path, client: Client) -> anyhow::Result<()> {
     let mut writer = object_store::buffered::BufWriter::with_capacity(
-        client,
+        client.store,
         path.clone(),
         10 * 1024 * 1024
     )
@@ -88,19 +87,15 @@ async fn multipart_put(slice: &[u8], path: &Path, client: Arc<dyn ObjectStore>)
     };
 }
 
-pub(crate) async fn handle_get(slice: &mut [u8], path: &Path, config: &Config) -> anyhow::Result<usize> {
-    let (client, _) = clients()
-        .try_get_with(config.get_hash(), dyn_connect(config)).await
-        .map_err(|e| anyhow!(e))?;
-
+pub(crate) async fn handle_get(client: Client, slice: &mut [u8], path: &Path) -> anyhow::Result<usize> {
     // Multipart Get
     if slice.len() > static_config().multipart_get_threshold as usize {
         let accum = multipart_get(slice, path, &client).await?;
         return Ok(accum);
     }
 
     // Single part Get
-    let body = client.get(path).await?;
+    let body = client.store.get(path).await?;
     let mut batch_stream = body.into_stream().chunks(8);
 
     let mut received_bytes = 0;
@@ -129,27 +124,19 @@ pub(crate) async fn handle_get(slice: &mut [u8], path: &Path, config: &Config) -
     Ok(received_bytes)
 }
 
-pub(crate) async fn handle_put(slice: &'static [u8], path: &Path, config: &Config) -> anyhow::Result<usize> {
-    let (client, _) = clients()
-        .try_get_with(config.get_hash(), dyn_connect(config)).await
-        .map_err(|e| anyhow!(e))?;
-
+pub(crate) async fn handle_put(client: Client, slice: &'static [u8], path: &Path) -> anyhow::Result<usize> {
     let len = slice.len();
     if len < static_config().multipart_put_threshold as usize {
-        let _ = client.put(path, slice.into()).await?;
+        let _ = client.store.put(path, slice.into()).await?;
     } else {
         let _ = multipart_put(slice, path, client).await?;
     }
 
     Ok(len)
 }
 
-pub(crate) async fn handle_delete(path: &Path, config: &Config) -> anyhow::Result<usize> {
-    let (client, _) = clients()
-        .try_get_with(config.get_hash(), dyn_connect(config)).await
-        .map_err(|e| anyhow!(e))?;
-
-    client.delete(path).await?;
+pub(crate) async fn handle_delete(client: Client, path: &Path) -> anyhow::Result<usize> {
+    client.store.delete(path).await?;
 
     Ok(0)
 }
@@ -159,7 +146,7 @@ pub extern "C" fn get(
     path: *const c_char,
     buffer: *mut u8,
     size: usize,
-    config: *const Config,
+    config: *const RawConfig,
     response: *mut Response,
     handle: *const c_void
 ) -> CResult {
@@ -195,7 +182,7 @@ pub extern "C" fn put(
     path: *const c_char,
     buffer: *const u8,
     size: usize,
-    config: *const Config,
+    config: *const RawConfig,
     response: *mut Response,
     handle: *const c_void
 ) -> CResult {
@@ -229,7 +216,7 @@ pub extern "C" fn put(
 #[no_mangle]
 pub extern "C" fn delete(
     path: *const c_char,
-    config: *const Config,
+    config: *const RawConfig,
     response: *mut Response,
     handle: *const c_void
 ) -> CResult {

src/crud_ops.rs

@@ -1,7 +1,7 @@
 use std::time::Duration;
 use backoff::backoff::Backoff;
+use object_store::RetryConfig;
 use once_cell::sync::Lazy;
-use crate::ConfigMeta;
 use std::error::Error as StdError;
 use anyhow::anyhow;
 
@@ -122,15 +122,15 @@ pub(crate) fn extract_error_info(error: &anyhow::Error) -> ErrorInfo {
 pub(crate) struct RetryState {
     pub(crate) start: std::time::Instant,
     pub(crate) attempts: Vec<ErrorInfo>,
-    pub(crate) config_meta: ConfigMeta
+    pub(crate) retry_config: RetryConfig
 }
 
 impl RetryState {
-    pub(crate) fn new(config_meta: ConfigMeta) -> RetryState {
+    pub(crate) fn new(retry_config: RetryConfig) -> RetryState {
         RetryState {
             start: std::time::Instant::now(),
             attempts: vec![],
-            config_meta
+            retry_config
         }
     }
 
@@ -145,8 +145,8 @@ impl RetryState {
         // We try to use the same settings as the object_store backoff but the implementation is
         // different so this is best effort.
         let mut backoff = backoff::ExponentialBackoff {
-            initial_interval: self.config_meta.retry_config.backoff.init_backoff,
-            max_interval: self.config_meta.retry_config.backoff.max_backoff,
+            initial_interval: self.retry_config.backoff.init_backoff,
+            max_interval: self.retry_config.backoff.max_backoff,
             ..Default::default()
         };
 
@@ -155,16 +155,16 @@ impl RetryState {
             let _ = backoff.next_backoff();
         }
 
-        backoff.next_backoff().unwrap_or(self.config_meta.retry_config.backoff.max_backoff)
+        backoff.next_backoff().unwrap_or(self.retry_config.backoff.max_backoff)
     }
 
     fn log_attempt(&mut self, info: ErrorInfo) {
         self.attempts.push(info);
     }
 
     pub(crate) fn should_retry_logic(&self) -> bool {
-        let max_retries = self.config_meta.retry_config.max_retries;
-        let retry_timeout = self.config_meta.retry_config.retry_timeout;
+        let max_retries = self.retry_config.max_retries;
+        let retry_timeout = self.retry_config.retry_timeout;
         let elapsed = self.start.elapsed();
         let all_retries = self.retries();