From 587090e7065bf5ce117f697f5c2c2ef65d28edb4 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Mon, 4 Nov 2024 17:32:31 -0500 Subject: [PATCH] Simplify Python binding of `StorageConfig` (#371) --- .../python/icechunk/_icechunk_python.pyi | 41 ++---- icechunk-python/src/storage.rs | 138 +++++++----------- 2 files changed, 69 insertions(+), 110 deletions(-) diff --git a/icechunk-python/python/icechunk/_icechunk_python.pyi b/icechunk-python/python/icechunk/_icechunk_python.pyi index 161ce842..e0a73340 100644 --- a/icechunk-python/python/icechunk/_icechunk_python.pyi +++ b/icechunk-python/python/icechunk/_icechunk_python.pyi @@ -105,27 +105,7 @@ class StorageConfig: storage_config = StorageConfig.s3_from_config("bucket", "prefix", ...) ``` """ - class Memory: - """Config for an in-memory storage backend""" - prefix: str - - class Filesystem: - """Config for a local filesystem storage backend""" - - root: str - - class S3: - """Config for an S3 Object Storage compatible storage backend""" - - bucket: str - prefix: str - credentials: S3Credentials | None - endpoint_url: str | None - allow_http: bool | None - region: str | None - - def __init__(self, storage: Memory | Filesystem | S3): ... @classmethod def memory(cls, prefix: str) -> StorageConfig: """Create a StorageConfig object for an in-memory storage backend with the given prefix""" @@ -137,7 +117,14 @@ class StorageConfig: ... @classmethod - def s3_from_env(cls, bucket: str, prefix: str) -> StorageConfig: + def s3_from_env( + cls, + bucket: str, + prefix: str, + endpoint_url: str | None, + allow_http: bool = False, + region: str | None = None, + ) -> StorageConfig: """Create a StorageConfig object for an S3 Object Storage compatible storage backend with the given bucket and prefix @@ -158,7 +145,7 @@ class StorageConfig: prefix: str, credentials: S3Credentials, endpoint_url: str | None, - allow_http: bool | None = None, + allow_http: bool = False, region: str | None = None, ) -> StorageConfig: """Create a StorageConfig object for an S3 Object Storage compatible storage @@ -175,7 +162,7 @@ class StorageConfig: bucket: str, prefix: str, endpoint_url: str | None, - allow_http: bool | None = None, + allow_http: bool = False, region: str | None = None, ) -> StorageConfig: """Create a StorageConfig object for an S3 Object Storage compatible storage @@ -269,7 +256,7 @@ class StoreConfig: inline_chunk_threshold_bytes: int | None = None, unsafe_overwrite_refs: bool | None = None, virtual_ref_config: VirtualRefConfig | None = None, - ): + ): """Create a StoreConfig object with the given configuration options Parameters @@ -284,24 +271,22 @@ class StoreConfig: Whether to allow overwriting refs in the store. Default is False. Experimental. virtual_ref_config: VirtualRefConfig | None Configurations for virtual references such as credentials and endpoints - + Returns ------- StoreConfig A StoreConfig object with the given configuration options - """ + """ ... async def async_pyicechunk_store_exists(storage: StorageConfig) -> bool: ... def pyicechunk_store_exists(storage: StorageConfig) -> bool: ... - async def async_pyicechunk_store_create( storage: StorageConfig, config: StoreConfig | None ) -> PyIcechunkStore: ... def pyicechunk_store_create( storage: StorageConfig, config: StoreConfig | None ) -> PyIcechunkStore: ... - async def async_pyicechunk_store_open_existing( storage: StorageConfig, read_only: bool, config: StoreConfig | None ) -> PyIcechunkStore: ... diff --git a/icechunk-python/src/storage.rs b/icechunk-python/src/storage.rs index 2b590739..2a4bb48f 100644 --- a/icechunk-python/src/storage.rs +++ b/icechunk-python/src/storage.rs @@ -1,6 +1,3 @@ -#![allow(clippy::too_many_arguments)] -// TODO: we only need that allow for PyStorageConfig, but i don't know how to set it - use std::path::PathBuf; use icechunk::{ @@ -33,6 +30,22 @@ impl From<&PyS3Credentials> for StaticS3Credentials { } } +impl From for StaticS3Credentials { + fn from(credentials: PyS3Credentials) -> Self { + StaticS3Credentials { + access_key_id: credentials.access_key_id, + secret_access_key: credentials.secret_access_key, + session_token: credentials.session_token, + } + } +} + +impl From for S3Credentials { + fn from(value: PyS3Credentials) -> Self { + S3Credentials::Static(value.into()) + } +} + #[pymethods] impl PyS3Credentials { #[new] @@ -51,35 +64,19 @@ impl PyS3Credentials { } #[pyclass(name = "StorageConfig")] -pub enum PyStorageConfig { - Memory { - prefix: Option, - }, - Filesystem { - root: String, - }, - S3 { - bucket: String, - prefix: String, - anon: bool, - credentials: Option, - endpoint_url: Option, - allow_http: Option, - region: Option, - }, -} +pub struct PyStorageConfig(StorageConfig); #[pymethods] impl PyStorageConfig { #[classmethod] #[pyo3(signature = (prefix = None))] fn memory(_cls: &Bound<'_, PyType>, prefix: Option) -> Self { - PyStorageConfig::Memory { prefix } + Self(StorageConfig::InMemory { prefix }) } #[classmethod] - fn filesystem(_cls: &Bound<'_, PyType>, root: String) -> Self { - PyStorageConfig::Filesystem { root } + fn filesystem(_cls: &Bound<'_, PyType>, root: PathBuf) -> Self { + Self(StorageConfig::LocalFileSystem { root }) } #[classmethod] @@ -87,7 +84,7 @@ impl PyStorageConfig { bucket, prefix, endpoint_url = None, - allow_http = None, + allow_http = false, region = None, ))] fn s3_from_env( @@ -95,18 +92,16 @@ impl PyStorageConfig { bucket: String, prefix: String, endpoint_url: Option, - allow_http: Option, + allow_http: bool, region: Option, ) -> Self { - PyStorageConfig::S3 { - bucket, - prefix, - anon: false, - credentials: None, - endpoint_url, - allow_http, + let config = S3Config { region, - } + endpoint: endpoint_url, + allow_http, + credentials: mk_credentials(None, false), + }; + Self(StorageConfig::S3ObjectStore { bucket, prefix, config: Some(config) }) } #[classmethod] @@ -115,7 +110,7 @@ impl PyStorageConfig { prefix, credentials, endpoint_url = None, - allow_http = None, + allow_http = false, region = None, ))] fn s3_from_config( @@ -124,18 +119,16 @@ impl PyStorageConfig { prefix: String, credentials: PyS3Credentials, endpoint_url: Option, - allow_http: Option, + allow_http: bool, region: Option, ) -> Self { - PyStorageConfig::S3 { - bucket, - prefix, - anon: false, - credentials: Some(credentials), - endpoint_url, - allow_http, + let config = S3Config { region, - } + endpoint: endpoint_url, + allow_http, + credentials: credentials.into(), + }; + Self(StorageConfig::S3ObjectStore { bucket, prefix, config: Some(config) }) } #[classmethod] @@ -143,7 +136,7 @@ impl PyStorageConfig { bucket, prefix, endpoint_url = None, - allow_http = None, + allow_http = false, region = None, ))] fn s3_anonymous( @@ -151,18 +144,16 @@ impl PyStorageConfig { bucket: String, prefix: String, endpoint_url: Option, - allow_http: Option, + allow_http: bool, region: Option, ) -> Self { - PyStorageConfig::S3 { - bucket, - prefix, - anon: true, - credentials: None, - endpoint_url, - allow_http, + let config = S3Config { region, - } + endpoint: endpoint_url, + allow_http, + credentials: mk_credentials(None, true), + }; + Self(StorageConfig::S3ObjectStore { bucket, prefix, config: Some(config) }) } } @@ -177,38 +168,21 @@ fn mk_credentials(config: Option<&PyS3Credentials>, anon: bool) -> S3Credentials } } +impl From for StorageConfig { + fn from(storage: PyStorageConfig) -> Self { + storage.0 + } +} + impl From<&PyStorageConfig> for StorageConfig { fn from(storage: &PyStorageConfig) -> Self { - match storage { - PyStorageConfig::Memory { prefix } => { - StorageConfig::InMemory { prefix: prefix.clone() } - } - PyStorageConfig::Filesystem { root } => { - StorageConfig::LocalFileSystem { root: PathBuf::from(root.clone()) } - } - PyStorageConfig::S3 { - bucket, - prefix, - anon, - credentials, - endpoint_url, - allow_http, - region, - } => { - let s3_config = S3Config { - region: region.clone(), - credentials: mk_credentials(credentials.as_ref(), *anon), - endpoint: endpoint_url.clone(), - allow_http: allow_http.unwrap_or(false), - }; + storage.0.clone() + } +} - StorageConfig::S3ObjectStore { - bucket: bucket.clone(), - prefix: prefix.clone(), - config: Some(s3_config), - } - } - } +impl AsRef for PyStorageConfig { + fn as_ref(&self) -> &StorageConfig { + &self.0 } }