diff --git a/.ci/cloudbuild.yaml b/.ci/cloudbuild.yaml index 4c40d3261..e31e8bebb 100644 --- a/.ci/cloudbuild.yaml +++ b/.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" diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 755ad5690..8e5d0b002 100644 --- a/.github/workflows/tests.yml +++ b/.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 }}" diff --git a/tests/system/test_asyncpg_connection.py b/tests/system/test_asyncpg_connection.py index d7602e537..1aae1f7e6 100644 --- a/tests/system/test_asyncpg_connection.py +++ b/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"] diff --git a/tests/system/test_pg8000_connection.py b/tests/system/test_pg8000_connection.py index 42a8d9894..c7074b0cf 100644 --- a/tests/system/test_pg8000_connection.py +++ b/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()