Skip to content

feat: add SQLCommenter observability to database drivers (fixes #2899)#2913

Open
tanujbolisetty wants to merge 5 commits intogoogleapis:mainfrom
tanujbolisetty:feat/sqlcommenter-integration
Open

feat: add SQLCommenter observability to database drivers (fixes #2899)#2913
tanujbolisetty wants to merge 5 commits intogoogleapis:mainfrom
tanujbolisetty:feat/sqlcommenter-integration

Conversation

@tanujbolisetty
Copy link
Copy Markdown

🚀 Overview

This PR introduces SQLCommenter observability across all supported database drivers in the GenAI Toolbox. It automatically injects rich contextual metadata (tool name, framework, agent controller, database driver, and OpenTelemetry traceparent) into every query executed by the toolbox.

By standardizing this telemetry pattern, downstream database administrators and cloud monitoring tools (like Cloud SQL Insights) can seamlessly map database loads back to specific AI agents and MCP tools without any manual parsing.

🛠 Features & Modifications

1. Comprehensive Driver Coverage

We have instrumented the RunSQL method on 100% of SQL-capable database sources.
The metadata is injected safely as a block comment (/* ... */) before the trailing semicolon (where applicable) to preserve syntax parsing.

  • Instrumented Standard SQL Drivers: alloydbpg, clickhouse, cockroachdb, cloudsqlmssql, cloudsqlmysql, cloudsqlpg, firebird, mindsdb, mssql, mysql, oceanbase, oracle, postgres, singlestore, snowflake, spanner, sqlite, tidb, trino, yugabytedb.
  • Instrumented Query/Graph Drivers: cassandra (CQL) and couchbase (N1QL). For Couchbase, the Source RunSQL signature and corresponding Tool wrapper Invoke interface were refactored to explicitly accept context.Context to enable propagation.
  • Native Labeling: bigquery is instrumented via native BigQuery Job Labels instead of SQL Comments, leveraging the same contextual map.
  • Exclusions (Safety): Object/NoSQL stores (mongodb, firestore, redis, bigtable, dgraph) were explicitly skipped as they do not support standard SQL comment blocks, preventing catastrophic parsing crashes.

2. Context Injection Pipeline

  • Created internal/sqlcommenter/commenter.go to handle context reading, deterministic key sorting, formatting (RFC 3986 URL encoding), and OpenTelemetry span extraction.

  • MCP Auto-Context Pipeline: internal/server/server.go was updated to trap the clientInfo.name payload from the MCP initialize handshake in a concurrent mcpClientNames sync map. api.go extracts this dynamically during tool invocation to populate the controller field automatically (e.g. associating queries directly with "Mosa" or "Claude Desktop").

  • Fallback configurations TOOLBOX_APP_NAME and TOOLBOX_AGENT_NAME were established to manually override the context if needed.

3. Core Framework & Build Architecture Upgrades

These changes resolve significant localized testing and binary compilation blockers encountered during development:

  • API UI Autostart: Modified internal/server/server.go configuration booting logic so that passing the --ui config implicitly turns on the REST API endpoints (if cfg.EnableAPI || cfg.UI).
  • Oracle CGO Decoupling: Separated the Oracle driver implementations into two new files (oci_cgo.go and oci_pure.go). By leveraging Go build tags (//go:build cgo), the toolbox now dynamically compiles using the lightweight go-ora/v2 Pure Go driver by default, entirely sidestepping notoriously complex godror C compiler dependencies native to Oracle cross-compilation.

🧪 Verification

  • Full test matrix suite completely passing: go test ./internal/sqlcommenter/....
  • Tested the append function string manipulation algorithm safely by mocking the WithDBDriver and WithToolName context decorators directly, avoiding the need for live mocked database instances in CI.
  • go build ./... successfully compiles across pure GO targets natively.

📝 Documentation

  • Overhauled docs/SQLCOMMENTER_README.md and explicitly mapped out the full driver support matrix, environment overrides, and NoSQL exclusions.

  • Recommendation: Ensure DEVELOPER.md or the primary README.md is updated to highlight that the default codebase cross-compiles with Oracle's PureGo driver to unblock new developers!

@tanujbolisetty tanujbolisetty requested review from a team as code owners March 31, 2026 19:28
@google-cla
Copy link
Copy Markdown

google-cla bot commented Mar 31, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request implements the SQLCommenter observability layer across all SQL-capable database drivers, enabling query tracking by tool name, agent identity, and trace metadata. It introduces the internal/sqlcommenter package and updates the MCP server to persist agent names across sessions. Feedback identifies a memory leak in the session name map, suggests refactoring duplicated RunSQL logic in the CockroachDB source, and recommends using %20 for space encoding to ensure strict adherence to the SQLCommenter specification.

@tanujbolisetty tanujbolisetty force-pushed the feat/sqlcommenter-integration branch from 8f3c9c0 to 9836bd1 Compare March 31, 2026 19:35
@tanujbolisetty tanujbolisetty force-pushed the feat/sqlcommenter-integration branch from 9836bd1 to b24f64f Compare March 31, 2026 19:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants