(arrow) branch main updated: GH-45705: [Python] Add support for SAS token in AzureFileSystem (#45706)

[ko...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

kou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new c3e399af43 GH-45705: [Python] Add support for SAS token in AzureFileSystem (#45706)
c3e399af43 is described below

commit c3e399af43c4a2e384a41fad0619589baa9045f0
Author: Mike <ja...@users.noreply.github.com>
AuthorDate: Wed Mar 12 21:49:20 2025 -0400

    GH-45705: [Python] Add support for SAS token in AzureFileSystem (#45706)
    
    ### Rationale for this change
    
    Python bindings could use the ability to use SAS tokens just like C++
    
    ### What changes are included in this PR?
    
    AzureFileSystem now includes a sas_token parameter
    
    ### Are these changes tested?
    
    Yes.
    
    ### Are there any user-facing changes?
    
    Additional parameter to AzureFileSystem
    
    * GitHub Issue: #45705
    
    Lead-authored-by: Mike Barry <gi...@barry.io>
    Co-authored-by: Mike <ja...@users.noreply.github.com>
    Co-authored-by: Sutou Kouhei <ko...@cozmixng.org>
    Signed-off-by: Sutou Kouhei <ko...@clear-code.com>
---
 python/pyarrow/_azurefs.pyx             | 20 +++++++++++++++++---
 python/pyarrow/includes/libarrow_fs.pxd |  1 +
 python/pyarrow/tests/test_fs.py         | 10 ++++++++++
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/python/pyarrow/_azurefs.pyx b/python/pyarrow/_azurefs.pyx
index 60c326d280..188d18a27a 100644
--- a/python/pyarrow/_azurefs.pyx
+++ b/python/pyarrow/_azurefs.pyx
@@ -47,7 +47,9 @@ cdef class AzureFileSystem(FileSystem):
         Azure Blob Storage account name. This is the globally unique identifier for the 
         storage account.
     account_key : str, default None
-        Account key of the storage account. Pass None to use default credential. 
+        Account key of the storage account. If sas_token and account_key are None the 
+        default credential will be used. The parameters account_key and sas_token are
+        mutually exclusive.
     blob_storage_authority : str, default None
         hostname[:port] of the Blob Service. Defaults to `.blob.core.windows.net`. Useful
         for connecting to a local emulator, like Azurite.
@@ -60,6 +62,10 @@ cdef class AzureFileSystem(FileSystem):
     dfs_storage_scheme : str, default None
         Either `http` or `https`. Defaults to `https`. Useful for connecting to a local 
         emulator, like Azurite.
+    sas_token : str, default None
+        SAS token for the storage account, used as an alternative to account_key. If sas_token
+        and account_key are None the default credential will be used. The parameters 
+        account_key and sas_token are mutually exclusive.
 
     Examples
     --------
@@ -79,10 +85,11 @@ cdef class AzureFileSystem(FileSystem):
     cdef:
         CAzureFileSystem* azurefs
         c_string account_key
+        c_string sas_token
 
     def __init__(self, account_name, *, account_key=None, blob_storage_authority=None,
                  dfs_storage_authority=None, blob_storage_scheme=None,
-                 dfs_storage_scheme=None):
+                 dfs_storage_scheme=None, sas_token=None):
         cdef:
             CAzureOptions options
             shared_ptr[CAzureFileSystem] wrapped
@@ -97,9 +104,15 @@ cdef class AzureFileSystem(FileSystem):
         if dfs_storage_scheme:
             options.dfs_storage_scheme = tobytes(dfs_storage_scheme)
 
+        if account_key and sas_token:
+            raise ValueError("Cannot specify both account_key and sas_token.")
+
         if account_key:
             options.ConfigureAccountKeyCredential(tobytes(account_key))
             self.account_key = tobytes(account_key)
+        elif sas_token:
+            options.ConfigureSASCredential(tobytes(sas_token))
+            self.sas_token = tobytes(sas_token)
         else:
             options.ConfigureDefaultCredential()
 
@@ -127,5 +140,6 @@ cdef class AzureFileSystem(FileSystem):
                 blob_storage_authority=frombytes(opts.blob_storage_authority),
                 dfs_storage_authority=frombytes(opts.dfs_storage_authority),
                 blob_storage_scheme=frombytes(opts.blob_storage_scheme),
-                dfs_storage_scheme=frombytes(opts.dfs_storage_scheme)
+                dfs_storage_scheme=frombytes(opts.dfs_storage_scheme),
+                sas_token=frombytes(self.sas_token)
             ),))
diff --git a/python/pyarrow/includes/libarrow_fs.pxd b/python/pyarrow/includes/libarrow_fs.pxd
index cc260b80c7..65c3b4e0bb 100644
--- a/python/pyarrow/includes/libarrow_fs.pxd
+++ b/python/pyarrow/includes/libarrow_fs.pxd
@@ -252,6 +252,7 @@ cdef extern from "arrow/filesystem/api.h" namespace "arrow::fs" nogil:
         c_bool Equals(const CAzureOptions& other)
         CStatus ConfigureDefaultCredential()
         CStatus ConfigureAccountKeyCredential(c_string account_key)
+        CStatus ConfigureSASCredential(c_string sas_token)
 
     cdef cppclass CAzureFileSystem "arrow::fs::AzureFileSystem":
         @staticmethod
diff --git a/python/pyarrow/tests/test_fs.py b/python/pyarrow/tests/test_fs.py
index 99d82c3cbf..417a1dcd91 100644
--- a/python/pyarrow/tests/test_fs.py
+++ b/python/pyarrow/tests/test_fs.py
@@ -1463,6 +1463,16 @@ def test_azurefs_options(pickle_module):
     assert pickle_module.loads(pickle_module.dumps(fs3)) == fs3
     assert fs3 != fs2
 
+    fs4 = AzureFileSystem(account_name='fake-account-name',
+                          sas_token='fakesastoken')
+    assert isinstance(fs4, AzureFileSystem)
+    assert pickle_module.loads(pickle_module.dumps(fs4)) == fs4
+    assert fs4 != fs3
+
+    with pytest.raises(ValueError):
+        AzureFileSystem(account_name='fake-account-name', account_key='fakeaccount',
+                        sas_token='fakesastoken')
+
     with pytest.raises(TypeError):
         AzureFileSystem()