chore: import upstream snapshot with attribution

tests/cloudsqlpg/cloud_sql_pg_vectorassist_test.go

@@ -0,0 +1,583 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cloudsqlpg
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"regexp"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/googleapis/mcp-toolbox/internal/testutils"
+	"github.com/googleapis/mcp-toolbox/tests"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+func createPostgresExtension(t *testing.T, ctx context.Context, pool *pgxpool.Pool, extensionName string) func() {
+	createExtensionCmd := fmt.Sprintf("CREATE EXTENSION IF NOT EXISTS %s", extensionName)
+	_, err := pool.Exec(ctx, createExtensionCmd)
+	if err != nil {
+		t.Fatalf("failed to create extension: %v", err)
+	}
+	return func() {
+		dropExtensionCmd := fmt.Sprintf("DROP EXTENSION IF EXISTS %s", extensionName)
+		_, err := pool.Exec(ctx, dropExtensionCmd)
+		if err != nil {
+			t.Fatalf("failed to drop extension: %v", err)
+		}
+	}
+}
+
+// setupVectorAssistTable prepares the database extensions and test data needed
+// to test the definespec, modifyspec, applyspec, and generatequery tools.
+func setupVectorAssistTable(t *testing.T, ctx context.Context, pool *pgxpool.Pool) (string, func(t *testing.T), func()) {
+	// Install necessary extensions for VectorAssist
+	dropExtensionFunc := createPostgresExtension(t, ctx, pool, "vector_assist")
+
+	uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
+	uniqueID = uniqueID[:10]
+	tableName := "vector_assist_test_" + uniqueID
+
+	// Create a table with vector data for defining/modifying/applying specs
+	createStmt := fmt.Sprintf(`
+		CREATE TABLE %s (
+			name TEXT,
+			category TEXT,
+			content TEXT,
+			embedding vector(3)
+		);
+	`, tableName)
+
+	_, err := pool.Exec(ctx, createStmt)
+	if err != nil {
+		t.Fatalf("failed to create vector assist test table: %v", err)
+	}
+
+	// Insert sample data to generate queries against
+	insertDataStmt := fmt.Sprintf(`
+		INSERT INTO %s (name, category, content, embedding)
+		VALUES
+		('Item 1', 'Document', 'Sample text document about AI', array_fill(0.1, ARRAY[3])::vector),
+		('Item 2', 'Document', 'Sample text document about databases', array_fill(0.2, ARRAY[3])::vector);
+	`, tableName)
+
+	_, err = pool.Exec(ctx, insertDataStmt)
+	if err != nil {
+		t.Fatalf("failed to insert data into vector assist table: %v", err)
+	}
+
+	// Return teardown function
+	teardown := func(t *testing.T) {
+		_, err := pool.Exec(context.Background(), fmt.Sprintf("DROP TABLE IF EXISTS %s;", tableName))
+		if err != nil {
+			t.Errorf("failed to drop vector assist table %s: %v", tableName, err)
+		}
+	}
+
+	return tableName, teardown, dropExtensionFunc
+}
+
+// TODO: Remove the test from this file and follow the existing test pattern
+// by calling the tests from cloudsqlpg_integration_test.go
+func TestVectorAssistIntegration(t *testing.T) {
+	// temporarily skip vector assist test since "Vector assist is temporarily
+	// disabled for all Cloud SQL for PostgreSQL instances"
+	// ref: https://docs.cloud.google.com/sql/docs/postgres/vector-assist-overview
+	t.Skip("Temporarily skipping vector assist test. Google had temporarily disabled vector assist for all CloudSQL for PostgreSQL instances.")
+
+	sourceConfig := getCloudSQLPgVars(t)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
+	args := []string{"--enable-api"}
+
+	pool, err := initCloudSQLPgConnectionPool(CloudSQLPostgresProject, CloudSQLPostgresRegion, CloudSQLPostgresInstance, "public", CloudSQLPostgresUser, CloudSQLPostgresPass, CloudSQLPostgresDatabase)
+	if err != nil {
+		t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
+	}
+
+	// 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
+
+	// set up data for param tool
+	createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetPostgresSQLParamToolInfo(tableNameParam)
+	teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
+	defer teardownTable1(t)
+
+	// set up data for auth tool
+	createAuthTableStmt, insertAuthTableStmt, authToolStmt, authTestParams := tests.GetPostgresSQLAuthToolInfo(tableNameAuth)
+	teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
+	defer teardownTable2(t)
+
+	// // Set up data for vector assist tools
+	vectorAssistTableName, teardownVectorAssistTable, dropExtension := setupVectorAssistTable(t, ctx, pool)
+	defer teardownVectorAssistTable(t)
+	defer dropExtension()
+
+	// Write config into a file and pass it to command
+	toolsFile := tests.GetToolsConfig(sourceConfig, CloudSQLPostgresToolType, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
+	toolsFile = tests.AddExecuteSqlConfig(t, toolsFile, "postgres-execute-sql")
+	tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
+	toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CloudSQLPostgresToolType, tmplSelectCombined, tmplSelectFilterCombined, "")
+
+	// Add vector assist tools to the configuration
+	toolsFile = AddVectorAssistConfig(t, toolsFile, "my-instance")
+
+	toolsFile = tests.AddPostgresPrebuiltConfig(t, toolsFile)
+	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
+	if err != nil {
+		t.Fatalf("command initialization returned an error: %s", err)
+	}
+	defer cleanup()
+
+	waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
+	if err != nil {
+		t.Logf("toolbox command logs: \n%s", out)
+		t.Fatalf("toolbox didn't start successfully: %s", err)
+	}
+
+	// Run vectorassist tool tests
+	specID := "va_spec_001"
+	RunVectorAssistDefineSpecToolInvokeTest(t, ctx, pool, vectorAssistTableName, specID)
+	RunVectorAssistModifySpecToolInvokeTest(t, ctx, pool, specID)
+	RunVectorAssistApplySpecToolInvokeTest(t, ctx, pool, specID)
+	RunVectorAssistGenerateQueryToolInvokeTest(t, ctx, pool, specID)
+	RunVectorAssistImproveQueryRecallToolInvokeTest(t, ctx, pool, vectorAssistTableName)
+	RunVectorAssistListSpecsToolInvokeTest(t, ctx, pool, vectorAssistTableName)
+	RunVectorAssistGetSpecToolInvokeTest(t, ctx, pool, specID)
+	RunVectorAssistDeleteSpecToolInvokeTest(t, ctx, pool, specID)
+}
+
+// AddVectorAssistConfig appends the vector assist tool configurations to the given tools file.
+func AddVectorAssistConfig(t *testing.T, config map[string]any, sourceName string) map[string]any {
+	tools, ok := config["tools"].(map[string]any)
+	if !ok {
+		t.Fatalf("unable to get tools from config")
+	}
+
+	tools["define_spec"] = map[string]any{
+		"type":   "vector-assist-define-spec",
+		"source": sourceName,
+	}
+	tools["modify_spec"] = map[string]any{
+		"type":   "vector-assist-modify-spec",
+		"source": sourceName,
+	}
+	tools["apply_spec"] = map[string]any{
+		"type":   "vector-assist-apply-spec",
+		"source": sourceName,
+	}
+	tools["generate_query"] = map[string]any{
+		"type":   "vector-assist-generate-query",
+		"source": sourceName,
+	}
+	tools["improve_query_recall"] = map[string]any{
+		"type":   "vector-assist-improve-query-recall",
+		"source": sourceName,
+	}
+	tools["list_specs"] = map[string]any{
+		"type":   "vector-assist-list-specs",
+		"source": sourceName,
+	}
+	tools["get_spec"] = map[string]any{
+		"type":   "vector-assist-get-spec",
+		"source": sourceName,
+	}
+	tools["delete_spec"] = map[string]any{
+		"type":   "vector-assist-delete-spec",
+		"source": sourceName,
+	}
+	config["tools"] = tools
+	return config
+}
+
+func RunVectorAssistImproveQueryRecallToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tableName string) {
+	validPayload := fmt.Sprintf(`{
+    "schema_name": "public",
+		"table_name": "%s",
+    "vector_column_name": "embedding",
+    "index_name":  "%s_embedding_idx",
+    "top_k": 10,
+    "target_recall": 0.95,
+    "distance_func": "cosine"
+  }`,
+		tableName, tableName)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "improve query recall valid parameters",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/improve_query_recall/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`ef_search`,
+			},
+		},
+		{
+			name:           "improve query recall missing required table_name",
+			requestBody:    bytes.NewBuffer([]byte(`{"schema_name": "public"}`)),
+			api:            "http://127.0.0.1:5000/api/tool/improve_query_recall/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistListSpecsToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tableName string) {
+	validPayload := fmt.Sprintf(`{
+    "table_name": "%s",
+    "column_name": "content"
+  }`,
+		tableName)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "list specs valid parameters",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/list_specs/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`va_spec_001`,
+			},
+		},
+		{
+			name:           "list specs missing required table_name",
+			requestBody:    bytes.NewBuffer([]byte(`{"column_name": "content"}`)),
+			api:            "http://127.0.0.1:5000/api/tool/list_specs/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistGetSpecToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, specID string) {
+	validPayload := fmt.Sprintf(`{
+    "spec_id": "%s"
+  }`,
+		specID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "get spec valid parameters",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/get_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				specID,
+			},
+		},
+		{
+			name:           "get spec missing spec_id",
+			requestBody:    bytes.NewBuffer([]byte(`{}`)),
+			api:            "http://127.0.0.1:5000/api/tool/get_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistDeleteSpecToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, specID string) {
+	validPayload := fmt.Sprintf(`{
+    "spec_id": "%s"
+  }`,
+		specID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "delete spec valid parameters",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/delete_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`delete_spec`,
+			},
+		},
+		{
+			name:           "delete spec missing spec_id",
+			requestBody:    bytes.NewBuffer([]byte(`{}`)),
+			api:            "http://127.0.0.1:5000/api/tool/delete_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistDefineSpecToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tableName string, specID string) {
+	validPayload := fmt.Sprintf(`{
+		"table_name": "%s",
+		"schema_name": "public",
+		"spec_id": "%s",
+		"vector_column_name": "embedding",
+		"text_column_name": "content",
+		"vector_index_type": "hnsw",
+		"embeddings_available": true,
+		"num_vectors": 2,
+		"dimensionality": 3,
+		"embedding_model": "textembedding-gecko",
+		"prefilter_column_names": ["category"],
+		"distance_func": "cosine",
+		"quantization": "halfvec",
+		"memory_budget_kb": 1024,
+		"target_recall": 0.95,
+		"target_top_k": 10,
+		"tune_vector_index": true
+	}`, tableName, specID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "invoke define_spec with all valid parameters",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/define_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			// Check for key identifiers instead of the entire JSON string
+			wantContains: []string{
+				`"vector_spec_id":"va_spec_001"`,
+				fmt.Sprintf(`"table_name":"%s"`, tableName),
+				`"recommendation_id"`, // Ensure a recommendation was generated
+			},
+		},
+		{
+			name:           "invoke define_spec with missing required table_name",
+			requestBody:    bytes.NewBuffer([]byte(`{"schema_name": "public", "spec_id": "va_spec_002"}`)),
+			api:            "http://127.0.0.1:5000/api/tool/define_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistModifySpecToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, specID string) {
+	validPayload := fmt.Sprintf(`{
+        "spec_id": "%s",
+        "memory_budget_kb": 2048,
+        "target_recall": 0.99
+    }`, specID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "modify existing spec with new constraints",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/modify_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"recommendation_id"`,
+			},
+		},
+		{
+			name:           "modify existing spec without required spec id",
+			requestBody:    bytes.NewBuffer([]byte(`{"target_recall": 0.99}`)),
+			api:            "http://127.0.0.1:5000/api/tool/modify_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistApplySpecToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, recommendationID string) {
+	validPayload := fmt.Sprintf(`{
+        "spec_id": "%s"
+    }`, recommendationID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "apply recommendation to database",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/apply_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`{"apply_spec":true}`,
+			},
+		},
+		{
+			name:           "apply recommendation to database without spec id",
+			requestBody:    bytes.NewBuffer([]byte(`{"schema_name": "public"}`)),
+			api:            "http://127.0.0.1:5000/api/tool/apply_spec/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"error"`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func RunVectorAssistGenerateQueryToolInvokeTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool, specID string) {
+	validPayload := fmt.Sprintf(`{
+        "spec_id": "%s",
+        "search_text": "What is the capital of France?",
+        "top_k": 5
+    }`, specID)
+
+	tcs := []struct {
+		name           string
+		requestBody    io.Reader
+		api            string
+		wantStatusCode int
+		wantContains   []string
+	}{
+		{
+			name:           "generate SQL for vector search",
+			requestBody:    bytes.NewBuffer([]byte(validPayload)),
+			api:            "http://127.0.0.1:5000/api/tool/generate_query/invoke",
+			wantStatusCode: http.StatusOK,
+			wantContains: []string{
+				`"generate_query"`,
+				`LIMIT 5`,
+			},
+		},
+	}
+
+	for _, tc := range tcs {
+		runVectorAssistToolInvokeTest(t, tc)
+	}
+}
+
+func runVectorAssistToolInvokeTest(t *testing.T, tc struct {
+	name           string
+	requestBody    io.Reader
+	api            string
+	wantStatusCode int
+	wantContains   []string
+}) {
+	t.Run(tc.name, func(t *testing.T) {
+		resp, body := tests.RunRequest(t, http.MethodPost, tc.api, tc.requestBody, nil)
+
+		if resp.StatusCode != tc.wantStatusCode {
+			t.Fatalf("tool %s: wrong status code: got %d, want %d, body: %s", tc.api, resp.StatusCode, tc.wantStatusCode, string(body))
+		}
+
+		if tc.wantStatusCode != http.StatusOK {
+			return
+		}
+
+		// Unmarshal the standard response wrapper
+		var bodyWrapper struct {
+			Result json.RawMessage `json:"result"`
+		}
+		if err := json.Unmarshal(body, &bodyWrapper); err != nil {
+			t.Fatalf("error decoding response wrapper: %v", err)
+		}
+
+		// Handle the double-unmarshal logic for stringified results
+		var resultString string
+		if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
+			resultString = string(bodyWrapper.Result)
+		}
+
+		// Verification loop
+		for _, expectedSubstr := range tc.wantContains {
+			if !strings.Contains(resultString, expectedSubstr) {
+				t.Errorf("Expected result to contain %q, but it did not.\nFull result: %s", expectedSubstr, resultString)
+			}
+		}
+	})
+}