fix(python): Load _expiry_time from botocore Credentials in CredentialProviderAWS (#23753)

Author: nameexhaustion

py-polars/polars/dependencies.py

@@ -155,6 +155,7 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]:
     import subprocess
 
     import altair
+    import boto3
     import deltalake
     import fsspec
     import gevent
@@ -179,6 +180,7 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]:
 
     # heavy/optional third party libs
     altair, _ALTAIR_AVAILABLE = _lazy_import("altair")
+    boto3, _BOTO3_AVAILABLE = _lazy_import("boto3")
     deltalake, _DELTALAKE_AVAILABLE = _lazy_import("deltalake")
     fsspec, _FSSPEC_AVAILABLE = _lazy_import("fsspec")
     gevent, _GEVENT_AVAILABLE = _lazy_import("gevent")
@@ -315,6 +317,7 @@ def import_optional(
     "subprocess",
     # lazy-load third party libs
     "altair",
+    "boto3",
     "deltalake",
     "fsspec",
     "gevent",

py-polars/polars/io/cloud/credential_provider/_providers.py

@@ -22,6 +22,8 @@
 from polars.io.cloud._utils import NoPickleOption
 
 if TYPE_CHECKING:
+    from polars.dependencies import boto3
+
     if sys.version_info >= (3, 10):
         from typing import TypeAlias
     else:
@@ -171,11 +173,17 @@ def retrieve_credentials_impl(self) -> CredentialProviderFunctionReturn:
             msg = "did not receive any credentials from boto3.Session.get_credentials()"
             raise self.EmptyCredentialError(msg)
 
+        expiry = (
+            int(expiry.timestamp())
+            if isinstance(expiry := getattr(creds, "_expiry_time", None), datetime)
+            else None
+        )
+
         return {
             "aws_access_key_id": creds.access_key,
             "aws_secret_access_key": creds.secret_key,
             **({"aws_session_token": creds.token} if creds.token is not None else {}),
-        }, None
+        }, expiry
 
     def _finish_assume_role(self, session: Any) -> CredentialProviderFunctionReturn:
         client = session.client("sts")
@@ -245,7 +253,7 @@ def _can_use_as_provider(self) -> bool:
 
         return True
 
-    def _session(self) -> Any:
+    def _session(self) -> boto3.Session:
         # Note: boto3 automatically sources the AWS_PROFILE env var
         import boto3
 

py-polars/tests/unit/io/cloud/test_credential_provider.py

@@ -1,5 +1,7 @@
 import io
 import pickle
+import sys
+from datetime import datetime, timezone
 from functools import lru_cache
 from pathlib import Path
 from typing import Any
@@ -539,3 +541,68 @@ def cache(value: ZeroHashWrap[Any]) -> int:
     assert cache(ZeroHashWrap(3)) == 3
     assert cache(ZeroHashWrap(7)) == 3
     assert cache(ZeroHashWrap("A")) == 3
+
+
+@pytest.mark.write_disk
+def test_credential_provider_aws_expiry(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    credential_file_path = tmp_path / "credentials.json"
+
+    credential_file_path.write_text(
+        """\
+{
+    "Version": 1,
+    "AccessKeyId": "123",
+    "SecretAccessKey": "456",
+    "SessionToken": "789",
+    "Expiration": "2099-01-01T00:00:00+00:00"
+}
+"""
+    )
+
+    cfg_file_path = tmp_path / "config"
+
+    credential_file_path_str = str(credential_file_path).replace("\\", "/")
+
+    cfg_file_path.write_text(f"""\
+[profile cred_process]
+credential_process = "{sys.executable}" -c "from pathlib import Path; print(Path('{credential_file_path_str}').read_text())"
+""")
+
+    monkeypatch.setenv("AWS_CONFIG_FILE", str(cfg_file_path))
+
+    creds, expiry = pl.CredentialProviderAWS(profile_name="cred_process")()
+
+    assert creds == {
+        "aws_access_key_id": "123",
+        "aws_secret_access_key": "456",
+        "aws_session_token": "789",
+    }
+
+    assert expiry is not None
+
+    assert datetime.fromtimestamp(expiry, tz=timezone.utc) == datetime.fromisoformat(
+        "2099-01-01T00:00:00+00:00"
+    )
+
+    credential_file_path.write_text(
+        """\
+{
+    "Version": 1,
+    "AccessKeyId": "...",
+    "SecretAccessKey": "...",
+    "SessionToken": "..."
+}
+"""
+    )
+
+    creds, expiry = pl.CredentialProviderAWS(profile_name="cred_process")()
+
+    assert creds == {
+        "aws_access_key_id": "...",
+        "aws_secret_access_key": "...",
+        "aws_session_token": "...",
+    }
+
+    assert expiry is None