feat: allow specifying ip address type (#79)

* feat: allow specifying ip address type as public and/or private when connecting

* make IPTypes enum public

* address review comments

* set default IP type to public

* add tests for specifying IP types and flag to skip private IP

* Add info about specifying IP types to README

* linting fixes

* allow only 1 IP type to be passed into connection args

* update README

* hard code default IP type

Author: shubha-rajan

README.md

@@ -56,6 +56,19 @@ connector.connect(
 ```
 **Note for SQL Server users**: If your SQL Server instance requires SSL, you need to download the CA certificate for your instance and include `cafile={path to downloaded certificate}` and `validate_host=False`. This is a workaround for a [known issue](https://issuetracker.google.com/184867147).
 
+### Specifying Public or Private IP
+The Cloud SQL Connector for Python can be used to connect to Cloud SQL instances using both public and private IP addresses. To specify which IP address to use to connect, set the `ip_type` keyword argument Possible values are `IPTypes.PUBLIC` and `IPTypes.PRIVATE`. 
+Example:
+```
+connector.connect(
+    "your:connection:string:", 
+    "pymysql",
+    ip_types=IPTypes.PRIVATE # Prefer private IP
+... insert other kwargs ...
+)
+```
+
+Note: If specifying Private IP, your application must already be in the same VPC network as your Cloud SQL Instance. 
 ### Setup for development
 
 Tests can be run with `nox`. Change directory into the `cloud-sql-python-connector` and just run `nox` to run the tests.

google/cloud/sql/connector/__init__.py

@@ -17,10 +17,10 @@
 from typing import Iterable
 
 from .connector import connect
-from .instance_connection_manager import CloudSQLConnectionError
+from .instance_connection_manager import CloudSQLConnectionError, IPTypes
 
 
-__ALL__ = [connect, CloudSQLConnectionError]
+__ALL__ = [connect, CloudSQLConnectionError, IPTypes]
 
 try:
     import pkg_resources

google/cloud/sql/connector/connector.py

@@ -17,13 +17,13 @@
 import concurrent
 from google.cloud.sql.connector.instance_connection_manager import (
     InstanceConnectionManager,
+    IPTypes,
 )
 from google.cloud.sql.connector.utils import generate_keys
 
 from threading import Thread
 from typing import Any, Dict, Optional
 
-
 # This thread is used to background processing
 _thread: Optional[Thread] = None
 _loop: Optional[asyncio.AbstractEventLoop] = None
@@ -48,7 +48,12 @@ def _get_keys(loop: asyncio.AbstractEventLoop) -> concurrent.futures.Future:
     return _keys
 
 
-def connect(instance_connection_string: str, driver: str, **kwargs: Any) -> Any:
+def connect(
+    instance_connection_string: str,
+    driver: str,
+    ip_types: IPTypes = IPTypes.PUBLIC,
+    **kwargs: Any
+) -> Any:
     """Prepares and returns a database connection object and starts a
     background thread to refresh the certificates and metadata.
 
@@ -59,6 +64,15 @@ def connect(instance_connection_string: str, driver: str, **kwargs: Any) -> Any:
 
         Example: example-proj:example-region-us6:example-instance
 
+    :type driver: str
+    :param: driver:
+        A string representing the driver to connect with. Supported drivers are
+        pymysql, pg8000, and pytds.
+
+    :type ip_types: IPTypes
+        The IP type (public or private)  used to connect. IP types
+        can be either IPTypes.PUBLIC or IPTypes.PRIVATE.
+
     :param kwargs:
         Pass in any driver-specific arguments needed to connect to the Cloud
         SQL instance.
@@ -86,10 +100,10 @@ def connect(instance_connection_string: str, driver: str, **kwargs: Any) -> Any:
         _instances[instance_connection_string] = icm
 
     if "timeout" in kwargs:
-        return icm.connect(driver, **kwargs)
+        return icm.connect(driver, ip_types, **kwargs)
     elif "connect_timeout" in kwargs:
         timeout = kwargs["connect_timeout"]
     else:
         timeout = 30  # 30s
 
-    return icm.connect(driver, timeout, **kwargs)
+    return icm.connect(driver, ip_types, timeout, **kwargs)

google/cloud/sql/connector/instance_connection_manager.py

@@ -23,6 +23,7 @@
 import asyncio
 import aiohttp
 import concurrent
+from enum import Enum
 import google.auth
 from google.auth.credentials import Credentials
 import google.auth.transport.requests
@@ -33,6 +34,7 @@
     Any,
     Awaitable,
     Coroutine,
+    Dict,
     Optional,
     TYPE_CHECKING,
     Union,
@@ -55,6 +57,11 @@
 _delay: int = 55 * 60
 
 
+class IPTypes(Enum):
+    PUBLIC: str = "PRIMARY"
+    PRIVATE: str = "PRIVATE"
+
+
 class ConnectionSSLContext(ssl.SSLContext):
     """Subclass of ssl.SSLContext with added request_ssl attribute. This is
     required for compatibility with pg8000 driver.
@@ -65,18 +72,37 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super(ConnectionSSLContext, self).__init__(*args, **kwargs)
 
 
+class CloudSQLConnectionError(Exception):
+    """
+    Raised when the provided connection string is not formatted
+    correctly.
+    """
+
+    def __init__(self, *args: Any) -> None:
+        super(CloudSQLConnectionError, self).__init__(self, *args)
+
+
+class CloudSQLIPTypeError(Exception):
+    """
+    Raised when IP address for the preferred IP type is not found.
+    """
+
+    def __init__(self, *args: Any) -> None:
+        super(CloudSQLIPTypeError, self).__init__(self, *args)
+
+
 class InstanceMetadata:
-    ip_address: str
+    ip_addrs: Dict[str, Any]
     context: ssl.SSLContext
 
     def __init__(
         self,
         ephemeral_cert: str,
-        ip_address: str,
+        ip_addrs: Dict[str, Any],
         private_key: bytes,
         server_ca_cert: str,
     ) -> None:
-        self.ip_address = ip_address
+        self.ip_addrs = ip_addrs
         self.context = ConnectionSSLContext()
 
         # tmpdir and its contents are automatically deleted after the CA cert
@@ -89,15 +115,16 @@ def __init__(
             self.context.load_cert_chain(cert_filename, keyfile=key_filename)
             self.context.load_verify_locations(cafile=ca_filename)
 
-
-class CloudSQLConnectionError(Exception):
-    """
-    Raised when the provided connection string is not formatted
-    correctly.
-    """
-
-    def __init__(self, *args: Any) -> None:
-        super(CloudSQLConnectionError, self).__init__(self, *args)
+    def get_preferred_ip(self, ip_type: IPTypes) -> str:
+        """Returns the first IP address for the instance, according to the preference
+        supplied by ip_type. If no IP addressess with the given preference are found,
+        an error is raised."""
+        if ip_type.value in self.ip_addrs:
+            return self.ip_addrs[ip_type.value]
+        raise CloudSQLIPTypeError(
+            "Cloud SQL instance does not have any IP addresses matching "
+            f"preference: {ip_type.value})"
+        )
 
 
 class InstanceConnectionManager:
@@ -241,7 +268,7 @@ async def _get_instance_data(self) -> InstanceMetadata:
 
         return InstanceMetadata(
             ephemeral_cert,
-            metadata["ip_addresses"]["PRIMARY"],
+            metadata["ip_addresses"],
             priv_key,
             metadata["server_ca_cert"],
         )
@@ -296,7 +323,13 @@ async def _schedule_refresh(self, delay: int) -> asyncio.Task:
 
         return await self._perform_refresh()
 
-    def connect(self, driver: str, timeout: int, **kwargs: Any) -> Any:
+    def connect(
+        self,
+        driver: str,
+        ip_type: IPTypes,
+        timeout: int,
+        **kwargs: Any,
+    ) -> Any:
         """A method that returns a DB-API connection to the database.
 
         :type driver: str
@@ -310,7 +343,7 @@ def connect(self, driver: str, timeout: int, **kwargs: Any) -> Any:
         """
 
         connect_future: concurrent.futures.Future = asyncio.run_coroutine_threadsafe(
-            self._connect(driver, **kwargs), self._loop
+            self._connect(driver, ip_type, **kwargs), self._loop
         )
 
         try:
@@ -321,7 +354,12 @@ def connect(self, driver: str, timeout: int, **kwargs: Any) -> Any:
         else:
             return connection
 
-    async def _connect(self, driver: str, **kwargs: Any) -> Any:
+    async def _connect(
+        self,
+        driver: str,
+        ip_type: IPTypes,
+        **kwargs: Any,
+    ) -> Any:
         """A method that returns a DB-API connection to the database.
 
         :type driver: str
@@ -344,14 +382,15 @@ async def _connect(self, driver: str, **kwargs: Any) -> Any:
         }
 
         instance_data: InstanceMetadata = await self._current
+        ip_address: str = instance_data.get_preferred_ip(ip_type)
 
         try:
             connector = connect_func[driver]
         except KeyError:
             raise KeyError("Driver {} is not supported.".format(driver))
 
         connect_partial = partial(
-            connector, instance_data.ip_address, instance_data.context, **kwargs
+            connector, ip_address, instance_data.context, **kwargs
         )
 
         return await self._loop.run_in_executor(None, connect_partial)

tests/unit/conftest.py tests/conftest.pytests/unit/conftest.py renamed to tests/conftest.py

@@ -15,12 +15,36 @@
 """
 import os
 import threading
-from typing import Generator
+from typing import Any, Generator
 
 import asyncio
 import pytest  # noqa F401 Needed to run the tests
 
 
+def pytest_addoption(parser: Any) -> None:
+    parser.addoption(
+        "--run_private_ip",
+        action="store_true",
+        default=False,
+        help="run tests that need to be running in VPC network",
+    )
+
+
+def pytest_configure(config: Any) -> None:
+    config.addinivalue_line(
+        "markers", "private_ip: mark test as requiring private IP access"
+    )
+
+
+def pytest_collection_modifyitems(config: Any, items: Any) -> None:
+    if config.getoption("--run_private_ip"):
+        return
+    skip_private_ip = pytest.mark.skip(reason="need --run_private_ip option to run")
+    for item in items:
+        if "private_ip" in item.keywords:
+            item.add_marker(skip_private_ip)
+
+
 @pytest.fixture
 def async_loop() -> Generator:
     """

tests/system/test_ip_types.py

@@ -0,0 +1,63 @@
+""""
+Copyright 2021 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
+
+  https://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 logging
+import os
+import uuid
+
+import pymysql
+import pytest
+import sqlalchemy
+from google.cloud.sql.connector import connector, IPTypes
+
+table_name = f"books_{uuid.uuid4().hex}"
+
+
+def init_connection_engine(ip_types: IPTypes) -> sqlalchemy.engine.Engine:
+    def getconn() -> pymysql.connections.Connection:
+        conn: pymysql.connections.Connection = connector.connect(
+            os.environ["MYSQL_CONNECTION_NAME"],
+            "pymysql",
+            ip_types=ip_types,
+            user=os.environ["MYSQL_USER"],
+            password=os.environ["MYSQL_PASS"],
+            db=os.environ["MYSQL_DB"],
+        )
+        return conn
+
+    engine = sqlalchemy.create_engine(
+        "mysql+pymysql://",
+        creator=getconn,
+    )
+    return engine
+
+
+def test_public_ip() -> None:
+    try:
+        pool = init_connection_engine(IPTypes.PUBLIC)
+    except Exception as e:
+        logging.exception("Failed to initialize pool with public IP", e)
+    with pool.connect() as conn:
+        conn.execute("SELECT 1")
+
+
+@pytest.mark.private_ip
+def test_private_ip() -> None:
+    try:
+        pool = init_connection_engine(IPTypes.PRIVATE)
+    except Exception as e:
+        logging.exception("Failed to initialize pool with private IP", e)
+    with pool.connect() as conn:
+        conn.execute("SELECT 1")

tests/unit/test_connector.py

@@ -45,5 +45,9 @@ async def timeout_stub(*args: Any, **kwargs: Any) -> None:
     mock_instances[connect_string] = icm
     with patch.dict(connector._instances, mock_instances):
         pytest.raises(
-            TimeoutError, connector.connect, connect_string, "pymysql", timeout=timeout
+            TimeoutError,
+            connector.connect,
+            connect_string,
+            "pymysql",
+            timeout=timeout,
         )