From 61379d8da6c1c52d0cc7fec718c01073bd9f0dc7 Mon Sep 17 00:00:00 2001 From: jackwotherspoon Date: Fri, 4 Apr 2025 13:27:39 +0000 Subject: [PATCH 1/6] test: add postgres MCP test --- .ci/cloudbuild.yaml | 109 +++++++++++++++---------- .github/workflows/tests.yml | 4 + tests/system/test_pg8000_connection.py | 19 +++++ 3 files changed, 91 insertions(+), 41 deletions(-) diff --git a/.ci/cloudbuild.yaml b/.ci/cloudbuild.yaml index 9ed3add58..4bdbd01c8 100644 --- a/.ci/cloudbuild.yaml +++ b/.ci/cloudbuild.yaml @@ -16,9 +16,32 @@ steps: - id: run integration tests name: python:${_VERSION} entrypoint: bash - env: + env: - "IP_TYPE=${_IP_TYPE}" - secretEnv: ["MYSQL_CONNECTION_NAME", "MYSQL_USER", "MYSQL_IAM_USER", "MYSQL_PASS", "MYSQL_DB", "POSTGRES_CONNECTION_NAME", "POSTGRES_USER", "POSTGRES_IAM_USER", "POSTGRES_PASS", "POSTGRES_DB", "POSTGRES_CAS_CONNECTION_NAME", "POSTGRES_CAS_PASS", "POSTGRES_CUSTOMER_CAS_CONNECTION_NAME", "POSTGRES_CUSTOMER_CAS_PASS", "POSTGRES_CUSTOMER_CAS_PASS_VALID_DOMAIN_NAME","SQLSERVER_CONNECTION_NAME", "SQLSERVER_USER", "SQLSERVER_PASS", "SQLSERVER_DB"] + secretEnv: + [ + "MYSQL_CONNECTION_NAME", + "MYSQL_USER", + "MYSQL_IAM_USER", + "MYSQL_PASS", + "MYSQL_DB", + "POSTGRES_CONNECTION_NAME", + "POSTGRES_USER", + "POSTGRES_IAM_USER", + "POSTGRES_PASS", + "POSTGRES_DB", + "POSTGRES_CAS_CONNECTION_NAME", + "POSTGRES_CAS_PASS", + "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", + "SQLSERVER_DB", + ] args: - "-c" - | @@ -26,48 +49,52 @@ steps: nox -s system-${_VERSION} availableSecrets: secretManager: - - versionName: 'projects/$PROJECT_ID/secrets/MYSQL_CONNECTION_NAME/versions/latest' - env: 'MYSQL_CONNECTION_NAME' - - versionName: 'projects/$PROJECT_ID/secrets/MYSQL_USER/versions/latest' - env: 'MYSQL_USER' - - versionName: 'projects/$PROJECT_ID/secrets/CLOUD_BUILD_MYSQL_IAM_USER/versions/latest' - env: 'MYSQL_IAM_USER' - - versionName: 'projects/$PROJECT_ID/secrets/MYSQL_PASS/versions/latest' - env: 'MYSQL_PASS' - - versionName: 'projects/$PROJECT_ID/secrets/MYSQL_DB/versions/latest' - env: 'MYSQL_DB' - - versionName: 'projects/$PROJECT_ID/secrets/POSTGRES_CONNECTION_NAME/versions/latest' - env: 'POSTGRES_CONNECTION_NAME' - - versionName: 'projects/$PROJECT_ID/secrets/POSTGRES_USER/versions/latest' - env: 'POSTGRES_USER' - - versionName: 'projects/$PROJECT_ID/secrets/CLOUD_BUILD_POSTGRES_IAM_USER/versions/latest' - env: 'POSTGRES_IAM_USER' - - versionName: 'projects/$PROJECT_ID/secrets/POSTGRES_PASS/versions/latest' - env: 'POSTGRES_PASS' - - versionName: 'projects/$PROJECT_ID/secrets/POSTGRES_DB/versions/latest' - env: 'POSTGRES_DB' - - versionName: 'projects/$PROJECT_ID/secrets/POSTGRES_CAS_CONNECTION_NAME/versions/latest' - env: 'POSTGRES_CAS_CONNECTION_NAME' - - versionName: 'projects/$PROJECT_ID/secrets/POSTGRES_CAS_PASS/versions/latest' - env: 'POSTGRES_CAS_PASS' - - versionName: 'projects/$PROJECT_ID/secrets/POSTGRES_CUSTOMER_CAS_CONNECTION_NAME/versions/latest' - env: 'POSTGRES_CUSTOMER_CAS_CONNECTION_NAME' - - versionName: 'projects/$PROJECT_ID/secrets/POSTGRES_CUSTOMER_CAS_PASS/versions/latest' - 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/SQLSERVER_CONNECTION_NAME/versions/latest' - env: 'SQLSERVER_CONNECTION_NAME' - - versionName: 'projects/$PROJECT_ID/secrets/SQLSERVER_USER/versions/latest' - env: 'SQLSERVER_USER' - - versionName: 'projects/$PROJECT_ID/secrets/SQLSERVER_PASS/versions/latest' - env: 'SQLSERVER_PASS' - - versionName: 'projects/$PROJECT_ID/secrets/SQLSERVER_DB/versions/latest' - env: 'SQLSERVER_DB' + - versionName: "projects/$PROJECT_ID/secrets/MYSQL_CONNECTION_NAME/versions/latest" + env: "MYSQL_CONNECTION_NAME" + - versionName: "projects/$PROJECT_ID/secrets/MYSQL_USER/versions/latest" + env: "MYSQL_USER" + - versionName: "projects/$PROJECT_ID/secrets/CLOUD_BUILD_MYSQL_IAM_USER/versions/latest" + env: "MYSQL_IAM_USER" + - versionName: "projects/$PROJECT_ID/secrets/MYSQL_PASS/versions/latest" + env: "MYSQL_PASS" + - versionName: "projects/$PROJECT_ID/secrets/MYSQL_DB/versions/latest" + env: "MYSQL_DB" + - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_CONNECTION_NAME/versions/latest" + env: "POSTGRES_CONNECTION_NAME" + - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_USER/versions/latest" + env: "POSTGRES_USER" + - versionName: "projects/$PROJECT_ID/secrets/CLOUD_BUILD_POSTGRES_IAM_USER/versions/latest" + env: "POSTGRES_IAM_USER" + - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_PASS/versions/latest" + env: "POSTGRES_PASS" + - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_DB/versions/latest" + env: "POSTGRES_DB" + - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_CAS_CONNECTION_NAME/versions/latest" + env: "POSTGRES_CAS_CONNECTION_NAME" + - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_CAS_PASS/versions/latest" + env: "POSTGRES_CAS_PASS" + - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_CUSTOMER_CAS_CONNECTION_NAME/versions/latest" + env: "POSTGRES_CUSTOMER_CAS_CONNECTION_NAME" + - versionName: "projects/$PROJECT_ID/secrets/POSTGRES_CUSTOMER_CAS_PASS/versions/latest" + 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" + env: "SQLSERVER_USER" + - versionName: "projects/$PROJECT_ID/secrets/SQLSERVER_PASS/versions/latest" + env: "SQLSERVER_PASS" + - versionName: "projects/$PROJECT_ID/secrets/SQLSERVER_DB/versions/latest" + env: "SQLSERVER_DB" substitutions: _VERSION: ${_VERSION} _IP_TYPE: ${_IP_TYPE} - + options: dynamicSubstitutions: true pool: diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index b8e6eb58d..698fdd6be 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -82,6 +82,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 @@ -104,6 +106,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_pg8000_connection.py b/tests/system/test_pg8000_connection.py index 42a8d9894..a7bc49631 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.""" + 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() From 12f08b266846d674fbdb39306227ebcbc866443b Mon Sep 17 00:00:00 2001 From: jackwotherspoon Date: Fri, 4 Apr 2025 13:40:07 +0000 Subject: [PATCH 2/6] test: add IAM test --- tests/system/test_pg8000_iam_auth.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/system/test_pg8000_iam_auth.py b/tests/system/test_pg8000_iam_auth.py index 38ee76612..3b73ae6e8 100644 --- a/tests/system/test_pg8000_iam_auth.py +++ b/tests/system/test_pg8000_iam_auth.py @@ -113,3 +113,19 @@ def test_lazy_pg8000_iam_authn_connection() -> None: curr_time = time[0] assert type(curr_time) is datetime connector.close() + + +def test_MCP_pg8000_iam_authn_connection() -> None: + """Basic test to get time from database.""" + inst_conn_name = os.environ["POSTGRES_MCP_CONNECTION_NAME"] + user = os.environ["POSTGRES_IAM_USER"] + db = os.environ["POSTGRES_DB"] + ip_type = os.environ.get("IP_TYPE", "public") + + engine, connector = create_sqlalchemy_engine(inst_conn_name, user, 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() From d5fa0e8b2171711ce52b77ce6eec4e8023e6c9ce Mon Sep 17 00:00:00 2001 From: jackwotherspoon Date: Fri, 4 Apr 2025 13:47:27 +0000 Subject: [PATCH 3/6] chore: update tests --- tests/system/test_pg8000_connection.py | 2 +- tests/system/test_pg8000_iam_auth.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/system/test_pg8000_connection.py b/tests/system/test_pg8000_connection.py index a7bc49631..c7074b0cf 100644 --- a/tests/system/test_pg8000_connection.py +++ b/tests/system/test_pg8000_connection.py @@ -194,7 +194,7 @@ def test_custom_SAN_with_dns_pg8000_connection() -> None: def test_MCP_pg8000_connection() -> None: - """Basic test to get time from database.""" + """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"] diff --git a/tests/system/test_pg8000_iam_auth.py b/tests/system/test_pg8000_iam_auth.py index 3b73ae6e8..44ff42228 100644 --- a/tests/system/test_pg8000_iam_auth.py +++ b/tests/system/test_pg8000_iam_auth.py @@ -116,7 +116,7 @@ def test_lazy_pg8000_iam_authn_connection() -> None: def test_MCP_pg8000_iam_authn_connection() -> None: - """Basic test to get time from database.""" + """Basic test to get time from database using MCP enabled instance.""" inst_conn_name = os.environ["POSTGRES_MCP_CONNECTION_NAME"] user = os.environ["POSTGRES_IAM_USER"] db = os.environ["POSTGRES_DB"] From e1c94c750834b2a32436a5ef533c52c7d74f0055 Mon Sep 17 00:00:00 2001 From: jackwotherspoon Date: Fri, 4 Apr 2025 13:51:03 +0000 Subject: [PATCH 4/6] test: add asyncpg MCP test --- tests/system/test_asyncpg_connection.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/system/test_asyncpg_connection.py b/tests/system/test_asyncpg_connection.py index d7602e537..7389aff34 100644 --- a/tests/system/test_asyncpg_connection.py +++ b/tests/system/test_asyncpg_connection.py @@ -220,6 +220,25 @@ 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 + ) + + 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"] From fc6acbfbe49a287a21c7e6c9a331e07221580771 Mon Sep 17 00:00:00 2001 From: jackwotherspoon Date: Fri, 4 Apr 2025 17:43:20 +0000 Subject: [PATCH 5/6] chore: remove auto iam authn test --- tests/system/test_pg8000_iam_auth.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/tests/system/test_pg8000_iam_auth.py b/tests/system/test_pg8000_iam_auth.py index 44ff42228..38ee76612 100644 --- a/tests/system/test_pg8000_iam_auth.py +++ b/tests/system/test_pg8000_iam_auth.py @@ -113,19 +113,3 @@ def test_lazy_pg8000_iam_authn_connection() -> None: curr_time = time[0] assert type(curr_time) is datetime connector.close() - - -def test_MCP_pg8000_iam_authn_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_IAM_USER"] - db = os.environ["POSTGRES_DB"] - ip_type = os.environ.get("IP_TYPE", "public") - - engine, connector = create_sqlalchemy_engine(inst_conn_name, user, 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() From fc20a02c357b5470c2357b14afa1165cdf08ee41 Mon Sep 17 00:00:00 2001 From: jackwotherspoon Date: Fri, 4 Apr 2025 18:06:07 +0000 Subject: [PATCH 6/6] chore: update asyncpg MCP test --- tests/system/test_asyncpg_connection.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/system/test_asyncpg_connection.py b/tests/system/test_asyncpg_connection.py index 7389aff34..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"}, ) @@ -229,7 +231,12 @@ async def test_MCP_sqlalchemy_connection_with_asyncpg() -> None: ip_type = os.environ.get("IP_TYPE", "public") pool, connector = await create_sqlalchemy_engine( - inst_conn_name, user, password, db, ip_type + inst_conn_name, + user, + password, + db, + ip_type, + statement_cache_size=0, ) async with pool.connect() as conn: