Skip to content

Commit 3e1b274

Browse files
AjmeraParth132anubhav756
AjmeraParth132
authored and
anubhav756
committed
feat: setup SQLCommenter and allow client metadata
1 parent 87a600a commit 3e1b274

14 files changed

Lines changed: 429 additions & 28 deletions

File tree

‎cmd/internal/flags.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ func PersistentFlags(parentCmd *cobra.Command, opts *ToolboxOptions) {
3434
persistentFlags.BoolVar(&opts.Cfg.TelemetryGCP, "telemetry-gcp", false, "Enable exporting directly to Google Cloud Monitoring.")
3535
persistentFlags.StringVar(&opts.Cfg.TelemetryOTLP, "telemetry-otlp", "", "Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318')")
3636
persistentFlags.StringVar(&opts.Cfg.TelemetryServiceName, "telemetry-service-name", "toolbox", "Sets the value of the service.name resource attribute for telemetry data.")
37+
persistentFlags.BoolVar(&opts.Cfg.SQLCommenter, "sql-commenter", false, "Enable appending SQLCommenter-format comments to SQL statements.")
3738
persistentFlags.StringSliceVar(&opts.Cfg.UserAgentMetadata, "user-agent-metadata", []string{}, "Appends additional metadata to the User-Agent.")
3839
}
3940

‎internal/server/config.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,8 @@ type ServerConfig struct {
6868
TelemetryOTLP string
6969
// TelemetryServiceName defines the value of service.name resource attribute.
7070
TelemetryServiceName string
71+
// SQLCommenter enables appending SQLCommenter-format comments to SQL statements.
72+
SQLCommenter bool
7173
// Stdio indicates if Toolbox is listening via MCP stdio.
7274
Stdio bool
7375
// DisableReload indicates if the user has disabled dynamic reloading for Toolbox.

‎internal/server/mcp.go‎

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -145,14 +145,15 @@ func (c traceContextCarrier) Keys() []string {
145145
return keys
146146
}
147147

148-
// extractTraceContext extracts W3C Trace Context from params._meta
149-
func extractTraceContext(ctx context.Context, body []byte) context.Context {
150-
// Try to parse the request to extract _meta
148+
// extractMeta parses params._meta from the request body in a single pass,
149+
// extracting both W3C Trace Context and client telemetry attributes.
150+
func extractMeta(ctx context.Context, body []byte) context.Context {
151151
var req struct {
152152
Params struct {
153153
Meta struct {
154-
Traceparent string `json:"traceparent,omitempty"`
155-
Tracestate string `json:"tracestate,omitempty"`
154+
Traceparent string `json:"traceparent,omitempty"`
155+
Tracestate string `json:"tracestate,omitempty"`
156+
TelemetryAttrs map[string]string `json:"dev.mcp-toolbox/telemetry,omitempty"`
156157
} `json:"_meta,omitempty"`
157158
} `json:"params,omitempty"`
158159
}
@@ -161,15 +162,27 @@ func extractTraceContext(ctx context.Context, body []byte) context.Context {
161162
return ctx
162163
}
163164

164-
// If traceparent is present, extract the context
165+
// Extract W3C Trace Context
165166
if req.Params.Meta.Traceparent != "" {
166167
carrier := traceContextCarrier{
167168
"traceparent": req.Params.Meta.Traceparent,
168169
}
169170
if req.Params.Meta.Tracestate != "" {
170171
carrier["tracestate"] = req.Params.Meta.Tracestate
171172
}
172-
return otel.GetTextMapPropagator().Extract(ctx, carrier)
173+
ctx = otel.GetTextMapPropagator().Extract(ctx, carrier)
174+
}
175+
176+
// Extract client telemetry attributes
177+
if attrs := req.Params.Meta.TelemetryAttrs; len(attrs) > 0 {
178+
ta := &util.TelemetryAttributes{
179+
ClientName: attrs["client.name"],
180+
ClientVersion: attrs["client.version"],
181+
ClientModel: attrs["client.model"],
182+
ClientUserID: attrs["client.user.id"],
183+
ClientAgentID: attrs["client.agent.id"],
184+
}
185+
ctx = util.WithTelemetryAttributes(ctx, ta)
173186
}
174187

175188
return ctx
@@ -191,6 +204,8 @@ func (s *stdioSession) Start(ctx context.Context) error {
191204
// readInputStream reads requests/notifications from MCP clients through stdin
192205
func (s *stdioSession) readInputStream(ctx context.Context) error {
193206
sessionStart := time.Now()
207+
ctx = util.WithUserAgent(ctx, s.server.version)
208+
ctx = util.WithSQLCommenterEnabled(ctx, s.server.sqlCommenterEnabled)
194209

195210
// Define attributes for session metrics
196211
// Note: mcp.protocol.version is added dynamically after protocol negotiation
@@ -238,7 +253,7 @@ func (s *stdioSession) readInputStream(ctx context.Context) error {
238253

239254
if err := func() error {
240255
// This ensures the transport span becomes a child of the client span
241-
msgCtx := extractTraceContext(ctx, []byte(line))
256+
msgCtx := extractMeta(ctx, []byte(line))
242257

243258
// Create span for STDIO transport
244259
msgCtx, span := s.server.instrumentation.Tracer.Start(msgCtx, "toolbox/server/mcp/stdio",
@@ -469,6 +484,9 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
469484
w.Header().Set("Content-Type", "application/json")
470485

471486
ctx := r.Context()
487+
ctx = util.WithLogger(ctx, s.logger)
488+
ctx = util.WithUserAgent(ctx, s.version)
489+
ctx = util.WithSQLCommenterEnabled(ctx, s.sqlCommenterEnabled)
472490

473491
// Read body first so we can extract trace context
474492
body, err := io.ReadAll(r.Body)
@@ -481,7 +499,7 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
481499
}
482500

483501
// This ensures the transport span becomes a child of the client span
484-
ctx = extractTraceContext(ctx, body)
502+
ctx = extractMeta(ctx, body)
485503

486504
// Create span for HTTP transport
487505
ctx, span := s.instrumentation.Tracer.Start(ctx, "toolbox/server/mcp/http",

‎internal/server/server.go‎

Lines changed: 21 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -51,16 +51,17 @@ import (
5151

5252
// Server contains info for running an instance of Toolbox. Should be instantiated with NewServer().
5353
type Server struct {
54-
version string
55-
toolboxUrl string
56-
srv *http.Server
57-
listener net.Listener
58-
root chi.Router
59-
logger log.Logger
60-
instrumentation *telemetry.Instrumentation
61-
sseManager *sseManager
62-
ResourceMgr *resources.ResourceManager
63-
mcpPrmFile string
54+
version string
55+
sqlCommenterEnabled bool
56+
toolboxUrl string
57+
srv *http.Server
58+
listener net.Listener
59+
root chi.Router
60+
logger log.Logger
61+
instrumentation *telemetry.Instrumentation
62+
sseManager *sseManager
63+
ResourceMgr *resources.ResourceManager
64+
mcpPrmFile string
6465
}
6566

6667
func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
@@ -380,15 +381,16 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
380381
resourceManager := resources.NewResourceManager(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
381382

382383
s := &Server{
383-
version: cfg.Version,
384-
srv: srv,
385-
root: r,
386-
logger: l,
387-
instrumentation: instrumentation,
388-
sseManager: sseManager,
389-
ResourceMgr: resourceManager,
390-
toolboxUrl: cfg.ToolboxUrl,
391-
mcpPrmFile: cfg.McpPrmFile,
384+
version: cfg.Version,
385+
sqlCommenterEnabled: cfg.SQLCommenter,
386+
srv: srv,
387+
root: r,
388+
logger: l,
389+
instrumentation: instrumentation,
390+
sseManager: sseManager,
391+
ResourceMgr: resourceManager,
392+
toolboxUrl: cfg.ToolboxUrl,
393+
mcpPrmFile: cfg.McpPrmFile,
392394
}
393395

394396
// cors

‎internal/sources/alloydbpg/alloydb_pg.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
"cloud.google.com/go/alloydbconn"
2424
"github.com/goccy/go-yaml"
2525
"github.com/googleapis/mcp-toolbox/internal/sources"
26+
"github.com/googleapis/mcp-toolbox/internal/sources/sqlcommenter"
2627
"github.com/googleapis/mcp-toolbox/internal/util"
2728
"github.com/googleapis/mcp-toolbox/internal/util/orderedmap"
2829
"github.com/jackc/pgx/v5/pgxpool"
@@ -103,6 +104,7 @@ func (s *Source) PostgresPool() *pgxpool.Pool {
103104
}
104105

105106
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
107+
statement = sqlcommenter.AppendComment(ctx, statement, SourceType)
106108
results, err := s.Pool.Query(ctx, statement, params...)
107109
if err != nil {
108110
return nil, fmt.Errorf("unable to execute query: %w", err)

‎internal/sources/cloudsqlmssql/cloud_sql_mssql.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"cloud.google.com/go/cloudsqlconn/sqlserver/mssql"
2525
"github.com/goccy/go-yaml"
2626
"github.com/googleapis/mcp-toolbox/internal/sources"
27+
"github.com/googleapis/mcp-toolbox/internal/sources/sqlcommenter"
2728
"github.com/googleapis/mcp-toolbox/internal/util"
2829
"github.com/googleapis/mcp-toolbox/internal/util/orderedmap"
2930
"go.opentelemetry.io/otel/trace"
@@ -108,6 +109,7 @@ func (s *Source) MSSQLDB() *sql.DB {
108109
}
109110

110111
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
112+
statement = sqlcommenter.AppendComment(ctx, statement, SourceType)
111113
results, err := s.MSSQLDB().QueryContext(ctx, statement, params...)
112114
if err != nil {
113115
return nil, fmt.Errorf("unable to execute query: %w", err)

‎internal/sources/cloudsqlmysql/cloud_sql_mysql.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"cloud.google.com/go/cloudsqlconn/mysql/mysql"
2525
"github.com/goccy/go-yaml"
2626
"github.com/googleapis/mcp-toolbox/internal/sources"
27+
"github.com/googleapis/mcp-toolbox/internal/sources/sqlcommenter"
2728
"github.com/googleapis/mcp-toolbox/internal/tools/mysql/mysqlcommon"
2829
"github.com/googleapis/mcp-toolbox/internal/util"
2930
"github.com/googleapis/mcp-toolbox/internal/util/orderedmap"
@@ -107,6 +108,7 @@ func (s *Source) MySQLDatabase() string {
107108
}
108109

109110
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
111+
statement = sqlcommenter.AppendComment(ctx, statement, SourceType)
110112
results, err := s.MySQLPool().QueryContext(ctx, statement, params...)
111113
if err != nil {
112114
return nil, fmt.Errorf("unable to execute query: %w", err)

‎internal/sources/cloudsqlpg/cloud_sql_pg.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"cloud.google.com/go/cloudsqlconn"
2323
"github.com/goccy/go-yaml"
2424
"github.com/googleapis/mcp-toolbox/internal/sources"
25+
"github.com/googleapis/mcp-toolbox/internal/sources/sqlcommenter"
2526
"github.com/googleapis/mcp-toolbox/internal/util"
2627
"github.com/googleapis/mcp-toolbox/internal/util/orderedmap"
2728
"github.com/jackc/pgx/v5/pgxpool"
@@ -109,6 +110,7 @@ func (s *Source) PostgresPool() *pgxpool.Pool {
109110
}
110111

111112
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
113+
statement = sqlcommenter.AppendComment(ctx, statement, SourceType)
112114
results, err := s.PostgresPool().Query(ctx, statement, params...)
113115
if err != nil {
114116
return nil, fmt.Errorf("unable to execute query: %w", err)

‎internal/sources/mssql/mssql.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222

2323
"github.com/goccy/go-yaml"
2424
"github.com/googleapis/mcp-toolbox/internal/sources"
25+
"github.com/googleapis/mcp-toolbox/internal/sources/sqlcommenter"
2526
"github.com/googleapis/mcp-toolbox/internal/util"
2627
"github.com/googleapis/mcp-toolbox/internal/util/orderedmap"
2728
_ "github.com/microsoft/go-mssqldb"
@@ -106,6 +107,7 @@ func (s *Source) MSSQLDB() *sql.DB {
106107
}
107108

108109
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
110+
statement = sqlcommenter.AppendComment(ctx, statement, SourceType)
109111
results, err := s.MSSQLDB().QueryContext(ctx, statement, params...)
110112
if err != nil {
111113
return nil, fmt.Errorf("unable to execute query: %w", err)

‎internal/sources/mysql/mysql.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
driver "github.com/go-sql-driver/mysql"
2424
"github.com/goccy/go-yaml"
2525
"github.com/googleapis/mcp-toolbox/internal/sources"
26+
"github.com/googleapis/mcp-toolbox/internal/sources/sqlcommenter"
2627
"github.com/googleapis/mcp-toolbox/internal/tools/mysql/mysqlcommon"
2728
"github.com/googleapis/mcp-toolbox/internal/util"
2829
"github.com/googleapis/mcp-toolbox/internal/util/orderedmap"
@@ -106,6 +107,7 @@ func (s *Source) MySQLDatabase() string {
106107
}
107108

108109
func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
110+
statement = sqlcommenter.AppendComment(ctx, statement, SourceType)
109111
results, err := s.MySQLPool().QueryContext(ctx, statement, params...)
110112
if err != nil {
111113
return nil, fmt.Errorf("unable to execute query: %w", err)

0 commit comments

Comments
 (0)