Skip to content

Commit f98efab

Browse files
committed
chore: add check for closed connector
1 parent 31ec9ca commit f98efab

File tree

2 files changed

+46
-16
lines changed

2 files changed

+46
-16
lines changed

google/cloud/sql/connector/connector.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,7 @@ def __init__(
153153
# connection name string and enable_iam_auth boolean flag
154154
self._cache: dict[tuple[str, bool], MonitoredCache] = {}
155155
self._client: Optional[CloudSQLClient] = None
156+
self._closed: bool = False
156157

157158
# initialize credentials
158159
scopes = ["https://www.googleapis.com/auth/sqlservice.admin"]
@@ -279,7 +280,11 @@ async def connect_async(
279280
and then subsequent attempt with IAM database authentication.
280281
KeyError: Unsupported database driver Must be one of pymysql, asyncpg,
281282
pg8000, and pytds.
283+
RuntimeError: Connector has been closed. Cannot connect using a closed
284+
Connector.
282285
"""
286+
if self._closed:
287+
raise RuntimeError("Cannot connect using a closed Connector.")
283288
if self._keys is None:
284289
self._keys = asyncio.create_task(generate_keys())
285290
if self._client is None:
@@ -462,13 +467,15 @@ def close(self) -> None:
462467
self._loop.call_soon_threadsafe(self._loop.stop)
463468
# wait for thread to finish closing (i.e. loop to stop)
464469
self._thread.join()
470+
self._closed = True
465471

466472
async def close_async(self) -> None:
467473
"""Helper function to cancel the cache's tasks
468474
and close aiohttp.ClientSession."""
469475
await asyncio.gather(*[cache.close() for cache in self._cache.values()])
470476
if self._client:
471477
await self._client.close()
478+
self._closed = True
472479

473480

474481
async def create_async_connector(

tests/unit/test_connector.py

Lines changed: 39 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,3 @@
1-
"""
2-
Copyright 2021 Google LLC
3-
4-
Licensed under the Apache License, Version 2.0 (the "License");
5-
you may not use this file except in compliance with the License.
6-
You may obtain a copy of the License at
7-
8-
https://www.apache.org/licenses/LICENSE-2.0
9-
10-
Unless required by applicable law or agreed to in writing, software
11-
distributed under the License is distributed on an "AS IS" BASIS,
12-
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13-
See the License for the specific language governing permissions and
14-
limitations under the License.
15-
"""
16-
171
import asyncio
182
import os
193
from typing import Union
@@ -468,3 +452,42 @@ def test_configured_quota_project_env_var(
468452
assert connector._quota_project == quota_project
469453
# unset env var
470454
del os.environ["GOOGLE_CLOUD_QUOTA_PROJECT"]
455+
456+
457+
@pytest.mark.asyncio
458+
async def test_connect_async_closed_connector(
459+
fake_credentials: Credentials, fake_client: CloudSQLClient
460+
) -> None:
461+
"""Test that calling connect_async() on a closed connector raises an error."""
462+
async with Connector(
463+
credentials=fake_credentials, loop=asyncio.get_running_loop()
464+
) as connector:
465+
connector._client = fake_client
466+
await connector.close_async()
467+
with pytest.raises(RuntimeError) as exc_info:
468+
await connector.connect_async(
469+
"test-project:test-region:test-instance",
470+
"asyncpg",
471+
user="my-user",
472+
password="my-pass",
473+
db="my-db",
474+
)
475+
assert exc_info.value.args[0] == "Cannot connect using a closed Connector."
476+
477+
478+
def test_connect_closed_connector(
479+
fake_credentials: Credentials, fake_client: CloudSQLClient
480+
) -> None:
481+
"""Test that calling connect() on a closed connector raises an error."""
482+
with Connector(credentials=fake_credentials) as connector:
483+
connector._client = fake_client
484+
connector.close()
485+
with pytest.raises(RuntimeError) as exc_info:
486+
connector.connect(
487+
"test-project:test-region:test-instance",
488+
"pg8000",
489+
user="my-user",
490+
password="my-pass",
491+
db="my-db",
492+
)
493+
assert exc_info.value.args[0] == "Cannot connect using a closed Connector."

0 commit comments

Comments
 (0)