From 454dfbcda3ac9b65aafb4bfe266c9d87c219e30d Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Fri, 8 Nov 2024 12:50:56 -0500 Subject: [PATCH] WIP: requester pays --- pyo3-object_store/Cargo.toml | 1 + pyo3-object_store/src/api.rs | 5 +- pyo3-object_store/src/client.rs | 79 +++++++++++++++++++++++ pyo3-object_store/type-hints/__init__.pyi | 1 + pyo3-object_store/type-hints/_aws.pyi | 4 +- pyo3-object_store/type-hints/_client.pyi | 5 +- tests/test_get.py | 19 +++++- 7 files changed, 109 insertions(+), 5 deletions(-) diff --git a/pyo3-object_store/Cargo.toml b/pyo3-object_store/Cargo.toml index a119dae..484e59c 100644 --- a/pyo3-object_store/Cargo.toml +++ b/pyo3-object_store/Cargo.toml @@ -15,6 +15,7 @@ include = ["src", "type-hints", "README.md", "LICENSE"] [dependencies] futures = "0.3" +http = "1.0.0" object_store = { version = "0.11", features = ["aws", "azure", "gcp", "http"] } pyo3 = { version = "0.22", features = ["chrono", "indexmap"] } pyo3-async-runtimes = { version = "0.22", features = ["tokio-runtime"] } diff --git a/pyo3-object_store/src/api.rs b/pyo3-object_store/src/api.rs index af8d555..0f27284 100644 --- a/pyo3-object_store/src/api.rs +++ b/pyo3-object_store/src/api.rs @@ -2,7 +2,9 @@ use pyo3::intern; use pyo3::prelude::*; use crate::error::*; -use crate::{PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store}; +use crate::{ + PyAzureStore, PyClientOptions, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store, +}; /// Export the default Python API as a submodule named `store` within the given parent module /// @@ -50,6 +52,7 @@ pub fn register_store_module( child_module.add_class::()?; child_module.add_class::()?; child_module.add_class::()?; + child_module.add_class::()?; parent_module.add_submodule(&child_module)?; diff --git a/pyo3-object_store/src/client.rs b/pyo3-object_store/src/client.rs index 65e7811..379cf67 100644 --- a/pyo3-object_store/src/client.rs +++ b/pyo3-object_store/src/client.rs @@ -1,9 +1,11 @@ use std::collections::HashMap; use std::str::FromStr; +use http::{HeaderMap, HeaderName, HeaderValue}; use object_store::{ClientConfigKey, ClientOptions}; use pyo3::prelude::*; use pyo3::pybacked::PyBackedStr; +use pyo3::types::PyDict; use crate::error::PyObjectStoreError; @@ -37,10 +39,15 @@ impl<'py> FromPyObject<'py> for PyClientConfigValue { /// A wrapper around `ClientOptions` that implements [`FromPyObject`]. #[derive(Debug)] +#[pyclass(name = "ClientOptions", frozen)] pub struct PyClientOptions(ClientOptions); impl<'py> FromPyObject<'py> for PyClientOptions { fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + if let Ok(options) = ob.downcast::() { + return Ok(Self(options.get().0.clone())); + } + let py_input = ob.extract::>()?; let mut options = ClientOptions::new(); for (key, value) in py_input.into_iter() { @@ -55,3 +62,75 @@ impl From for ClientOptions { value.0 } } + +#[pymethods] +impl PyClientOptions { + #[new] + #[pyo3(signature = (*, default_headers = None, **kwargs))] + // TODO: add kwargs + fn py_new( + default_headers: Option, + kwargs: Option<&Bound>, + ) -> PyResult { + let mut options = ClientOptions::default(); + if let Some(default_headers) = default_headers { + options = options.with_default_headers(default_headers.0); + } + if let Some(kwargs) = kwargs { + let kwargs = kwargs.extract::>()?; + for (key, value) in kwargs.into_iter() { + options = options.with_config(key.0, value); + } + } + + Ok(Self(options)) + } +} + +// use pyo3::prelude::*; +// use pyo3::types::PyDict; + +// #[pyfunction] +// #[pyo3(signature = (**kwds))] +// fn num_kwds(kwds: Option<&Bound<'_, PyDict>>) -> usize { +// kwds.map_or(0, |dict| dict.len()) +// } + +// #[pymodule] +// fn module_with_functions(m: &Bound<'_, PyModule>) -> PyResult<()> { +// m.add_function(wrap_pyfunction!(num_kwds, m)?) +// } + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct PyHeaderName(HeaderName); + +impl<'py> FromPyObject<'py> for PyHeaderName { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + // TODO: check that this works on both str and bytes input + Ok(Self(HeaderName::from_bytes(ob.extract()?).unwrap())) + } +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct PyHeaderValue(HeaderValue); + +impl<'py> FromPyObject<'py> for PyHeaderValue { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + Ok(Self( + HeaderValue::from_str(ob.extract::()?.as_ref()).unwrap(), + )) + } +} + +pub struct PyHeaderMap(HeaderMap); + +impl<'py> FromPyObject<'py> for PyHeaderMap { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + let py_input = ob.extract::>()?; + let mut header = HeaderMap::with_capacity(py_input.len()); + for (key, value) in py_input.into_iter() { + header.insert(key.0, value.0); + } + Ok(Self(header)) + } +} diff --git a/pyo3-object_store/type-hints/__init__.pyi b/pyo3-object_store/type-hints/__init__.pyi index 4c1b7b0..7c23ae6 100644 --- a/pyo3-object_store/type-hints/__init__.pyi +++ b/pyo3-object_store/type-hints/__init__.pyi @@ -6,6 +6,7 @@ from ._aws import S3Store as S3Store from ._azure import AzureConfigKey as AzureConfigKey from ._azure import AzureStore as AzureStore from ._client import ClientConfigKey as ClientConfigKey +from ._client import ClientOptions as ClientOptions from ._gcs import GCSConfigKey as GCSConfigKey from ._gcs import GCSStore as GCSStore from ._http import HTTPStore as HTTPStore diff --git a/pyo3-object_store/type-hints/_aws.pyi b/pyo3-object_store/type-hints/_aws.pyi index 97fef44..ff9d394 100644 --- a/pyo3-object_store/type-hints/_aws.pyi +++ b/pyo3-object_store/type-hints/_aws.pyi @@ -4,7 +4,7 @@ import boto3 import botocore import botocore.session -from ._client import ClientConfigKey +from ._client import ClientConfigKey, ClientOptions from ._retry import RetryConfig S3ConfigKey = Literal[ @@ -164,7 +164,7 @@ class S3Store: bucket: str, *, config: Dict[S3ConfigKey | str, str] | None = None, - client_options: Dict[ClientConfigKey, str | bool] | None = None, + client_options: Dict[ClientConfigKey, str | bool] | ClientOptions | None = None, retry_config: RetryConfig | None = None, ) -> S3Store: """Construct a new S3Store with credentials inferred from a boto3 Session diff --git a/pyo3-object_store/type-hints/_client.pyi b/pyo3-object_store/type-hints/_client.pyi index 4424758..9007388 100644 --- a/pyo3-object_store/type-hints/_client.pyi +++ b/pyo3-object_store/type-hints/_client.pyi @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Dict, Literal ClientConfigKey = Literal[ "allow_http", @@ -61,3 +61,6 @@ Either lower case or upper case strings are accepted. response body has finished. - `"user_agent"`: User-Agent header to be used by this client. """ + +class ClientOptions: + def __init__(self, *, default_headers: Dict[str, str], **kwargs) -> None: ... diff --git a/tests/test_get.py b/tests/test_get.py index 7bef9b2..e454636 100644 --- a/tests/test_get.py +++ b/tests/test_get.py @@ -1,7 +1,24 @@ import pytest import obstore as obs -from obstore.store import MemoryStore +import obstore + +dir(obstore) +import obstore.store + +dir(obstore.store) + +from obstore.store import MemoryStore, ClientOptions, S3Store + +import boto3 + +session = boto3.Session() + +options = ClientOptions(default_headers={b"x-amz-request-payer": "requester"}) +store = S3Store.from_session(session, "naip-visualization", client_options=options) +result = obs.list(store, "ny/2022/60cm/rgb/40073/") +out = result.collect() +len(out) def test_stream_sync():