chore: merge main

Author: jackwotherspoon

google/cloud/sql/connector/connection_name.py

@@ -19,6 +19,10 @@
 # Additionally, we have to support legacy "domain-scoped" projects
 # (e.g. "google.com:PROJECT")
 CONN_NAME_REGEX = re.compile(("([^:]+(:[^:]+)?):([^:]+):([^:]+)"))
+# The domain name pattern in accordance with RFC 1035, RFC 1123 and RFC 2181.
+DOMAIN_NAME_REGEX = re.compile(
+    r"^(?:[_a-z0-9](?:[_a-z0-9-]{0,61}[a-z0-9])?\.)+(?:[a-z](?:[a-z0-9-]{0,61}[a-z0-9])?)?$"
+)
 
 
 @dataclass
@@ -43,6 +47,12 @@ def get_connection_string(self) -> str:
         return f"{self.project}:{self.region}:{self.instance_name}"
 
 
+def _is_valid_domain(domain_name: str) -> bool:
+    if DOMAIN_NAME_REGEX.fullmatch(domain_name) is None:
+        return False
+    return True
+
+
 def _parse_connection_name(connection_name: str) -> ConnectionName:
     return _parse_connection_name_with_domain_name(connection_name, "")
 

google/cloud/sql/connector/resolver.py

@@ -17,6 +17,7 @@
 from google.cloud.sql.connector.connection_name import (
     _parse_connection_name_with_domain_name,
 )
+from google.cloud.sql.connector.connection_name import _is_valid_domain
 from google.cloud.sql.connector.connection_name import _parse_connection_name
 from google.cloud.sql.connector.connection_name import ConnectionName
 from google.cloud.sql.connector.exceptions import DnsResolutionError
@@ -40,8 +41,16 @@ async def resolve(self, dns: str) -> ConnectionName:  # type: ignore
             conn_name = _parse_connection_name(dns)
         except ValueError:
             # The connection name was not project:region:instance format.
-            # Attempt to query a TXT record to get connection name.
-            conn_name = await self.query_dns(dns)
+            # Check if connection name is a valid DNS domain name
+            if _is_valid_domain(dns):
+                # Attempt to query a TXT record to get connection name.
+                conn_name = await self.query_dns(dns)
+            else:
+                raise ValueError(
+                    "Arg `instance_connection_string` must have "
+                    "format: PROJECT:REGION:INSTANCE or be a valid DNS domain "
+                    f"name, got {dns}."
+                )
         return conn_name
 
     async def query_dns(self, dns: str) -> ConnectionName:

tests/system/test_asyncpg_connection.py

@@ -16,13 +16,15 @@
 
 import asyncio
 import os
-from typing import Any
+from typing import Any, Union
 
 import asyncpg
 import sqlalchemy
 import sqlalchemy.ext.asyncio
 
 from google.cloud.sql.connector import Connector
+from google.cloud.sql.connector import DefaultResolver
+from google.cloud.sql.connector import DnsResolver
 
 
 async def create_sqlalchemy_engine(
@@ -31,6 +33,7 @@ async def create_sqlalchemy_engine(
     password: str,
     db: str,
     refresh_strategy: str = "background",
+    resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
 ) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, Connector]:
     """Creates a connection pool for a Cloud SQL instance and returns the pool
     and the connector. Callers are responsible for closing the pool and the
@@ -64,9 +67,16 @@ async def create_sqlalchemy_engine(
             Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
             or "background". For serverless environments use "lazy" to avoid
             errors resulting from CPU being throttled.
+        resolver (Optional[google.cloud.sql.connector.DefaultResolver]):
+            Resolver class for resolving instance connection name. Use
+            google.cloud.sql.connector.DnsResolver when resolving DNS domain
+            names or google.cloud.sql.connector.DefaultResolver for regular
+            instance connection names ("my-project:my-region:my-instance").
     """
     loop = asyncio.get_running_loop()
-    connector = Connector(loop=loop, refresh_strategy=refresh_strategy)
+    connector = Connector(
+        loop=loop, refresh_strategy=refresh_strategy, resolver=resolver
+    )
 
     async def getconn() -> asyncpg.Connection:
         conn: asyncpg.Connection = await connector.connect_async(
@@ -183,6 +193,24 @@ async def test_lazy_sqlalchemy_connection_with_asyncpg() -> None:
     await connector.close_async()
 
 
+async def test_custom_SAN_with_dns_sqlalchemy_connection_with_asyncpg() -> None:
+    """Basic test to get time from database."""
+    inst_conn_name = os.environ["POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME"]
+    user = os.environ["POSTGRES_USER"]
+    password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
+    db = os.environ["POSTGRES_DB"]
+
+    pool, connector = await create_sqlalchemy_engine(
+        inst_conn_name, user, password, db, resolver=DnsResolver
+    )
+
+    async with pool.connect() as conn:
+        res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone()
+        assert res[0] == 1
+
+    await connector.close_async()
+
+
 async def test_connection_with_asyncpg() -> None:
     """Basic test to get time from database."""
     inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]

tests/system/test_pg8000_connection.py

@@ -34,7 +34,7 @@ def create_sqlalchemy_engine(
     password: str,
     db: str,
     refresh_strategy: str = "background",
-    resolver: Union[DefaultResolver, DnsResolver] = DefaultResolver,
+    resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
 ) -> tuple[sqlalchemy.engine.Engine, Connector]:
     """Creates a connection pool for a Cloud SQL instance and returns the pool
     and the connector. Callers are responsible for closing the pool and the
@@ -69,11 +69,11 @@ def create_sqlalchemy_engine(
             Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
             or "background". For serverless environments use "lazy" to avoid
             errors resulting from CPU being throttled.
-        resolver (Optional[google.cloud.sql.connector.DefaultResolver | google.cloud.sql.connector.DnsResolver])
-            Resolver class for the Cloud SQL Connector. Can be one of
-            DefaultResolver (default) or DnsResolver. The resolver tells the
-            connector whether to resolve the 'instance_connection_name' as a
-            Cloud SQL instance connection name or as a domain name.
+        resolver (Optional[google.cloud.sql.connector.DefaultResolver]):
+            Resolver class for resolving instance connection name. Use
+            google.cloud.sql.connector.DnsResolver when resolving DNS domain
+            names or google.cloud.sql.connector.DefaultResolver for regular
+            instance connection names ("my-project:my-region:my-instance").
     """
     connector = Connector(refresh_strategy=refresh_strategy, resolver=resolver)
 
@@ -165,15 +165,15 @@ def test_customer_managed_CAS_pg8000_connection() -> None:
     connector.close()
 
 
-def test_domain_name_pg8000_connection() -> None:
-    """Basic test to get time from database using domain name to connect."""
-    domain_name = os.environ["POSTGRES_CUSTOMER_CAS_DOMAIN_NAME"]
+def test_custom_SAN_with_dns_pg8000_connection() -> None:
+    """Basic test to get time from database."""
+    inst_conn_name = os.environ["POSTGRES_CUSTOMER_CAS_DOMAIN_NAME"]
     user = os.environ["POSTGRES_USER"]
     password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
     db = os.environ["POSTGRES_DB"]
 
     engine, connector = create_sqlalchemy_engine(
-        domain_name, user, password, db, "lazy", DnsResolver
+        inst_conn_name, user, password, db, resolver=DnsResolver
     )
     with engine.connect() as conn:
         time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()

tests/unit/test_connection_name.py

@@ -17,6 +17,7 @@
 from google.cloud.sql.connector.connection_name import (
     _parse_connection_name_with_domain_name,
 )
+from google.cloud.sql.connector.connection_name import _is_valid_domain
 from google.cloud.sql.connector.connection_name import _parse_connection_name
 from google.cloud.sql.connector.connection_name import ConnectionName
 
@@ -100,3 +101,40 @@ def test_parse_connection_name_with_domain_name(
     assert expected == _parse_connection_name_with_domain_name(
         connection_name, domain_name
     )
+
+
+@pytest.mark.parametrize(
+    "domain_name, expected",
+    [
+        (
+            "prod-db.mycompany.example.com",
+            True,
+        ),
+        (
+            "example.com.",  # trailing dot
+            True,
+        ),
+        (
+            "-example.com.",  # leading hyphen
+            False,
+        ),
+        (
+            "example",  # missing TLD
+            False,
+        ),
+        (
+            "127.0.0.1",  # IPv4 address
+            False,
+        ),
+        (
+            "0:0:0:0:0:0:0:1",  # IPv6 address
+            False,
+        ),
+    ],
+)
+def test_is_valid_domain(domain_name: str, expected: bool) -> None:
+    """
+    Test that _is_valid_domain works correctly for
+    parsing domain names.
+    """
+    assert expected == _is_valid_domain(domain_name)