feat: add support for PSC (#291)

.github/workflows/tests.yaml

@@ -169,6 +169,7 @@ jobs:
             ALLOYDB_CLUSTER_PASS:${{ secrets.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_CLUSTER_PASS
             ALLOYDB_IAM_USER:${{ secrets.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_PYTHON_IAM_USER
             ALLOYDB_INSTANCE_IP:${{ secrets.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_INSTANCE_IP
+            ALLOYDB_PSC_INSTANCE_URI:${{ secrets.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_PSC_INSTANCE_URI
 
       - name: Run tests
         env:
@@ -178,6 +179,7 @@ jobs:
           ALLOYDB_IAM_USER: '${{ steps.secrets.outputs.ALLOYDB_IAM_USER }}'
           ALLOYDB_INSTANCE_IP: '${{ steps.secrets.outputs.ALLOYDB_INSTANCE_IP }}'
           ALLOYDB_INSTANCE_URI: '${{ steps.secrets.outputs.ALLOYDB_INSTANCE_URI }}'
+          ALLOYDB_PSC_INSTANCE_URI: '${{ steps.secrets.outputs.ALLOYDB_PSC_INSTANCE_URI }}'
         run: nox -s system-${{ matrix.python-version }}
 
       - name: FlakyBot (Linux)

README.md

@@ -95,7 +95,8 @@ connections. These functions are used with your database driver to connect to
 your AlloyDB instance.
 
 AlloyDB supports network connectivity through public IP addresses and private,
-internal IP addresses. By default this package will attempt to connect over a
+internal IP addresses, as well as [Private Service Connect][psc] (PSC). 
+By default this package will attempt to connect over a
 private IP connection. When doing so, this package must be run in an
 environment that is connected to the [VPC Network][vpc] that hosts your
 AlloyDB private IP address.
@@ -104,6 +105,7 @@ Please see [Configuring AlloyDB Connectivity][alloydb-connectivity] for more det
 
 [vpc]: https://cloud.google.com/vpc/docs/vpc
 [alloydb-connectivity]: https://cloud.google.com/alloydb/docs/configure-connectivity
+[psc]: https://cloud.google.com/vpc/docs/private-service-connect
 
 ### Synchronous Driver Usage
 
@@ -384,10 +386,13 @@ connector.connect(
 
 The AlloyDB Python Connector by default will attempt to establish connections
 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()`.
+over a public IP address or Private Service Connect (PSC), set the `ip_type`
+keyword argument when initializing a `Connector()` or when calling 
+`connector.connect()`.
+
+Possible values for `ip_type` are `"PRIVATE"` (default value), `"PUBLIC"`, 
+and `"PSC"`.
 
-Possible values for `ip_type` are `"PRIVATE"` (default value), and `"PUBLIC"`.
 Example:
 
 ```python

google/cloud/alloydb/connector/client.py

@@ -129,9 +129,15 @@ async def _get_metadata(
         resp = await self._client.get(url, headers=headers, raise_for_status=True)
         resp_dict = await resp.json()
 
+        # Remove trailing period from PSC DNS name.
+        psc_dns = resp_dict.get("pscDnsName")
+        if psc_dns:
+            psc_dns = psc_dns.rstrip(".")
+
         return {
             "PRIVATE": resp_dict.get("ipAddress"),
             "PUBLIC": resp_dict.get("publicIpAddress"),
+            "PSC": psc_dns,
         }
 
     async def _get_client_certificate(

google/cloud/alloydb/connector/instance.py

@@ -48,6 +48,7 @@ class IPTypes(Enum):
 
     PUBLIC: str = "PUBLIC"
     PRIVATE: str = "PRIVATE"
+    PSC: str = "PSC"
 
     @classmethod
     def _missing_(cls, value: object) -> None:

google/cloud/alloydb/connector/refresh.py

@@ -86,7 +86,8 @@ def __init__(
         self.ip_addrs = ip_addrs
         # create TLS context
         self.context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
-        # update ssl.PROTOCOL_TLS_CLIENT default
+        # TODO: Set check_hostname to True to verify the identity in the
+        # certificate once PSC DNS is populated in all existing clusters.
         self.context.check_hostname = False
         # force TLSv1.3
         self.context.minimum_version = ssl.TLSVersion.TLSv1_3

tests/system/test_asyncpg_psc.py

@@ -0,0 +1,98 @@
+# Copyright 2024 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 os
+from typing import Tuple
+
+import asyncpg
+import pytest
+import sqlalchemy
+import sqlalchemy.ext.asyncio
+
+from google.cloud.alloydb.connector import AsyncConnector
+
+
+async def create_sqlalchemy_engine(
+    inst_uri: str,
+    user: str,
+    password: str,
+    db: str,
+) -> Tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, AsyncConnector]:
+    """Creates a connection pool for an AlloyDB instance and returns the pool
+    and the connector. Callers are responsible for closing the pool and the
+    connector.
+
+    A sample invocation looks like:
+
+        engine, connector = await create_sqlalchemy_engine(
+                inst_uri,
+                user,
+                password,
+                db,
+        )
+        async with engine.connect() as conn:
+            time = await conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
+            curr_time = time[0]
+            # do something with query result
+            await connector.close()
+
+    Args:
+        instance_uri (str):
+            The instance URI specifies the instance relative to the project,
+            region, and cluster. For example:
+            "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance"
+        user (str):
+            The database user name, e.g., postgres
+        password (str):
+            The database user's password, e.g., secret-password
+        db_name (str):
+            The name of the database, e.g., mydb
+    """
+    connector = AsyncConnector()
+
+    async def getconn() -> asyncpg.Connection:
+        conn: asyncpg.Connection = await connector.connect(
+            inst_uri,
+            "asyncpg",
+            user=user,
+            password=password,
+            db=db,
+            ip_type="PSC",
+        )
+        return conn
+
+    # create SQLAlchemy connection pool
+    engine = sqlalchemy.ext.asyncio.create_async_engine(
+        "postgresql+asyncpg://",
+        async_creator=getconn,
+        execution_options={"isolation_level": "AUTOCOMMIT"},
+    )
+    return engine, connector
+
+
+@pytest.mark.asyncio
+async def test_connection_with_asyncpg() -> None:
+    """Basic test to get time from database."""
+    inst_uri = os.environ["ALLOYDB_PSC_INSTANCE_URI"]
+    user = os.environ["ALLOYDB_USER"]
+    password = os.environ["ALLOYDB_PASS"]
+    db = os.environ["ALLOYDB_DB"]
+
+    pool, connector = await create_sqlalchemy_engine(inst_uri, user, password, db)
+
+    async with pool.connect() as conn:
+        res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone()
+        assert res[0] == 1
+
+    await connector.close()

tests/system/test_pg8000_psc.py

@@ -0,0 +1,96 @@
+# Copyright 2024 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
+#
+#     http://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.
+
+from datetime import datetime
+import os
+from typing import Tuple
+
+import pg8000
+import sqlalchemy
+
+from google.cloud.alloydb.connector import Connector
+
+
+def create_sqlalchemy_engine(
+    inst_uri: str,
+    user: str,
+    password: str,
+    db: str,
+) -> Tuple[sqlalchemy.engine.Engine, Connector]:
+    """Creates a connection pool for an AlloyDB instance and returns the pool
+    and the connector. Callers are responsible for closing the pool and the
+    connector.
+
+    A sample invocation looks like:
+
+        engine, connector = create_sqlalchemy_engine(
+                inst_uri,
+                user,
+                password,
+                db,
+        )
+        with engine.connect() as conn:
+            time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
+            conn.commit()
+            curr_time = time[0]
+            # do something with query result
+            connector.close()
+
+    Args:
+        instance_uri (str):
+            The instance URI specifies the instance relative to the project,
+            region, and cluster. For example:
+            "projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance"
+        user (str):
+            The database user name, e.g., postgres
+        password (str):
+            The database user's password, e.g., secret-password
+        db_name (str):
+            The name of the database, e.g., mydb
+    """
+    connector = Connector()
+
+    def getconn() -> pg8000.dbapi.Connection:
+        conn: pg8000.dbapi.Connection = connector.connect(
+            inst_uri,
+            "pg8000",
+            user=user,
+            password=password,
+            db=db,
+            ip_type="PSC",
+        )
+        return conn
+
+    # create SQLAlchemy connection pool
+    engine = sqlalchemy.create_engine(
+        "postgresql+pg8000://",
+        creator=getconn,
+    )
+    return engine, connector
+
+
+def test_pg8000_time() -> None:
+    """Basic test to get time from database."""
+    inst_uri = os.environ["ALLOYDB_PSC_INSTANCE_URI"]
+    user = os.environ["ALLOYDB_USER"]
+    password = os.environ["ALLOYDB_PASS"]
+    db = os.environ["ALLOYDB_DB"]
+
+    engine, connector = create_sqlalchemy_engine(inst_uri, user, password, db)
+    with engine.connect() as conn:
+        time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
+        conn.commit()
+        curr_time = time[0]
+        assert type(curr_time) is datetime
+    connector.close()

tests/unit/mocks.py

@@ -16,6 +16,7 @@
 from datetime import datetime
 from datetime import timedelta
 from datetime import timezone
+import ipaddress
 import ssl
 import struct
 from typing import Any, Callable, Dict, List, Optional, Tuple
@@ -60,14 +61,15 @@ def valid(self) -> bool:
 
 
 def generate_cert(
-    common_name: str, expires_in: int = 60
+    common_name: str, expires_in: int = 60, server_cert: bool = False
 ) -> Tuple[x509.CertificateBuilder, rsa.RSAPrivateKey]:
     """
     Generate a private key and cert object to be used in testing.
 
     Args:
         common_name (str): The Common Name for the certificate.
         expires_in (int): Time in minutes until expiry of certificate.
+        server_cert (bool): Whether it is a server certificate.
 
     Returns:
         Tuple[x509.CertificateBuilder, rsa.RSAPrivateKey]
@@ -97,6 +99,17 @@ def generate_cert(
         .not_valid_before(now)
         .not_valid_after(expiration)
     )
+    if server_cert:
+        cert = cert.add_extension(
+            x509.SubjectAlternativeName(
+                general_names=[
+                    x509.IPAddress(ipaddress.ip_address("127.0.0.1")),
+                    x509.IPAddress(ipaddress.ip_address("10.0.0.1")),
+                    x509.DNSName("x.y.alloydb.goog."),
+                ]
+            ),
+            critical=False,
+        )
     return cert, key
 
 
@@ -112,6 +125,7 @@ def __init__(
         ip_addrs: Dict = {
             "PRIVATE": "127.0.0.1",
             "PUBLIC": "0.0.0.0",
+            "PSC": "x.y.alloydb.goog",
         },
         server_name: str = "00000000-0000-0000-0000-000000000000.server.alloydb",
         cert_before: datetime = datetime.now(timezone.utc),
@@ -137,7 +151,9 @@ def __init__(
             self.root_key, hashes.SHA256()
         )
         # build server cert
-        self.server_cert, self.server_key = generate_cert(self.server_name)
+        self.server_cert, self.server_key = generate_cert(
+            self.server_name, server_cert=True
+        )
         # create server cert signed by root cert
         self.server_cert = self.server_cert.sign(self.root_key, hashes.SHA256())
 

tests/unit/test_async_connector.py

@@ -69,6 +69,18 @@ async def test_AsyncConnector_init(credentials: FakeCredentials) -> None:
             IPTypes.PUBLIC,
             IPTypes.PUBLIC,
         ),
+        (
+            "psc",
+            IPTypes.PSC,
+        ),
+        (
+            "PSC",
+            IPTypes.PSC,
+        ),
+        (
+            IPTypes.PSC,
+            IPTypes.PSC,
+        ),
     ],
 )
 async def test_AsyncConnector_init_ip_type(
@@ -90,7 +102,7 @@ async def test_AsyncConnector_init_bad_ip_type(credentials: FakeCredentials) ->
         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'."
+        == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'PUBLIC', 'PRIVATE', 'PSC'."
     )
 
 
@@ -276,5 +288,5 @@ async def test_async_connect_bad_ip_type(
             )
         assert (
             exc_info.value.args[0]
-            == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'PUBLIC', 'PRIVATE'."
+            == f"Incorrect value for ip_type, got '{bad_ip_type}'. Want one of: 'PUBLIC', 'PRIVATE', 'PSC'."
         )