diff --git a/google/cloud/alloydbconnector/async_connector.py b/google/cloud/alloydbconnector/async_connector.py index 7aac297e..1ba50861 100644 --- a/google/cloud/alloydbconnector/async_connector.py +++ b/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: diff --git a/google/cloud/alloydbconnector/connector.py b/google/cloud/alloydbconnector/connector.py index 0b100a82..af888456 100644 --- a/google/cloud/alloydbconnector/connector.py +++ b/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 diff --git a/tests/unit/mocks.py b/tests/unit/mocks.py index 2dfbd9a4..5bf6df57 100644 --- a/tests/unit/mocks.py +++ b/tests/unit/mocks.py @@ -21,7 +21,7 @@ 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 @@ -29,6 +29,7 @@ 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]: diff --git a/tests/unit/test_async_connector.py b/tests/unit/test_async_connector.py index fb5076bd..92a345b9 100644 --- a/tests/unit/test_async_connector.py +++ b/tests/unit/test_async_connector.py @@ -13,7 +13,7 @@ # 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 @@ -21,6 +21,7 @@ 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: """ diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py index 349120c6..b950de61 100644 --- a/tests/unit/test_connector.py +++ b/tests/unit/test_connector.py @@ -21,6 +21,7 @@ from mock import patch from mocks import FakeAlloyDBClient from mocks import FakeCredentials +from mocks import FakeCredentialsRequiresScopes from mocks import write_static_info import pytest @@ -43,10 +44,40 @@ def test_Connector_init(credentials: FakeCredentials) -> None: assert connector._alloydb_api_endpoint == "alloydb.googleapis.com" assert connector._client is None assert connector._credentials == credentials + assert connector._db_credentials == credentials assert connector._closed is False connector.close() +def test_Connector_init_db_credentials(credentials: FakeCredentials) -> None: + """ + Test to check whether the __init__ method of Connector + properly sets db_credentials when specified. + """ + db_credentials = FakeCredentials() + connector = Connector(credentials, db_credentials) + assert connector._db_credentials == db_credentials + connector.close() + + +def test_Connector_init_scopes() -> None: + """ + Test to check whether the __init__ method of Connector + properly sets the credential's scopes. + """ + credentials = FakeCredentialsRequiresScopes() + connector = Connector(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" + ] + connector.close() + + def test_Connector_init_bad_ip_type(credentials: FakeCredentials) -> None: """Test that Connector errors due to bad ip_type str.""" bad_ip_type = "BAD-IP-TYPE" @@ -149,6 +180,7 @@ def test_Connector_context_manager(credentials: FakeCredentials) -> None: assert connector._alloydb_api_endpoint == "alloydb.googleapis.com" assert connector._client is None assert connector._credentials == credentials + assert connector._db_credentials == credentials def test_Connector_close(credentials: FakeCredentials) -> None: @@ -170,7 +202,8 @@ def test_Connector_close(credentials: FakeCredentials) -> None: @pytest.mark.usefixtures("proxy_server") def test_connect(credentials: FakeCredentials, fake_client: FakeAlloyDBClient) -> None: """ - Test that connector.connect returns connection object. + Test that connector.connect returns connection object and refreshes + credentials. """ client = fake_client with Connector(credentials) as connector: @@ -187,6 +220,35 @@ def test_connect(credentials: FakeCredentials, fake_client: FakeAlloyDBClient) - ) # check connection is returned assert connection is True + # check DB authentication refreshed the credentials + assert connector._credentials.token + assert connector._db_credentials.token + + +@pytest.mark.usefixtures("proxy_server") +def test_connect_db_credentials( + credentials: FakeCredentials, fake_client: FakeAlloyDBClient +) -> None: + """ + Test that connector.connect refreshes only the DB credentials when + specified. + """ + client = fake_client + db_credentials = FakeCredentials() + with Connector(credentials, db_credentials) as connector: + connector._client = client + # patch db connection creation + with patch("google.cloud.alloydbconnector.pg8000.connect"): + connector.connect( + "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance", + "pg8000", + 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 def test_connect_bad_ip_type(