Skip to content

Commit ae7db85

Browse files
authored
chore: add check for closed connector (#1269)
1 parent b33295d commit ae7db85

File tree

3 files changed

+69
-1
lines changed

3 files changed

+69
-1
lines changed

google/cloud/sql/connector/connector.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
from google.cloud.sql.connector.enums import DriverMapping
3535
from google.cloud.sql.connector.enums import IPTypes
3636
from google.cloud.sql.connector.enums import RefreshStrategy
37+
from google.cloud.sql.connector.exceptions import ClosedConnectorError
3738
from google.cloud.sql.connector.exceptions import ConnectorLoopError
3839
from google.cloud.sql.connector.instance import RefreshAheadCache
3940
from google.cloud.sql.connector.lazy import LazyRefreshCache
@@ -155,6 +156,7 @@ def __init__(
155156
# connection name string and enable_iam_auth boolean flag
156157
self._cache: dict[tuple[str, bool], MonitoredCache] = {}
157158
self._client: Optional[CloudSQLClient] = None
159+
self._closed: bool = False
158160

159161
# initialize credentials
160162
scopes = ["https://www.googleapis.com/auth/sqlservice.admin"]
@@ -244,6 +246,12 @@ def connect(
244246
# connect runs sync database connections on background thread.
245247
# Async database connections should call 'connect_async' directly to
246248
# avoid hanging indefinitely.
249+
250+
# Check if the connector is closed before attempting to connect.
251+
if self._closed:
252+
raise ClosedConnectorError(
253+
"Connection attempt failed because the connector has already been closed."
254+
)
247255
connect_future = asyncio.run_coroutine_threadsafe(
248256
self.connect_async(instance_connection_string, driver, **kwargs),
249257
self._loop,
@@ -281,7 +289,13 @@ async def connect_async(
281289
and then subsequent attempt with IAM database authentication.
282290
KeyError: Unsupported database driver Must be one of pymysql, asyncpg,
283291
pg8000, and pytds.
292+
RuntimeError: Connector has been closed. Cannot connect using a closed
293+
Connector.
284294
"""
295+
if self._closed:
296+
raise ClosedConnectorError(
297+
"Connection attempt failed because the connector has already been closed."
298+
)
285299
# check if event loop is running in current thread
286300
if self._loop != asyncio.get_running_loop():
287301
raise ConnectorLoopError(
@@ -477,9 +491,10 @@ def close(self) -> None:
477491
async def close_async(self) -> None:
478492
"""Helper function to cancel the cache's tasks
479493
and close aiohttp.ClientSession."""
480-
await asyncio.gather(*[cache.close() for cache in self._cache.values()])
494+
self._closed = True
481495
if self._client:
482496
await self._client.close()
497+
await asyncio.gather(*[cache.close() for cache in self._cache.values()])
483498

484499

485500
async def create_async_connector(

google/cloud/sql/connector/exceptions.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,3 +84,10 @@ class CacheClosedError(Exception):
8484
Exception to be raised when a ConnectionInfoCache can not be accessed after
8585
it is closed.
8686
"""
87+
88+
89+
class ClosedConnectorError(Exception):
90+
"""
91+
Exception to be raised when a Connector is closed and connect method is
92+
called on it.
93+
"""

tests/unit/test_connector.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from google.cloud.sql.connector import IPTypes
3030
from google.cloud.sql.connector.client import CloudSQLClient
3131
from google.cloud.sql.connector.connection_name import ConnectionName
32+
from google.cloud.sql.connector.exceptions import ClosedConnectorError
3233
from google.cloud.sql.connector.exceptions import CloudSQLIPTypeError
3334
from google.cloud.sql.connector.exceptions import ConnectorLoopError
3435
from google.cloud.sql.connector.exceptions import IncompatibleDriverError
@@ -502,3 +503,48 @@ def test_configured_quota_project_env_var(
502503
assert connector._quota_project == quota_project
503504
# unset env var
504505
del os.environ["GOOGLE_CLOUD_QUOTA_PROJECT"]
506+
507+
508+
@pytest.mark.asyncio
509+
async def test_connect_async_closed_connector(
510+
fake_credentials: Credentials, fake_client: CloudSQLClient
511+
) -> None:
512+
"""Test that calling connect_async() on a closed connector raises an error."""
513+
async with Connector(
514+
credentials=fake_credentials, loop=asyncio.get_running_loop()
515+
) as connector:
516+
connector._client = fake_client
517+
await connector.close_async()
518+
with pytest.raises(ClosedConnectorError) as exc_info:
519+
await connector.connect_async(
520+
"test-project:test-region:test-instance",
521+
"asyncpg",
522+
user="my-user",
523+
password="my-pass",
524+
db="my-db",
525+
)
526+
assert (
527+
exc_info.value.args[0]
528+
== "Connection attempt failed because the connector has already been closed."
529+
)
530+
531+
532+
def test_connect_closed_connector(
533+
fake_credentials: Credentials, fake_client: CloudSQLClient
534+
) -> None:
535+
"""Test that calling connect() on a closed connector raises an error."""
536+
with Connector(credentials=fake_credentials) as connector:
537+
connector._client = fake_client
538+
connector.close()
539+
with pytest.raises(ClosedConnectorError) as exc_info:
540+
connector.connect(
541+
"test-project:test-region:test-instance",
542+
"pg8000",
543+
user="my-user",
544+
password="my-pass",
545+
db="my-db",
546+
)
547+
assert (
548+
exc_info.value.args[0]
549+
== "Connection attempt failed because the connector has already been closed."
550+
)

0 commit comments

Comments
 (0)