feat: add support for PSC

.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,102 @@
+# 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
+
+# [START alloydb_sqlalchemy_connect_async_connector_psc]
+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
+
+
+# [END alloydb_sqlalchemy_connect_async_connector_psc]
+
+
+@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,100 @@
+# 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
+
+# [START alloydb_sqlalchemy_connect_connector_psc]
+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
+
+
+# [END alloydb_sqlalchemy_connect_connector_psc]
+
+
+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'."
  )