Skip to content

Commit 94a1d99

Browse files
AjmeraParth132anubhav756
AjmeraParth132
authored and
anubhav756
committed
feat: setup SQLCommenter and allow client metadata
1 parent 6632d96 commit 94a1d99

15 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

‎docs/en/reference/cli.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ description: >
2222
| | `--telemetry-gcp` | Enable exporting directly to Google Cloud Monitoring. | |
2323
| | `--telemetry-otlp` | Enable exporting using OpenTelemetry Protocol (OTLP) to the specified endpoint (e.g. 'http://127.0.0.1:4318') | |
2424
| | `--telemetry-service-name` | Sets the value of the service.name resource attribute for telemetry data. | `toolbox` |
25+
| | `--sql-commenter` | Enable appending SQLCommenter-format comments to SQL statements. | `false` |
2526
| | `--config` | File path specifying the tool configuration. Cannot be used with --configs or --config-folder. | |
2627
| | `--configs` | Multiple file paths specifying tool configurations. Files will be merged. Cannot be used with --config or --config-folder. | |
2728
| | `--config-folder` | Directory path containing YAML tool configuration files. All .yaml and .yml files in the directory will be loaded and merged. Cannot be used with --config or --configs. | |

‎internal/server/config.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@ type ServerConfig struct {
6464
TelemetryOTLP string
6565
// TelemetryServiceName defines the value of service.name resource attribute.
6666
TelemetryServiceName string
67+
// SQLCommenter enables appending SQLCommenter-format comments to SQL statements.
68+
SQLCommenter bool
6769
// Stdio indicates if Toolbox is listening via MCP stdio.
6870
Stdio bool
6971
// DisableReload indicates if the user has disabled dynamic reloading for Toolbox.

‎internal/server/mcp.go‎

Lines changed: 26 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",
@@ -464,6 +479,8 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
464479

465480
ctx := r.Context()
466481
ctx = util.WithLogger(ctx, s.logger)
482+
ctx = util.WithUserAgent(ctx, s.version)
483+
ctx = util.WithSQLCommenterEnabled(ctx, s.sqlCommenterEnabled)
467484

468485
// Read body first so we can extract trace context
469486
body, err := io.ReadAll(r.Body)
@@ -476,7 +493,7 @@ func httpHandler(s *Server, w http.ResponseWriter, r *http.Request) {
476493
}
477494

478495
// This ensures the transport span becomes a child of the client span
479-
ctx = extractTraceContext(ctx, body)
496+
ctx = extractMeta(ctx, body)
480497

481498
// Create span for HTTP transport
482499
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
@@ -49,16 +49,17 @@ import (
4949

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

6465
func InitializeConfigs(ctx context.Context, cfg ServerConfig) (
@@ -378,15 +379,16 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
378379
resourceManager := resources.NewResourceManager(sourcesMap, authServicesMap, embeddingModelsMap, toolsMap, toolsetsMap, promptsMap, promptsetsMap)
379380

380381
s := &Server{
381-
version: cfg.Version,
382-
srv: srv,
383-
root: r,
384-
logger: l,
385-
instrumentation: instrumentation,
386-
sseManager: sseManager,
387-
ResourceMgr: resourceManager,
388-
toolboxUrl: cfg.ToolboxUrl,
389-
mcpPrmFile: cfg.McpPrmFile,
382+
version: cfg.Version,
383+
sqlCommenterEnabled: cfg.SQLCommenter,
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,
390392
}
391393

392394
// 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)

0 commit comments

Comments
 (0)