From b58883a0059ec99006f083609c159ab112e81e1b Mon Sep 17 00:00:00 2001 From: Vincent Chan Date: Tue, 12 Aug 2025 12:44:42 -0700 Subject: [PATCH 1/2] Introduce GCP Postgers integration tests --- .circleci/continue_config.yml | 2 ++ .circleci/manage-test-db.sh | 26 +++++++++++++++++++ Makefile | 3 +++ sqlmesh/core/config/connection.py | 6 ----- .../engine_adapter/integration/__init__.py | 1 + .../engine_adapter/integration/config.yaml | 11 ++++++++ .../integration/test_integration.py | 6 ++++- 7 files changed, 48 insertions(+), 7 deletions(-) diff --git a/.circleci/continue_config.yml b/.circleci/continue_config.yml index b93caf482e..04135574a9 100644 --- a/.circleci/continue_config.yml +++ b/.circleci/continue_config.yml @@ -246,6 +246,7 @@ jobs: echo "export SNOWFLAKE_DATABASE='$TEST_DB_NAME'" >> "$BASH_ENV" echo "export DATABRICKS_CATALOG='$TEST_DB_NAME'" >> "$BASH_ENV" echo "export REDSHIFT_DATABASE='$TEST_DB_NAME'" >> "$BASH_ENV" + echo "export GCP_POSTGRES_DATABASE='$TEST_DB_NAME'" >> "$BASH_ENV" - run: name: Create test database command: ./.circleci/manage-test-db.sh << parameters.engine >> "$TEST_DB_NAME" up @@ -303,6 +304,7 @@ workflows: - bigquery - clickhouse-cloud - athena + - gcp-postgres filters: branches: only: diff --git a/.circleci/manage-test-db.sh b/.circleci/manage-test-db.sh index de8922c988..a9dd8a6d5a 100755 --- a/.circleci/manage-test-db.sh +++ b/.circleci/manage-test-db.sh @@ -109,6 +109,32 @@ clickhouse-cloud_init() { echo "Clickhouse Cloud instance $CLICKHOUSE_CLOUD_HOST is up and running" } +# GCP Postgres +gcp-postgres_init() { + # Download and start Cloud SQL Proxy + curl -fsSL -o cloud-sql-proxy https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v2.18.0/cloud-sql-proxy.linux.amd64 + chmod +x cloud-sql-proxy + echo $GCP_POSTGRES_KEYFILE_JSON > /tmp/keyfile.json + ./cloud-sql-proxy --credentials-file /tmp/keyfile.json $GCP_POSTGRES_INSTANCE_CONNECTION_STRING & + + # Wait for proxy to start + sleep 5 +} + +gcp-postgres_exec() { + PGPASSWORD=$GCP_POSTGRES_PASSWORD psql -h 127.0.0.1 -U $GCP_POSTGRES_USER -c "$1" postgres +} + +gcp-postgres_up() { + gcp-postgres_exec "create database $1" +} + +gcp-postgres_down() { + gcp-postgres_exec "drop database $1" +} + + + INIT_FUNC="${ENGINE}_init" UP_FUNC="${ENGINE}_up" DOWN_FUNC="${ENGINE}_down" diff --git a/Makefile b/Makefile index 0a89bba437..855d866c84 100644 --- a/Makefile +++ b/Makefile @@ -173,6 +173,9 @@ clickhouse-cloud-test: guard-CLICKHOUSE_CLOUD_HOST guard-CLICKHOUSE_CLOUD_USERNA athena-test: guard-AWS_ACCESS_KEY_ID guard-AWS_SECRET_ACCESS_KEY guard-ATHENA_S3_WAREHOUSE_LOCATION engine-athena-install pytest -n auto -m "athena" --retries 3 --junitxml=test-results/junit-athena.xml +gcp-postgres-test: guard-GCP_POSTGRES_INSTANCE_CONNECTION_STRING guard-GCP_POSTGRES_USER guard-GCP_POSTGRES_PASSWORD guard-GCP_POSTGRES_KEYFILE_JSON engine-gcppostgres-install + pytest -n auto -m "gcp_postgres" --retries 3 --junitxml=test-results/junit-gcp-postgres.xml + vscode_settings: mkdir -p .vscode cp -r ./tooling/vscode/*.json .vscode/ diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 415e916365..4f289262ea 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -1210,12 +1210,6 @@ def _validate_auth_method(cls, data: t.Any) -> t.Any: password = data.get("password") enable_iam_auth = data.get("enable_iam_auth") - if password and enable_iam_auth: - raise ConfigError( - "Invalid GCP Postgres connection configuration - both password and" - " enable_iam_auth set. Use password when connecting to a postgres" - " user and enable_iam_auth 'True' when connecting to an IAM user." - ) if not password and not enable_iam_auth: raise ConfigError( "GCP Postgres connection configuration requires either password set" diff --git a/tests/core/engine_adapter/integration/__init__.py b/tests/core/engine_adapter/integration/__init__.py index 7e35b832be..15339eeaa6 100644 --- a/tests/core/engine_adapter/integration/__init__.py +++ b/tests/core/engine_adapter/integration/__init__.py @@ -82,6 +82,7 @@ def pytest_marks(self) -> t.List[MarkDecorator]: IntegrationTestEngine("bigquery", native_dataframe_type="bigframe", cloud=True), IntegrationTestEngine("databricks", native_dataframe_type="pyspark", cloud=True), IntegrationTestEngine("snowflake", native_dataframe_type="snowpark", cloud=True), + IntegrationTestEngine("gcp_postgres", cloud=True), ] ENGINES_BY_NAME = {e.engine: e for e in ENGINES} diff --git a/tests/core/engine_adapter/integration/config.yaml b/tests/core/engine_adapter/integration/config.yaml index d18ea5366f..4aee4640a3 100644 --- a/tests/core/engine_adapter/integration/config.yaml +++ b/tests/core/engine_adapter/integration/config.yaml @@ -186,5 +186,16 @@ gateways: state_connection: type: duckdb + inttest_gcp_postgres: + connection: + type: gcp_postgres + instance_connection_string: {{ env_var("GCP_POSTGRES_INSTANCE_CONNECTION_STRING") }} + user: {{ env_var("GCP_POSTGRES_USER") }} + password: {{ env_var("GCP_POSTGRES_PASSWORD") }} + keyfile_json: {{ env_var("GCP_POSTGRES_KEYFILE_JSON", "") }} + db: {{ env_var("GCP_POSTGRES_DATABASE") }} + enable_iam_auth: true + check_import: false + model_defaults: dialect: duckdb diff --git a/tests/core/engine_adapter/integration/test_integration.py b/tests/core/engine_adapter/integration/test_integration.py index eb593c273b..0739256d6e 100644 --- a/tests/core/engine_adapter/integration/test_integration.py +++ b/tests/core/engine_adapter/integration/test_integration.py @@ -1587,7 +1587,11 @@ def _normalize_snowflake(name: str, prefix_regex: str = "(sqlmesh__)(.*)"): k: [_normalize_snowflake(name) for name in v] for k, v in object_names.items() } - init_example_project(tmp_path, ctx.mark.split("_")[0], schema_name=schema_name) + if ctx.mark.startswith("gcp_postgres"): + engine_type = "gcp_postgres" + else: + engine_type = ctx.mark.split("_")[0] + init_example_project(tmp_path, engine_type, schema_name=schema_name) config = load_config_from_paths( Config, From e0b6d552bf97619ea86f66bf866146b9a050df95 Mon Sep 17 00:00:00 2001 From: Vincent Chan Date: Tue, 12 Aug 2025 16:53:09 -0700 Subject: [PATCH 2/2] PR feedback --- .circleci/manage-test-db.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/manage-test-db.sh b/.circleci/manage-test-db.sh index a9dd8a6d5a..80ca075912 100755 --- a/.circleci/manage-test-db.sh +++ b/.circleci/manage-test-db.sh @@ -114,7 +114,7 @@ gcp-postgres_init() { # Download and start Cloud SQL Proxy curl -fsSL -o cloud-sql-proxy https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v2.18.0/cloud-sql-proxy.linux.amd64 chmod +x cloud-sql-proxy - echo $GCP_POSTGRES_KEYFILE_JSON > /tmp/keyfile.json + echo "$GCP_POSTGRES_KEYFILE_JSON" > /tmp/keyfile.json ./cloud-sql-proxy --credentials-file /tmp/keyfile.json $GCP_POSTGRES_INSTANCE_CONNECTION_STRING & # Wait for proxy to start