From 8ce5c65170b3370ed8eded9ba5c5d4ee2ab492e2 Mon Sep 17 00:00:00 2001 From: duwenxin Date: Sun, 8 Mar 2026 23:19:32 -0400 Subject: [PATCH 1/6] ci: Change CockroachDB test source --- go.mod | 1 + go.sum | 2 + .../cockroachdb_integration_test.go | 68 +++++++++---------- 3 files changed, 37 insertions(+), 34 deletions(-) diff --git a/go.mod b/go.mod index f9e9ad96fbf2..669fd24963ee 100644 --- a/go.mod +++ b/go.mod @@ -49,6 +49,7 @@ require ( github.com/snowflakedb/gosnowflake v1.18.1 github.com/spf13/cobra v1.10.1 github.com/testcontainers/testcontainers-go v0.40.0 + github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.40.0 github.com/thlib/go-timezone-local v0.0.7 github.com/trinodb/trino-go-client v0.330.0 github.com/valkey-io/valkey-go v1.0.68 diff --git a/go.sum b/go.sum index dc0c0ebb2f70..641157ecdb38 100644 --- a/go.sum +++ b/go.sum @@ -1357,6 +1357,8 @@ github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= github.com/testcontainers/testcontainers-go v0.40.0 h1:pSdJYLOVgLE8YdUY2FHQ1Fxu+aMnb6JfVz1mxk7OeMU= github.com/testcontainers/testcontainers-go v0.40.0/go.mod h1:FSXV5KQtX2HAMlm7U3APNyLkkap35zNLxukw9oBi/MY= +github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.40.0 h1:UNYfrnFV9mkO93Sw6hqRA5KbE9DsAvDeYKD4GDConiE= +github.com/testcontainers/testcontainers-go/modules/cockroachdb v0.40.0/go.mod h1:O8By1J/1y726YYk7obTIXxfv2OzonVe+ORq9Z+K+fDg= github.com/thlib/go-timezone-local v0.0.7 h1:fX8zd3aJydqLlTs/TrROrIIdztzsdFV23OzOQx31jII= github.com/thlib/go-timezone-local v0.0.7/go.mod h1:/Tnicc6m/lsJE0irFMA0LfIwTBo4QP7A8IfyIv4zZKI= github.com/tidwall/pretty v1.0.0/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk= diff --git a/tests/cockroachdb/cockroachdb_integration_test.go b/tests/cockroachdb/cockroachdb_integration_test.go index 464b82238e41..71c83392f1ba 100644 --- a/tests/cockroachdb/cockroachdb_integration_test.go +++ b/tests/cockroachdb/cockroachdb_integration_test.go @@ -17,8 +17,6 @@ package cockroachdb import ( "context" "fmt" - "net/url" - "os" "regexp" "strings" "testing" @@ -28,37 +26,23 @@ import ( "github.com/googleapis/genai-toolbox/internal/testutils" "github.com/googleapis/genai-toolbox/tests" "github.com/jackc/pgx/v5/pgxpool" + tccockroachdb "github.com/testcontainers/testcontainers-go/modules/cockroachdb" ) var ( CockroachDBSourceType = "cockroachdb" CockroachDBToolType = "cockroachdb-sql" - CockroachDBDatabase = getEnvOrDefault("COCKROACHDB_DATABASE", "defaultdb") - CockroachDBHost = getEnvOrDefault("COCKROACHDB_HOST", "localhost") - CockroachDBPort = getEnvOrDefault("COCKROACHDB_PORT", "26257") - CockroachDBUser = getEnvOrDefault("COCKROACHDB_USER", "root") - CockroachDBPass = getEnvOrDefault("COCKROACHDB_PASS", "") + CockroachDBDatabase = "defaultdb" ) -func getEnvOrDefault(key, defaultValue string) string { - if value := os.Getenv(key); value != "" { - return value - } - return defaultValue -} - -func getCockroachDBVars(t *testing.T) map[string]any { - if CockroachDBHost == "" { - t.Skip("COCKROACHDB_HOST not set, skipping CockroachDB integration test") - } - +func getCockroachDBVars(host, port string) map[string]any { return map[string]any{ "type": CockroachDBSourceType, - "host": CockroachDBHost, - "port": CockroachDBPort, + "host": host, + "port": port, "database": CockroachDBDatabase, - "user": CockroachDBUser, - "password": CockroachDBPass, + "user": "root", + "password": "", "maxRetries": 5, "retryBaseDelay": "500ms", "queryParams": map[string]string{ @@ -67,15 +51,8 @@ func getCockroachDBVars(t *testing.T) map[string]any { } } -func initCockroachDBConnectionPool(host, port, user, pass, dbname string) (*pgxpool.Pool, error) { - connURL := &url.URL{ - Scheme: "postgres", - User: url.UserPassword(user, pass), - Host: fmt.Sprintf("%s:%s", host, port), - Path: dbname, - RawQuery: "sslmode=disable&application_name=cockroachdb-integration-test", - } - pool, err := pgxpool.New(context.Background(), connURL.String()) +func initCockroachDBConnectionPool(connString string) (*pgxpool.Pool, error) { + pool, err := pgxpool.New(context.Background(), connString) if err != nil { return nil, fmt.Errorf("unable to create connection pool: %w", err) } @@ -84,13 +61,36 @@ func initCockroachDBConnectionPool(host, port, user, pass, dbname string) (*pgxp } func TestCockroachDB(t *testing.T) { - sourceConfig := getCockroachDBVars(t) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) defer cancel() + tccockroachdbContainer, err := tccockroachdb.Run(ctx, "cockroachdb/cockroach:latest-v23.1") + if err != nil { + t.Fatalf("failed to start container: %s", err) + } + defer tccockroachdbContainer.Terminate(ctx) + + host, err := tccockroachdbContainer.Host(ctx) + if err != nil { + t.Fatalf("failed to get host: %s", err) + } + + port, err := tccockroachdbContainer.MappedPort(ctx, "26257/tcp") + if err != nil { + t.Fatalf("failed to get port: %s", err) + } + + connString, err := tccockroachdbContainer.ConnectionString(ctx) + if err != nil { + t.Fatalf("failed to get connection string: %s", err) + } + + connString += "?sslmode=disable" + + sourceConfig := getCockroachDBVars(host, port.Port()) var args []string - pool, err := initCockroachDBConnectionPool(CockroachDBHost, CockroachDBPort, CockroachDBUser, CockroachDBPass, CockroachDBDatabase) + pool, err := initCockroachDBConnectionPool(connString) if err != nil { t.Fatalf("unable to create cockroachdb connection pool: %s", err) } From 443cdb5b6ba608e3e31844b76a9cdf19eee06879 Mon Sep 17 00:00:00 2001 From: duwenxin Date: Sun, 8 Mar 2026 23:51:18 -0400 Subject: [PATCH 2/6] fix lint --- .ci/integration.cloudbuild.yaml | 9 +-------- tests/cockroachdb/cockroachdb_integration_test.go | 8 ++++++-- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/.ci/integration.cloudbuild.yaml b/.ci/integration.cloudbuild.yaml index 519c218301be..be34659b32bb 100644 --- a/.ci/integration.cloudbuild.yaml +++ b/.ci/integration.cloudbuild.yaml @@ -359,10 +359,8 @@ steps: entrypoint: /bin/bash env: - "GOPATH=/gopath" - - "COCKROACHDB_DATABASE=$_DATABASE_NAME" - - "COCKROACHDB_PORT=$_COCKROACHDB_PORT" - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL" - secretEnv: ["COCKROACHDB_USER", "COCKROACHDB_HOST","CLIENT_ID"] + secretEnv: ["CLIENT_ID"] volumes: - name: "go" path: "/gopath" @@ -1173,11 +1171,6 @@ availableSecrets: env: MARIADB_HOST - versionName: projects/$PROJECT_ID/secrets/mongodb_uri/versions/latest env: MONGODB_URI - - versionName: projects/$PROJECT_ID/secrets/cockroachdb_user/versions/latest - env: COCKROACHDB_USER - - versionName: projects/$PROJECT_ID/secrets/cockroachdb_host/versions/latest - env: COCKROACHDB_HOST - options: logging: CLOUD_LOGGING_ONLY diff --git a/tests/cockroachdb/cockroachdb_integration_test.go b/tests/cockroachdb/cockroachdb_integration_test.go index 71c83392f1ba..c7b5d281e831 100644 --- a/tests/cockroachdb/cockroachdb_integration_test.go +++ b/tests/cockroachdb/cockroachdb_integration_test.go @@ -68,7 +68,11 @@ func TestCockroachDB(t *testing.T) { if err != nil { t.Fatalf("failed to start container: %s", err) } - defer tccockroachdbContainer.Terminate(ctx) + defer func() { + if err := tccockroachdbContainer.Terminate(ctx); err != nil { + t.Logf("failed to terminate testcontainer: %s", err) + } + }() host, err := tccockroachdbContainer.Host(ctx) if err != nil { @@ -84,7 +88,7 @@ func TestCockroachDB(t *testing.T) { if err != nil { t.Fatalf("failed to get connection string: %s", err) } - + connString += "?sslmode=disable" sourceConfig := getCockroachDBVars(host, port.Port()) From 86ead517cba092b68847603d59d305efea344451 Mon Sep 17 00:00:00 2001 From: duwenxin Date: Mon, 9 Mar 2026 00:18:39 -0400 Subject: [PATCH 3/6] fix conn string --- tests/cockroachdb/cockroachdb_integration_test.go | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/cockroachdb/cockroachdb_integration_test.go b/tests/cockroachdb/cockroachdb_integration_test.go index c7b5d281e831..07e835792cdd 100644 --- a/tests/cockroachdb/cockroachdb_integration_test.go +++ b/tests/cockroachdb/cockroachdb_integration_test.go @@ -89,7 +89,13 @@ func TestCockroachDB(t *testing.T) { t.Fatalf("failed to get connection string: %s", err) } - connString += "?sslmode=disable" + if !strings.Contains(connString, "sslmode=") { + if strings.Contains(connString, "?") { + connString += "&sslmode=disable" + } else { + connString += "?sslmode=disable" + } + } sourceConfig := getCockroachDBVars(host, port.Port()) var args []string From 5ad681041c9f89d74ad608c17c11100718ad150b Mon Sep 17 00:00:00 2001 From: duwenxin Date: Mon, 9 Mar 2026 00:45:43 -0400 Subject: [PATCH 4/6] fix conn str --- tests/cockroachdb/cockroachdb_integration_test.go | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/tests/cockroachdb/cockroachdb_integration_test.go b/tests/cockroachdb/cockroachdb_integration_test.go index 07e835792cdd..61f459d2a1ad 100644 --- a/tests/cockroachdb/cockroachdb_integration_test.go +++ b/tests/cockroachdb/cockroachdb_integration_test.go @@ -84,18 +84,7 @@ func TestCockroachDB(t *testing.T) { t.Fatalf("failed to get port: %s", err) } - connString, err := tccockroachdbContainer.ConnectionString(ctx) - if err != nil { - t.Fatalf("failed to get connection string: %s", err) - } - - if !strings.Contains(connString, "sslmode=") { - if strings.Contains(connString, "?") { - connString += "&sslmode=disable" - } else { - connString += "?sslmode=disable" - } - } + connString := fmt.Sprintf("postgres://root@%s:%s/defaultdb?sslmode=disable", host, port.Port()) sourceConfig := getCockroachDBVars(host, port.Port()) var args []string From 0549d52389f24385bbe70b56aa5a6b342fbc423d Mon Sep 17 00:00:00 2001 From: duwenxin Date: Mon, 9 Mar 2026 00:52:27 -0400 Subject: [PATCH 5/6] add insecure flag --- tests/cockroachdb/cockroachdb_integration_test.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/cockroachdb/cockroachdb_integration_test.go b/tests/cockroachdb/cockroachdb_integration_test.go index 61f459d2a1ad..c691a3b53f39 100644 --- a/tests/cockroachdb/cockroachdb_integration_test.go +++ b/tests/cockroachdb/cockroachdb_integration_test.go @@ -26,6 +26,7 @@ import ( "github.com/googleapis/genai-toolbox/internal/testutils" "github.com/googleapis/genai-toolbox/tests" "github.com/jackc/pgx/v5/pgxpool" + "github.com/testcontainers/testcontainers-go" tccockroachdb "github.com/testcontainers/testcontainers-go/modules/cockroachdb" ) @@ -64,7 +65,9 @@ func TestCockroachDB(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) defer cancel() - tccockroachdbContainer, err := tccockroachdb.Run(ctx, "cockroachdb/cockroach:latest-v23.1") + tccockroachdbContainer, err := tccockroachdb.Run(ctx, "cockroachdb/cockroach:latest-v23.1", + testcontainers.WithCmd("start-single-node", "--insecure"), + ) if err != nil { t.Fatalf("failed to start container: %s", err) } From 3189b35c0c53443a6abbccd1d9eb581ca6b3cf0f Mon Sep 17 00:00:00 2001 From: duwenxin Date: Mon, 9 Mar 2026 01:10:28 -0400 Subject: [PATCH 6/6] fix cleanup issue --- tests/cockroachdb/cockroachdb_integration_test.go | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/tests/cockroachdb/cockroachdb_integration_test.go b/tests/cockroachdb/cockroachdb_integration_test.go index c691a3b53f39..a83bc3aded1f 100644 --- a/tests/cockroachdb/cockroachdb_integration_test.go +++ b/tests/cockroachdb/cockroachdb_integration_test.go @@ -71,11 +71,11 @@ func TestCockroachDB(t *testing.T) { if err != nil { t.Fatalf("failed to start container: %s", err) } - defer func() { - if err := tccockroachdbContainer.Terminate(ctx); err != nil { + t.Cleanup(func() { + if err := tccockroachdbContainer.Terminate(context.Background()); err != nil { t.Logf("failed to terminate testcontainer: %s", err) } - }() + }) host, err := tccockroachdbContainer.Host(ctx) if err != nil { @@ -114,11 +114,6 @@ func TestCockroachDB(t *testing.T) { // Generate a unique ID uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "") - // This will execute after all tool tests complete (success, fail, or t.Fatal) - t.Cleanup(func() { - tests.CleanupPostgresTables(t, context.Background(), pool, uniqueID) - }) - //Create table names using the UUID tableNameParam := "param_table_" + uniqueID tableNameAuth := "auth_table_" + uniqueID