Support async refreshable credential callbacks in Python bindings

icechunk-python/python/icechunk/credentials.py

@@ -1,7 +1,9 @@
+import asyncio
+import inspect
 import pickle
-from collections.abc import Callable, Mapping
+from collections.abc import Awaitable, Callable, Coroutine, Mapping
 from datetime import datetime
-from typing import cast
+from typing import Any, TypeVar, cast
 
 from icechunk._icechunk_python import (
     AzureCredentials,
@@ -46,16 +48,50 @@
 
 AnyCredential = Credentials.S3 | Credentials.Gcs | Credentials.Azure
 
+S3CredentialProvider = Callable[[], S3StaticCredentials | Awaitable[S3StaticCredentials]]
+GcsCredentialProvider = Callable[[], GcsBearerCredential | Awaitable[GcsBearerCredential]]
+CredentialType = TypeVar("CredentialType")
+
+
+def _resolve_callback_once(
+    get_credentials: Callable[[], CredentialType | Awaitable[CredentialType]],
+) -> CredentialType:
+    value = get_credentials()
+    if inspect.isawaitable(value):
+        if inspect.iscoroutine(value):
+            coroutine = cast(Coroutine[Any, Any, CredentialType], value)
+        else:
+            awaitable = cast(Awaitable[CredentialType], value)
+
+            async def _as_coroutine() -> CredentialType:
+                return await awaitable
+
+            coroutine = _as_coroutine()
+
+        try:
+            return asyncio.run(coroutine)
+        except RuntimeError as err:
+            coroutine.close()
+            if "asyncio.run() cannot be called from a running event loop" in str(err):
+                raise ValueError(
+                    "scatter_initial_credentials=True cannot eagerly evaluate async "
+                    "credential callbacks while an event loop is running. "
+                    "Set scatter_initial_credentials=False in async contexts."
+                ) from err
+            raise
+
+    return value
+
 
 def s3_refreshable_credentials(
-    get_credentials: Callable[[], S3StaticCredentials],
+    get_credentials: S3CredentialProvider,
     scatter_initial_credentials: bool = False,
 ) -> S3Credentials.Refreshable:
     """Create refreshable credentials for S3 and S3 compatible object stores.
 
     Parameters
     ----------
-    get_credentials: Callable[[], S3StaticCredentials]
+    get_credentials: Callable[[], S3StaticCredentials | Awaitable[S3StaticCredentials]]
         Use this function to get and refresh the credentials. The function must be pickable.
     scatter_initial_credentials: bool, optional
         Immediately call and store the value returned by get_credentials. This is useful if the
@@ -64,7 +100,9 @@ def s3_refreshable_credentials(
         set of credentials has expired, the cached value is no longer used. Notice that credentials
         obtained are stored, and they can be sent over the network if you pickle the session/repo.
     """
-    current = get_credentials() if scatter_initial_credentials else None
+    current = (
+        _resolve_callback_once(get_credentials) if scatter_initial_credentials else None
+    )
     return S3Credentials.Refreshable(pickle.dumps(get_credentials), current)
 
 
@@ -116,7 +154,7 @@ def s3_credentials(
     expires_after: datetime | None = None,
     anonymous: bool | None = None,
     from_env: bool | None = None,
-    get_credentials: Callable[[], S3StaticCredentials] | None = None,
+    get_credentials: S3CredentialProvider | None = None,
     scatter_initial_credentials: bool = False,
 ) -> AnyS3Credential:
     """Create credentials for S3 and S3 compatible object stores.
@@ -137,7 +175,7 @@ def s3_credentials(
         If set to True requests to the object store will not be signed
     from_env: bool | None
         Fetch credentials from the operative system environment
-    get_credentials: Callable[[], S3StaticCredentials] | None
+    get_credentials: Callable[[], S3StaticCredentials | Awaitable[S3StaticCredentials]] | None
         Use this function to get and refresh object store credentials
     scatter_initial_credentials: bool, optional
         Immediately call and store the value returned by get_credentials. This is useful if the
@@ -218,14 +256,14 @@ def gcs_static_credentials(
 
 
 def gcs_refreshable_credentials(
-    get_credentials: Callable[[], GcsBearerCredential],
+    get_credentials: GcsCredentialProvider,
     scatter_initial_credentials: bool = False,
 ) -> GcsCredentials.Refreshable:
     """Create refreshable credentials for Google Cloud Storage object store.
 
     Parameters
     ----------
-    get_credentials: Callable[[], S3StaticCredentials]
+    get_credentials: Callable[[], GcsBearerCredential | Awaitable[GcsBearerCredential]]
         Use this function to get and refresh the credentials. The function must be pickable.
     scatter_initial_credentials: bool, optional
         Immediately call and store the value returned by get_credentials. This is useful if the
@@ -235,7 +273,9 @@ def gcs_refreshable_credentials(
         obtained are stored, and they can be sent over the network if you pickle the session/repo.
     """
 
-    current = get_credentials() if scatter_initial_credentials else None
+    current = (
+        _resolve_callback_once(get_credentials) if scatter_initial_credentials else None
+    )
     return GcsCredentials.Refreshable(pickle.dumps(get_credentials), current)
 
 
@@ -257,7 +297,7 @@ def gcs_credentials(
     bearer_token: str | None = None,
     from_env: bool | None = None,
     anonymous: bool | None = None,
-    get_credentials: Callable[[], GcsBearerCredential] | None = None,
+    get_credentials: GcsCredentialProvider | None = None,
     scatter_initial_credentials: bool = False,
 ) -> AnyGcsCredential:
     """Create credentials Google Cloud Storage object store.

icechunk-python/python/icechunk/storage.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from datetime import datetime
 
 from icechunk._icechunk_python import (
@@ -125,7 +125,9 @@ def s3_storage(
     expires_after: datetime | None = None,
     anonymous: bool | None = None,
     from_env: bool | None = None,
-    get_credentials: Callable[[], S3StaticCredentials] | None = None,
+    get_credentials: (
+        Callable[[], S3StaticCredentials | Awaitable[S3StaticCredentials]] | None
+    ) = None,
     scatter_initial_credentials: bool = False,
     force_path_style: bool = False,
     network_stream_timeout_seconds: int = 60,
@@ -157,7 +159,7 @@ def s3_storage(
         If set to True requests to the object store will not be signed
     from_env: bool | None
         Fetch credentials from the operative system environment
-    get_credentials: Callable[[], S3StaticCredentials] | None
+    get_credentials: Callable[[], S3StaticCredentials | Awaitable[S3StaticCredentials]] | None
         Use this function to get and refresh object store credentials
     scatter_initial_credentials: bool, optional
         Immediately call and store the value returned by get_credentials. This is useful if the
@@ -254,7 +256,9 @@ def tigris_storage(
     expires_after: datetime | None = None,
     anonymous: bool | None = None,
     from_env: bool | None = None,
-    get_credentials: Callable[[], S3StaticCredentials] | None = None,
+    get_credentials: (
+        Callable[[], S3StaticCredentials | Awaitable[S3StaticCredentials]] | None
+    ) = None,
     scatter_initial_credentials: bool = False,
     network_stream_timeout_seconds: int = 60,
 ) -> Storage:
@@ -288,7 +292,7 @@ def tigris_storage(
         If set to True requests to the object store will not be signed
     from_env: bool | None
         Fetch credentials from the operative system environment
-    get_credentials: Callable[[], S3StaticCredentials] | None
+    get_credentials: Callable[[], S3StaticCredentials | Awaitable[S3StaticCredentials]] | None
         Use this function to get and refresh object store credentials
     scatter_initial_credentials: bool, optional
         Immediately call and store the value returned by get_credentials. This is useful if the
@@ -340,7 +344,9 @@ def r2_storage(
     expires_after: datetime | None = None,
     anonymous: bool | None = None,
     from_env: bool | None = None,
-    get_credentials: Callable[[], S3StaticCredentials] | None = None,
+    get_credentials: (
+        Callable[[], S3StaticCredentials | Awaitable[S3StaticCredentials]] | None
+    ) = None,
     scatter_initial_credentials: bool = False,
     network_stream_timeout_seconds: int = 60,
 ) -> Storage:
@@ -374,7 +380,7 @@ def r2_storage(
         If set to True requests to the object store will not be signed
     from_env: bool | None
         Fetch credentials from the operative system environment
-    get_credentials: Callable[[], S3StaticCredentials] | None
+    get_credentials: Callable[[], S3StaticCredentials | Awaitable[S3StaticCredentials]] | None
         Use this function to get and refresh object store credentials
     scatter_initial_credentials: bool, optional
         Immediately call and store the value returned by get_credentials. This is useful if the
@@ -436,7 +442,9 @@ def gcs_storage(
     anonymous: bool | None = None,
     from_env: bool | None = None,
     config: dict[str, str] | None = None,
-    get_credentials: Callable[[], GcsBearerCredential] | None = None,
+    get_credentials: (
+        Callable[[], GcsBearerCredential | Awaitable[GcsBearerCredential]] | None
+    ) = None,
     scatter_initial_credentials: bool = False,
 ) -> Storage:
     """Create a Storage instance that saves data in Google Cloud Storage object store.
@@ -461,7 +469,7 @@ def gcs_storage(
         Fetch credentials from the operative system environment
     config: dict[str, str] | None
         A dictionary of options for the Google Cloud Storage object store. See https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html#variants for a list of possible configuration keys.
-    get_credentials: Callable[[], GcsBearerCredential] | None
+    get_credentials: Callable[[], GcsBearerCredential | Awaitable[GcsBearerCredential]] | None
         Use this function to get and refresh object store credentials
     scatter_initial_credentials: bool, optional
         Immediately call and store the value returned by get_credentials. This is useful if the

icechunk-python/src/config.rs

@@ -29,7 +29,7 @@ use icechunk::{
 };
 use pyo3::{
     Bound, FromPyObject, Py, PyErr, PyResult, Python, pyclass, pymethods,
-    types::{PyAnyMethods, PyModule, PyType},
+    types::{PyAny, PyAnyMethods, PyModule, PyType},
 };
 
 use crate::errors::PyIcechunkStoreError;
@@ -141,34 +141,101 @@ pub(crate) fn datetime_repr(d: &DateTime<Utc>) -> String {
 struct PythonCredentialsFetcher<CredType> {
     pub pickled_function: Vec<u8>,
     pub initial: Option<CredType>,
+    // Intentionally skipped during serialization — on deserialization,
+    // `await_python_callback` calls `current_task_locals()` to get
+    // fresh locals from whatever event loop is active at that point.
+    #[serde(skip, default)]
+    pub task_locals: Option<pyo3_async_runtimes::TaskLocals>,
 }
 
 impl<CredType> PythonCredentialsFetcher<CredType> {
     fn new(pickled_function: Vec<u8>) -> Self {
-        PythonCredentialsFetcher { pickled_function, initial: None }
+        PythonCredentialsFetcher {
+            pickled_function,
+            initial: None,
+            task_locals: current_task_locals(),
+        }
     }
 
     fn new_with_initial<C>(pickled_function: Vec<u8>, current: C) -> Self
     where
         C: Into<CredType>,
     {
-        PythonCredentialsFetcher { pickled_function, initial: Some(current.into()) }
+        PythonCredentialsFetcher {
+            pickled_function,
+            initial: Some(current.into()),
+            task_locals: current_task_locals(),
+        }
     }
 }
 
+fn current_task_locals() -> Option<pyo3_async_runtimes::TaskLocals> {
+    Python::attach(|py| pyo3_async_runtimes::tokio::get_current_locals(py).ok())
+}
+
+enum PythonCallbackResult<PyCred> {
+    Value(PyCred),
+    Awaitable(Py<PyAny>),
+}
+
+fn is_awaitable(py: Python<'_>, value: &Bound<'_, PyAny>) -> Result<bool, PyErr> {
+    let inspect = PyModule::import(py, "inspect")?;
+    inspect.getattr("isawaitable")?.call1((value,))?.extract()
+}
+
 fn call_pickled<PyCred>(
     py: Python<'_>,
-    pickled_function: Vec<u8>,
-) -> Result<PyCred, PyErr>
+    pickled_function: &[u8],
+) -> Result<PythonCallbackResult<PyCred>, PyErr>
 where
     PyCred: for<'a, 'py> FromPyObject<'a, 'py>,
     for<'a, 'py> <PyCred as FromPyObject<'a, 'py>>::Error: Into<PyErr>,
 {
     let pickle_module = PyModule::import(py, "pickle")?;
     let loads_function = pickle_module.getattr("loads")?;
     let fetcher = loads_function.call1((pickled_function,))?;
-    let creds: PyCred = fetcher.call0()?.extract().map_err(Into::into)?;
-    Ok(creds)
+    let value = fetcher.call0()?;
+    let value_for_extract = value.clone();
+    let extract_err = match value_for_extract.extract::<PyCred>() {
+        Ok(creds) => return Ok(PythonCallbackResult::Value(creds)),
+        Err(extract_err) => extract_err,
+    };
+
+    if is_awaitable(py, &value)? {
+        Ok(PythonCallbackResult::Awaitable(value.unbind()))
+    } else {
+        Err(extract_err.into())
+    }
+}
+
+async fn await_python_callback<PyCred>(
+    awaitable: Py<PyAny>,
+    task_locals: Option<pyo3_async_runtimes::TaskLocals>,
+) -> Result<PyCred, PyErr>
+where
+    PyCred: for<'a, 'py> FromPyObject<'a, 'py>,
+    for<'a, 'py> <PyCred as FromPyObject<'a, 'py>>::Error: Into<PyErr>,
+{
+    if let Some(task_locals) = current_task_locals().or(task_locals) {
+        let fut = Python::attach(|py| {
+            pyo3_async_runtimes::into_future_with_locals(
+                &task_locals,
+                awaitable.bind(py).clone(),
+            )
+        })?;
+        let value = fut.await?;
+        return Python::attach(|py| value.bind(py).extract().map_err(Into::into));
+    }
+
+    // No task locals available — no Python event loop is currently running.
+    // Fall back to asyncio.run() which creates a temporary event loop to
+    // drive the awaitable. This is the correct path when an async credential
+    // callback is used from a sync context (e.g. Repository.open).
+    Python::attach(|py| {
+        let asyncio = PyModule::import(py, "asyncio")?;
+        let value = asyncio.getattr("run")?.call1((awaitable.bind(py),))?;
+        value.extract().map_err(Into::into)
+    })
 }
 
 #[async_trait]
@@ -190,11 +257,23 @@ impl S3CredentialsFetcher for PythonCredentialsFetcher<S3StaticCredentials> {
                 _ => {}
             }
         }
-        Python::attach(|py| {
-            call_pickled::<PyS3StaticCredentials>(py, self.pickled_function.clone())
-                .map(|c| c.into())
-        })
-        .map_err(|e: PyErr| e.to_string())
+        let callback_result = Python::attach(|py| {
+            call_pickled::<PyS3StaticCredentials>(py, &self.pickled_function)
+        });
+
+        let creds = match callback_result {
+            Ok(PythonCallbackResult::Value(creds)) => Ok(creds),
+            Ok(PythonCallbackResult::Awaitable(awaitable)) => {
+                await_python_callback::<PyS3StaticCredentials>(
+                    awaitable,
+                    self.task_locals.clone(),
+                )
+                .await
+            }
+            Err(err) => Err(err),
+        };
+
+        creds.map(Into::into).map_err(|e: PyErr| e.to_string())
     }
 }
 
@@ -217,11 +296,23 @@ impl GcsCredentialsFetcher for PythonCredentialsFetcher<GcsBearerCredential> {
                 _ => {}
             }
         }
-        Python::attach(|py| {
-            call_pickled::<PyGcsBearerCredential>(py, self.pickled_function.clone())
-                .map(|c| c.into())
-        })
-        .map_err(|e: PyErr| e.to_string())
+        let callback_result = Python::attach(|py| {
+            call_pickled::<PyGcsBearerCredential>(py, &self.pickled_function)
+        });
+
+        let creds = match callback_result {
+            Ok(PythonCallbackResult::Value(creds)) => Ok(creds),
+            Ok(PythonCallbackResult::Awaitable(awaitable)) => {
+                await_python_callback::<PyGcsBearerCredential>(
+                    awaitable,
+                    self.task_locals.clone(),
+                )
+                .await
+            }
+            Err(err) => Err(err),
+        };
+
+        creds.map(Into::into).map_err(|e: PyErr| e.to_string())
     }
 }