test: add Postgres MCP tests (GoogleCloudPlatform#1266)

Add integration tests for Postgres for Cloud SQL's Managed Connection Pooling (MCP).

Author: jackwotherspoon

.ci/cloudbuild.yaml

@@ -37,6 +37,8 @@ steps:
         "POSTGRES_CUSTOMER_CAS_CONNECTION_NAME",
         "POSTGRES_CUSTOMER_CAS_PASS",
         "POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME",
+        "POSTGRES_MCP_CONNECTION_NAME",
+        "POSTGRES_MCP_PASS",
         "SQLSERVER_CONNECTION_NAME",
         "SQLSERVER_USER",
         "SQLSERVER_PASS",
@@ -83,6 +85,10 @@ availableSecrets:
       env: "POSTGRES_CUSTOMER_CAS_PASS"
     - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME/versions/latest"
       env: "POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME"
+    - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_MCP_CONNECTION_NAME/versions/latest"
+      env: "POSTGRES_MCP_CONNECTION_NAME"
+    - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_MCP_PASS/versions/latest"
+      env: "POSTGRES_MCP_PASS"
     - versionName: "projects/$PROJECT_ID/secrets/SQLSERVER_CONNECTION_NAME/versions/latest"
       env: "SQLSERVER_CONNECTION_NAME"
     - versionName: "projects/$PROJECT_ID/secrets/SQLSERVER_USER/versions/latest"

.github/workflows/tests.yml

@@ -84,6 +84,8 @@ jobs:
             POSTGRES_CUSTOMER_CAS_CONNECTION_NAME:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_CUSTOMER_CAS_CONNECTION_NAME
             POSTGRES_CUSTOMER_CAS_PASS:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_CUSTOMER_CAS_PASS
             POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME
+            POSTGRES_MCP_CONNECTION_NAME:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_MCP_CONNECTION_NAME
+            POSTGRES_MCP_PASS:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_MCP_PASS
             SQLSERVER_CONNECTION_NAME:${{ vars.GOOGLE_CLOUD_PROJECT }}/SQLSERVER_CONNECTION_NAME
             SQLSERVER_USER:${{ vars.GOOGLE_CLOUD_PROJECT }}/SQLSERVER_USER
             SQLSERVER_PASS:${{ vars.GOOGLE_CLOUD_PROJECT }}/SQLSERVER_PASS
@@ -108,6 +110,8 @@ jobs:
           POSTGRES_CUSTOMER_CAS_CONNECTION_NAME: "${{ steps.secrets.outputs.POSTGRES_CUSTOMER_CAS_CONNECTION_NAME }}"
           POSTGRES_CUSTOMER_CAS_PASS: "${{ steps.secrets.outputs.POSTGRES_CUSTOMER_CAS_PASS }}"
           POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME: "${{ steps.secrets.outputs.POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME }}"
+          POSTGRES_MCP_CONNECTION_NAME: "${{ steps.secrets.outputs.POSTGRES_MCP_CONNECTION_NAME }}"
+          POSTGRES_MCP_PASS: "${{ steps.secrets.outputs.POSTGRES_MCP_PASS }}"
           SQLSERVER_CONNECTION_NAME: "${{ steps.secrets.outputs.SQLSERVER_CONNECTION_NAME }}"
           SQLSERVER_USER: "${{ steps.secrets.outputs.SQLSERVER_USER }}"
           SQLSERVER_PASS: "${{ steps.secrets.outputs.SQLSERVER_PASS }}"

tests/system/test_asyncpg_connection.py

@@ -35,6 +35,7 @@ async def create_sqlalchemy_engine(
     ip_type: str = "public",
     refresh_strategy: str = "background",
     resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
+    **kwargs: Any,
 ) -> 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
@@ -92,6 +93,7 @@ async def create_sqlalchemy_engine(
             password=password,
             db=db,
             ip_type=ip_type,  # can be "public", "private" or "psc"
+            **kwargs,  # additional asyncpg connection args
         ),
         execution_options={"isolation_level": "AUTOCOMMIT"},
     )
@@ -220,6 +222,30 @@ async def test_custom_SAN_with_dns_sqlalchemy_connection_with_asyncpg() -> None:
     await connector.close_async()
 
 
+async def test_MCP_sqlalchemy_connection_with_asyncpg() -> None:
+    """Basic test to get time from database using MCP enabled instance."""
+    inst_conn_name = os.environ["POSTGRES_MCP_CONNECTION_NAME"]
+    user = os.environ["POSTGRES_USER"]
+    password = os.environ["POSTGRES_MCP_PASS"]
+    db = os.environ["POSTGRES_DB"]
+    ip_type = os.environ.get("IP_TYPE", "public")
+
+    pool, connector = await create_sqlalchemy_engine(
+        inst_conn_name,
+        user,
+        password,
+        db,
+        ip_type,
+        statement_cache_size=0,
+    )
+
+    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

@@ -191,3 +191,22 @@ def test_custom_SAN_with_dns_pg8000_connection() -> None:
         curr_time = time[0]
         assert type(curr_time) is datetime
     connector.close()
+
+
+def test_MCP_pg8000_connection() -> None:
+    """Basic test to get time from database using MCP enabled instance."""
+    inst_conn_name = os.environ["POSTGRES_MCP_CONNECTION_NAME"]
+    user = os.environ["POSTGRES_USER"]
+    password = os.environ["POSTGRES_MCP_PASS"]
+    db = os.environ["POSTGRES_DB"]
+    ip_type = os.environ.get("IP_TYPE", "public")
+
+    engine, connector = create_sqlalchemy_engine(
+        inst_conn_name, user, password, db, ip_type
+    )
+    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()