feat: add support for a custom user agent (#233)

google/cloud/alloydb/connector/async_connector.py

@@ -54,12 +54,14 @@ def __init__(
         quota_project: Optional[str] = None,
         alloydb_api_endpoint: str = "https://alloydb.googleapis.com",
         enable_iam_auth: bool = False,
+        user_agent: Optional[str] = None,
     ) -> None:
         self._instances: Dict[str, Instance] = {}
         # initialize default params
         self._quota_project = quota_project
         self._alloydb_api_endpoint = alloydb_api_endpoint
         self._enable_iam_auth = enable_iam_auth
+        self._user_agent = user_agent
         # initialize credentials
         scopes = ["https://www.googleapis.com/auth/cloud-platform"]
         if credentials:
@@ -108,6 +110,7 @@ async def connect(
                 self._alloydb_api_endpoint,
                 self._quota_project,
                 self._credentials,
+                user_agent=self._user_agent,
                 driver=driver,
             )
 

google/cloud/alloydb/connector/client.py

@@ -31,6 +31,20 @@
 logger = logging.getLogger(name=__name__)
 
 
+def _format_user_agent(
+    driver: Optional[str],
+    custom_user_agent: Optional[str],
+) -> str:
+    """
+    Appends user-defined user agents to the base default agent.
+    """
+    agent = f"{USER_AGENT}+{driver}" if driver else USER_AGENT
+    if custom_user_agent:
+        agent = f"{agent} {custom_user_agent}"
+
+    return agent
+
+
 class AlloyDBClient:
     def __init__(
         self,
@@ -39,6 +53,7 @@ def __init__(
         credentials: Credentials,
         client: Optional[aiohttp.ClientSession] = None,
         driver: Optional[str] = None,
+        user_agent: Optional[str] = None,
     ) -> None:
         """
         Establish the client to be used for AlloyDB Admin API requests.
@@ -58,7 +73,7 @@ def __init__(
                 Optional, defaults to None and creates new client.
             driver (str): Database driver to be used by the client.
         """
-        user_agent = f"{USER_AGENT}+{driver}" if driver else USER_AGENT
+        user_agent = _format_user_agent(driver, user_agent)
         headers = {
             "x-goog-api-client": user_agent,
             "User-Agent": user_agent,

google/cloud/alloydb/connector/connector.py

@@ -64,6 +64,7 @@ def __init__(
         quota_project: Optional[str] = None,
         alloydb_api_endpoint: str = "https://alloydb.googleapis.com",
         enable_iam_auth: bool = False,
+        user_agent: Optional[str] = None,
     ) -> None:
         # create event loop and start it in background thread
         self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
@@ -74,6 +75,7 @@ def __init__(
         self._quota_project = quota_project
         self._alloydb_api_endpoint = alloydb_api_endpoint
         self._enable_iam_auth = enable_iam_auth
+        self._user_agent = user_agent
         # initialize credentials
         scopes = ["https://www.googleapis.com/auth/cloud-platform"]
         if credentials:
@@ -137,6 +139,7 @@ async def connect_async(self, instance_uri: str, driver: str, **kwargs: Any) ->
                 self._alloydb_api_endpoint,
                 self._quota_project,
                 self._credentials,
+                user_agent=self._user_agent,
                 driver=driver,
             )
         enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth)

tests/unit/test_client.py

@@ -104,6 +104,26 @@ async def test_AlloyDBClient_init_(credentials: FakeCredentials) -> None:
     await client.close()
 
 
+@pytest.mark.asyncio
+async def test_AlloyDBClient_init_custom_user_agent(
+    credentials: FakeCredentials,
+) -> None:
+    """
+    Test to check that custom user agents are included in HTTP requests.
+    """
+    client = AlloyDBClient(
+        "www.test-endpoint.com",
+        "my-quota-project",
+        credentials,
+        user_agent="custom-agent/v1.0.0 other-agent/v2.0.0",
+    )
+    assert (
+        client._client.headers["User-Agent"]
+        == f"alloydb-python-connector/{version} custom-agent/v1.0.0 other-agent/v2.0.0"
+    )
+    await client.close()
+
+
 @pytest.mark.parametrize(
     "driver",
     [None, "pg8000", "asyncpg"],