feat: add support for specifying separate credentials for authenticating with DB (#510)

This PR does the following:

- When authenticating with the DB, the alloydb.login scope is used instead of the cloud-platform scope. Authentication with the AlloyDB Admin API continues to use the cloud-platform scope.
- The user can specify separate credentials, with the db_credentials argument passed into the Connector() and AsyncConnector(). When specified, only the db_credentials will be used to authenticate with the DB. If not specified, the existing behavior is preserved with credentials being used to authenticate with the DB.

Author: rhatgadkar-goog

google/cloud/alloydbconnector/async_connector.py

@@ -48,6 +48,15 @@ class AsyncConnector:
         credentials (google.auth.credentials.Credentials):
             A credentials object created from the google-auth Python library.
             If not specified, Application Default Credentials are used.
+            These are the credentials used for authenticating with the AlloyDB
+            Admin API.
+        db_credentials (google.auth.credentials.Credentials):
+            A credentials object created from the google-auth Python library.
+            This is only used when Auto IAM AuthN is enabled.
+            If not specified, the credentials used for authenticating with the
+            AlloyDB Admin API will also be used to authenticate with the DB.
+            If specified, the credential's scope should be
+            "https://www.googleapis.com/auth/alloydb.login".
         quota_project (str): The Project ID for an existing Google Cloud
             project. The project specified is used for quota and
             billing purposes.
@@ -67,6 +76,7 @@ class AsyncConnector:
     def __init__(
         self,
         credentials: Optional[Credentials] = None,
+        db_credentials: Optional[Credentials] = None,
         quota_project: Optional[str] = None,
         alloydb_api_endpoint: str = "alloydb.googleapis.com",
         enable_iam_auth: bool = False,
@@ -88,13 +98,23 @@ def __init__(
             refresh_strategy = RefreshStrategy(refresh_strategy.upper())
         self._refresh_strategy = refresh_strategy
         self._user_agent = user_agent
-        # initialize credentials
+        # initialize credentials for authenticating with AlloyDB Admin API
         scopes = ["https://www.googleapis.com/auth/cloud-platform"]
         if credentials:
             self._credentials = with_scopes_if_required(credentials, scopes=scopes)
         # otherwise use application default credentials
         else:
             self._credentials, _ = google.auth.default(scopes=scopes)
+        # initialize credentials for authenticating with the DB
+        if db_credentials:
+            self._db_credentials = db_credentials
+        # otherwise use the same credentials as the one for authenticating with
+        # AlloyDB Admin API
+        else:
+            scopes = ["https://www.googleapis.com/auth/alloydb.login"]
+            self._db_credentials = with_scopes_if_required(
+                self._credentials, scopes=scopes
+            )
 
         # check if AsyncConnector is being initialized with event loop running
         # Otherwise we will lazy init keys
@@ -196,13 +216,13 @@ async def connect(
         logger.debug(f"['{instance_uri}']: Connecting to {ip_address}:5433")
 
         # callable to be used for auto IAM authn
-        def get_authentication_token() -> str:
+        async def get_authentication_token() -> str:
             """Get OAuth2 access token to be used for IAM database authentication"""
             # refresh credentials if expired
-            if not self._credentials.valid:
+            if not self._db_credentials.valid:
                 request = google.auth.transport.requests.Request()
-                self._credentials.refresh(request)
-            return self._credentials.token
+                await asyncio.to_thread(self._db_credentials.refresh, request)
+            return self._db_credentials.token
 
         # if enable_iam_auth is set, use auth token as database password
         if enable_iam_auth:

google/cloud/alloydbconnector/connector.py

@@ -62,6 +62,14 @@ class Connector:
         credentials (google.auth.credentials.Credentials):
             A credentials object created from the google-auth Python library.
             If not specified, Application Default Credentials are used.
+            These are the credentials used for authenticating with the AlloyDB
+            Admin API.
+        db_credentials (google.auth.credentials.Credentials):
+            A credentials object created from the google-auth Python library.
+            If not specified, the credentials used for authenticating with the
+            AlloyDB Admin API will also be used to authenticate with the DB.
+            If specified, the credential's scope should be
+            "https://www.googleapis.com/auth/alloydb.login".
         quota_project (str): The Project ID for an existing Google Cloud
             project. The project specified is used for quota and
             billing purposes.
@@ -87,6 +95,7 @@ class Connector:
     def __init__(
         self,
         credentials: Optional[Credentials] = None,
+        db_credentials: Optional[Credentials] = None,
         quota_project: Optional[str] = None,
         alloydb_api_endpoint: str = "alloydb.googleapis.com",
         enable_iam_auth: bool = False,
@@ -113,13 +122,23 @@ def __init__(
             refresh_strategy = RefreshStrategy(refresh_strategy.upper())
         self._refresh_strategy = refresh_strategy
         self._user_agent = user_agent
-        # initialize credentials
+        # initialize credentials for authenticating with AlloyDB Admin API
         scopes = ["https://www.googleapis.com/auth/cloud-platform"]
         if credentials:
             self._credentials = with_scopes_if_required(credentials, scopes=scopes)
         # otherwise use application default credentials
         else:
             self._credentials, _ = default(scopes=scopes)
+        # initialize credentials for authenticating with the DB
+        if db_credentials:
+            self._db_credentials = db_credentials
+        # otherwise use the same credentials as the one for authenticating with
+        # AlloyDB Admin API
+        else:
+            scopes = ["https://www.googleapis.com/auth/alloydb.login"]
+            self._db_credentials = with_scopes_if_required(
+                self._credentials, scopes=scopes
+            )
         self._keys = asyncio.wrap_future(
             asyncio.run_coroutine_threadsafe(generate_keys(), self._loop),
             loop=self._loop,
@@ -296,14 +315,14 @@ def metadata_exchange(
             auth_type = connectorspb.MetadataExchangeRequest.AUTO_IAM
 
         # Ensure the credentials are in fact valid before proceeding.
-        if not self._credentials.token_state == TokenState.FRESH:
-            self._credentials.refresh(requests.Request())
+        if not self._db_credentials.token_state == TokenState.FRESH:
+            self._db_credentials.refresh(requests.Request())
 
         # form metadata exchange request
         req = connectorspb.MetadataExchangeRequest(
             user_agent=f"{self._client._user_agent}",  # type: ignore
             auth_type=auth_type,
-            oauth2_token=self._credentials.token,
+            oauth2_token=self._db_credentials.token,
         )
 
         # set I/O timeout

tests/unit/mocks.py

@@ -21,14 +21,15 @@
 import json
 import ssl
 import struct
-from typing import Any, Callable, Literal, Optional
+from typing import Any, Callable, Literal, Optional, Sequence
 
 from cryptography import x509
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
 from cryptography.x509.oid import NameOID
 from google.auth.credentials import _helpers
+from google.auth.credentials import Scoped
 from google.auth.credentials import TokenState
 from google.auth.transport import requests
 
@@ -57,6 +58,11 @@ def expired(self) -> bool:
         """
         return False if not self.expiry else True
 
+    @property
+    def valid(self) -> bool:
+        """Checks if the credentials are valid."""
+        return self.token is not None and not self.expired
+
     @property
     def token_state(
         self,
@@ -87,6 +93,26 @@ def token_state(
         return TokenState.FRESH
 
 
+class FakeCredentialsRequiresScopes(Scoped):
+    def requires_scopes(self) -> bool:
+        """
+        Overrides the requires_scopes() method of the Scoped class to require
+        scopes for these credentials.
+        """
+        return True
+
+    def with_scopes(
+        self, scopes: Sequence[str], default_scopes: Optional[Sequence[str]] = None
+    ) -> "FakeCredentialsRequiresScopes":
+        """
+        Overrides the with_scopes() method of the Scoped class to create a
+        copy of these credentials with the specified scopes.
+        """
+        f = FakeCredentialsRequiresScopes()
+        f._scopes = scopes
+        return f
+
+
 def generate_cert(
     common_name: str, expires_in: int = 60, server_cert: bool = False
 ) -> tuple[x509.CertificateBuilder, rsa.RSAPrivateKey]:

tests/unit/test_async_connector.py

@@ -13,14 +13,15 @@
 # limitations under the License.
 
 import asyncio
-from typing import Union
+from typing import Any, Union
 
 from google.api_core.exceptions import RetryError
 from google.api_core.retry.retry_unary_async import AsyncRetry
 from mock import patch
 from mocks import FakeAlloyDBClient
 from mocks import FakeConnectionInfo
 from mocks import FakeCredentials
+from mocks import FakeCredentialsRequiresScopes
 import pytest
 
 from google.cloud.alloydbconnector import AsyncConnector
@@ -44,11 +45,42 @@ async def test_AsyncConnector_init(credentials: FakeCredentials) -> None:
     assert connector._alloydb_api_endpoint == ALLOYDB_API_ENDPOINT
     assert connector._client is None
     assert connector._credentials == credentials
+    assert connector._db_credentials == credentials
     assert connector._enable_iam_auth is False
     assert connector._closed is False
     await connector.close()
 
 
+@pytest.mark.asyncio
+async def test_AsyncConnector_init_db_credentials(credentials: FakeCredentials) -> None:
+    """
+    Test to check whether the __init__ method of AsyncConnector
+    properly sets db_credentials when specified.
+    """
+    db_credentials = FakeCredentials()
+    connector = AsyncConnector(credentials, db_credentials)
+    assert connector._db_credentials == db_credentials
+    await connector.close()
+
+
+async def test_AsyncConnector_init_scopes() -> None:
+    """
+    Test to check whether the __init__ method of AsyncConnector
+    properly sets the credential's scopes.
+    """
+    credentials = FakeCredentialsRequiresScopes()
+    connector = AsyncConnector(credentials)
+    assert connector._credentials != credentials
+    assert connector._credentials._scopes == [
+        "https://www.googleapis.com/auth/cloud-platform"
+    ]
+    assert connector._db_credentials != credentials
+    assert connector._db_credentials._scopes == [
+        "https://www.googleapis.com/auth/alloydb.login"
+    ]
+    await connector.close()
+
+
 @pytest.mark.parametrize(
     "ip_type, expected",
     [
@@ -154,6 +186,7 @@ async def test_AsyncConnector_context_manager(
         assert connector._alloydb_api_endpoint == ALLOYDB_API_ENDPOINT
         assert connector._client is None
         assert connector._credentials == credentials
+        assert connector._db_credentials == credentials
         assert connector._enable_iam_auth is False
 
 
@@ -197,6 +230,66 @@ async def test_connect_and_close(credentials: FakeCredentials) -> None:
         assert connection.result() is True
 
 
+@pytest.mark.asyncio
+async def test_connect_iam_authn(credentials: FakeCredentials) -> None:
+    """
+    Test that connector.connect, with IAM authentication, refreshes credentials.
+    """
+    async with AsyncConnector(credentials, enable_iam_auth=True) as connector:
+        connector._client = FakeAlloyDBClient()
+
+        async def custom_connect(*_: Any, **kwargs: Any) -> bool:
+            passwd = kwargs.pop("password")
+            await passwd()
+
+        # patch db connection creation
+        with patch(
+            "google.cloud.alloydbconnector.asyncpg.connect", side_effect=custom_connect
+        ):
+            await connector.connect(
+                TEST_INSTANCE_NAME,
+                "asyncpg",
+                user="test-user",
+                password="test-password",
+                db="test-db",
+            )
+        # check DB authentication refreshed the credentials
+        assert connector._credentials.token
+        assert connector._db_credentials.token
+
+
+@pytest.mark.asyncio
+async def test_connect_db_credentials_iam_authn(credentials: FakeCredentials) -> None:
+    """
+    Test that connector.connect, with IAM authentication, refreshes only the DB
+    credentials when specified.
+    """
+    db_credentials = FakeCredentials()
+    async with AsyncConnector(
+        credentials, db_credentials, enable_iam_auth=True
+    ) as connector:
+        connector._client = FakeAlloyDBClient()
+
+        async def custom_connect(*_: Any, **kwargs: Any) -> bool:
+            passwd = kwargs.pop("password")
+            await passwd()
+
+        # patch db connection creation
+        with patch(
+            "google.cloud.alloydbconnector.asyncpg.connect", side_effect=custom_connect
+        ):
+            await connector.connect(
+                TEST_INSTANCE_NAME,
+                "asyncpg",
+                user="test-user",
+                password="test-password",
+                db="test-db",
+            )
+        # check DB authentication refreshed only the DB credential's token
+        assert not connector._credentials.token
+        assert connector._db_credentials.token
+
+
 @pytest.mark.asyncio
 async def test_force_refresh(credentials: FakeCredentials) -> None:
     """