fix(io): pass GCS token expiry to gcsfs

pyiceberg/io/fsspec.py

@@ -66,6 +66,7 @@
     GCS_SERVICE_HOST,
     GCS_SESSION_KWARGS,
     GCS_TOKEN,
+    GCS_TOKEN_EXPIRES_AT_MS,
     GCS_VERSION_AWARE,
     HF_ENDPOINT,
     HF_TOKEN,
@@ -95,6 +96,7 @@
 )
 from pyiceberg.typedef import Properties
 from pyiceberg.types import strtobool
+from pyiceberg.utils.datetime import millis_to_datetime
 from pyiceberg.utils.properties import get_first_property_value, get_header_properties, property_as_bool
 
 logger = logging.getLogger(__name__)
@@ -243,11 +245,16 @@ def _s3(properties: Properties) -> AbstractFileSystem:
 def _gs(properties: Properties) -> AbstractFileSystem:
     # https://gcsfs.readthedocs.io/en/latest/api.html#gcsfs.core.GCSFileSystem
     from gcsfs import GCSFileSystem
+    from google.oauth2.credentials import Credentials
+
+    token: str | Credentials | None = properties.get(GCS_TOKEN)
+    if token and (expires_at := properties.get(GCS_TOKEN_EXPIRES_AT_MS)):
+        token = Credentials(token, expiry=millis_to_datetime(int(expires_at)))
 
     return GCSFileSystem(
         project=properties.get(GCS_PROJECT_ID),
         access=properties.get(GCS_ACCESS, "full_control"),
-        token=properties.get(GCS_TOKEN),
+        token=token,
         consistency=properties.get(GCS_CONSISTENCY, "none"),
         cache_timeout=properties.get(GCS_CACHE_TIMEOUT),
         requester_pays=property_as_bool(properties, GCS_REQUESTER_PAYS, False),

tests/io/test_fsspec.py

@@ -20,12 +20,14 @@
 import tempfile
 import threading
 import uuid
+from datetime import datetime
 from unittest import mock
 
 import pytest
 from botocore.awsrequest import AWSRequest
 from fsspec.implementations.local import LocalFileSystem
 from fsspec.spec import AbstractFileSystem
+from google.oauth2.credentials import Credentials
 from requests_mock import Mocker
 
 from pyiceberg.catalog.rest.auth import AUTH_MANAGER
@@ -725,6 +727,29 @@ def test_adls_account_name_not_overridden_when_in_properties() -> None:
         )
 
 
+def test_fsspec_gcs_token_with_expiry() -> None:
+    """Test that a GCS token with an expiry is passed to gcsfs as google-auth credentials."""
+    properties: Properties = {"gcs.oauth2.token": "token", "gcs.oauth2.token-expires-at": "1700000000000"}
+
+    with mock.patch("gcsfs.GCSFileSystem") as mock_gcsfs:
+        FsspecFileIO(properties=properties).new_input(location="gs://bucket/path/file.parquet")
+
+    token = mock_gcsfs.call_args.kwargs["token"]
+    assert isinstance(token, Credentials)
+    assert token.token == "token"
+    assert token.expiry == datetime(2023, 11, 14, 22, 13, 20)
+
+
+def test_fsspec_gcs_token_without_expiry() -> None:
+    """Test that a GCS token without an expiry is passed to gcsfs unchanged."""
+    properties: Properties = {"gcs.oauth2.token": "token"}
+
+    with mock.patch("gcsfs.GCSFileSystem") as mock_gcsfs:
+        FsspecFileIO(properties=properties).new_input(location="gs://bucket/path/file.parquet")
+
+    assert mock_gcsfs.call_args.kwargs["token"] == "token"
+
+
 @pytest.mark.gcs
 def test_fsspec_new_input_file_gcs(fsspec_fileio_gcs: FsspecFileIO) -> None:
     """Test creating a new input file from a fsspec file-io"""