refactor(sources): deduplicate identical RunSQL patterns across pgx abstraction drivers

Author: tanujbolisetty

internal/sources/alloydbpg/alloydb_pg.go

@@ -23,9 +23,7 @@ import (
 	"cloud.google.com/go/alloydbconn"
 	"github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/internal/sources"
-	"github.com/googleapis/genai-toolbox/internal/sqlcommenter"
 	"github.com/googleapis/genai-toolbox/internal/util"
-	"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"go.opentelemetry.io/otel/trace"
 )
@@ -104,35 +102,7 @@ func (s *Source) PostgresPool() *pgxpool.Pool {
 }
 
 func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
-	// Inject the database driver into the context for SQLCommenter
-	ctx = sqlcommenter.WithDBDriver(ctx, "pgx")
-	// Decorate the statement with SQLCommenter metadata from the context
-	statement = sqlcommenter.AppendComment(ctx, statement)
-
-	results, err := s.Pool.Query(ctx, statement, params...)
-	if err != nil {
-		return nil, fmt.Errorf("unable to execute query: %w", err)
-	}
-	defer results.Close()
-
-	fields := results.FieldDescriptions()
-	var out []any
-	for results.Next() {
-		v, err := results.Values()
-		if err != nil {
-			return nil, fmt.Errorf("unable to parse row: %w", err)
-		}
-		row := orderedmap.Row{}
-		for i, f := range fields {
-			row.Add(f.Name, v[i])
-		}
-		out = append(out, row)
-	}
-	// this will catch actual query execution errors
-	if err := results.Err(); err != nil {
-		return nil, fmt.Errorf("unable to execute query: %w", err)
-	}
-	return out, nil
+	return sources.RunSQLWithPgxQueryer(ctx, s.PostgresPool(), statement, params, "pgx")
 }
 
 func getOpts(ipType, userAgent string, useIAM bool) ([]alloydbconn.Option, error) {

internal/sources/cloudsqlpg/cloud_sql_pg.go

@@ -22,9 +22,7 @@ import (
 	"cloud.google.com/go/cloudsqlconn"
 	"github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/internal/sources"
-	"github.com/googleapis/genai-toolbox/internal/sqlcommenter"
 	"github.com/googleapis/genai-toolbox/internal/util"
-	"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"go.opentelemetry.io/otel/trace"
 )
@@ -102,35 +100,7 @@ func (s *Source) PostgresPool() *pgxpool.Pool {
 }
 
 func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
-	// Inject the database driver into the context for SQLCommenter
-	ctx = sqlcommenter.WithDBDriver(ctx, "pgx")
-	// Decorate the statement with SQLCommenter metadata from the context
-	statement = sqlcommenter.AppendComment(ctx, statement)
-
-	results, err := s.PostgresPool().Query(ctx, statement, params...)
-	if err != nil {
-		return nil, fmt.Errorf("unable to execute query: %w", err)
-	}
-	defer results.Close()
-
-	fields := results.FieldDescriptions()
-	var out []any
-	for results.Next() {
-		values, err := results.Values()
-		if err != nil {
-			return nil, fmt.Errorf("unable to parse row: %w", err)
-		}
-		row := orderedmap.Row{}
-		for i, f := range fields {
-			row.Add(f.Name, values[i])
-		}
-		out = append(out, row)
-	}
-	// this will catch actual query execution errors
-	if err := results.Err(); err != nil {
-		return nil, fmt.Errorf("unable to execute query: %w", err)
-	}
-	return out, nil
+	return sources.RunSQLWithPgxQueryer(ctx, s.PostgresPool(), statement, params, "pgx")
 }
 
 func getConnectionConfig(ctx context.Context, user, pass, dbname string) (string, bool, error) {

internal/sources/cockroachdb/cockroachdb.go

@@ -28,9 +28,7 @@ import (
 	"github.com/cockroachdb/cockroach-go/v2/crdb/crdbpgxv5"
 	"github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/internal/sources"
-	"github.com/googleapis/genai-toolbox/internal/sqlcommenter"
 	"github.com/googleapis/genai-toolbox/internal/util"
-	"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"go.opentelemetry.io/otel/trace"
@@ -147,34 +145,7 @@ func (s *Source) ExecuteTxWithRetry(ctx context.Context, fn func(pgx.Tx) error)
 }
 
 func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
-	// Inject the database driver into the context for SQLCommenter
-	ctx = sqlcommenter.WithDBDriver(ctx, "pgx")
-	// Decorate the statement with SQLCommenter metadata from the context
-	statement = sqlcommenter.AppendComment(ctx, statement)
-
-	rows, err := s.Query(ctx, statement, params...)
-	if err != nil {
-		return nil, err
-	}
-	defer rows.Close()
-
-	fields := rows.FieldDescriptions()
-	var out []any
-	for rows.Next() {
-		v, err := rows.Values()
-		if err != nil {
-			return nil, fmt.Errorf("unable to parse row: %w", err)
-		}
-		row := orderedmap.Row{}
-		for i, f := range fields {
-			row.Add(f.Name, v[i])
-		}
-		out = append(out, row)
-	}
-	if err := rows.Err(); err != nil {
-		return nil, fmt.Errorf("unable to execute query: %w", err)
-	}
-	return out, nil
+	return sources.RunSQLWithPgxQueryer(ctx, s, statement, params, "pgx")
 }
 
 // Query executes a query using the connection pool with MCP security enforcement.

internal/sources/pgx_util.go

@@ -0,0 +1,63 @@
+// Copyright 2025 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 sources
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/googleapis/genai-toolbox/internal/sqlcommenter"
+	"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
+	"github.com/jackc/pgx/v5"
+)
+
+// PgxQueryer abstracts connection pools and wrapper classes that can execute native Pgx queries.
+type PgxQueryer interface {
+	Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
+}
+
+// RunSQLWithPgxQueryer executes a standard SQL statement with SQLCommenter telemetry
+// across any driver natively supporting the pgx execution interface.
+func RunSQLWithPgxQueryer(ctx context.Context, queryer PgxQueryer, statement string, params []any, driver string) (any, error) {
+	// Inject the database driver into the context for SQLCommenter
+	ctx = sqlcommenter.WithDBDriver(ctx, driver)
+	// Decorate the statement with SQLCommenter metadata from the context
+	statement = sqlcommenter.AppendComment(ctx, statement)
+
+	results, err := queryer.Query(ctx, statement, params...)
+	if err != nil {
+		return nil, fmt.Errorf("unable to execute query: %w", err)
+	}
+	defer results.Close()
+
+	fields := results.FieldDescriptions()
+	var out []any
+	for results.Next() {
+		values, err := results.Values()
+		if err != nil {
+			return nil, fmt.Errorf("unable to parse row: %w", err)
+		}
+		row := orderedmap.Row{}
+		for i, f := range fields {
+			row.Add(f.Name, values[i])
+		}
+		out = append(out, row)
+	}
+	// this will catch actual query execution errors
+	if err := results.Err(); err != nil {
+		return nil, fmt.Errorf("unable to execute query: %w", err)
+	}
+	return out, nil
+}

internal/sources/postgres/postgres.go

@@ -22,9 +22,7 @@ import (
 
 	"github.com/goccy/go-yaml"
 	"github.com/googleapis/genai-toolbox/internal/sources"
-	"github.com/googleapis/genai-toolbox/internal/sqlcommenter"
 	"github.com/googleapis/genai-toolbox/internal/util"
-	"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"go.opentelemetry.io/otel/trace"
@@ -103,35 +101,7 @@ func (s *Source) PostgresPool() *pgxpool.Pool {
 }
 
 func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
-	// Inject the database driver into the context for SQLCommenter
-	ctx = sqlcommenter.WithDBDriver(ctx, "pgx")
-	// Decorate the statement with SQLCommenter metadata from the context
-	statement = sqlcommenter.AppendComment(ctx, statement)
-
-	results, err := s.PostgresPool().Query(ctx, statement, params...)
-	if err != nil {
-		return nil, fmt.Errorf("unable to execute query: %w", err)
-	}
-	defer results.Close()
-
-	fields := results.FieldDescriptions()
-	var out []any
-	for results.Next() {
-		values, err := results.Values()
-		if err != nil {
-			return nil, fmt.Errorf("unable to parse row: %w", err)
-		}
-		row := orderedmap.Row{}
-		for i, f := range fields {
-			row.Add(f.Name, values[i])
-		}
-		out = append(out, row)
-	}
-	// this will catch actual query execution errors
-	if err := results.Err(); err != nil {
-		return nil, fmt.Errorf("unable to execute query: %w", err)
-	}
-	return out, nil
+	return sources.RunSQLWithPgxQueryer(ctx, s.PostgresPool(), statement, params, "pgx")
 }
 
 func initPostgresConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string, queryParams map[string]string, queryExecMode string) (*pgxpool.Pool, error) {

internal/sources/yugabytedb/yugabytedb.go

@@ -126,7 +126,6 @@ func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (an
 		out = append(out, vMap)
 	}
 
-	// this will catch actual query execution errors
 	if err := results.Err(); err != nil {
 		return nil, fmt.Errorf("unable to execute query: %w", err)
 	}