Repository: GoogleCloudPlatform/cloud-sql-python-connector
Pull request #995 (2024-01-31 17:29:52)

chore: refactor credentials to be initialized in Connector
Refactor the `credentials` from being initialized at the `Instance` level to the `Connector` level.

 Improvement already made as a new credentials object was being used for each `Instance`, now done once and shared across all instances.

Part of the work to move to using one client session per Connector #873 


Patch:
diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py
index 81044f0..11765df 100755
--- a/google/cloud/sql/connector/connector.py
+++ b/google/cloud/sql/connector/connector.py
@@ -21,7 +21,11 @@ import logging
 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 @@ import google.cloud.sql.connector.pytds as pytds
 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 @@ class Connector:
             )
         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(
diff --git a/google/cloud/sql/connector/exceptions.py b/google/cloud/sql/connector/exceptions.py
index 33b9fda..fa56bdf 100644
--- a/google/cloud/sql/connector/exceptions.py
+++ b/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
diff --git a/google/cloud/sql/connector/instance.py b/google/cloud/sql/connector/instance.py
index e51d9d5..6b64741 100644
--- a/google/cloud/sql/connector/instance.py
+++ b/google/cloud/sql/connector/instance.py
@@ -28,14 +28,12 @@ from google.auth.credentials import Credentials
 
 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 @@ class Instance:
             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 @@ class Instance:
         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 @@ class Instance:
         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
         )
diff --git a/google/cloud/sql/connector/utils.py b/google/cloud/sql/connector/utils.py
index 08f5145..175e763 100755
--- a/google/cloud/sql/connector/utils.py
+++ b/google/cloud/sql/connector/utils.py
@@ -13,14 +13,11 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 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


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 9b971eb..c1f0c3b 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -23,7 +23,6 @@ from aioresponses import aioresponses
 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,32 +145,31 @@ 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
@@ -179,45 +177,47 @@ async def connector(fake_credentials: Credentials) -> AsyncGenerator[Connector,
     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()
diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py
index 0a59be1..f402fe0 100644
--- a/tests/unit/test_connector.py
+++ b/tests/unit/test_connector.py
@@ -15,6 +15,7 @@ limitations under the License.
 """
 import asyncio
 
+from google.auth.credentials import Credentials
 from mock import patch
 from mocks import MockInstance
 import pytest  # noqa F401 Needed to run the tests
@@ -25,7 +26,7 @@ from google.cloud.sql.connector import IPTypes
 from google.cloud.sql.connector.exceptions import ConnectorLoopError
 
 
-def test_connect_enable_iam_auth_error() -> None:
+def test_connect_enable_iam_auth_error(fake_credentials: Credentials) -> None:
     """Test that calling connect() with different enable_iam_auth
     argument values throws error."""
     connect_string = "my-project:my-region:my-instance"
@@ -34,7 +35,7 @@ def test_connect_enable_iam_auth_error() -> None:
     mock_instances = {}
     mock_instances[connect_string] = instance
     # init Connector
-    connector = Connector()
+    connector = Connector(credentials=fake_credentials)
     with patch.dict(connector._instances, mock_instances):
         # try to connect using enable_iam_auth=True, should raise error
         pytest.raises(
@@ -60,11 +61,11 @@ def test_connect_with_unsupported_driver(connector: Connector) -> None:
 
 
 @pytest.mark.asyncio
-async def test_connect_ConnectorLoopError() -> None:
+async def test_connect_ConnectorLoopError(fake_credentials: Credentials) -> None:
     """Test that ConnectorLoopError is thrown when Connector.connect
     is called with event loop running in current thread."""
     current_loop = asyncio.get_running_loop()
-    connector = Connector(loop=current_loop)
+    connector = Connector(credentials=fake_credentials, loop=current_loop)
     # try to connect using current thread's loop, should raise error
     pytest.raises(
         ConnectorLoopError,
@@ -74,35 +75,57 @@ async def test_connect_ConnectorLoopError() -> None:
     )
 
 
-def test_Connector_Init() -> None:
+def test_Connector_Init(fake_credentials: Credentials) -> None:
     """Test that Connector __init__ sets default properties properly."""
-    connector = Connector()
-    assert connector._ip_type == IPTypes.PUBLIC
-    assert connector._enable_iam_auth is False
-    assert connector._timeout == 30
-    assert connector._credentials is None
-    connector.close()
+    with patch("google.auth.default") as mock_auth:
+        mock_auth.return_value = fake_credentials, None
+        connector = Connector()
+        assert connector._ip_type == IPTypes.PUBLIC
+        assert connector._enable_iam_auth is False
+        assert connector._timeout == 30
+        assert connector._credentials == fake_credentials
+        mock_auth.assert_called_once()
+        connector.close()
+
+
+def test_Connector_Init_with_credentials(fake_credentials: Credentials) -> None:
+    """Test that Connector uses custom credentials when given them."""
+    with patch(
+        "google.cloud.sql.connector.connector.with_scopes_if_required"
+    ) as mock_auth:
+        mock_auth.return_value = fake_credentials
+        connector = Connector(credentials=fake_credentials)
+        assert connector._credentials == fake_credentials
+        mock_auth.assert_called_once()
+        connector.close()
+
+
+def test_Connector_Init_with_bad_credentials_type() -> None:
+    """Test that Connector with bad custom credentials type throws error."""
+    pytest.raises(TypeError, Connector, credentials="bad creds")
 
 
-def test_Connector_Init_context_manager() -> None:
+def test_Connector_Init_context_manager(fake_credentials: Credentials) -> None:
     """Test that Connector as context manager sets default properties properly."""
-    with Connector() as connector:
+    with Connector(credentials=fake_credentials) as connector:
         assert connector._ip_type == IPTypes.PUBLIC
         assert connector._enable_iam_auth is False
         assert connector._timeout == 30
-        assert connector._credentials is None
+        assert connector._credentials == fake_credentials
 
 
 @pytest.mark.asyncio
-async def test_Connector_Init_async_context_manager() -> None:
+async def test_Connector_Init_async_context_manager(
+    fake_credentials: Credentials,
+) -> None:
     """Test that Connector as async context manager sets default properties
     properly."""
     loop = asyncio.get_running_loop()
-    async with Connector(loop=loop) as connector:
+    async with Connector(credentials=fake_credentials, loop=loop) as connector:
         assert connector._ip_type == IPTypes.PUBLIC
         assert connector._enable_iam_auth is False
         assert connector._timeout == 30
-        assert connector._credentials is None
+        assert connector._credentials == fake_credentials
         assert connector._loop == loop
 
 
@@ -120,18 +143,18 @@ def test_Connector_connect(connector: Connector) -> None:
 
 
 @pytest.mark.asyncio
-async def test_create_async_connector() -> None:
+async def test_create_async_connector(fake_credentials: Credentials) -> None:
     """Test that create_async_connector properly initializes connector
     object using current thread's event loop"""
-    connector = await create_async_connector()
+    connector = await create_async_connector(credentials=fake_credentials)
     assert connector._loop == asyncio.get_running_loop()
     await connector.close_async()
 
 
-def test_Connector_close_kills_thread() -> None:
+def test_Connector_close_kills_thread(fake_credentials: Credentials) -> None:
     """Test that Connector.close kills background threads."""
     # open and close Connector object
-    connector = Connector()
+    connector = Connector(credentials=fake_credentials)
     # verify background thread exists
     assert connector._thread
     connector.close()
@@ -139,10 +162,10 @@ def test_Connector_close_kills_thread() -> None:
     assert connector._thread.is_alive() is False
 
 
-def test_Connector_close_called_multiple_times() -> None:
+def test_Connector_close_called_multiple_times(fake_credentials: Credentials) -> None:
     """Test that Connector.close can be called multiple times."""
     # open and close Connector object
-    connector = Connector()
+    connector = Connector(credentials=fake_credentials)
     # verify background thread exists
     assert connector._thread
     connector.close()
diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py
index 1bc1d7c..6a3fd62 100644
--- a/tests/unit/test_instance.py
+++ b/tests/unit/test_instance.py
@@ -27,7 +27,6 @@ import pytest  # noqa F401 Needed to run the tests
 
 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.instance import _parse_instance_connection_name
 from google.cloud.sql.connector.instance import ConnectionInfo
 from google.cloud.sql.connector.instance import Instance
@@ -83,15 +82,14 @@ async def test_Instance_init(
     keys = asyncio.wrap_future(
         asyncio.run_coroutine_threadsafe(generate_keys(), event_loop), loop=event_loop
     )
-    with patch("google.cloud.sql.connector.utils.default") as mock_auth:
-        mock_auth.return_value = fake_credentials, None
-        instance = Instance(
-            connect_string,
-            "pymysql",
-            keys,
-            event_loop,
-            user_agent="custom/v1.0.0",
-        )
+    instance = Instance(
+        connect_string,
+        "pymysql",
+        keys,
+        event_loop,
+        credentials=fake_credentials,
+        user_agent="custom/v1.0.0",
+    )
     project_result = instance._project
     region_result = instance._region
     instance_result = instance._instance
@@ -108,22 +106,6 @@ async def test_Instance_init(
     await instance.close()
 
 
-@pytest.mark.asyncio
-async def test_Instance_init_bad_credentials() -> None:
-    """
-    Test to check whether the __init__ method of Instance
-    throws proper error for bad credentials arg type.
-    """
-    event_loop = asyncio.get_running_loop()
-    connect_string = "test-project:test-region:test-instance"
-    keys = asyncio.wrap_future(
-        asyncio.run_coroutine_threadsafe(generate_keys(), event_loop), loop=event_loop
-    )
-    with pytest.raises(CredentialsTypeError):
-        instance = Instance(connect_string, "pymysql", keys, event_loop, credentials=1)
-        await instance.close()
-
-
 @pytest.mark.asyncio
 async def test_schedule_refresh_replaces_result(
     instance: Instance, test_rate_limiter: AsyncRateLimiter
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index d7c384d..fe190ce 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -13,8 +13,6 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-from google.auth.credentials import Credentials
-from mock import patch
 import pytest  # noqa F401 Needed to run the tests
 
 from google.cloud.sql.connector import utils
@@ -76,29 +74,3 @@ def test_format_database_user_mysql() -> None:
     user2 = utils.format_database_user("MYSQL_8_0", "test")
     assert user == "test"
     assert user2 == "test"
-
-
-def test_auth_init_with_credentials_object(fake_credentials: Credentials) -> None:
-    """
-    Test that _auth_init initializes credentials with scopes
-    when passed a google.auth.credentials.Credentials object.
-    """
-    with patch("google.cloud.sql.connector.utils.with_scopes_if_required") as mock_auth:
-        mock_auth.return_value = fake_credentials
-        credentials = utils._auth_init(credentials=fake_credentials)
-        assert isinstance(credentials, Credentials)
-        mock_auth.assert_called_once()
-
-
-def test_auth_init_with_default_credentials(
-    fake_credentials: Credentials,
-) -> None:
-    """
-    Test that _auth_init initializes _credentials
-    with application default credentials when credentials are not specified.
-    """
-    with patch("google.cloud.sql.connector.utils.default") as mock_auth:
-        mock_auth.return_value = fake_credentials, None
-        credentials = utils._auth_init(credentials=None)
-        assert isinstance(credentials, Credentials)
-        mock_auth.assert_called_once()