Repository: GoogleCloudPlatform/cloud-sql-python-connector
Pull request #1029 (2024-03-04 12:54:15)

feat: support `ip_type` as str
Adding support for `ip_type` as `str` type with options as `"public"`, `"private"` or `"psc"`.

If users want to configure a Connector for a specific IP type they currently need to import `IPTypes` first and then set `ip_type=IPTypes.PUBLIC` in the Connector config or during their call to connect.

```python
from google.cloud.sql.connector import IPTypes
```

This is an extra step that is unnecessary and causes friction to users, we should support users passing in the IPTypes enum or string value. This will allow users to not have to import the additional class.




Patch:
diff --git a/README.md b/README.md
index 1d9b28e..4700ea1 100644
--- a/README.md
+++ b/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
 
diff --git a/google/cloud/sql/connector/connector.py b/google/cloud/sql/connector/connector.py
index df27192..760db2c 100755
--- a/google/cloud/sql/connector/connector.py
+++ b/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 @@ class Connector:
         # 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 @@ class Connector:
             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
diff --git a/google/cloud/sql/connector/instance.py b/google/cloud/sql/connector/instance.py
index dda5eb9..5d8f412 100644
--- a/google/cloud/sql/connector/instance.py
+++ b/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]


Test patch:
diff --git a/tests/system/test_pg8000_connection.py b/tests/system/test_pg8000_connection.py
index 77c17ab..9cf56aa 100644
--- a/tests/system/test_pg8000_connection.py
+++ b/tests/system/test_pg8000_connection.py
@@ -42,6 +42,7 @@ def init_connection_engine() -> sqlalchemy.engine.Engine:
                 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
 
diff --git a/tests/system/test_pymysql_connection.py b/tests/system/test_pymysql_connection.py
index 5258be7..fe5ad94 100644
--- a/tests/system/test_pymysql_connection.py
+++ b/tests/system/test_pymysql_connection.py
@@ -42,6 +42,7 @@ def init_connection_engine() -> sqlalchemy.engine.Engine:
                 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
 
diff --git a/tests/system/test_pytds_connection.py b/tests/system/test_pytds_connection.py
index c21a8e0..54e2ecc 100644
--- a/tests/system/test_pytds_connection.py
+++ b/tests/system/test_pytds_connection.py
@@ -43,6 +43,7 @@ def init_connection_engine() -> sqlalchemy.engine.Engine:
                 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
 
diff --git a/tests/unit/test_connector.py b/tests/unit/test_connector.py
index a9c57de..08e9aa5 100644
--- a/tests/unit/test_connector.py
+++ b/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