feat: access sqlite dbs from object storage

crates/datafusion_ext/Cargo.toml

@@ -15,19 +15,20 @@ unicode_expressions = []
 ioutil = { path = "../ioutil" }
 telemetry = { path = "../telemetry" }
 catalog = { path = "../catalog" }
+decimal = { path = "../decimal" }
+protogen = { path = "../protogen" }
+pgrepr = { path = "../pgrepr" }
 serde_json = { workspace = true }
 datafusion = { workspace = true }
 async-trait = { workspace = true }
+object_store = { workspace = true }
+tracing = { workspace = true }
+thiserror.workspace = true
+futures = { workspace = true }
 async-recursion = "1.0.4"
 uuid = { version = "1.7.0", features = ["v4", "fast-rng", "macro-diagnostics"] }
 regex = "1.10"
 once_cell = "1.19.0"
-tracing = { workspace = true }
-thiserror.workspace = true
-decimal = { path = "../decimal" }
-protogen = { path = "../protogen" }
-pgrepr = { path = "../pgrepr" }
-futures = { workspace = true }
 parking_lot = "0.12.1"
 bson = "2.9.0"
 

crates/datafusion_ext/src/errors.rs

@@ -42,6 +42,12 @@ pub enum ExtensionError {
     #[error("object store: {0}")]
     ObjectStore(String),
 
+    #[error(transparent)]
+    ObjectStoreCrate(#[from] object_store::Error),
+
+    #[error(transparent)]
+    ObjectStorePath(#[from] object_store::path::Error),
+
     #[error("{0}")]
     String(String),
 }

crates/datasources/src/common/errors.rs

@@ -35,3 +35,9 @@ pub enum DatasourceCommonError {
 }
 
 pub type Result<T, E = DatasourceCommonError> = std::result::Result<T, E>;
+
+impl From<DatasourceCommonError> for datafusion_ext::errors::ExtensionError {
+    fn from(value: DatasourceCommonError) -> Self {
+        datafusion_ext::errors::ExtensionError::access(value)
+    }
+}

crates/datasources/src/common/url.rs

@@ -149,6 +149,7 @@ impl TryFrom<&str> for DatasourceUrl {
     }
 }
 
+
 impl TryFrom<FuncParamValue> for DatasourceUrl {
     type Error = ExtensionError;
 
@@ -161,6 +162,12 @@ impl TryFrom<FuncParamValue> for DatasourceUrl {
     }
 }
 
+impl From<PathBuf> for DatasourceUrl {
+    fn from(path: PathBuf) -> Self {
+        DatasourceUrl::File(path)
+    }
+}
+
 impl TryFrom<DatasourceUrl> for ObjectStoreUrl {
     type Error = DataFusionError;
 

crates/datasources/src/json/errors.rs

@@ -1,8 +1,6 @@
 use datafusion::error::DataFusionError;
 use datafusion_ext::errors::ExtensionError;
 
-use crate::object_store::errors::ObjectStoreSourceError;
-
 #[derive(Debug, thiserror::Error)]
 pub enum JsonError {
     #[error("Unsupported json type: {0}")]
@@ -18,7 +16,7 @@ pub enum JsonError {
     SendAlreadyInProgress,
 
     #[error(transparent)]
-    ObjectStoreSource(#[from] ObjectStoreSourceError),
+    ObjectStoreSource(#[from] crate::object_store::errors::ObjectStoreSourceError),
 
     #[error(transparent)]
     ObjectStore(#[from] object_store::Error),

crates/datasources/src/sqlite/errors.rs

@@ -14,10 +14,13 @@ pub enum SqliteError {
     MpscSendError(String),
 
     #[error(transparent)]
-    FmtError(#[from] std::fmt::Error),
+    Fmt(#[from] std::fmt::Error),
 
     #[error(transparent)]
-    DatasourceCommonError(#[from] crate::common::errors::DatasourceCommonError),
+    DatasourceCommon(#[from] crate::common::errors::DatasourceCommonError),
+
+    #[error("Unimplemented: {0}")]
+    Unimplemented(&'static str),
 
     #[error("Missing data for column {0}")]
     MissingDataForColumn(usize),
@@ -28,11 +31,41 @@ pub enum SqliteError {
         to: datafusion::arrow::datatypes::DataType,
     },
 
+    #[error("found {num} objects matching specification '{url}'")]
+    NoMatchingObjectFound {
+        url: crate::common::url::DatasourceUrl,
+        num: usize,
+    },
+
     #[error(transparent)]
     ArrowError(#[from] datafusion::arrow::error::ArrowError),
 
     #[error(transparent)]
     ReprError(#[from] repr::error::ReprError),
+
+    #[error(transparent)]
+    IoError(#[from] std::io::Error),
+
+    #[error(transparent)]
+    ObjectStoreSource(#[from] crate::object_store::errors::ObjectStoreSourceError),
+
+    #[error(transparent)]
+    ObjectStoreError(#[from] object_store::Error),
+
+    #[error(transparent)]
+    ObjectStorePath(#[from] object_store::path::Error),
+
+    #[error(transparent)]
+    LakeStorageOptions(#[from] crate::lake::LakeStorageOptionsError),
+
+    #[error(transparent)]
+    ExtensionError(#[from] datafusion_ext::errors::ExtensionError),
 }
 
 pub type Result<T, E = SqliteError> = std::result::Result<T, E>;
+
+impl From<SqliteError> for datafusion_ext::errors::ExtensionError {
+    fn from(value: SqliteError) -> Self {
+        datafusion_ext::errors::ExtensionError::access(value)
+    }
+}

crates/datasources/src/sqlite/mod.rs

@@ -34,21 +34,101 @@ use datafusion_ext::errors::ExtensionError;
 use datafusion_ext::functions::VirtualLister;
 use datafusion_ext::metrics::DataSourceMetricsStreamAdapter;
 use futures::{StreamExt, TryStreamExt};
+use object_store::ObjectStore;
+use protogen::metastore::types::options::StorageOptions;
+use uuid::Uuid;
 
-use self::errors::Result;
+use self::errors::{Result, SqliteError};
 use self::wrapper::SqliteAsyncClient;
+use crate::common::url::DatasourceUrl;
 use crate::common::util::{self, COUNT_SCHEMA};
+use crate::lake::storage_options_into_store_access;
+use crate::object_store::ObjStoreAccessor;
 
 type DataFusionResult<T> = Result<T, DataFusionError>;
 
 #[derive(Debug, Clone)]
 pub struct SqliteAccess {
     pub db: PathBuf,
+    pub cache: Option<Arc<tempfile::TempDir>>,
 }
 
 impl SqliteAccess {
+    pub async fn new(url: DatasourceUrl, opts: Option<StorageOptions>) -> Result<Self> {
+        match url {
+            DatasourceUrl::File(ref location) => {
+                if !location.try_exists()? {
+                    Err(SqliteError::NoMatchingObjectFound {
+                        url: url.clone(),
+                        num: 0,
+                    })
+                } else {
+                    Ok(Self {
+                        db: location.clone(),
+                        cache: None,
+                    })
+                }
+            }
+            DatasourceUrl::Url(_) => {
+                let storage_options = match opts {
+                    Some(v) => v,
+                    None => {
+                        return Err(SqliteError::Internal(
+                            "storage options are required".to_string(),
+                        ))
+                    }
+                };
+                let store_access = storage_options_into_store_access(&url, &storage_options)?;
+
+                let accessor = ObjStoreAccessor::new(store_access)?;
+                let mut list = accessor.list_globbed(url.path()).await?;
+                if list.len() != 1 {
+                    return Err(SqliteError::NoMatchingObjectFound {
+                        url,
+                        num: list.len(),
+                    });
+                }
+
+                let store = accessor.into_object_store();
+
+                let obj = list.pop().unwrap().location;
+                let payload = store.get(&obj).await?.bytes().await?;
+
+                let tmpdir = Arc::new(
+                    tempfile::Builder::new()
+                        .prefix(
+                            storage_options
+                                .inner
+                                .get("__tmp_prefix")
+                                .map(|i| i.to_owned())
+                                .unwrap_or_else(|| Uuid::new_v4().to_string())
+                                .as_str(),
+                        )
+                        .rand_bytes(8)
+                        .tempdir()?,
+                );
+
+                let tmpdir_path = tmpdir.path();
+                let local_store =
+                    object_store::local::LocalFileSystem::new_with_prefix(tmpdir_path)?;
+
+                let local_path =
+                    object_store::path::Path::parse(obj.filename().unwrap_or("sqlite"))?;
+
+                local_store.put(&local_path, payload).await?;
+
+                let db = tmpdir_path.join(local_path.filename().unwrap());
+
+                Ok(Self {
+                    db,
+                    cache: Some(tmpdir.clone()),
+                })
+            }
+        }
+    }
+
     pub async fn connect(&self) -> Result<SqliteAccessState> {
-        let client = SqliteAsyncClient::new(self.db.to_path_buf()).await?;
+        let client = SqliteAsyncClient::new(self.db.to_path_buf(), self.cache.clone()).await?;
         Ok(SqliteAccessState { client })
     }
 
@@ -64,12 +144,17 @@ impl SqliteAccess {
     }
 }
 
+
 #[derive(Clone, Debug)]
 pub struct SqliteAccessState {
     client: SqliteAsyncClient,
 }
 
 impl SqliteAccessState {
+    pub fn is_local_file(&self) -> bool {
+        self.client.is_local_file()
+    }
+
     async fn validate_table_access(&self, table: &str) -> Result<()> {
         let query = format!("SELECT * FROM {table} WHERE FALSE");
         let _ = self.client.query_all(query).await?;
@@ -301,6 +386,12 @@ impl TableProvider for SqliteTableProvider {
             return Err(DataFusionError::Execution("cannot overwrite".to_string()));
         }
 
+        if !self.state.is_local_file() {
+            return Err(DataFusionError::Execution(
+                "cannot write remote file".to_string(),
+            ));
+        }
+
         Ok(Arc::new(SqliteInsertExec {
             input,
             table: self.table.to_string(),

crates/datasources/src/sqlite/wrapper.rs

@@ -1,6 +1,7 @@
 use std::fmt;
 use std::path::PathBuf;
 use std::pin::Pin;
+use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use async_sqlite::rusqlite;
@@ -10,6 +11,7 @@ use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::error::DataFusionError;
 use datafusion::physical_plan::RecordBatchStream;
 use futures::{Future, FutureExt, Stream};
+use tempfile;
 use tokio::sync::mpsc;
 
 use super::convert::Converter;
@@ -19,6 +21,8 @@ use crate::sqlite::errors::Result;
 pub struct SqliteAsyncClient {
     path: PathBuf,
     inner: async_sqlite::Client,
+    // we're just tying the lifetime of the tempdir to this connection
+    cache: Option<Arc<tempfile::TempDir>>,
 }
 
 impl fmt::Debug for SqliteAsyncClient {
@@ -28,12 +32,13 @@ impl fmt::Debug for SqliteAsyncClient {
 }
 
 impl SqliteAsyncClient {
-    pub async fn new(path: PathBuf) -> Result<Self> {
+    pub async fn new(path: PathBuf, cache: Option<Arc<tempfile::TempDir>>) -> Result<Self> {
         let inner = async_sqlite::ClientBuilder::new()
             .path(&path)
             .open()
             .await?;
-        Ok(Self { path, inner })
+
+        Ok(Self { path, inner, cache })
     }
 
     /// Query and return a RecordBatchStream for sqlite data.
@@ -118,6 +123,10 @@ impl SqliteAsyncClient {
             })
             .await?)
     }
+
+    pub fn is_local_file(&self) -> bool {
+        self.cache.is_none()
+    }
 }
 
 #[derive(Debug, Clone)]