diff --git a/obstore/python/obstore/_store/__init__.pyi b/obstore/python/obstore/_store/__init__.pyi index 39eb60ec..6bc3ee0a 100644 --- a/obstore/python/obstore/_store/__init__.pyi +++ b/obstore/python/obstore/_store/__init__.pyi @@ -16,6 +16,10 @@ from ._azure import AzureCredentialProvider as AzureCredentialProvider from ._azure import AzureSASToken as AzureSASToken from ._azure import AzureStore as AzureStore from ._client import ClientConfig as ClientConfig +from ._client import ClientFactory as ClientFactory +from ._client import HttpRequest as HttpRequest +from ._client import HttpResponse as HttpResponse +from ._client import HttpService as HttpService from ._gcs import GCSConfig as GCSConfig from ._gcs import GCSCredential as GCSCredential from ._gcs import GCSCredentialProvider as GCSCredentialProvider @@ -42,6 +46,7 @@ def from_url( client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, credential_provider: S3CredentialProvider | None = None, + client_factory: ClientFactory | None = None, **kwargs: Unpack[S3Config], ) -> ObjectStore: ... @overload @@ -52,6 +57,7 @@ def from_url( client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, credential_provider: GCSCredentialProvider | None = None, + client_factory: ClientFactory | None = None, **kwargs: Unpack[GCSConfig], ) -> ObjectStore: ... @overload @@ -62,6 +68,7 @@ def from_url( client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, credential_provider: AzureCredentialProvider | None = None, + client_factory: ClientFactory | None = None, **kwargs: Unpack[AzureConfig], ) -> ObjectStore: ... @overload @@ -73,6 +80,7 @@ def from_url( retry_config: None = None, automatic_cleanup: bool = False, mkdir: bool = False, + client_factory: ClientFactory | None = None, ) -> ObjectStore: ... def from_url( # type: ignore[misc] # docstring in pyi file url: str, @@ -81,6 +89,7 @@ def from_url( # type: ignore[misc] # docstring in pyi file client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, credential_provider: Callable | None = None, + client_factory: ClientFactory | None = None, **kwargs: Any, ) -> ObjectStore: """Easy construction of store by URL, identifying the relevant store. @@ -120,6 +129,7 @@ def from_url( # type: ignore[misc] # docstring in pyi file client_options: HTTP Client options. Defaults to None. retry_config: Retry configuration. Defaults to None. credential_provider: A callback to provide custom credentials to the underlying store classes. + client_factory: A custom HTTP client factory to use for requests. Defaults to None, which uses the Rust `reqwest` library to handle HTTP requests. kwargs: per-store configuration passed down to store-specific builders. """ diff --git a/obstore/python/obstore/_store/_aws.pyi b/obstore/python/obstore/_store/_aws.pyi index 3efa8bb7..d85eaaea 100644 --- a/obstore/python/obstore/_store/_aws.pyi +++ b/obstore/python/obstore/_store/_aws.pyi @@ -3,7 +3,7 @@ from collections.abc import Coroutine from datetime import datetime from typing import Any, Literal, Protocol, TypedDict -from ._client import ClientConfig +from ._client import ClientConfig, ClientFactory from ._retry import RetryConfig if sys.version_info >= (3, 10): @@ -476,6 +476,7 @@ class S3Store: client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, credential_provider: S3CredentialProvider | None = None, + client_factory: ClientFactory | None = None, **kwargs: Unpack[S3Config], # type: ignore # noqa: PGH003 (bucket key overlaps with positional arg) ) -> None: """Create a new S3Store. @@ -489,6 +490,7 @@ class S3Store: client_options: HTTP Client options. Defaults to None. retry_config: Retry configuration. Defaults to None. credential_provider: A callback to provide custom S3 credentials. + client_factory: A custom HTTP client factory to use for requests. Defaults to None, which uses the Rust `reqwest` library to handle HTTP requests. kwargs: AWS configuration values. Supports the same values as `config`, but as named keyword args. Returns: @@ -504,6 +506,7 @@ class S3Store: client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, credential_provider: S3CredentialProvider | None = None, + client_factory: ClientFactory | None = None, **kwargs: Unpack[S3Config], ) -> Self: """Parse available connection info from a well-known storage URL. @@ -529,6 +532,7 @@ class S3Store: client_options: HTTP Client options. Defaults to None. retry_config: Retry configuration. Defaults to None. credential_provider: A callback to provide custom S3 credentials. + client_factory: A custom HTTP client factory to use for requests. Defaults to None, which uses the Rust `reqwest` library to handle HTTP requests. kwargs: AWS configuration values. Supports the same values as `config`, but as named keyword args. diff --git a/obstore/python/obstore/_store/_azure.pyi b/obstore/python/obstore/_store/_azure.pyi index c2742155..993b5d70 100644 --- a/obstore/python/obstore/_store/_azure.pyi +++ b/obstore/python/obstore/_store/_azure.pyi @@ -3,7 +3,7 @@ from collections.abc import Coroutine from datetime import datetime from typing import Any, Protocol, TypedDict -from ._client import ClientConfig +from ._client import ClientConfig, ClientFactory from ._retry import RetryConfig if sys.version_info >= (3, 10): @@ -332,6 +332,7 @@ class AzureStore: client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, credential_provider: AzureCredentialProvider | None = None, + client_factory: ClientFactory | None = None, **kwargs: Unpack[AzureConfig], # type: ignore # noqa: PGH003 (container_name key overlaps with positional arg) ) -> None: """Construct a new AzureStore. @@ -345,6 +346,7 @@ class AzureStore: client_options: HTTP Client options. Defaults to None. retry_config: Retry configuration. Defaults to None. credential_provider: A callback to provide custom Azure credentials. + client_factory: A custom HTTP client factory to use for requests. Defaults to None, which uses the Rust `reqwest` library to handle HTTP requests. kwargs: Azure configuration values. Supports the same values as `config`, but as named keyword args. Returns: @@ -362,6 +364,7 @@ class AzureStore: client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, credential_provider: AzureCredentialProvider | None = None, + client_factory: ClientFactory | None = None, **kwargs: Unpack[AzureConfig], ) -> Self: """Construct a new AzureStore with values populated from a well-known storage URL. @@ -396,6 +399,7 @@ class AzureStore: client_options: HTTP Client options. Defaults to None. retry_config: Retry configuration. Defaults to None. credential_provider: A callback to provide custom Azure credentials. + client_factory: A custom HTTP client factory to use for requests. Defaults to None, which uses the Rust `reqwest` library to handle HTTP requests. kwargs: Azure configuration values. Supports the same values as `config`, but as named keyword args. Returns: diff --git a/obstore/python/obstore/_store/_client.pyi b/obstore/python/obstore/_store/_client.pyi index 626d4854..1dfe41a0 100644 --- a/obstore/python/obstore/_store/_client.pyi +++ b/obstore/python/obstore/_store/_client.pyi @@ -1,5 +1,8 @@ +from collections.abc import AsyncIterable, Buffer, Iterable from datetime import timedelta -from typing import TypedDict +from http import HTTPMethod, HTTPStatus +from typing import Literal, Protocol, TypedDict +from urllib.parse import ParseResult class ClientConfig(TypedDict, total=False): """HTTP client configuration. @@ -85,3 +88,25 @@ class ClientConfig(TypedDict, total=False): """ user_agent: str """User-Agent header to be used by this client.""" + +class HttpRequest(TypedDict): + method: HTTPMethod + uri: ParseResult + version: Literal["0.9", "1.0", "1.1", "2.0", "3.0"] + headers: Iterable[tuple[str, bytes]] + body: Buffer | None + +class HttpResponse(Protocol): + status: int | HTTPStatus + version: Literal["0.9", "1.0", "1.1", "2.0", "3.0"] + headers: Iterable[tuple[str, str | bytes]] + # TODO: not sure yet what body will look like + body: AsyncIterable[Buffer] + +# This maps to what is called PyHttpConnector in Rust +class ClientFactory(Protocol): + def connect(self, options: ClientConfig) -> HttpService: ... + +class HttpService(Protocol): + async def __call__(self, req: HttpRequest) -> HttpResponse: + """Perform the given `HttpRequest`, returning an `HttpResponse`.""" diff --git a/obstore/python/obstore/_store/_gcs.pyi b/obstore/python/obstore/_store/_gcs.pyi index e9813a94..e1e5cc6f 100644 --- a/obstore/python/obstore/_store/_gcs.pyi +++ b/obstore/python/obstore/_store/_gcs.pyi @@ -3,7 +3,7 @@ from collections.abc import Coroutine from datetime import datetime from typing import Any, Protocol, TypedDict -from ._client import ClientConfig +from ._client import ClientConfig, ClientFactory from ._retry import RetryConfig if sys.version_info >= (3, 11): @@ -146,6 +146,7 @@ class GCSStore: client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, credential_provider: GCSCredentialProvider | None = None, + client_factory: ClientFactory | None = None, **kwargs: Unpack[GCSConfig], # type: ignore # noqa: PGH003 (bucket key overlaps with positional arg) ) -> None: """Construct a new GCSStore. @@ -159,6 +160,7 @@ class GCSStore: client_options: HTTP Client options. Defaults to None. retry_config: Retry configuration. Defaults to None. credential_provider: A callback to provide custom Google credentials. + client_factory: A custom HTTP client factory to use for requests. Defaults to None, which uses the Rust `reqwest` library to handle HTTP requests. kwargs: GCS configuration values. Supports the same values as `config`, but as named keyword args. Returns: @@ -176,6 +178,7 @@ class GCSStore: client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, credential_provider: GCSCredentialProvider | None = None, + client_factory: ClientFactory | None = None, **kwargs: Unpack[GCSConfig], ) -> Self: """Construct a new GCSStore with values populated from a well-known storage URL. @@ -198,6 +201,7 @@ class GCSStore: client_options: HTTP Client options. Defaults to None. retry_config: Retry configuration. Defaults to None. credential_provider: A callback to provide custom Google credentials. + client_factory: A custom HTTP client factory to use for requests. Defaults to None, which uses the Rust `reqwest` library to handle HTTP requests. kwargs: GCS configuration values. Supports the same values as `config`, but as named keyword args. Returns: diff --git a/obstore/python/obstore/_store/_http.pyi b/obstore/python/obstore/_store/_http.pyi index 3558709d..eab72195 100644 --- a/obstore/python/obstore/_store/_http.pyi +++ b/obstore/python/obstore/_store/_http.pyi @@ -1,6 +1,6 @@ import sys -from ._client import ClientConfig +from ._client import ClientConfig, ClientFactory from ._retry import RetryConfig if sys.version_info >= (3, 11): @@ -17,6 +17,7 @@ class HTTPStore: *, client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, + client_factory: ClientFactory | None = None, ) -> None: """Construct a new HTTPStore from a URL. @@ -31,6 +32,7 @@ class HTTPStore: Keyword Args: client_options: HTTP Client options. Defaults to None. retry_config: Retry configuration. Defaults to None. + client_factory: A custom HTTP client factory to use for requests. Defaults to None, which uses the Rust `reqwest` library to handle HTTP requests. Returns: HTTPStore @@ -44,6 +46,7 @@ class HTTPStore: *, client_options: ClientConfig | None = None, retry_config: RetryConfig | None = None, + client_factory: ClientFactory | None = None, ) -> Self: """Construct a new HTTPStore from a URL. diff --git a/obstore/python/obstore/client/__init__.py b/obstore/python/obstore/client/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/obstore/python/obstore/client/aiohttp.py b/obstore/python/obstore/client/aiohttp.py new file mode 100644 index 00000000..27930a0c --- /dev/null +++ b/obstore/python/obstore/client/aiohttp.py @@ -0,0 +1,102 @@ +"""Aiohttp client implementation.""" + +# ruff: noqa: PLR2004 + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal +from weakref import WeakSet + +from aiohttp import ClientSession +from multidict import MultiDict + +if TYPE_CHECKING: + from collections.abc import AsyncIterable, Iterable + + from aiohttp import ClientResponse + + from obstore.store import ClientConfig, ClientFactory, HttpRequest, HttpResponse + + +class AiohttpClientFactory: + """A client factory for Aiohttp.""" + + _sessions: WeakSet[ClientSession] + + def __init__(self) -> None: + """Create a new AiohttpClientFactory.""" + self._sessions = WeakSet() + + def connect(self, options: ClientConfig) -> _AiohttpService: # noqa: ARG002 + """Create a new HTTP Client.""" + session = ClientSession() + self._sessions.add(session) + return _AiohttpService(session) + + async def close_all(self) -> None: + """Close all generated aiohttp ClientSession instances.""" + futs = [session.close() for session in self._sessions] + await asyncio.gather(*futs) + + +class _AiohttpService: + session: ClientSession + + def __init__(self, session: ClientSession) -> None: + self.session = session + + async def __call__(self, req: HttpRequest) -> HttpResponse: + method = req["method"] + url = req["uri"].geturl() + + headers: MultiDict[str] = MultiDict() + for header_name, header_value in req["headers"]: + # TODO: it seems like aiohttp only allows string header values? + headers.add(header_name, str(header_value)) + + async with self.session.request(method, url, headers=headers) as resp: + version = _get_http_version_from_response(resp) + return _AiohttpResponse( + status=resp.status, + version=version, + headers=resp.headers.items(), + body=resp.content, + ) + + +def _get_http_version_from_response( + resp: ClientResponse, +) -> Literal["0.9", "1.0", "1.1", "2.0", "3.0"]: + v = resp.version + if v is not None: + if v.major == 0 and v.minor == 9: + return "0.9" + if v.major == 1 and v.minor == 0: + return "1.0" + if v.major == 1 and v.minor == 1: + return "1.1" + if v.major == 2 and v.minor == 0: + return "2.0" + if v.major == 3 and v.minor == 0: + return "3.0" + + return "1.1" + + +@dataclass +class _AiohttpResponse: + status: int + version: Literal["0.9", "1.0", "1.1", "2.0", "3.0"] + headers: Iterable[tuple[str, str | bytes]] + body: AsyncIterable + + +if TYPE_CHECKING: + # Just for testing + def _accepts_factory(factory: ClientFactory) -> None: + pass + + aiohttp_factory = AiohttpClientFactory() + _accepts_factory(aiohttp_factory) diff --git a/obstore/python/obstore/store.py b/obstore/python/obstore/store.py index 4d6823ec..e3f2ff0d 100644 --- a/obstore/python/obstore/store.py +++ b/obstore/python/obstore/store.py @@ -47,9 +47,13 @@ AzureSASToken, # noqa: TC004 BackoffConfig, # noqa: TC004 ClientConfig, # noqa: TC004 + ClientFactory, # noqa: TC004 GCSConfig, # noqa: TC004 GCSCredential, # noqa: TC004 GCSCredentialProvider, # noqa: TC004 + HttpRequest, # noqa: TC004 + HttpResponse, # noqa: TC004 + HttpService, # noqa: TC004 RetryConfig, # noqa: TC004 S3Config, # noqa: TC004 S3Credential, # noqa: TC004 @@ -82,11 +86,15 @@ "AzureStore", "BackoffConfig", "ClientConfig", + "ClientFactory", "GCSConfig", "GCSCredential", "GCSCredentialProvider", "GCSStore", "HTTPStore", + "HttpRequest", + "HttpResponse", + "HttpService", "LocalStore", "MemoryStore", "RetryConfig", diff --git a/pyo3-object_store/src/aws/store.rs b/pyo3-object_store/src/aws/store.rs index 8cbb0535..9af87924 100644 --- a/pyo3-object_store/src/aws/store.rs +++ b/pyo3-object_store/src/aws/store.rs @@ -17,7 +17,7 @@ use crate::error::{GenericError, ParseUrlError, PyObjectStoreError, PyObjectStor use crate::path::PyPath; use crate::prefix::MaybePrefixedStore; use crate::retry::PyRetryConfig; -use crate::PyUrl; +use crate::{PyHttpConnector, PyUrl}; #[derive(Debug, Clone, PartialEq)] struct S3Config { @@ -26,6 +26,7 @@ struct S3Config { client_options: Option, retry_config: Option, credential_provider: Option, + client_factory: Option, } impl S3Config { @@ -54,6 +55,9 @@ impl S3Config { if let Some(credential_provider) = &self.credential_provider { kwargs.set_item("credential_provider", credential_provider)?; } + if let Some(client_factory) = &self.client_factory { + kwargs.set_item("client_factory", client_factory)?; + } PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) } @@ -85,7 +89,8 @@ impl PyS3Store { impl PyS3Store { // Create from parameters #[new] - #[pyo3(signature = (bucket=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + #[expect(clippy::too_many_arguments)] + #[pyo3(signature = (bucket=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, client_factory=None, **kwargs))] fn new( bucket: Option, prefix: Option, @@ -93,6 +98,7 @@ impl PyS3Store { client_options: Option, retry_config: Option, credential_provider: Option, + client_factory: Option, kwargs: Option, ) -> PyObjectStoreResult { let mut builder = AmazonS3Builder::from_env(); @@ -124,6 +130,10 @@ impl PyS3Store { builder = builder.with_credentials(Arc::new(credential_provider)); } + if let Some(client_factory) = client_factory.clone() { + builder = builder.with_http_connector(client_factory); + } + builder = combined_config.clone().apply_config(builder); Ok(Self { @@ -134,12 +144,14 @@ impl PyS3Store { client_options, retry_config, credential_provider, + client_factory, }, }) } #[classmethod] - #[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + #[expect(clippy::too_many_arguments)] + #[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, client_factory=None, **kwargs))] pub(crate) fn from_url<'py>( cls: &Bound<'py, PyType>, url: PyUrl, @@ -147,6 +159,7 @@ impl PyS3Store { client_options: Option, retry_config: Option, credential_provider: Option, + client_factory: Option, kwargs: Option, ) -> PyObjectStoreResult> { // We manually parse the URL to find the prefix because `with_url` does not apply the @@ -168,6 +181,7 @@ impl PyS3Store { kwargs.set_item("client_options", client_options)?; kwargs.set_item("retry_config", retry_config)?; kwargs.set_item("credential_provider", credential_provider)?; + kwargs.set_item("client_factory", client_factory)?; Ok(cls.call((), Some(&kwargs))?) } diff --git a/pyo3-object_store/src/azure/store.rs b/pyo3-object_store/src/azure/store.rs index edfd2b57..b651fd2b 100644 --- a/pyo3-object_store/src/azure/store.rs +++ b/pyo3-object_store/src/azure/store.rs @@ -15,7 +15,7 @@ use crate::config::PyConfigValue; use crate::error::{GenericError, ParseUrlError, PyObjectStoreError, PyObjectStoreResult}; use crate::path::PyPath; use crate::retry::PyRetryConfig; -use crate::{MaybePrefixedStore, PyUrl}; +use crate::{MaybePrefixedStore, PyHttpConnector, PyUrl}; #[derive(Debug, Clone, PartialEq)] struct AzureConfig { @@ -24,6 +24,7 @@ struct AzureConfig { client_options: Option, retry_config: Option, credential_provider: Option, + client_factory: Option, } impl AzureConfig { @@ -60,6 +61,9 @@ impl AzureConfig { if let Some(credential_provider) = &self.credential_provider { kwargs.set_item("credential_provider", credential_provider)?; } + if let Some(client_factory) = &self.client_factory { + kwargs.set_item("client_factory", client_factory)?; + } PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) } @@ -91,7 +95,7 @@ impl PyAzureStore { impl PyAzureStore { // Create from parameters #[new] - #[pyo3(signature = (container_name=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + #[pyo3(signature = (container_name=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, client_factory=None, **kwargs))] fn new( container_name: Option, mut prefix: Option, @@ -99,6 +103,7 @@ impl PyAzureStore { client_options: Option, retry_config: Option, credential_provider: Option, + client_factory: Option, kwargs: Option, ) -> PyObjectStoreResult { let mut builder = MicrosoftAzureBuilder::from_env(); @@ -141,6 +146,10 @@ impl PyAzureStore { builder = builder.with_credentials(Arc::new(credential_provider)); } + if let Some(client_factory) = client_factory.clone() { + builder = builder.with_http_connector(client_factory); + } + builder = combined_config.clone().apply_config(builder); Ok(Self { @@ -151,12 +160,13 @@ impl PyAzureStore { client_options, retry_config, credential_provider, + client_factory, }, }) } #[classmethod] - #[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + #[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, client_factory=None, **kwargs))] pub(crate) fn from_url<'py>( cls: &Bound<'py, PyType>, url: PyUrl, @@ -164,6 +174,7 @@ impl PyAzureStore { client_options: Option, retry_config: Option, credential_provider: Option, + client_factory: Option, kwargs: Option, ) -> PyObjectStoreResult> { // We manually parse the URL to find the prefix because `parse_url` does not apply the @@ -185,6 +196,7 @@ impl PyAzureStore { kwargs.set_item("client_options", client_options)?; kwargs.set_item("retry_config", retry_config)?; kwargs.set_item("credential_provider", credential_provider)?; + kwargs.set_item("client_factory", client_factory)?; Ok(cls.call((), Some(&kwargs))?) } diff --git a/pyo3-object_store/src/client/connector.rs b/pyo3-object_store/src/client/connector.rs new file mode 100644 index 00000000..56309670 --- /dev/null +++ b/pyo3-object_store/src/client/connector.rs @@ -0,0 +1,301 @@ +use async_trait::async_trait; +use object_store::client::{ + HttpClient, HttpConnector, HttpError, HttpRequest, HttpResponse, HttpResponseBody, HttpService, +}; +use object_store::{ClientOptions, Result}; +use pyo3::exceptions::PyValueError; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedStr; +use pyo3::types::{PyDict, PyString, PyTuple}; + +use crate::client::options::PyHeaderMap; +use crate::PyClientOptions; + +/// An [HttpConnector] defined from Python. +#[derive(Debug)] +pub struct PyHttpConnector(Py); + +impl PyHttpConnector { + fn equals(&self, py: Python, other: &Self) -> PyResult { + self.0 + .call_method1(py, "__eq__", PyTuple::new(py, vec![&other.0])?)? + .extract(py) + } +} + +impl HttpConnector for PyHttpConnector { + fn connect(&self, options: &ClientOptions) -> Result { + let py_options = PyClientOptions::from(options.clone()); + let http_service = Python::attach(|py| { + self.0 + .call_method1(py, intern!(py, "connect"), (py_options,)) + }) + .expect("httpconnector.connect"); + let client = HttpClient::new(PyHttpService(http_service)); + Ok(client) + } +} + +impl Clone for PyHttpConnector { + fn clone(&self) -> Self { + Python::attach(|py| Self(self.0.clone_ref(py))) + } +} + +impl PartialEq for PyHttpConnector { + fn eq(&self, other: &Self) -> bool { + Python::attach(|py| self.equals(py, other)).unwrap_or(false) + } +} + +impl<'py> FromPyObject<'_, 'py> for PyHttpConnector { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, PyAny>) -> std::result::Result { + let py = obj.py(); + if !obj.hasattr(intern!(py, "connect"))? { + Err(PyValueError::new_err( + "client_factory must have a method named `connect`.", + )) + } else { + Ok(Self(obj.as_unbound().clone_ref(py))) + } + } +} + +impl<'py> IntoPyObject<'py> for PyHttpConnector { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + (&self).into_pyobject(py) + } +} + +impl<'py> IntoPyObject<'py> for &PyHttpConnector { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(self.0.bind(py).clone()) + } +} + +/// An [HttpService] defined from Python. +#[derive(Debug)] +pub struct PyHttpService(Py); + +#[async_trait] +impl HttpService for PyHttpService { + /// Perform [`HttpRequest`] returning [`HttpResponse`] + async fn call(&self, req: HttpRequest) -> Result { + let py_req = PyHttpRequest(req); + let py_resp = Python::attach(|py| { + self.0 + .call1(py, (py_req,)) + .expect("httpservice.call") + .extract::(py) + .expect("py http response extraction") + }); + Ok(py_resp.0) + } +} + +pub struct PyHttpRequest(HttpRequest); + +impl<'py> IntoPyObject<'py> for PyHttpRequest { + type Target = PyDict; + type Output = Bound<'py, PyDict>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let (parts, body) = self.0.into_parts(); + + let dict = PyDict::new(py); + dict.set_item( + intern!(py, "method"), + PyHttpMethod(parts.method).into_pyobject(py)?, + )?; + dict.set_item(intern!(py, "uri"), PyUri(parts.uri).into_pyobject(py)?)?; + dict.set_item( + intern!(py, "version"), + PyHttpVersion(parts.version).into_pyobject(py)?, + )?; + dict.set_item( + intern!(py, "headers"), + PyHeaderMap(parts.headers).into_pyobject(py)?, + )?; + + // TODO: body doesn't currently offer a way to access the underlying PutPayload for + // Inner::PutPayload. + dict.set_item(intern!(py, "body"), body.as_bytes().map(|v| v.as_ref()))?; + + Ok(dict) + } +} + +pub struct PyHttpMethod(http::Method); + +impl<'py> IntoPyObject<'py> for PyHttpMethod { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + // TODO: in the future we could cache these http method constants so we aren't continually + // accessing them + let py_http_mod = py + .import(intern!(py, "http"))? + .get_item(intern!(py, "HTTPMethod"))?; + py_http_mod.call1((self.0.as_str(),)) + } +} + +pub struct PyUri(http::Uri); + +impl<'py> IntoPyObject<'py> for PyUri { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let uri_parts = self.0.into_parts(); + + let mut scheme = ""; + let mut netloc = ""; + let mut path = ""; + // TODO: upstream doesn't have a way to access params? + let params = ""; + let mut query = ""; + // TODO: upstream doesn't have a way to access fragment? + let fragment = ""; + + if let Some(s) = &uri_parts.scheme { + scheme = s.as_str(); + } + + if let Some(path_and_query) = &uri_parts.path_and_query { + path = path_and_query.path(); + if let Some(q) = path_and_query.query() { + query = q; + } + } + + if let Some(authority) = &uri_parts.authority { + netloc = authority.as_str(); + } + + let kwargs = PyDict::new(py); + kwargs.set_item(intern!(py, "scheme"), PyString::new(py, scheme))?; + kwargs.set_item(intern!(py, "netloc"), PyString::new(py, netloc))?; + kwargs.set_item(intern!(py, "path"), path)?; + kwargs.set_item(intern!(py, "params"), params)?; + kwargs.set_item(intern!(py, "query"), query)?; + kwargs.set_item(intern!(py, "fragment"), fragment)?; + + let urllib_parse_mod = py.import(intern!(py, "urllib.parse"))?; + let parse_result_cls = urllib_parse_mod.getattr(intern!(py, "ParseResult"))?; + parse_result_cls.call((), Some(&kwargs)) + } +} + +pub struct PyHttpVersion(http::Version); + +impl<'py> IntoPyObject<'py> for PyHttpVersion { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + use http::Version; + + let s = match self.0 { + Version::HTTP_09 => intern!(py, "0.9"), + Version::HTTP_10 => intern!(py, "1.0"), + Version::HTTP_11 => intern!(py, "1.1"), + Version::HTTP_2 => intern!(py, "2.0"), + Version::HTTP_3 => intern!(py, "3.0"), + _ => unimplemented!("Unknown http version"), + }; + Ok(s.clone()) + } +} + +impl<'py> FromPyObject<'_, 'py> for PyHttpVersion { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, PyAny>) -> std::result::Result { + use http::Version; + + let version_input = obj.extract::()?; + let http_version = match version_input.as_ref() { + "0.9" => Version::HTTP_09, + "1.0" => Version::HTTP_10, + "1.1" => Version::HTTP_11, + "2.0" => Version::HTTP_2, + "3.0" => Version::HTTP_3, + _ => panic!("Unsupported HTTP version"), + }; + Ok(Self(http_version)) + } +} + +pub struct PyHttpResponse(HttpResponse); + +impl<'py> FromPyObject<'_, 'py> for PyHttpResponse { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, PyAny>) -> std::result::Result { + let py = obj.py(); + let status = obj + .getattr(intern!(py, "status"))? + .extract::()?; + let version = obj + .getattr(intern!(py, "version"))? + .extract::()?; + let headers = obj + .getattr(intern!(py, "headers"))? + .extract::()?; + + // TODO: construct body. This probably will have to be a Python object that we poll? + + let resp = http::Response::new(body); + let (mut parts, body) = resp.into_parts(); + + parts.status = status.0; + parts.version = version.0; + parts.headers = headers.0; + + // There's also a `http::response::Builder` API but I can't figure out how I'd convert that + // `Builder` to `Response`?? + Ok(Self(http::Response::from_parts(parts, body))) + } +} + +pub struct PyHttpStatusCode(http::StatusCode); + +impl<'py> FromPyObject<'_, 'py> for PyHttpStatusCode { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, PyAny>) -> std::result::Result { + let code = obj.extract::()?; + Ok(PyHttpStatusCode( + http::StatusCode::from_u16(code).expect("invalid http status code"), + )) + } +} + +pub struct PyHttpResponseBody(HttpResponseBody); + +pub struct PyHttpError(HttpError); + +impl<'py> FromPyObject<'_, 'py> for PyHttpError { + type Error = PyErr; + + fn extract(_obj: Borrowed<'_, 'py, PyAny>) -> std::result::Result { + todo!("create http error kind") + } +} diff --git a/pyo3-object_store/src/client/mod.rs b/pyo3-object_store/src/client/mod.rs new file mode 100644 index 00000000..b35d86a1 --- /dev/null +++ b/pyo3-object_store/src/client/mod.rs @@ -0,0 +1,5 @@ +mod connector; +mod options; + +pub use connector::{PyHttpConnector, PyHttpService}; +pub use options::{PyClientConfigKey, PyClientOptions}; diff --git a/pyo3-object_store/src/client.rs b/pyo3-object_store/src/client/options.rs similarity index 91% rename from pyo3-object_store/src/client.rs rename to pyo3-object_store/src/client/options.rs index 9fe159fb..4b82d227 100644 --- a/pyo3-object_store/src/client.rs +++ b/pyo3-object_store/src/client/options.rs @@ -6,7 +6,7 @@ use object_store::{ClientConfigKey, ClientOptions}; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::pybacked::{PyBackedBytes, PyBackedStr}; -use pyo3::types::{PyDict, PyString}; +use pyo3::types::{PyDict, PyString, PyTuple}; use crate::config::PyConfigValue; use crate::error::PyObjectStoreError; @@ -125,7 +125,7 @@ impl From for ClientOptions { } #[derive(Clone, Debug, PartialEq)] -struct PyHeaderMap(HeaderMap); +pub(crate) struct PyHeaderMap(pub(crate) HeaderMap); impl<'py> FromPyObject<'_, 'py> for PyHeaderMap { type Error = PyErr; @@ -152,29 +152,27 @@ impl<'py> FromPyObject<'_, 'py> for PyHeaderMap { } impl<'py> IntoPyObject<'py> for PyHeaderMap { - type Target = PyDict; - type Output = Bound<'py, PyDict>; + type Target = PyTuple; + type Output = Bound<'py, PyTuple>; type Error = PyErr; fn into_pyobject(self, py: Python<'py>) -> Result { - let dict = PyDict::new(py); - for (key, value) in self.0.iter() { - dict.set_item(key.as_str(), value.as_bytes())?; - } - Ok(dict) + (&self).into_pyobject(py) } } impl<'py> IntoPyObject<'py> for &PyHeaderMap { - type Target = PyDict; - type Output = Bound<'py, PyDict>; + type Target = PyTuple; + type Output = Bound<'py, PyTuple>; type Error = PyErr; fn into_pyobject(self, py: Python<'py>) -> Result { - let dict = PyDict::new(py); - for (key, value) in self.0.iter() { - dict.set_item(key.as_str(), value.as_bytes())?; + let mut headers = vec![]; + + for (header_name, header_value) in self.0.iter() { + headers.push((header_name.as_str(), header_value.as_bytes()).into_pyobject(py)?) } - Ok(dict) + + PyTuple::new(py, headers) } } diff --git a/pyo3-object_store/src/gcp/store.rs b/pyo3-object_store/src/gcp/store.rs index 6e31f9c8..4c53b170 100644 --- a/pyo3-object_store/src/gcp/store.rs +++ b/pyo3-object_store/src/gcp/store.rs @@ -15,7 +15,7 @@ use crate::error::{GenericError, ParseUrlError, PyObjectStoreError, PyObjectStor use crate::gcp::credentials::PyGcpCredentialProvider; use crate::path::PyPath; use crate::retry::PyRetryConfig; -use crate::{MaybePrefixedStore, PyUrl}; +use crate::{MaybePrefixedStore, PyHttpConnector, PyUrl}; #[derive(Debug, Clone, PartialEq)] struct GCSConfig { @@ -24,6 +24,7 @@ struct GCSConfig { client_options: Option, retry_config: Option, credential_provider: Option, + client_factory: Option, } impl GCSConfig { @@ -52,6 +53,9 @@ impl GCSConfig { if let Some(credential_provider) = &self.credential_provider { kwargs.set_item("credential_provider", credential_provider)?; } + if let Some(client_factory) = &self.client_factory { + kwargs.set_item("client_factory", client_factory)?; + } PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) } @@ -83,7 +87,7 @@ impl PyGCSStore { impl PyGCSStore { // Create from parameters #[new] - #[pyo3(signature = (bucket=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + #[pyo3(signature = (bucket=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, client_factory=None, **kwargs))] fn new( bucket: Option, prefix: Option, @@ -91,6 +95,7 @@ impl PyGCSStore { client_options: Option, retry_config: Option, credential_provider: Option, + client_factory: Option, kwargs: Option, ) -> PyObjectStoreResult { let mut builder = GoogleCloudStorageBuilder::from_env(); @@ -111,6 +116,10 @@ impl PyGCSStore { if let Some(credential_provider) = credential_provider.clone() { builder = builder.with_credentials(Arc::new(credential_provider)); } + if let Some(client_factory) = client_factory.clone() { + builder = builder.with_http_connector(client_factory); + } + Ok(Self { store: Arc::new(MaybePrefixedStore::new(builder.build()?, prefix.clone())), config: GCSConfig { @@ -119,12 +128,13 @@ impl PyGCSStore { client_options, retry_config, credential_provider, + client_factory, }, }) } #[classmethod] - #[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + #[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, client_factory=None, **kwargs))] pub(crate) fn from_url<'py>( cls: &Bound<'py, PyType>, url: PyUrl, @@ -132,6 +142,7 @@ impl PyGCSStore { client_options: Option, retry_config: Option, credential_provider: Option, + client_factory: Option, kwargs: Option, ) -> PyObjectStoreResult> { // We manually parse the URL to find the prefix because `parse_url` does not apply the @@ -153,6 +164,7 @@ impl PyGCSStore { kwargs.set_item("client_options", client_options)?; kwargs.set_item("retry_config", retry_config)?; kwargs.set_item("credential_provider", credential_provider)?; + kwargs.set_item("client_factory", client_factory)?; Ok(cls.call((), Some(&kwargs))?) } diff --git a/pyo3-object_store/src/http.rs b/pyo3-object_store/src/http.rs index 6c0f5f7f..ba3f9e08 100644 --- a/pyo3-object_store/src/http.rs +++ b/pyo3-object_store/src/http.rs @@ -7,13 +7,14 @@ use pyo3::{intern, IntoPyObjectExt}; use crate::error::PyObjectStoreResult; use crate::retry::PyRetryConfig; -use crate::{PyClientOptions, PyUrl}; +use crate::{PyClientOptions, PyHttpConnector, PyUrl}; #[derive(Debug, Clone, PartialEq)] struct HTTPConfig { url: PyUrl, client_options: Option, retry_config: Option, + client_factory: Option, } impl HTTPConfig { @@ -27,6 +28,9 @@ impl HTTPConfig { if let Some(retry_config) = &self.retry_config { kwargs.set_item(intern!(py, "retry_config"), retry_config.clone())?; } + if let Some(client_factory) = &self.client_factory { + kwargs.set_item("client_factory", client_factory)?; + } PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) } @@ -59,11 +63,12 @@ impl PyHttpStore { #[pymethods] impl PyHttpStore { #[new] - #[pyo3(signature = (url, *, client_options=None, retry_config=None))] + #[pyo3(signature = (url, *, client_options=None, retry_config=None, client_factory=None))] fn new( url: PyUrl, client_options: Option, retry_config: Option, + client_factory: Option, ) -> PyObjectStoreResult { let mut builder = HttpBuilder::new().with_url(url.clone()); if let Some(client_options) = client_options.clone() { @@ -72,24 +77,30 @@ impl PyHttpStore { if let Some(retry_config) = retry_config.clone() { builder = builder.with_retry(retry_config.into()) } + if let Some(client_factory) = client_factory.clone() { + builder = builder.with_http_connector(client_factory); + } + Ok(Self { store: Arc::new(builder.build()?), config: HTTPConfig { url, client_options, retry_config, + client_factory, }, }) } #[classmethod] - #[pyo3(signature = (url, *, client_options=None, retry_config=None))] + #[pyo3(signature = (url, *, client_options=None, retry_config=None, client_factory=None))] pub(crate) fn from_url<'py>( cls: &Bound<'py, PyType>, py: Python<'py>, url: PyUrl, client_options: Option, retry_config: Option, + client_factory: Option, ) -> PyObjectStoreResult> { // Note: we pass **back** through Python so that if cls is a subclass, we instantiate the // subclass @@ -97,6 +108,7 @@ impl PyHttpStore { kwargs.set_item("url", url)?; kwargs.set_item("client_options", client_options)?; kwargs.set_item("retry_config", retry_config)?; + kwargs.set_item("client_factory", client_factory)?; Ok(cls.call((), Some(&kwargs))?) } diff --git a/pyo3-object_store/src/lib.rs b/pyo3-object_store/src/lib.rs index 0d7d7fdf..fd90f013 100644 --- a/pyo3-object_store/src/lib.rs +++ b/pyo3-object_store/src/lib.rs @@ -22,7 +22,7 @@ mod url; pub use api::{register_exceptions_module, register_store_module}; pub use aws::PyS3Store; pub use azure::PyAzureStore; -pub use client::{PyClientConfigKey, PyClientOptions}; +pub use client::{PyClientConfigKey, PyClientOptions, PyHttpConnector, PyHttpService}; pub use error::{PyObjectStoreError, PyObjectStoreResult}; pub use gcp::PyGCSStore; pub use http::PyHttpStore; diff --git a/pyo3-object_store/src/simple.rs b/pyo3-object_store/src/simple.rs index 1966756b..cf29f35d 100644 --- a/pyo3-object_store/src/simple.rs +++ b/pyo3-object_store/src/simple.rs @@ -10,8 +10,8 @@ use crate::error::GenericError; use crate::retry::PyRetryConfig; use crate::url::PyUrl; use crate::{ - PyAzureStore, PyClientOptions, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, - PyObjectStoreResult, PyS3Store, + PyAzureStore, PyClientOptions, PyGCSStore, PyHttpConnector, PyHttpStore, PyLocalStore, + PyMemoryStore, PyObjectStoreResult, PyS3Store, }; /// Simple construction of stores by url. @@ -19,7 +19,8 @@ use crate::{ // AWS/Azure/Google config keys could overlap. And so we don't want to accidentally parse a config // as an AWS config before knowing that the URL scheme is AWS. #[pyfunction] -#[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] +#[expect(clippy::too_many_arguments)] +#[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, client_factory=None, **kwargs))] pub fn from_url<'py>( py: Python<'py>, url: PyUrl, @@ -27,6 +28,7 @@ pub fn from_url<'py>( client_options: Option, retry_config: Option, credential_provider: Option>, + client_factory: Option, kwargs: Option>, ) -> PyObjectStoreResult> { let (scheme, _) = ObjectStoreScheme::parse(url.as_ref()).map_err(object_store::Error::from)?; @@ -38,6 +40,7 @@ pub fn from_url<'py>( client_options, retry_config, credential_provider.map(|x| x.extract()).transpose()?, + client_factory, kwargs.map(|x| x.extract()).transpose()?, ), ObjectStoreScheme::GoogleCloudStorage => PyGCSStore::from_url( @@ -47,6 +50,7 @@ pub fn from_url<'py>( client_options, retry_config, credential_provider.map(|x| x.extract()).transpose()?, + client_factory, kwargs.map(|x| x.extract()).transpose()?, ), ObjectStoreScheme::MicrosoftAzure => PyAzureStore::from_url( @@ -56,6 +60,7 @@ pub fn from_url<'py>( client_options, retry_config, credential_provider.map(|x| x.extract()).transpose()?, + client_factory, kwargs.map(|x| x.extract()).transpose()?, ), ObjectStoreScheme::Http => { @@ -66,6 +71,7 @@ pub fn from_url<'py>( url, client_options, retry_config, + client_factory, ) } ObjectStoreScheme::Local => {