feat: support ip_type as str (#1029)

Author: jackwotherspoon

README.md

@@ -188,7 +188,7 @@ from google.cloud.sql.connector import Connector, IPTypes
 
 # Note: all parameters below are optional
 connector = Connector(
-    ip_type=IPTypes.PUBLIC,
+    ip_type="public",  # can also be "private" or "psc"
     enable_iam_auth=False,
     timeout=30,
     credentials=custom_creds # google.auth.credentials.Credentials
@@ -261,18 +261,16 @@ using both public and private IP addresses, as well as
 with, set the `ip_type` keyword argument when initializing a `Connector()` or when
 calling `connector.connect()`.
 
-Possible values for `ip_type` are `IPTypes.PUBLIC` (default value),
-`IPTypes.PRIVATE`, and `IPTypes.PSC`.
+Possible values for `ip_type` are `"public"` (default value),
+`"private"`, and `"psc"`.
 
 Example:
 
 ```python
-from google.cloud.sql.connector import IPTypes
-
 conn = connector.connect(
     "project:region:instance",
     "pymysql",
-    ip_type=IPTypes.PRIVATE # use private IP
+    ip_type="private"  # use private IP
 ... insert other kwargs ...
 )
 ```
@@ -333,7 +331,7 @@ conn = connector.connect(
     db="my-db-name",
     active_directory_auth=True,
     server_name="private.[instance].[location].[project].cloudsql.[domain]",
-    ip_type=IPTypes.PRIVATE
+    ip_type="private"
 )
 ```
 
@@ -358,7 +356,7 @@ your web application through the following:
 ```python
 from flask import Flask
 from flask_sqlalchemy import SQLAlchemy
-from google.cloud.sql.connector import Connector, IPTypes
+from google.cloud.sql.connector import Connector
 
 
 # initialize Python Connector object
@@ -372,7 +370,7 @@ def getconn():
         user="my-user",
         password="my-password",
         db="my-database",
-        ip_type= IPTypes.PUBLIC  # IPTypes.PRIVATE for private IP
+        ip_type="public"  # "private" for private IP
     )
     return conn
 
@@ -407,7 +405,7 @@ from sqlalchemy import create_engine
 from sqlalchemy.engine import Engine
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker
-from google.cloud.sql.connector import Connector, IPTypes
+from google.cloud.sql.connector import Connector
 
 # helper function to return SQLAlchemy connection pool
 def init_connection_pool(connector: Connector) -> Engine:
@@ -419,7 +417,7 @@ def init_connection_pool(connector: Connector) -> Engine:
             user="my-user",
             password="my-password",
             db="my-database",
-            ip_type= IPTypes.PUBLIC  # IPTypes.PRIVATE for private IP
+            ip_type="public"  # "private" for private IP
         )
         return conn
 

google/cloud/sql/connector/connector.py

@@ -87,7 +87,7 @@ class Connector:
 
     def __init__(
         self,
-        ip_type: IPTypes = IPTypes.PUBLIC,
+        ip_type: str | IPTypes = IPTypes.PUBLIC,
         enable_iam_auth: bool = False,
         timeout: int = 30,
         credentials: Optional[Credentials] = None,
@@ -130,10 +130,13 @@ def __init__(
         # 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._user_agent = user_agent
+        # if ip_type is str, convert to IPTypes enum
+        if isinstance(ip_type, str):
+            ip_type = IPTypes._get_ip_type_from_str(ip_type)
+        self._ip_type = ip_type
 
     def connect(
         self, instance_connection_string: str, driver: str, **kwargs: Any
@@ -252,6 +255,9 @@ async def connect_async(
             raise KeyError(f"Driver '{driver}' is not supported.")
 
         ip_type = kwargs.pop("ip_type", self._ip_type)
+        # if ip_type is str, convert to IPTypes enum
+        if isinstance(ip_type, str):
+            ip_type = IPTypes._get_ip_type_from_str(ip_type)
         kwargs["timeout"] = kwargs.get("timeout", self._timeout)
 
         # Host and ssl options come from the certificates and metadata, so we don't

google/cloud/sql/connector/instance.py

@@ -64,6 +64,21 @@ class IPTypes(Enum):
     PRIVATE: str = "PRIVATE"
     PSC: str = "PSC"
 
+    @staticmethod
+    def _get_ip_type_from_str(ip_type_str: str) -> IPTypes:
+        """Utility method to convert IP type from a str into IPTypes."""
+        if ip_type_str.lower() == "public":
+            ip_type = IPTypes.PUBLIC
+        elif ip_type_str.lower() == "private":
+            ip_type = IPTypes.PRIVATE
+        elif ip_type_str.lower() == "psc":
+            ip_type = IPTypes.PSC
+        else:
+            raise ValueError(
+                f"Incorrect value for ip_type, got '{ip_type_str}'. Want one of: 'public', 'private' or 'psc'."
+            )
+        return ip_type
+
 
 class ConnectionInfo:
     ip_addrs: Dict[str, Any]

tests/system/test_pg8000_connection.py

@@ -42,6 +42,7 @@ def getconn() -> pg8000.dbapi.Connection:
                 user=os.environ["POSTGRES_USER"],
                 password=os.environ["POSTGRES_PASS"],
                 db=os.environ["POSTGRES_DB"],
+                ip_type="public",  # can also be "private" or "psc"
             )
             return conn
 

tests/system/test_pymysql_connection.py

@@ -42,6 +42,7 @@ def getconn() -> pymysql.connections.Connection:
                 user=os.environ["MYSQL_USER"],
                 password=os.environ["MYSQL_PASS"],
                 db=os.environ["MYSQL_DB"],
+                ip_type="public",  # can also be "private" or "psc"
             )
             return conn
 

tests/system/test_pytds_connection.py

@@ -43,6 +43,7 @@ def getconn() -> pytds.Connection:
                 user=os.environ["SQLSERVER_USER"],
                 password=os.environ["SQLSERVER_PASS"],
                 db=os.environ["SQLSERVER_DB"],
+                ip_type="public",  # can also be "private" or "psc"
             )
             return conn
 

tests/unit/test_connector.py

@@ -130,6 +130,56 @@ async def test_Connector_Init_async_context_manager(
         assert connector._loop == loop
 
 
+@pytest.mark.parametrize(
+    "ip_type",
+    ["public", "private", "psc", "PUBLIC", "PRIVATE", "PSC"],
+)
+def test_Connector_Init_ip_type_str(
+    ip_type: str, fake_credentials: Credentials
+) -> None:
+    """Test that Connector properly sets ip_type when given str."""
+    with Connector(ip_type=ip_type, credentials=fake_credentials) as connector:
+        if ip_type.lower() == "public":
+            assert connector._ip_type == IPTypes.PUBLIC
+        if ip_type.lower() == "private":
+            assert connector._ip_type == IPTypes.PRIVATE
+        if ip_type.lower() == "psc":
+            assert connector._ip_type == IPTypes.PSC
+
+
+def test_Connector_Init_bad_ip_type(fake_credentials: Credentials) -> 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=fake_credentials)
+    assert (
+        exc_info.value.args[0]
+        == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'public', 'private' or 'psc'."
+    )
+
+
+def test_Connector_connect_bad_ip_type(
+    fake_credentials: Credentials, fake_client: CloudSQLClient
+) -> None:
+    """Test that Connector.connect errors due to bad ip_type str."""
+    with Connector(credentials=fake_credentials) as connector:
+        connector._client = fake_client
+        bad_ip_type = "bad-ip-type"
+        with pytest.raises(ValueError) as exc_info:
+            connector.connect(
+                "test-project:test-region:test-instance",
+                "pg8000",
+                user="my-user",
+                password="my-pass",
+                db="my-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' or 'psc'."
+        )
+
+
 @pytest.mark.asyncio
 async def test_Connector_connect_async(
     fake_credentials: Credentials, fake_client: CloudSQLClient