adlfs user agent

adlfs/spec.py

@@ -73,6 +73,12 @@
 _SOCKET_TIMEOUT_DEFAULT = object()
 
 
+def _get_user_agent():
+    from adlfs import __version__
+
+    return f"adlfs/{__version__}"
+
+
 # https://github.com/Azure/azure-sdk-for-python/issues/11419#issuecomment-628143480
 def make_callback(key, callback):
     if callback is None:
@@ -118,6 +124,41 @@ 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=_get_user_agent(),
+        )
+    elif location_mode is not None:
+        return AIOBlobServiceClient(
+            account_url=account_url,
+            credential=None,
+            _location_mode=location_mode,
+            user_agent=_get_user_agent(),
+        )
+    else:
+        return AIOBlobServiceClient(
+            account_url=account_url,
+            user_agent=_get_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=_get_user_agent(),
+    )
+
+
 class AzureBlobFileSystem(AsyncFileSystem):
     """
     Access Azure Datalake Gen2 and Azure Storage if it were a file system using Multiprotocol Access
@@ -473,8 +514,10 @@ 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 +530,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(
@@ -2047,27 +2089,28 @@ def connect_client(self):
             creds = [self.fs.sync_credential, self.fs.account_key, self.fs.credential]
             if any(creds):
                 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
-                ).get_container_client(self.container_name)
+                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
+                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
+                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

@@ -2,13 +2,15 @@
 import os
 import tempfile
 from unittest import mock
+from unittest.mock import patch
 
 import azure.storage.blob.aio
 import dask.dataframe as dd
 import fsspec
 import numpy as np
 import pandas as pd
 import pytest
+from azure.storage.blob.aio import BlobServiceClient as AIOBlobServiceClient
 from packaging.version import parse as parse_version
 from pandas.testing import assert_frame_equal
 
@@ -2045,3 +2047,100 @@ 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_blob_file_connection_str(
+    storage: azure.storage.blob.BlobServiceClient, mocker
+):
+    from adlfs import __version__
+
+    fs = AzureBlobFileSystem(
+        account_name=storage.account_name, connection_string=CONN_STR
+    )
+    mock_client = mocker.patch.object(
+        AIOBlobServiceClient, "from_connection_string", return_value=mocker.MagicMock()
+    )
+    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
+):
+    from adlfs import __version__
+
+    path = "root/a/file.txt"
+    mocker.patch("adlfs.spec.filter_blobs", [])
+
+    mock_details = mocker.PropertyMock(
+        return_value={
+            "name": path,
+            "metadata": {},
+            "size": 0,
+            "content_settings": {"content_type": ""},
+        }
+    )
+    mocker.patch.object(AzureBlobFile, "details", mock_details)
+    mock_client = mocker.patch("adlfs.spec.AIOBlobServiceClient", spec=True)
+
+    mock_container_client = mocker.MagicMock()
+    mock_blob_client = mocker.MagicMock()
+    mock_blob_client.get_blob_properties = mocker.AsyncMock(return_value={})
+    mock_blob_client.__aenter__ = mocker.AsyncMock(return_value=mock_blob_client)
+    mock_blob_client.__aexit__ = mocker.AsyncMock()
+    mock_blob_client.close = mocker.AsyncMock()
+    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,
+    )
+    AzureBlobFile(fs, "data/root/a/file.txt", mode="rb")
+
+    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__}"
+
+
+@patch("adlfs.spec.AIOBlobServiceClient")
+def test_user_agent_connection_str(
+    storage: azure.storage.blob.BlobServiceClient, mocker
+):
+    from adlfs import __version__
+
+    mock_client_instance = mocker.MagicMock()
+    mock_client_instance.close = mocker.AsyncMock()
+    mock_client = mocker.patch(
+        "adlfs.spec.AIOBlobServiceClient.from_connection_string",
+        return_value=mock_client_instance,
+    )
+
+    AzureBlobFileSystem(account_name=storage.account_name, connection_string=CONN_STR)
+    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):
+    from adlfs import __version__
+
+    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__}"