feat: add support for PSC connections (#766)

Author: jackwotherspoon

README.md

@@ -252,9 +252,17 @@ with Connector() as connector:
             print(row)
 ```
 
-### Specifying Public or Private IP
+### Specifying IP Address Type
+
+The Cloud SQL Python Connector can be used to connect to Cloud SQL instances
+using both public and private IP addresses, as well as
+[Private Service Connect][psc] (PSC). To specify which IP address type to connect
+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`.
 
-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:
 
 ```python
@@ -268,9 +276,14 @@ conn = connector.connect(
 )
 ```
 
-Note: If specifying Private IP, your application must already be in the same VPC network as your Cloud SQL Instance.
+Note: If specifying Private IP or Private Service Connect, your application must be
+attached to the proper VPC network to connect to your Cloud SQL instance. For most
+applications this will require the use of a [VPC Connector][vpc-connector].
+
+[psc]: https://cloud.google.com/vpc/docs/private-service-connect
+[vpc-connector]: https://cloud.google.com/vpc/docs/configure-serverless-vpc-access#create-connector
 
-### IAM Authentication
+### Automatic IAM Database Authentication
 
 Connections using [Automatic IAM database authentication](https://cloud.google.com/sql/docs/postgres/authentication#automatic) are supported when using Postgres or MySQL drivers.
 First, make sure to [configure your Cloud SQL Instance to allow IAM authentication](https://cloud.google.com/sql/docs/postgres/create-edit-iam-instances#configure-iam-db-instance)

google/cloud/sql/connector/connector.py

@@ -18,12 +18,16 @@
 import asyncio
 from functools import partial
 import logging
+import socket
 from threading import Thread
 from types import TracebackType
 from typing import Any, Dict, Optional, Type, TYPE_CHECKING
 
 import google.cloud.sql.connector.asyncpg as asyncpg
-from google.cloud.sql.connector.exceptions import ConnectorLoopError
+from google.cloud.sql.connector.exceptions import (
+    ConnectorLoopError,
+    DnsNameResolutionError,
+)
 from google.cloud.sql.connector.instance import (
     Instance,
     IPTypes,
@@ -238,6 +242,21 @@ async def connect_async(
         # attempt to make connection to Cloud SQL instance
         try:
             instance_data, ip_address = await instance.connect_info(ip_type)
+            # resolve DNS name into IP address for PSC
+            if ip_type.value == "PSC":
+                addr_info = await self._loop.getaddrinfo(
+                    ip_address, None, family=socket.AF_INET, type=socket.SOCK_STREAM
+                )
+                # getaddrinfo returns a list of 5-tuples that contain socket
+                # connection info in the form
+                # (family, type, proto, canonname, sockaddr), where sockaddr is a
+                # 2-tuple in the form (ip_address, port)
+                try:
+                    ip_address = addr_info[0][4][0]
+                except IndexError as e:
+                    raise DnsNameResolutionError(
+                        f"['{instance_connection_string}']: DNS name could not be resolved into IP address"
+                    ) from e
 
             # format `user` param for automatic IAM database authn
             if enable_iam_auth:

google/cloud/sql/connector/exceptions.py

@@ -63,3 +63,10 @@ class AutoIAMAuthNotSupported(Exception):
     """
 
     pass
+
+
+class DnsNameResolutionError(Exception):
+    """
+    Exception to be raised when the DnsName of a PSC connection to a
+    Cloud SQL instance can not be resolved to a proper IP address.
+    """

google/cloud/sql/connector/instance.py

@@ -60,6 +60,7 @@
 class IPTypes(Enum):
     PUBLIC: str = "PRIMARY"
     PRIVATE: str = "PRIVATE"
+    PSC: str = "PSC"
 
 
 class InstanceMetadata:

google/cloud/sql/connector/refresh_utils.py

@@ -108,8 +108,16 @@ async def _get_metadata(
             f'[{project}:{region}:{instance}]: Provided region was mismatched - got region {region}, expected {ret_dict["region"]}.'
         )
 
+    ip_addresses = (
+        {ip["type"]: ip["ipAddress"] for ip in ret_dict["ipAddresses"]}
+        if "ipAddresses" in ret_dict
+        else {}
+    )
+    if "dnsName" in ret_dict:
+        ip_addresses["PSC"] = ret_dict["dnsName"]
+
     metadata = {
-        "ip_addresses": {ip["type"]: ip["ipAddress"] for ip in ret_dict["ipAddresses"]},
+        "ip_addresses": ip_addresses,
         "server_ca_cert": ret_dict["serverCaCert"]["cert"],
         "database_version": ret_dict["databaseVersion"],
     }

tests/unit/mocks.py

@@ -219,6 +219,7 @@ def connect_settings(self, ip_addrs: Optional[Dict] = None) -> str:
                         datetime.datetime.utcnow() + datetime.timedelta(minutes=10)
                     ),
                 },
+                "dnsName": "abcde.12345.us-central1.sql.goog",
                 "ipAddresses": ip_addresses,
                 "region": self.region,
                 "databaseVersion": self.db_version,

tests/unit/test_instance.py

@@ -301,6 +301,11 @@ async def test_get_preferred_ip(instance: Instance) -> None:
     # verify private ip address is preferred
     assert ip_addr == "1.1.1.1"
 
+    # test PSC as preferred IP type for connection
+    ip_addr = instance_metadata.get_preferred_ip(IPTypes.PSC)
+    # verify PSC ip address is preferred
+    assert ip_addr == "abcde.12345.us-central1.sql.goog"
+
 
 @pytest.mark.asyncio
 async def test_get_preferred_ip_CloudSQLIPTypeError(instance: Instance) -> None:
@@ -319,6 +324,10 @@ async def test_get_preferred_ip_CloudSQLIPTypeError(instance: Instance) -> None:
     with pytest.raises(CloudSQLIPTypeError):
         instance_metadata.get_preferred_ip(IPTypes.PRIVATE)
 
+    # test error when PSC is missing
+    with pytest.raises(CloudSQLIPTypeError):
+        instance_metadata.get_preferred_ip(IPTypes.PSC)
+
 
 @pytest.mark.asyncio
 async def test_ClientResponseError(