Fix credential extraction; never try to await non-awaitable object (#345)

* Fix credential extraction

* elide

Author: kylebarron

pyo3-object_store/src/aws/credentials.rs

@@ -9,7 +9,7 @@ use pyo3::intern;
 use pyo3::prelude::*;
 
 use crate::aws::store::PyAmazonS3Config;
-use crate::credentials::{TemporaryToken, TokenCache};
+use crate::credentials::{is_awaitable, TemporaryToken, TokenCache};
 
 /// A wrapper around an [AwsCredential] that includes an optional expiry timestamp.
 struct PyAwsCredential {
@@ -138,10 +138,10 @@ impl PyCredentialProviderResult {
 
 impl<'py> FromPyObject<'py> for PyCredentialProviderResult {
     fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
-        if let Ok(credentials) = ob.extract() {
-            Ok(Self::Sync(credentials))
-        } else {
+        if is_awaitable(ob)? {
             Ok(Self::Async(ob.clone().unbind()))
+        } else {
+            Ok(Self::Sync(ob.extract()?))
         }
     }
 }
@@ -164,8 +164,8 @@ impl PyAWSCredentialProvider {
         let credential = self
             .call()
             .await
-            .map_err(|err| object_store::Error::Generic {
-                store: "External AWS credential provider",
+            .map_err(|err| object_store::Error::Unauthenticated {
+                path: "External AWS credential provider".to_string(),
                 source: Box::new(err),
             })?;
 

pyo3-object_store/src/azure/credentials.rs

@@ -11,7 +11,7 @@ use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
 
 use crate::azure::error::Error;
-use crate::credentials::{TemporaryToken, TokenCache};
+use crate::credentials::{is_awaitable, TemporaryToken, TokenCache};
 use crate::PyObjectStoreError;
 
 struct PyAzureAccessKey {
@@ -200,10 +200,10 @@ impl PyCredentialProviderResult {
 
 impl<'py> FromPyObject<'py> for PyCredentialProviderResult {
     fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
-        if let Ok(credentials) = ob.extract() {
-            Ok(Self::Sync(credentials))
-        } else {
+        if is_awaitable(ob)? {
             Ok(Self::Async(ob.clone().unbind()))
+        } else {
+            Ok(Self::Sync(ob.extract()?))
         }
     }
 }
@@ -223,8 +223,8 @@ impl PyAzureCredentialProvider {
         let credential = self
             .call()
             .await
-            .map_err(|err| object_store::Error::Generic {
-                store: "External Azure credential provider",
+            .map_err(|err| object_store::Error::Unauthenticated {
+                path: "External Azure credential provider".to_string(),
                 source: Box::new(err),
             })?;
 

pyo3-object_store/src/credentials.rs

@@ -1,5 +1,8 @@
 use chrono::Utc;
 use chrono::{DateTime, TimeDelta};
+use pyo3::intern;
+use pyo3::prelude::*;
+use pyo3::types::PyTuple;
 use std::future::Future;
 use tokio::sync::Mutex;
 
@@ -87,3 +90,12 @@ impl<T: Clone + Send> TokenCache<T> {
         Ok(token)
     }
 }
+
+/// Check whether a Python object is awaitable
+pub(crate) fn is_awaitable(ob: &Bound<PyAny>) -> PyResult<bool> {
+    let py = ob.py();
+    let inspect_mod = py.import(intern!(py, "inspect"))?;
+    inspect_mod
+        .call_method1(intern!(py, "isawaitable"), PyTuple::new(py, [ob])?)?
+        .extract::<bool>()
+}

pyo3-object_store/src/gcp/credentials.rs

@@ -8,7 +8,7 @@ use pyo3::exceptions::PyTypeError;
 use pyo3::intern;
 use pyo3::prelude::*;
 
-use crate::credentials::{TemporaryToken, TokenCache};
+use crate::credentials::{is_awaitable, TemporaryToken, TokenCache};
 
 /// Ref https://github.com/apache/arrow-rs/pull/6638
 const DEFAULT_GCP_MIN_TTL: TimeDelta = TimeDelta::minutes(4);
@@ -114,10 +114,10 @@ impl PyCredentialProviderResult {
 
 impl<'py> FromPyObject<'py> for PyCredentialProviderResult {
     fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
-        if let Ok(credentials) = ob.extract() {
-            Ok(Self::Sync(credentials))
-        } else {
+        if is_awaitable(ob)? {
             Ok(Self::Async(ob.clone().unbind()))
+        } else {
+            Ok(Self::Sync(ob.extract()?))
         }
     }
 }
@@ -140,8 +140,8 @@ impl PyGcpCredentialProvider {
         let credential = self
             .call()
             .await
-            .map_err(|err| object_store::Error::Generic {
-                store: "External GCP credential provider",
+            .map_err(|err| object_store::Error::Unauthenticated {
+                path: "External GCP credential provider".to_string(),
                 source: Box::new(err),
             })?;
 

tests/store/test_s3.py

@@ -1,11 +1,12 @@
 # ruff: noqa: PGH003
 
 import pickle
+from datetime import UTC, datetime
 
 import pytest
 
 import obstore as obs
-from obstore.exceptions import BaseError
+from obstore.exceptions import BaseError, UnauthenticatedError
 from obstore.store import S3Store, from_url
 
 
@@ -97,3 +98,27 @@ def test_config_round_trip():
     assert store.prefix == new_store.prefix
     assert store.client_options == new_store.client_options
     assert store.retry_config == new_store.retry_config
+
+
+def test_invalid_credential_provider():
+    """Test that passing an invalid synchronous credential provider raises an error.
+
+    instead of trying to await the value.
+    """
+
+    def credential_provider():
+        return {"access_key_id": "str", "expires_at": datetime.now(UTC)}
+
+    store = S3Store("bucket", credential_provider=credential_provider)  # type: ignore
+    with pytest.raises(UnauthenticatedError):
+        obs.list(store).collect()
+
+
+@pytest.mark.asyncio
+async def test_invalid_credential_provider_async():
+    async def credential_provider():
+        return {"access_key_id": "str", "expires_at": datetime.now(UTC)}
+
+    store = S3Store("bucket", credential_provider=credential_provider)  # type: ignore
+    with pytest.raises(UnauthenticatedError):
+        await obs.list(store).collect_async()