chore: refactor credentials to be initialized in Connector (#995)

Author: jackwotherspoon

google/cloud/sql/connector/connector.py

@@ -21,7 +21,11 @@
 import socket
 from threading import Thread
 from types import TracebackType
-from typing import Any, Dict, Optional, Type, TYPE_CHECKING
+from typing import Any, Dict, Optional, Type
+
+import google.auth
+from google.auth.credentials import Credentials
+from google.auth.credentials import with_scopes_if_required
 
 import google.cloud.sql.connector.asyncpg as asyncpg
 from google.cloud.sql.connector.exceptions import ConnectorLoopError
@@ -34,9 +38,6 @@
 from google.cloud.sql.connector.utils import format_database_user
 from google.cloud.sql.connector.utils import generate_keys
 
-if TYPE_CHECKING:
-    from google.auth.credentials import Credentials
-
 logger = logging.getLogger(name=__name__)
 
 ASYNC_DRIVERS = ["asyncpg"]
@@ -109,13 +110,26 @@ def __init__(
             )
         self._instances: Dict[str, Instance] = {}
 
+        # initialize credentials
+        scopes = ["https://www.googleapis.com/auth/sqlservice.admin"]
+        if credentials:
+            # verfiy custom credentials are proper type
+            # and atleast base class of google.auth.credentials
+            if not isinstance(credentials, Credentials):
+                raise TypeError(
+                    "credentials must be of type google.auth.credentials.Credentials,"
+                    f" got {type(credentials)}"
+                )
+            self._credentials = with_scopes_if_required(credentials, scopes=scopes)
+        # otherwise use application default credentials
+        else:
+            self._credentials, _ = google.auth.default(scopes=scopes)
         # set default params for connections
         self._timeout = timeout
         self._enable_iam_auth = enable_iam_auth
         self._ip_type = ip_type
         self._quota_project = quota_project
         self._sqladmin_api_endpoint = sqladmin_api_endpoint
-        self._credentials = credentials
         self._user_agent = user_agent
 
     def connect(

google/cloud/sql/connector/exceptions.py

@@ -48,14 +48,6 @@ class PlatformNotSupportedError(Exception):
     pass
 
 
-class CredentialsTypeError(Exception):
-    """
-    Raised when credentials parameter is not proper type.
-    """
-
-    pass
-
-
 class AutoIAMAuthNotSupported(Exception):
     """
     Exception to be raised when Automatic IAM Authentication is not

google/cloud/sql/connector/instance.py

@@ -28,14 +28,12 @@
 
 from google.cloud.sql.connector.exceptions import AutoIAMAuthNotSupported
 from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError
-from google.cloud.sql.connector.exceptions import CredentialsTypeError
 from google.cloud.sql.connector.exceptions import TLSVersionError
 from google.cloud.sql.connector.rate_limiter import AsyncRateLimiter
 from google.cloud.sql.connector.refresh_utils import _get_ephemeral
 from google.cloud.sql.connector.refresh_utils import _get_metadata
 from google.cloud.sql.connector.refresh_utils import _is_valid
 from google.cloud.sql.connector.refresh_utils import _seconds_until_refresh
-from google.cloud.sql.connector.utils import _auth_init
 from google.cloud.sql.connector.utils import write_to_file
 from google.cloud.sql.connector.version import __version__ as version
 
@@ -157,7 +155,6 @@ class Instance:
     :type credentials: google.auth.credentials.Credentials
     :param credentials
         Credentials object used to authenticate connections to Cloud SQL server.
-        If not specified, Application Default Credentials are used.
 
     :param enable_iam_auth
         Enables automatic IAM database authentication for Postgres or MySQL
@@ -206,7 +203,7 @@ def _client_session(self) -> aiohttp.ClientSession:
             self.__client_session = aiohttp.ClientSession(headers=headers)
         return self.__client_session
 
-    _credentials: Optional[Credentials] = None
+    _credentials: Credentials
     _keys: asyncio.Future
 
     _instance_connection_string: str
@@ -227,7 +224,7 @@ def __init__(
         driver_name: str,
         keys: asyncio.Future,
         loop: asyncio.AbstractEventLoop,
-        credentials: Optional[Credentials] = None,
+        credentials: Credentials,
         enable_iam_auth: bool = False,
         quota_project: Optional[str] = None,
         sqladmin_api_endpoint: str = "https://sqladmin.googleapis.com",
@@ -250,13 +247,7 @@ def __init__(
         self._sqladmin_api_endpoint = sqladmin_api_endpoint
         self._loop = loop
         self._keys = keys
-        # validate credentials type
-        if not isinstance(credentials, Credentials) and credentials is not None:
-            raise CredentialsTypeError(
-                "Arg credentials must be type 'google.auth.credentials.Credentials' "
-                "or None (to use Application Default Credentials)"
-            )
-        self._credentials = _auth_init(credentials)
+        self._credentials = credentials
         self._refresh_rate_limiter = AsyncRateLimiter(
             max_capacity=2, rate=1 / 30, loop=self._loop
         )

google/cloud/sql/connector/utils.py

@@ -13,14 +13,11 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-from typing import Optional, Tuple
+from typing import Tuple
 
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
-from google.auth import default
-from google.auth.credentials import Credentials
-from google.auth.credentials import with_scopes_if_required
 
 
 async def generate_keys() -> Tuple[bytes, str]:
@@ -104,23 +101,3 @@ def format_database_user(database_version: str, user: str) -> str:
         return user.split("@")[0]
 
     return user
-
-
-def _auth_init(credentials: Optional[Credentials]) -> Credentials:
-    """Creates google.auth credentials object with scopes required to make
-    calls to the Cloud SQL Admin APIs.
-
-    :type credentials: google.auth.credentials.Credentials
-    :param credentials
-        Credentials object used to authenticate connections to Cloud SQL server.
-        If not specified, Application Default Credentials are used.
-    """
-    scopes = ["https://www.googleapis.com/auth/sqlservice.admin"]
-    # if Credentials object is passed in, use for authentication
-    if isinstance(credentials, Credentials):
-        credentials = with_scopes_if_required(credentials, scopes=scopes)
-    # otherwise use application default credentials
-    else:
-        credentials, _ = default(scopes=scopes)
-
-    return credentials

tests/conftest.py

@@ -23,7 +23,6 @@
 from google.auth.credentials import Credentials
 from google.auth.credentials import with_scopes_if_required
 from google.oauth2 import service_account
-from mock import patch
 import pytest  # noqa F401 Needed to run the tests
 from unit.mocks import FakeCSQLInstance  # type: ignore
 
@@ -146,78 +145,79 @@ async def instance(
     keys = asyncio.create_task(generate_keys())
     _, client_key = await keys
 
-    with patch("google.cloud.sql.connector.utils.default") as mock_auth:
-        mock_auth.return_value = fake_credentials, None
-        # mock Cloud SQL Admin API calls
-        with aioresponses() as mocked:
-            mocked.get(
-                f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{mock_instance.project}/instances/{mock_instance.name}/connectSettings",
-                status=200,
-                body=mock_instance.connect_settings(),
-                repeat=True,
-            )
-            mocked.post(
-                f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{mock_instance.project}/instances/{mock_instance.name}:generateEphemeralCert",
-                status=200,
-                body=mock_instance.generate_ephemeral(client_key),
-                repeat=True,
-            )
-
-            instance = Instance(
-                f"{mock_instance.project}:{mock_instance.region}:{mock_instance.name}",
-                "pg8000",
-                keys,
-                loop,
-            )
-
-            yield instance
-            await instance.close()
+    # mock Cloud SQL Admin API calls
+    with aioresponses() as mocked:
+        mocked.get(
+            f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{mock_instance.project}/instances/{mock_instance.name}/connectSettings",
+            status=200,
+            body=mock_instance.connect_settings(),
+            repeat=True,
+        )
+        mocked.post(
+            f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{mock_instance.project}/instances/{mock_instance.name}:generateEphemeralCert",
+            status=200,
+            body=mock_instance.generate_ephemeral(client_key),
+            repeat=True,
+        )
+
+        instance = Instance(
+            f"{mock_instance.project}:{mock_instance.region}:{mock_instance.name}",
+            "pg8000",
+            keys,
+            loop,
+            fake_credentials,
+        )
+
+        yield instance
+        await instance.close()
 
 
 @pytest.fixture
 async def connector(fake_credentials: Credentials) -> AsyncGenerator[Connector, None]:
     instance_connection_name = "my-project:my-region:my-instance"
     project, region, instance_name = instance_connection_name.split(":")
     # initialize connector
-    connector = Connector()
-    with patch("google.cloud.sql.connector.utils.default") as mock_auth:
-        mock_auth.return_value = fake_credentials, None
-        # mock Cloud SQL Admin API calls
-        mock_instance = FakeCSQLInstance(project, region, instance_name)
-
-        async def wait_for_keys(future: asyncio.Future) -> Tuple[bytes, str]:
-            """
-            Helper method to await keys of Connector in tests prior to
-            initializing an Instance object.
-            """
-            return await future
-
-        # converting asyncio.Future into concurrent.Future
-        # await keys in background thread so that .result() is set
-        # required because keys are needed for mocks, but are not awaited
-        # in the code until Instance() is initialized
-        _, client_key = asyncio.run_coroutine_threadsafe(
-            wait_for_keys(connector._keys), connector._loop
-        ).result()
-        with aioresponses() as mocked:
-            mocked.get(
-                f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{project}/instances/{instance_name}/connectSettings",
-                status=200,
-                body=mock_instance.connect_settings(),
-                repeat=True,
-            )
-            mocked.post(
-                f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{project}/instances/{instance_name}:generateEphemeralCert",
-                status=200,
-                body=mock_instance.generate_ephemeral(client_key),
-                repeat=True,
-            )
-            # initialize Instance using mocked API calls
-            instance = Instance(
-                instance_connection_name, "pg8000", connector._keys, connector._loop
-            )
-
-            connector._instances[instance_connection_name] = instance
-
-            yield connector
-            connector.close()
+    connector = Connector(credentials=fake_credentials)
+    # mock Cloud SQL Admin API calls
+    mock_instance = FakeCSQLInstance(project, region, instance_name)
+
+    async def wait_for_keys(future: asyncio.Future) -> Tuple[bytes, str]:
+        """
+        Helper method to await keys of Connector in tests prior to
+        initializing an Instance object.
+        """
+        return await future
+
+    # converting asyncio.Future into concurrent.Future
+    # await keys in background thread so that .result() is set
+    # required because keys are needed for mocks, but are not awaited
+    # in the code until Instance() is initialized
+    _, client_key = asyncio.run_coroutine_threadsafe(
+        wait_for_keys(connector._keys), connector._loop
+    ).result()
+    with aioresponses() as mocked:
+        mocked.get(
+            f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{project}/instances/{instance_name}/connectSettings",
+            status=200,
+            body=mock_instance.connect_settings(),
+            repeat=True,
+        )
+        mocked.post(
+            f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{project}/instances/{instance_name}:generateEphemeralCert",
+            status=200,
+            body=mock_instance.generate_ephemeral(client_key),
+            repeat=True,
+        )
+        # initialize Instance using mocked API calls
+        instance = Instance(
+            instance_connection_name,
+            "pg8000",
+            connector._keys,
+            connector._loop,
+            fake_credentials,
+        )
+
+        connector._instances[instance_connection_name] = instance
+
+        yield connector
+        connector.close()