support adls.token

Author: kevinjqliu

pyiceberg/io/__init__.py

@@ -86,6 +86,7 @@
 ADLS_DFS_STORAGE_AUTHORITY = "adls.dfs-storage-authority"
 ADLS_BLOB_STORAGE_SCHEME = "adls.blob-storage-scheme"
 ADLS_DFS_STORAGE_SCHEME = "adls.dfs-storage-scheme"
+ADLS_TOKEN = "adls.token"
 GCS_TOKEN = "gcs.oauth2.token"
 GCS_TOKEN_EXPIRES_AT_MS = "gcs.oauth2.token-expires-at"
 GCS_PROJECT_ID = "gcs.project-id"

pyiceberg/io/fsspec.py

@@ -48,6 +48,7 @@
     ADLS_CREDENTIAL,
     ADLS_SAS_TOKEN,
     ADLS_TENANT_ID,
+    ADLS_TOKEN,
     AWS_ACCESS_KEY_ID,
     AWS_REGION,
     AWS_SECRET_ACCESS_KEY,
@@ -192,7 +193,11 @@ def _gs(properties: Properties) -> AbstractFileSystem:
 
 
 def _adls(properties: Properties) -> AbstractFileSystem:
+    import time
+
     from adlfs import AzureBlobFileSystem
+    from azure.core.credentials import AccessToken
+    from azure.core.credentials_async import AsyncTokenCredential
 
     for key, sas_token in {
         key.replace(f"{ADLS_SAS_TOKEN}.", ""): value for key, value in properties.items() if key.startswith(ADLS_SAS_TOKEN)
@@ -202,9 +207,29 @@ def _adls(properties: Properties) -> AbstractFileSystem:
         if ADLS_SAS_TOKEN not in properties:
             properties[ADLS_SAS_TOKEN] = sas_token
 
+    class StaticTokenCredential(AsyncTokenCredential):
+        def __init__(self, token_string: str) -> None:
+            self._token = token_string
+            # If no expiry provided, set 1 hour from now
+            self._expires_on = int(time.time()) + 3600
+
+        async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+            return AccessToken(self._token, self._expires_on)
+
+    if ADLS_TOKEN in properties:
+        token = properties.get(ADLS_TOKEN)
+        if token is not None:
+            credential = StaticTokenCredential(token)
+        else:
+            credential = None
+    elif ADLS_CREDENTIAL in properties:
+        credential = properties.get(ADLS_CREDENTIAL)
+    else:
+        credential = None
+
     return AzureBlobFileSystem(
         connection_string=properties.get(ADLS_CONNECTION_STRING),
-        credential=properties.get(ADLS_CREDENTIAL),
+        credential=credential,
         account_name=properties.get(ADLS_ACCOUNT_NAME),
         account_key=properties.get(ADLS_ACCOUNT_KEY),
         sas_token=properties.get(ADLS_SAS_TOKEN),