Skip to content

Commit

Permalink
Address PR comments
Browse files Browse the repository at this point in the history
  • Loading branch information
rhatgadkar-goog committed Feb 4, 2025
1 parent 1ff96c7 commit 5d9c473
Show file tree
Hide file tree
Showing 5 changed files with 33 additions and 46 deletions.
15 changes: 3 additions & 12 deletions google/cloud/alloydb/connector/async_connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,9 @@
from __future__ import annotations

import asyncio
import io
import logging
from types import TracebackType
from typing import Any, Optional, TYPE_CHECKING, Union
from typing import Any, Optional, TYPE_CHECKING

import google.auth
from google.auth.credentials import with_scopes_if_required
Expand All @@ -30,7 +29,7 @@
from google.cloud.alloydb.connector.enums import RefreshStrategy
from google.cloud.alloydb.connector.instance import RefreshAheadCache
from google.cloud.alloydb.connector.lazy import LazyRefreshCache
from google.cloud.alloydb.connector.static import StaticConnectionInfoCache
from google.cloud.alloydb.connector.types import CacheTypes
from google.cloud.alloydb.connector.utils import generate_keys

if TYPE_CHECKING:
Expand Down Expand Up @@ -61,9 +60,6 @@ class AsyncConnector:
of the following: RefreshStrategy.LAZY ("LAZY") or
RefreshStrategy.BACKGROUND ("BACKGROUND").
Default: RefreshStrategy.BACKGROUND
static_conn_info (io.TextIOBase): A file-like JSON object that contains
static connection info for the StaticConnectionInfoCache.
Defaults to None, which will not use the StaticConnectionInfoCache.
"""

def __init__(
Expand All @@ -75,9 +71,8 @@ def __init__(
ip_type: str | IPTypes = IPTypes.PRIVATE,
user_agent: Optional[str] = None,
refresh_strategy: str | RefreshStrategy = RefreshStrategy.BACKGROUND,
static_conn_info: Optional[io.TextIOBase] = None,
) -> None:
self._cache: dict[str, Union[RefreshAheadCache, LazyRefreshCache]] = {}
self._cache: dict[str, CacheTypes] = {}
# initialize default params
self._quota_project = quota_project
self._alloydb_api_endpoint = alloydb_api_endpoint
Expand Down Expand Up @@ -106,7 +101,6 @@ def __init__(
except RuntimeError:
self._keys = None
self._client: Optional[AlloyDBClient] = None
self._static_conn_info = static_conn_info

async def connect(
self,
Expand Down Expand Up @@ -147,11 +141,8 @@ async def connect(
enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth)

# use existing connection info if possible
cache: Union[RefreshAheadCache, LazyRefreshCache, StaticConnectionInfoCache]
if instance_uri in self._cache:
cache = self._cache[instance_uri]
elif self._static_conn_info:
cache = StaticConnectionInfoCache(instance_uri, self._static_conn_info)
else:
if self._refresh_strategy == RefreshStrategy.LAZY:
logger.debug(
Expand Down
9 changes: 3 additions & 6 deletions google/cloud/alloydb/connector/connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import struct
from threading import Thread
from types import TracebackType
from typing import Any, Optional, TYPE_CHECKING, Union
from typing import Any, Optional, TYPE_CHECKING

from google.auth import default
from google.auth.credentials import TokenState
Expand All @@ -35,7 +35,7 @@
from google.cloud.alloydb.connector.instance import RefreshAheadCache
from google.cloud.alloydb.connector.lazy import LazyRefreshCache
import google.cloud.alloydb.connector.pg8000 as pg8000
from google.cloud.alloydb.connector.static import StaticConnectionInfoCache
from google.cloud.alloydb.connector.types import CacheTypes
from google.cloud.alloydb.connector.utils import generate_keys
import google.cloud.alloydb_connectors_v1.proto.resources_pb2 as connectorspb

Expand Down Expand Up @@ -93,7 +93,7 @@ def __init__(
self._loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
self._thread = Thread(target=self._loop.run_forever, daemon=True)
self._thread.start()
self._cache: dict[str, Union[RefreshAheadCache, LazyRefreshCache]] = {}
self._cache: dict[str, CacheTypes] = {}
# initialize default params
self._quota_project = quota_project
self._alloydb_api_endpoint = alloydb_api_endpoint
Expand Down Expand Up @@ -176,11 +176,8 @@ async def connect_async(self, instance_uri: str, driver: str, **kwargs: Any) ->
)
enable_iam_auth = kwargs.pop("enable_iam_auth", self._enable_iam_auth)
# use existing connection info if possible
cache: Union[RefreshAheadCache, LazyRefreshCache, StaticConnectionInfoCache]
if instance_uri in self._cache:
cache = self._cache[instance_uri]
elif self._static_conn_info:
cache = StaticConnectionInfoCache(instance_uri, self._static_conn_info)
else:
if self._refresh_strategy == RefreshStrategy.LAZY:
logger.debug(
Expand Down
23 changes: 23 additions & 0 deletions google/cloud/alloydb/connector/types.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.

import typing

from google.cloud.alloydb.connector.instance import RefreshAheadCache
from google.cloud.alloydb.connector.lazy import LazyRefreshCache
from google.cloud.alloydb.connector.static import StaticConnectionInfoCache

CacheTypes = typing.Union[
RefreshAheadCache, LazyRefreshCache, StaticConnectionInfoCache
]
5 changes: 4 additions & 1 deletion tests/unit/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import socket
import ssl
from threading import Thread
from typing import Final

from aiofiles.tempfile import TemporaryDirectory
from mocks import FakeAlloyDBClient
Expand All @@ -26,6 +27,8 @@

from google.cloud.alloydb.connector.utils import _write_to_file

DELAY: Final[float] = 1.0


@pytest.fixture
def credentials() -> FakeCredentials:
Expand Down Expand Up @@ -86,4 +89,4 @@ def proxy_server(fake_instance: FakeInstance) -> None:
daemon=True,
)
thread.start()
thread.join(0.1) # wait 100ms to allow the proxy server to start
thread.join(DELAY) # add a delay to allow the proxy server to start
27 changes: 0 additions & 27 deletions tests/unit/test_async_connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
from mocks import FakeAlloyDBClient
from mocks import FakeConnectionInfo
from mocks import FakeCredentials
from mocks import write_static_info
import pytest

from google.cloud.alloydb.connector import AsyncConnector
Expand Down Expand Up @@ -334,29 +333,3 @@ async def test_Connector_remove_cached_no_ip_type(credentials: FakeCredentials)
await connector.connect(instance_uri, "asyncpg", ip_type="private")
# check that cache has been removed from dict
assert instance_uri not in connector._cache


async def test_Connector_static_connection_info(
credentials: FakeCredentials, fake_client: FakeAlloyDBClient
) -> None:
"""
Test that AsyncConnector.__init__() can specify a static connection info to
connect to an instance.
"""
static_info = write_static_info(fake_client.instance)
async with AsyncConnector(
credentials=credentials, static_conn_info=static_info
) as connector:
connector._client = fake_client
# patch db connection creation
with patch("google.cloud.alloydb.connector.asyncpg.connect") as mock_connect:
mock_connect.return_value = True
connection = await connector.connect(
fake_client.instance.uri(),
"asyncpg",
user="test-user",
password="test-password",
db="test-db",
)
# check connection is returned
assert connection is True

0 comments on commit 5d9c473

Please sign in to comment.