test(e2e): add end-to-end test for postgres connection with sqldata iptype

Author: hessjcg

.github/workflows/tests.yml

@@ -86,6 +86,10 @@ jobs:
             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
+            POSTGRES_AIDE_CONNECTION_NAME:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_AIDE_CONNECTION_NAME
+            POSTGRES_AIDE_USER:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_AIDE_USER
+            POSTGRES_AIDE_PASS:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_AIDE_PASS
+            POSTGRES_AIDE_DB:${{ vars.GOOGLE_CLOUD_PROJECT }}/POSTGRES_AIDE_DB
             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
@@ -112,6 +116,10 @@ jobs:
           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 }}"
+          POSTGRES_AIDE_CONNECTION_NAME: "${{ steps.secrets.outputs.POSTGRES_AIDE_CONNECTION_NAME }}"
+          POSTGRES_AIDE_USER: "${{ steps.secrets.outputs.POSTGRES_AIDE_USER }}"
+          POSTGRES_AIDE_PASS: "${{ steps.secrets.outputs.POSTGRES_AIDE_PASS }}"
+          POSTGRES_AIDE_DB: "${{ steps.secrets.outputs.POSTGRES_AIDE_DB }}"
           SQLSERVER_CONNECTION_NAME: "${{ steps.secrets.outputs.SQLSERVER_CONNECTION_NAME }}"
           SQLSERVER_USER: "${{ steps.secrets.outputs.SQLSERVER_USER }}"
           SQLSERVER_PASS: "${{ steps.secrets.outputs.SQLSERVER_PASS }}"

build.sh

@@ -120,6 +120,10 @@ function write_e2e_env(){
     POSTGRES_CUSTOMER_CAS_INVALID_DOMAIN_NAME=POSTGRES_CUSTOMER_CAS_INVALID_DOMAIN_NAME
     POSTGRES_MCP_CONNECTION_NAME=POSTGRES_MCP_CONNECTION_NAME
     POSTGRES_MCP_PASS=POSTGRES_MCP_PASS
+    POSTGRES_AIDE_CONNECTION_NAME=POSTGRES_AIDE_CONNECTION_NAME
+    POSTGRES_AIDE_USER=POSTGRES_AIDE_USER
+    POSTGRES_AIDE_PASS=POSTGRES_AIDE_PASS
+    POSTGRES_AIDE_DB=POSTGRES_AIDE_DB
     SQLSERVER_CONNECTION_NAME=SQLSERVER_CONNECTION_NAME
     SQLSERVER_USER=SQLSERVER_USER
     SQLSERVER_PASS=SQLSERVER_PASS

tests/system/test_asyncpg_connection.py

@@ -20,6 +20,7 @@
 from typing import Any
 
 import asyncpg
+import pytest
 import sqlalchemy
 import sqlalchemy.ext.asyncio
 
@@ -283,3 +284,52 @@ async def test_lazy_connection_with_asyncpg() -> None:
         assert res[0][0] == 1
 
     await connector.close_async()
+
+
+async def test_AIDE_sqlalchemy_connection_with_asyncpg() -> None:
+    """Basic test to get time from database using AIDE instance."""
+    if "POSTGRES_AIDE_CONNECTION_NAME" not in os.environ:
+        pytest.skip("POSTGRES_AIDE_CONNECTION_NAME not set")
+    inst_conn_name = os.environ["POSTGRES_AIDE_CONNECTION_NAME"]
+    user = os.environ.get("POSTGRES_AIDE_USER", os.environ.get("POSTGRES_USER", "postgres"))
+    password = os.environ.get("POSTGRES_AIDE_PASS", os.environ.get("POSTGRES_PASS", ""))
+    db = os.environ.get("POSTGRES_AIDE_DB", os.environ.get("POSTGRES_DB", "postgres"))
+
+    pool, connector = await create_sqlalchemy_engine(
+        inst_conn_name,
+        user,
+        password,
+        db,
+        ip_type="sqldata",
+    )
+
+    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_AIDE_connection_with_asyncpg() -> None:
+    """Basic test to get time from database using AIDE instance."""
+    if "POSTGRES_AIDE_CONNECTION_NAME" not in os.environ:
+        pytest.skip("POSTGRES_AIDE_CONNECTION_NAME not set")
+    inst_conn_name = os.environ["POSTGRES_AIDE_CONNECTION_NAME"]
+    user = os.environ.get("POSTGRES_AIDE_USER", os.environ.get("POSTGRES_USER", "postgres"))
+    password = os.environ.get("POSTGRES_AIDE_PASS", os.environ.get("POSTGRES_PASS", ""))
+    db = os.environ.get("POSTGRES_AIDE_DB", os.environ.get("POSTGRES_DB", "postgres"))
+
+    pool, connector = await create_asyncpg_pool(
+        inst_conn_name,
+        user,
+        password,
+        db,
+        ip_type="sqldata",
+    )
+
+    async with pool.acquire() as conn:
+        res = await conn.fetch("SELECT 1")
+        assert res[0][0] == 1
+
+    await connector.close_async()
+

tests/system/test_pg8000_connection.py

@@ -19,6 +19,7 @@
 import os
 
 # [START cloud_sql_connector_postgres_pg8000]
+import pytest
 import sqlalchemy
 
 from google.cloud.sql.connector import Connector
@@ -209,3 +210,28 @@ def test_MCP_pg8000_connection() -> None:
         curr_time = time[0]
         assert type(curr_time) is datetime
     connector.close()
+
+
+def test_AIDE_pg8000_connection() -> None:
+    """Basic test to get time from database using AIDE instance."""
+    if "POSTGRES_AIDE_CONNECTION_NAME" not in os.environ:
+        pytest.skip("POSTGRES_AIDE_CONNECTION_NAME not set")
+    inst_conn_name = os.environ["POSTGRES_AIDE_CONNECTION_NAME"]
+    user = os.environ.get("POSTGRES_AIDE_USER", os.environ.get("POSTGRES_USER", "postgres"))
+    password = os.environ.get("POSTGRES_AIDE_PASS", os.environ.get("POSTGRES_PASS", ""))
+    db = os.environ.get("POSTGRES_AIDE_DB", os.environ.get("POSTGRES_DB", "postgres"))
+
+    engine, connector = create_sqlalchemy_engine(
+        inst_conn_name,
+        user,
+        password,
+        db,
+        ip_type="sqldata",
+    )
+    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()
+