test

tests/common.go

@@ -207,18 +207,21 @@ func AddPostgresPrebuiltConfig(t *testing.T, config map[string]any) map[string]a
 		PostgresListQueryStatsToolType          = "postgres-list-query-stats"
 		PostgresGetColumnCardinalityToolType    = "postgres-get-column-cardinality"
 		PostgresListTableStats                  = "postgres-list-table-stats"
+
 		PostgresListPublicationTablesToolType   = "postgres-list-publication-tables"
 		PostgresListTablespacesToolType         = "postgres-list-tablespaces"
 		PostgresListPGSettingsToolType          = "postgres-list-pg-settings"
 		PostgresListDatabaseStatsToolType       = "postgres-list-database-stats"
 		PostgresListRolesToolType               = "postgres-list-roles"
 		PostgresListStoredProcedureToolType     = "postgres-list-stored-procedure"
+
 	)
 
 	tools, ok := config["tools"].(map[string]any)
 	if !ok {
 		t.Fatalf("unable to get tools from config")
 	}
+
 	tools["list_tables"] = map[string]any{
 		"type":        PostgresListTablesToolType,
 		"source":      "my-instance",
@@ -938,6 +941,8 @@ func TestCloudSQLMySQL_IPTypeParsingFromYAML(t *testing.T) {
 
 // Finds and drops all tables in a postgres database.
 func CleanupPostgresTables(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
+
+	t.Logf("in cleanupPostgrestTables");
 	query := `
 	SELECT table_name FROM information_schema.tables
 	WHERE table_schema = 'public' AND table_type = 'BASE TABLE';`
@@ -959,14 +964,31 @@ func CleanupPostgresTables(t *testing.T, ctx context.Context, pool *pgxpool.Pool
 	}
 
 	if len(tablesToDrop) == 0 {
+		t.Logf("No tables to drop in 'public' schema")
 		return
 	}
+	t.Logf("Tables to drop in 'public' schema: %s", strings.Join(tablesToDrop, ", "))
 	
 	dropQuery := fmt.Sprintf("DROP TABLE IF EXISTS %s CASCADE;", strings.Join(tablesToDrop, ", "))
 
 	if _, err := pool.Exec(ctx, dropQuery); err != nil {
 		t.Fatalf("Failed to drop all tables in 'public' schema: %v", err)
 	}
+
+	t.Logf("Dropped tables in 'public' schema: %s", strings.Join(tablesToDrop, ", "))
+
+
+	// // 1. Drop the entire public schema (this kills tables, views, types, etc.)
+    // dropSchema := "DROP SCHEMA public CASCADE;"
+    // // 2. Recreate the empty public schema
+    // createSchema := "CREATE SCHEMA public;"
+    // // 3. Grant permissions back (Postgres default)
+    // grantPublic := "GRANT ALL ON SCHEMA public TO public;"
+
+    // _, err := pool.Exec(ctx, dropSchema + createSchema + grantPublic)
+    // if err != nil {
+    //     t.Fatalf("Failed to nuclear-wipe the public schema: %v", err)
+    // }
 }
 
 // Finds and drops all tables in a mysql database.
@@ -1030,3 +1052,36 @@ func CleanupMSSQLTables(t *testing.T, ctx context.Context, pool *sql.DB) {
 	}
 
 }
+func CleanupSingleStoreTables(t *testing.T, ctx context.Context, pool *sql.DB) {
+
+	// SingleStore cleanup logic would go here
+	fmt.Logf("in cleanupSingleStoreTables");
+	rows, err := pool.QueryContext(ctx, "SHOW TABLES;")
+	if err != nil {
+		t.Fatalf("Failed to query for all SingleStore tables: %v", err)
+	}
+	defer rows.Close()
+
+	var tablesToDrop []string
+	for rows.Next() {
+		var tableName string
+		if err := rows.Scan(&tableName); err != nil {
+			t.Errorf("Failed to scan SingleStore table name: %v", err)
+			continue
+		}
+		tablesToDrop = append(tablesToDrop, fmt.Sprintf("`%s`", tableName))
+	}
+
+	for _, table := range tablesToDrop {
+		// Using fmt.Sprintf for table names is necessary as you cannot 
+		// use placeholders (?) for identifiers like table names in SQL.
+		// We use %q to wrap the name in double quotes for safety.
+		query := fmt.Sprintf("DROP TABLE IF EXISTS %s", table)
+		if _, err := pool.ExecContext(ctx, query); err != nil {
+			t.Logf("Warning: failed to drop table %s: %v", table, err)
+			// We use Logf instead of Fatalf here so one failure 
+			// doesn't stop the entire cleanup process.
+		}
+	}
+
+}

tests/postgres/postgres_integration_test.go

@@ -81,9 +81,28 @@ func initPostgresConnectionPool(host, port, user, pass, dbname string) (*pgxpool
 	return pool, nil
 }
 
+func CreateIsolatedSchema(t *testing.T, ctx context.Context, pool *pgxpool.Pool) (string, func()) {
+	
+	schemaName := "test_schema_" + strings.ReplaceAll(uuid.New().String(), "-", "")
+	
+	_, err := pool.Exec(ctx, fmt.Sprintf("CREATE SCHEMA %q;", schemaName))
+	if err != nil {
+		t.Fatalf("failed to create sandbox schema: %v", err)
+	}
+
+	cleanup := func() {
+		// Use Background context to ensure cleanup runs even if the test timed out
+		_, err := pool.Exec(context.Background(), fmt.Sprintf("DROP SCHEMA %q CASCADE;", schemaName))
+		if err != nil {
+			t.Logf("Cleanup warning: failed to drop schema %s: %v", schemaName, err)
+		}
+	}
+
+	return schemaName, cleanup
+}
 func TestPostgres(t *testing.T) {
 	sourceConfig := getPostgresVars(t)
-	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
 	defer cancel()
 
 	var args []string
@@ -92,10 +111,21 @@ func TestPostgres(t *testing.T) {
 	if err != nil {
 		t.Fatalf("unable to create postgres connection pool: %s", err)
 	}
+	defer pool.Close()
 
+	//this was conflicting and too slow
 	// cleanup test environment
 	tests.CleanupPostgresTables(t, ctx, pool)
 
+	// schemaName, cleanupSchema := CreateIsolatedSchema(t, ctx, pool)
+	// defer cleanupSchema()
+
+	// Force the pool to ONLY use this schema for this test run
+	// _, err = pool.Exec(ctx, fmt.Sprintf("SET search_path TO %q;", schemaName))
+	// if err != nil {
+	// 	t.Fatalf("failed to set search path: %v", err)
+	// }
+
 	// create table name with UUID
 	tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
 	tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
@@ -126,6 +156,7 @@ func TestPostgres(t *testing.T) {
 	insertStmt, searchStmt := tests.GetPostgresVectorSearchStmts(vectorTableName)
 	toolsFile = tests.AddSemanticSearchConfig(t, toolsFile, PostgresToolType, insertStmt, searchStmt)
 
+	// Start toolbox command
 	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
 	if err != nil {
 		t.Fatalf("command initialization returned an error: %s", err)

tests/singlestore/singlestore_integration_test.go

@@ -186,6 +186,8 @@ func TestSingleStoreToolEndpoints(t *testing.T) {
 
 	var args []string
 
+	tests.CleanupSingleStoreTables(t, ctx, sourceConfig)
+
 	pool, err := initSingleStoreConnectionPool(SingleStoreHost, SingleStorePort, SingleStoreUser, SingleStorePass, SingleStoreDatabase)
 	if err != nil {
 		t.Fatalf("unable to create SingleStore connection pool: %s", err)
@@ -218,6 +220,7 @@ func TestSingleStoreToolEndpoints(t *testing.T) {
 	}
 	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)