feat: support ip_type as str (#267)

.flake8

@@ -15,7 +15,7 @@
 # limitations under the License.
 
 [flake8]
-ignore = E203, E231, E266, E501, W503, ANN101, ANN401
+ignore = E203, E231, E266, E501, W503, ANN101, ANN102, ANN401
 exclude =
   # Exclude generated code.
   **/proto/**

README.md

@@ -381,12 +381,11 @@ to your instance's private IP. To change this, such as connecting to AlloyDB
 over a public IP address, set the `ip_type` keyword argument when initializing
 a `Connector()` or when calling `connector.connect()`.
 
-Possible values for `ip_type` are `IPTypes.PRIVATE` (default value), and
-`IPTypes.PUBLIC`.
+Possible values for `ip_type` are `"PRIVATE"` (default value), and `"PUBLIC"`.
 Example:
 
 ```python
-from google.cloud.alloydb.connector import Connector, IPTypes
+from google.cloud.alloydb.connector import Connector
 
 import sqlalchemy
 
@@ -401,7 +400,7 @@ def getconn():
     user="my-user",
     password="my-password",
     db="my-db-name",
-    ip_type=IPTypes.PUBLIC,  # use public IP
+    ip_type="PUBLIC",  # use public IP
   )
 
 # create connection pool

google/cloud/alloydb/connector/async_connector.py

@@ -47,8 +47,8 @@ class AsyncConnector:
         alloydb_api_endpoint (str): Base URL to use when calling
             the AlloyDB API endpoint. Defaults to "https://alloydb.googleapis.com".
         enable_iam_auth (bool): Enables automatic IAM database authentication.
-        ip_type (IPTypes): Default IP type for all AlloyDB connections.
-            Defaults to IPTypes.PRIVATE for private IP connections.
+        ip_type (str | IPTypes): Default IP type for all AlloyDB connections.
+            Defaults to IPTypes.PRIVATE ("PRIVATE") for private IP connections.
     """
 
     def __init__(
@@ -57,14 +57,17 @@ def __init__(
         quota_project: Optional[str] = None,
         alloydb_api_endpoint: str = "https://alloydb.googleapis.com",
         enable_iam_auth: bool = False,
-        ip_type: IPTypes = IPTypes.PRIVATE,
+        ip_type: str | IPTypes = IPTypes.PRIVATE,
         user_agent: Optional[str] = None,
     ) -> None:
         self._instances: Dict[str, Instance] = {}
         # initialize default params
         self._quota_project = quota_project
         self._alloydb_api_endpoint = alloydb_api_endpoint
         self._enable_iam_auth = enable_iam_auth
+        # if ip_type is str, convert to IPTypes enum
+        if isinstance(ip_type, str):
+            ip_type = IPTypes(ip_type.upper())
         self._ip_type = ip_type
         self._user_agent = user_agent
         # initialize credentials
@@ -144,7 +147,10 @@ async def connect(
         kwargs.pop("port", None)
 
         # get connection info for AlloyDB instance
-        ip_type: IPTypes = kwargs.pop("ip_type", self._ip_type)
+        ip_type: str | IPTypes = kwargs.pop("ip_type", self._ip_type)
+        # if ip_type is str, convert to IPTypes enum
+        if isinstance(ip_type, str):
+            ip_type = IPTypes(ip_type.upper())
         ip_address, context = await instance.connection_info(ip_type)
 
         # callable to be used for auto IAM authn

google/cloud/alloydb/connector/connector.py

@@ -57,8 +57,8 @@ class Connector:
         alloydb_api_endpoint (str): Base URL to use when calling
             the AlloyDB API endpoint. Defaults to "https://alloydb.googleapis.com".
         enable_iam_auth (bool): Enables automatic IAM database authentication.
-        ip_type (IPTypes): Default IP type for all AlloyDB connections.
-            Defaults to IPTypes.PRIVATE for private IP connections.
+        ip_type (str | IPTypes): Default IP type for all AlloyDB connections.
+            Defaults to IPTypes.PRIVATE ("PRIVATE") for private IP connections.
     """
 
     def __init__(
@@ -67,7 +67,7 @@ def __init__(
         quota_project: Optional[str] = None,
         alloydb_api_endpoint: str = "https://alloydb.googleapis.com",
         enable_iam_auth: bool = False,
-        ip_type: IPTypes = IPTypes.PRIVATE,
+        ip_type: str | IPTypes = IPTypes.PRIVATE,
         user_agent: Optional[str] = None,
     ) -> None:
         # create event loop and start it in background thread
@@ -79,6 +79,9 @@ def __init__(
         self._quota_project = quota_project
         self._alloydb_api_endpoint = alloydb_api_endpoint
         self._enable_iam_auth = enable_iam_auth
+        # if ip_type is str, convert to IPTypes enum
+        if isinstance(ip_type, str):
+            ip_type = IPTypes(ip_type.upper())
         self._ip_type = ip_type
         self._user_agent = user_agent
         # initialize credentials
@@ -171,7 +174,10 @@ async def connect_async(self, instance_uri: str, driver: str, **kwargs: Any) ->
         kwargs.pop("port", None)
 
         # get connection info for AlloyDB instance
-        ip_type: IPTypes = kwargs.pop("ip_type", self._ip_type)
+        ip_type: IPTypes | str = kwargs.pop("ip_type", self._ip_type)
+        # if ip_type is str, convert to IPTypes enum
+        if isinstance(ip_type, str):
+            ip_type = IPTypes(ip_type.upper())
         ip_address, context = await instance.connection_info(ip_type)
 
         # synchronous drivers are blocking and run using executor

google/cloud/alloydb/connector/instance.py

@@ -49,6 +49,13 @@ class IPTypes(Enum):
     PUBLIC: str = "PUBLIC"
     PRIVATE: str = "PRIVATE"
 
+    @classmethod
+    def _missing_(cls, value: object) -> None:
+        raise ValueError(
+            f"Incorrect value for ip_type, got '{value}'. Want one of: "
+            f"{', '.join([repr(m.value) for m in cls])}."
+        )
+
 
 def _parse_instance_uri(instance_uri: str) -> Tuple[str, str, str, str]:
     # should take form "projects/<PROJECT>/locations/<REGION>/clusters/<CLUSTER>/instances/<INSTANCE>"

tests/system/test_asyncpg_public_ip.py

@@ -22,7 +22,6 @@
 import sqlalchemy.ext.asyncio
 
 from google.cloud.alloydb.connector import AsyncConnector
-from google.cloud.alloydb.connector import IPTypes
 
 
 async def create_sqlalchemy_engine(
@@ -70,7 +69,7 @@ async def getconn() -> asyncpg.Connection:
             user=user,
             password=password,
             db=db,
-            ip_type=IPTypes.PUBLIC,
+            ip_type="PUBLIC",
         )
         return conn
 

tests/system/test_pg8000_public_ip.py

@@ -21,7 +21,6 @@
 import sqlalchemy
 
 from google.cloud.alloydb.connector import Connector
-from google.cloud.alloydb.connector import IPTypes
 
 
 def create_sqlalchemy_engine(
@@ -70,7 +69,7 @@ def getconn() -> pg8000.dbapi.Connection:
             user=user,
             password=password,
             db=db,
-            ip_type=IPTypes.PUBLIC,
+            ip_type="PUBLIC",
         )
         return conn
 

tests/unit/test_async_connector.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+from typing import Union
 
 from mock import patch
 from mocks import FakeAlloyDBClient
@@ -21,6 +22,7 @@
 import pytest
 
 from google.cloud.alloydb.connector import AsyncConnector
+from google.cloud.alloydb.connector import IPTypes
 
 ALLOYDB_API_ENDPOINT = "https://alloydb.googleapis.com"
 
@@ -40,6 +42,58 @@ async def test_AsyncConnector_init(credentials: FakeCredentials) -> None:
     await connector.close()
 
 
+@pytest.mark.parametrize(
+    "ip_type, expected",
+    [
+        (
+            "private",
+            IPTypes.PRIVATE,
+        ),
+        (
+            "PRIVATE",
+            IPTypes.PRIVATE,
+        ),
+        (
+            IPTypes.PRIVATE,
+            IPTypes.PRIVATE,
+        ),
+        (
+            "public",
+            IPTypes.PUBLIC,
+        ),
+        (
+            "PUBLIC",
+            IPTypes.PUBLIC,
+        ),
+        (
+            IPTypes.PUBLIC,
+            IPTypes.PUBLIC,
+        ),
+    ],
+)
+async def test_AsyncConnector_init_ip_type(
+    ip_type: Union[str, IPTypes], expected: IPTypes, credentials: FakeCredentials
+) -> None:
+    """
+    Test to check whether the __init__ method of AsyncConnector
+    properly sets ip_type.
+    """
+    connector = AsyncConnector(credentials=credentials, ip_type=ip_type)
+    assert connector._ip_type == expected
+    connector.close()
+
+
+async def test_AsyncConnector_init_bad_ip_type(credentials: FakeCredentials) -> None:
+    """Test that AsyncConnector errors due to bad ip_type str."""
+    bad_ip_type = "BAD-IP-TYPE"
+    with pytest.raises(ValueError) as exc_info:
+        AsyncConnector(ip_type=bad_ip_type, credentials=credentials)
+    assert (
+        exc_info.value.args[0]
+        == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'PUBLIC', 'PRIVATE'."
+    )
+
+
 @pytest.mark.asyncio
 async def test_AsyncConnector_context_manager(
     credentials: FakeCredentials,
@@ -202,3 +256,25 @@ def test_synchronous_init(credentials: FakeCredentials) -> None:
     """
     connector = AsyncConnector(credentials)
     assert connector._keys is None
+
+
+async def test_async_connect_bad_ip_type(
+    credentials: FakeCredentials, fake_client: FakeAlloyDBClient
+) -> None:
+    """Test that AyncConnector.connect errors due to bad ip_type str."""
+    async with AsyncConnector(credentials=credentials) as connector:
+        connector._client = fake_client
+        bad_ip_type = "BAD-IP-TYPE"
+        with pytest.raises(ValueError) as exc_info:
+            await connector.connect(
+                "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance",
+                "asyncpg",
+                user="test-user",
+                password="test-password",
+                db="test-db",
+                ip_type=bad_ip_type,
+            )
+        assert (
+            exc_info.value.args[0]
+            == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'PUBLIC', 'PRIVATE'."
+        )

tests/unit/test_connector.py

@@ -14,13 +14,15 @@
 
 import asyncio
 from threading import Thread
+from typing import Union
 
 from mock import patch
 from mocks import FakeAlloyDBClient
 from mocks import FakeCredentials
 import pytest
 
 from google.cloud.alloydb.connector import Connector
+from google.cloud.alloydb.connector import IPTypes
 
 
 def test_Connector_init(credentials: FakeCredentials) -> None:
@@ -36,6 +38,58 @@ def test_Connector_init(credentials: FakeCredentials) -> None:
     connector.close()
 
 
+def test_Connector_init_bad_ip_type(credentials: FakeCredentials) -> None:
+    """Test that Connector errors due to bad ip_type str."""
+    bad_ip_type = "BAD-IP-TYPE"
+    with pytest.raises(ValueError) as exc_info:
+        Connector(ip_type=bad_ip_type, credentials=credentials)
+    assert (
+        exc_info.value.args[0]
+        == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'PUBLIC', 'PRIVATE'."
+    )
+
+
+@pytest.mark.parametrize(
+    "ip_type, expected",
+    [
+        (
+            "private",
+            IPTypes.PRIVATE,
+        ),
+        (
+            "PRIVATE",
+            IPTypes.PRIVATE,
+        ),
+        (
+            IPTypes.PRIVATE,
+            IPTypes.PRIVATE,
+        ),
+        (
+            "public",
+            IPTypes.PUBLIC,
+        ),
+        (
+            "PUBLIC",
+            IPTypes.PUBLIC,
+        ),
+        (
+            IPTypes.PUBLIC,
+            IPTypes.PUBLIC,
+        ),
+    ],
+)
+def test_Connector_init_ip_type(
+    ip_type: Union[str, IPTypes], expected: IPTypes, credentials: FakeCredentials
+) -> None:
+    """
+    Test to check whether the __init__ method of Connector
+    properly sets ip_type.
+    """
+    connector = Connector(credentials=credentials, ip_type=ip_type)
+    assert connector._ip_type == expected
+    connector.close()
+
+
 def test_Connector_context_manager(credentials: FakeCredentials) -> None:
     """
     Test to check whether the __init__ method of Connector
@@ -84,6 +138,28 @@ def test_connect(credentials: FakeCredentials, fake_client: FakeAlloyDBClient) -
         assert connection is True
 
 
+def test_connect_bad_ip_type(
+    credentials: FakeCredentials, fake_client: FakeAlloyDBClient
+) -> None:
+    """Test that Connector.connect errors due to bad ip_type str."""
+    with Connector(credentials=credentials) as connector:
+        connector._client = fake_client
+        bad_ip_type = "BAD-IP-TYPE"
+        with pytest.raises(ValueError) as exc_info:
+            connector.connect(
+                "projects/test-project/locations/test-region/clusters/test-cluster/instances/test-instance",
+                "pg8000",
+                user="test-user",
+                password="test-password",
+                db="test-db",
+                ip_type=bad_ip_type,
+            )
+        assert (
+            exc_info.value.args[0]
+            == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'PUBLIC', 'PRIVATE'."
+        )
+
+
 def test_connect_unsupported_driver(credentials: FakeCredentials) -> None:
     """
     Test that connector.connect errors with unsupported database driver.