added adlfs user agent

Author: anjaliratnam-msft

adlfs/spec.py

@@ -45,6 +45,7 @@
     get_blob_metadata,
     match_blob_version,
 )
+from ._version import version as __version__
 
 logger = logging.getLogger(__name__)
 
@@ -71,6 +72,7 @@
 _DEFAULT_BLOCK_SIZE = 4 * 1024 * 1024
 
 _SOCKET_TIMEOUT_DEFAULT = object()
+_USER_AGENT = f"adlfs/{__version__}"
 
 
 # https://github.com/Azure/azure-sdk-for-python/issues/11419#issuecomment-628143480
@@ -118,6 +120,35 @@ def _coalesce_version_id(*args) -> Optional[str]:
         return version_ids.pop()
 
 
+def _create_aio_blob_service_client(account_url: str, location_mode: Optional[str] = None, credential: Optional[str] = None) -> AIOBlobServiceClient:
+    if credential is not None:
+        return AIOBlobServiceClient(
+            account_url=account_url,
+            credential=credential,
+            _location_mode=location_mode,
+            user_agent=_USER_AGENT,
+        )
+    elif location_mode is not None:
+        return AIOBlobServiceClient(
+            account_url=account_url,
+            credential=None,
+            _location_mode=location_mode,
+            user_agent=_USER_AGENT,
+        )
+    else:
+        return AIOBlobServiceClient(
+            account_url=account_url,
+            user_agent=_USER_AGENT,
+        )
+
+
+def _create_aio_blob_service_client_from_connection_string(connection_string: str) -> AIOBlobServiceClient:
+    return AIOBlobServiceClient.from_connection_string(
+        conn_str=connection_string,
+        user_agent=_USER_AGENT,
+    )
+
+
 class AzureBlobFileSystem(AsyncFileSystem):
     """
     Access Azure Datalake Gen2 and Azure Storage if it were a file system using Multiprotocol Access
@@ -473,8 +504,8 @@ def do_connect(self):
 
         try:
             if self.connection_string is not None:
-                self.service_client = AIOBlobServiceClient.from_connection_string(
-                    conn_str=self.connection_string
+                self.service_client = _create_aio_blob_service_client_from_connection_string(
+                    connection_string=self.connection_string,
                 )
             elif self.account_name is not None:
                 if hasattr(self, "account_host"):
@@ -487,26 +518,25 @@ def do_connect(self):
                 creds = [self.credential, self.account_key]
                 if any(creds):
                     self.service_client = [
-                        AIOBlobServiceClient(
+                        _create_aio_blob_service_client(
                             account_url=self.account_url,
+                            location_mode=self.location_mode,
                             credential=cred,
-                            _location_mode=self.location_mode,
                         )
                         for cred in creds
                         if cred is not None
                     ][0]
                 elif self.sas_token is not None:
                     if not self.sas_token.startswith("?"):
                         self.sas_token = f"?{self.sas_token}"
-                    self.service_client = AIOBlobServiceClient(
+                    self.service_client = _create_aio_blob_service_client(
                         account_url=self.account_url + self.sas_token,
-                        credential=None,
-                        _location_mode=self.location_mode,
+                        location_mode=self.location_mode,
                     )
                 else:
                     # Fall back to anonymous login, and assume public container
-                    self.service_client = AIOBlobServiceClient(
-                        account_url=self.account_url
+                    self.service_client = _create_aio_blob_service_client(
+                        account_url=self.account_url,
                     )
             else:
                 raise ValueError(
@@ -1950,6 +1980,7 @@ def __init__(
         self.end = None
         self.start = None
         self.closed = False
+        self._user_agent = f"adlfs/{__version__}"
 
         if cache_options is None:
             cache_options = {}
@@ -2045,29 +2076,33 @@ def connect_client(self):
                 )
 
             creds = [self.fs.sync_credential, self.fs.account_key, self.fs.credential]
+            print(f"creds: {creds}")
             if any(creds):
+                print("call")
                 self.container_client = [
-                    AIOBlobServiceClient(
+                    _create_aio_blob_service_client(
                         account_url=self.fs.account_url,
                         credential=cred,
-                        _location_mode=self.fs.location_mode,
+                        location_mode=self.fs.location_mode,
                     ).get_container_client(self.container_name)
                     for cred in creds
                     if cred is not None
                 ][0]
             elif self.fs.connection_string is not None:
-                self.container_client = AIOBlobServiceClient.from_connection_string(
-                    conn_str=self.fs.connection_string
+                print("call")
+                self.container_client = _create_aio_blob_service_client_from_connection_string(
+                    connection_string=self.fs.connection_string,
                 ).get_container_client(self.container_name)
             elif self.fs.sas_token is not None:
-                self.container_client = AIOBlobServiceClient(
-                    account_url=self.fs.account_url + self.fs.sas_token, credential=None
+                print("call")
+                self.container_client = _create_aio_blob_service_client(
+                    account_url=self.fs.account_url + self.fs.sas_token,
                 ).get_container_client(self.container_name)
             else:
-                self.container_client = AIOBlobServiceClient(
-                    account_url=self.fs.account_url
+                print("call")
+                self.container_client = _create_aio_blob_service_client(
+                    account_url=self.fs.account_url,
                 ).get_container_client(self.container_name)
-
         except Exception as e:
             raise ValueError(
                 f"Unable to fetch container_client with provided params for {e}!!"

adlfs/tests/test_spec.py

@@ -13,6 +13,10 @@
 from pandas.testing import assert_frame_equal
 
 from adlfs import AzureBlobFile, AzureBlobFileSystem
+from adlfs._version import version as __version__
+from azure.storage.blob.aio import BlobServiceClient as AIOBlobServiceClient
+import asyncio
+
 
 URL = "http://127.0.0.1:10000"
 ACCOUNT_NAME = "devstoreaccount1"
@@ -2045,3 +2049,64 @@ def test_open_file_x(storage: azure.storage.blob.BlobServiceClient, tmpdir):
         with fs.open("data/afile", "xb") as f:
             pass
     assert fs.cat_file("data/afile") == b"data"
+
+
+def test_user_agent_connection_str(storage: azure.storage.blob.BlobServiceClient, mocker):
+    fs = AzureBlobFileSystem(
+        account_name=storage.account_name, connection_string=CONN_STR
+    )
+    mock_client = mocker.patch.object(AIOBlobServiceClient, "from_connection_string", return_value=mocker.AsyncMock())
+   
+    fs.do_connect()
+    mock_client.assert_called_once()
+    assert "user_agent" in mock_client.call_args.kwargs
+    assert mock_client.call_args.kwargs["user_agent"] == f"adlfs/{__version__}"
+
+
+def test_user_agent_initializer(storage: azure.storage.blob.BlobServiceClient, mocker):
+    fs = AzureBlobFileSystem(
+        account_name=storage.account_name,
+    )
+    mock_client = mocker.patch("adlfs.spec.AIOBlobServiceClient", spec=True)
+
+    fs.do_connect()
+    mock_client.assert_called_once()
+    assert "user_agent" in mock_client.call_args.kwargs
+    assert mock_client.call_args.kwargs["user_agent"] == f"adlfs/{__version__}"
+
+
+def test_user_agent_blob_file_connection_str(storage: azure.storage.blob.BlobServiceClient, mocker):
+    fs = AzureBlobFileSystem(
+        account_name=storage.account_name, connection_string=CONN_STR
+    )
+    mock_client = mocker.patch.object(AIOBlobServiceClient, "from_connection_string", return_value=mocker.AsyncMock())
+    mock_client.return_value.get_container_client = mocker.MagicMock()
+
+    f = AzureBlobFile(fs, "data/root/a/file.txt", mode="rb")
+    f.container_name = "data"
+    f.connect_client()
+
+    mock_client.assert_called_once()
+    assert "user_agent" in mock_client.call_args.kwargs
+    assert mock_client.call_args.kwargs["user_agent"] == f"adlfs/{__version__}"
+
+
+def test_user_agent_blob_file_initializer(storage: azure.storage.blob.BlobServiceClient, mocker):
+    mock_client = mocker.patch("adlfs.spec.AIOBlobServiceClient", spec= True)
+    mock_container_client = mocker.MagicMock()
+    mock_blob_client = mocker.MagicMock()
+    mock_blob_client.__aenter__ = mocker.AsyncMock(return_value=mock_blob_client)
+    mock_blob_client.__aexit__ = mocker.AsyncMock()
+    mock_blob_client.get_blob_properties = mocker.AsyncMock(return_value={})
+    mock_container_client.get_blob_client = mocker.MagicMock(return_value=mock_blob_client)
+    mock_client.return_value.get_container_client = mocker.MagicMock(return_value=mock_container_client)
+    
+    fs = AzureBlobFileSystem(
+        account_name=storage.account_name,
+    )   
+    f = AzureBlobFile(fs, "data/root/a/file.txt", mode="rb")
+    f.container_name = "data"
+
+    mock_client.assert_called_once()
+    assert "user_agent" in mock_client.call_args.kwargs
+    assert mock_client.call_args.kwargs["user_agent"] == f"adlfs/{__version__}"