feat: Enable GCS, S3, R2 and MinIO as object stores for local runs

crates/datasources/src/native/access.rs

@@ -125,12 +125,7 @@ impl NativeTableStorage {
 
     pub async fn delete_table(&self, table: &TableEntry) -> Result<()> {
         let prefix = format!("databases/{}/tables/{}", self.db_id, table.meta.id);
-        let path: ObjectStorePath = match &self.conf {
-            StorageConfig::Gcs { bucket, .. } => format!("gs://{}/{}", bucket, prefix).into(),
-            StorageConfig::Memory => format!("memory://{}", prefix).into(),
-            _ => prefix.into(),
-        };
-        let mut x = self.store.list(Some(&path)).await?;
+        let mut x = self.store.list(Some(&prefix.into())).await?;
         while let Some(meta) = x.next().await {
             let meta = meta?;
             self.store.delete(&meta.location).await?
@@ -153,9 +148,21 @@ impl NativeTableStorage {
         let prefix = format!("databases/{}/tables/{}", self.db_id, table.meta.id);
 
         let url = match &self.conf {
-            StorageConfig::Gcs { bucket, .. } => {
-                Url::parse(&format!("gs://{}/{}", bucket, prefix.clone()))?
+            StorageConfig::S3 {
+                endpoint, bucket, ..
+            } => {
+                let mut s3_url = if let Some(endpoint) = endpoint {
+                    endpoint.clone()
+                } else {
+                    "s3://".to_string()
+                };
+
+                if let Some(bucket) = bucket {
+                    s3_url = format!("{s3_url}/{bucket}");
+                }
+                Url::parse(&format!("{s3_url}/{prefix}"))?
             }
+            StorageConfig::Gcs { bucket, .. } => Url::parse(&format!("gs://{bucket}/{prefix}"))?,
             StorageConfig::Local { path } => {
                 let path =
                     fs::canonicalize(path)
@@ -168,7 +175,7 @@ impl NativeTableStorage {
                 Url::from_file_path(path).map_err(|_| NativeError::Static("Path not absolute"))?
             }
             StorageConfig::Memory => {
-                let s = format!("memory://{}", prefix.clone());
+                let s = format!("memory://{prefix}");
                 Url::parse(&s)?
             }
         };

crates/glaredb/src/args/local.rs

@@ -35,6 +35,9 @@ pub struct LocalClientOpts {
     #[clap(short = 'c', long, value_parser)]
     pub cloud_url: Option<Url>,
 
+    #[clap(flatten)]
+    pub storage_config: StorageConfigArgs,
+
     /// Ignores the proxy and directly goes to the server for remote execution.
     ///
     /// (Internal)

crates/glaredb/src/args/mod.rs

@@ -1,3 +1,4 @@
+use anyhow::anyhow;
 use anyhow::Result;
 use clap::{Parser, ValueEnum};
 use std::fmt::Write as _;
@@ -81,3 +82,23 @@ pub struct PgProxyArgs {
     #[clap(long)]
     pub cloud_auth_code: String,
 }
+
+#[derive(Debug, Clone, Parser)]
+pub struct StorageConfigArgs {
+    /// URL of the object store in which to keep the data in.
+    #[clap(short, long)]
+    pub location: Option<String>,
+
+    /// Storage options for building the object store.
+    #[clap(short = 'o', long, requires = "location", value_parser=parse_key_value_pair)]
+    pub storage_options: Vec<(String, String)>,
+}
+
+fn parse_key_value_pair(key_value_pair: &str) -> Result<(String, String)> {
+    key_value_pair
+        .split_once('=')
+        .map(|(key, value)| (key.to_string(), value.to_string()))
+        .ok_or(anyhow!(
+            "Expected key-value pair delimited by an equals sign, got '{key_value_pair}'"
+        ))
+}

crates/glaredb/src/local.rs

@@ -1,4 +1,4 @@
-use crate::args::{LocalClientOpts, OutputMode};
+use crate::args::{LocalClientOpts, OutputMode, StorageConfigArgs};
 use crate::highlighter::{SQLHighlighter, SQLHinter, SQLValidator};
 use crate::prompt::SQLPrompt;
 use crate::util::MetastoreClientMode;
@@ -16,8 +16,11 @@ use datafusion::physical_plan::SendableRecordBatchStream;
 use futures::StreamExt;
 use pgrepr::format::Format;
 use reedline::{FileBackedHistory, Reedline, Signal};
+use std::collections::HashMap;
 
 use datafusion_ext::vars::SessionVars;
+use metastore::local::start_inprocess;
+use object_store_util::shared::SharedObjectStore;
 use sqlexec::engine::EngineStorageConfig;
 use sqlexec::engine::{Engine, SessionStorageConfig, TrackedSession};
 use sqlexec::parser;
@@ -48,15 +51,37 @@ pub struct LocalSession {
 impl LocalSession {
     pub async fn connect(opts: LocalClientOpts) -> Result<Self> {
         // Connect to metastore.
-        let mode = MetastoreClientMode::new_local(opts.data_dir.clone())?;
-        let metastore_client = mode.into_client().await?;
-        let tracker = Arc::new(Tracker::Nop);
+        let (storage_conf, metastore_client) = if let StorageConfigArgs {
+            location: Some(location),
+            storage_options,
+        } = &opts.storage_config
+        {
+            // TODO: try to consolidate with --data-dir option
+            let conf = EngineStorageConfig::try_from_options(
+                location,
+                HashMap::from_iter(storage_options.clone()),
+            )?;
+            let store = conf
+                .storage_config(&SessionStorageConfig::default())?
+                .new_object_store()?;
+            // Wrap up the store with a shared one, so that we get to use the non-atomic
+            // copy-if-not-exists that is defined there when initializing the lease
+            let store = SharedObjectStore::new(store);
+            let client = start_inprocess(Arc::new(store)).await?;
+            (conf, client)
+        } else {
+            let conf = match &opts.data_dir {
+                Some(path) => EngineStorageConfig::Local { path: path.clone() },
+                None => EngineStorageConfig::Memory,
+            };
 
-        let storage_conf = match &opts.data_dir {
-            Some(path) => EngineStorageConfig::Local { path: path.clone() },
-            None => EngineStorageConfig::Memory,
+            let mode = MetastoreClientMode::new_local(opts.data_dir.clone())?;
+            let client = mode.into_client().await?;
+            (conf, client)
         };
 
+        let tracker = Arc::new(Tracker::Nop);
+
         let engine = Engine::new(
             metastore_client,
             storage_conf,
@@ -131,9 +156,16 @@ impl LocalSession {
     }
 
     async fn run_interactive(&mut self) -> Result<()> {
-        let info = match &self.opts.data_dir {
-            Some(path) => format!("Persisting database at path: {}", path.display()),
-            None => "Using in-memory catalog".to_string(),
+        let info = match (&self.opts.storage_config, &self.opts.data_dir) {
+            (
+                StorageConfigArgs {
+                    location: Some(location),
+                    ..
+                },
+                _,
+            ) => format!("Persisting database at location: {location}"),
+            (_, Some(path)) => format!("Persisting database at path: {}", path.display()),
+            (_, None) => "Using in-memory catalog".to_string(),
         };
 
         println!("{info}");

crates/glaredb/src/server.rs

@@ -78,6 +78,7 @@ impl ComputeServer {
         let storage_conf = match (data_dir, service_account_key) {
             (None, Some(key)) => EngineStorageConfig::Gcs {
                 service_account_key: key,
+                bucket: None,
             },
             (Some(dir), None) => EngineStorageConfig::Local { path: dir },
             (None, None) => EngineStorageConfig::Memory,

crates/metastore/src/local.rs

@@ -27,8 +27,7 @@ pub async fn start_inprocess_local(
 
 /// Starts an in-process metastore service, returning a client for the service.
 ///
-/// Useful for some tests, as well as when running GlareDB locally for testing.
-/// This should never be used in production.
+/// Useful for tests, as well as when running GlareDB locally.
 pub async fn start_inprocess(
     store: Arc<dyn ObjectStore>,
 ) -> Result<MetastoreServiceClient<Channel>> {

crates/object_store_util/Cargo.toml

@@ -6,6 +6,7 @@ edition = { workspace = true }
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+anyhow = "1.0"
 logutil = { path = "../logutil" }
 object_store = { workspace = true }
 tempfile = "3"

crates/object_store_util/src/conf.rs

@@ -1,3 +1,4 @@
+use object_store::aws::{AmazonS3Builder, S3CopyIfNotExists};
 use object_store::{
     gcp::GoogleCloudStorageBuilder, local::LocalFileSystem, memory::InMemory,
     Error as ObjectStoreError, ObjectStore,
@@ -11,6 +12,13 @@ static IN_MEMORY_STORE: Lazy<Arc<InMemory>> = Lazy::new(|| Arc::new(InMemory::ne
 /// Configuration options for various types of storage we support.
 #[derive(Debug, Clone)]
 pub enum StorageConfig {
+    S3 {
+        access_key_id: String,
+        secret_access_key: String,
+        region: Option<String>,
+        endpoint: Option<String>,
+        bucket: Option<String>,
+    },
     Gcs {
         service_account_key: String,
         bucket: String,
@@ -25,6 +33,39 @@ impl StorageConfig {
     /// Create a new object store using this config.
     pub fn new_object_store(&self) -> Result<Arc<dyn ObjectStore>, ObjectStoreError> {
         Ok(match self {
+            StorageConfig::S3 {
+                access_key_id,
+                secret_access_key,
+                region,
+                endpoint,
+                bucket,
+            } => {
+                let mut builder = AmazonS3Builder::new()
+                    .with_access_key_id(access_key_id)
+                    .with_secret_access_key(secret_access_key)
+                    .with_region(region.clone().unwrap_or_default());
+
+                if let Some(endpoint) = endpoint {
+                    if endpoint.starts_with("http://") {
+                        builder = builder.with_allow_http(true);
+                    }
+                    builder = builder.with_endpoint(endpoint);
+                    if endpoint.contains("r2.cloudflarestorage.com") {
+                        // Ensure `ObjectStore::copy_if_not_exists` is enabled on the S3 client for
+                        // Cloudflare R2 with the adequate header
+                        builder = builder.with_copy_if_not_exists(S3CopyIfNotExists::Header(
+                            "cf-copy-destination-if-none-match".to_string(),
+                            "*".to_string(),
+                        ))
+                    }
+                }
+
+                if let Some(bucket) = bucket {
+                    builder = builder.with_bucket_name(bucket);
+                }
+
+                Arc::new(builder.build()?)
+            }
             StorageConfig::Gcs {
                 service_account_key,
                 bucket,

crates/object_store_util/src/shared.rs

@@ -1,7 +1,10 @@
+use anyhow::anyhow;
 use async_trait::async_trait;
 use bytes::Bytes;
 use futures::stream::BoxStream;
-use object_store::{path::Path, GetResult, ListResult, ObjectMeta, ObjectStore, Result};
+use object_store::{
+    path::Path, Error as ObjectStoreError, GetResult, ListResult, ObjectMeta, ObjectStore, Result,
+};
 use object_store::{GetOptions, MultipartId};
 use std::ops::Range;
 use std::sync::Arc;
@@ -86,14 +89,29 @@ impl ObjectStore for SharedObjectStore {
     }
 
     async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
-        self.inner.copy_if_not_exists(from, to).await
+        match self.inner.copy_if_not_exists(from, to).await {
+            Ok(_) => Ok(()),
+            Err(ObjectStoreError::NotSupported { .. }) => {
+                // Go with the poor man's copy-if-not-exists: try a regular copy if the path doesn't exist
+                match self.head(to).await {
+                    Ok(_) => {
+                        return Err(ObjectStoreError::AlreadyExists {
+                            path: to.to_string(),
+                            source: anyhow!(
+                            "Object at path {to} already exists, can't perform copy-if-not-exists"
+                        )
+                            .into(),
+                        })
+                    }
+                    Err(ObjectStoreError::NotFound { .. }) => self.copy(from, to).await,
+                    Err(e) => Err(e),
+                }
+            }
+            Err(e) => Err(e),
+        }
     }
 
     async fn rename(&self, from: &Path, to: &Path) -> Result<()> {
-        self.rename(from, to).await
-    }
-
-    async fn rename_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
-        self.inner.rename_if_not_exists(from, to).await
+        self.inner.rename(from, to).await
     }
 }