feat(tidb): add schema introspection tools and prebuilt config

cmd/internal/imports.go

@@ -224,6 +224,10 @@ import (
 	_ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqliteexecutesql"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqlitesql"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbexecutesql"
+	_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbgetqueryplan"
+	_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidblistactivequeries"
+	_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidblisttables"
+	_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidblisttiflashreplicas"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbsql"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/trino/trinoexecutesql"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/trino/trinosql"

internal/prebuiltconfigs/tools/tidb.yaml

@@ -0,0 +1,77 @@
+# Copyright 2026 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.
+
+# TiDB Prebuilt Configuration
+#
+# TiDB is a distributed SQL database that supports Hybrid Transactional
+# and Analytical Processing (HTAP) workloads. It is MySQL-compatible and
+# provides horizontal scalability, strong consistency, and high availability.
+#
+# Key TiDB-specific features exposed by these tools:
+# - TiFlash: Columnar storage for real-time analytics (requires TiDB 4.0+)
+# - Distributed query execution with transparent scaling
+#
+# Environment variables:
+#   TIDB_HOST     - TiDB server host (default: localhost)
+#   TIDB_PORT     - TiDB server port (default: 4000)
+#   TIDB_DATABASE - Database name (required)
+#   TIDB_USER     - Database user (required)
+#   TIDB_PASSWORD - Database password (required)
+#
+# For TiDB Cloud, SSL is automatically enabled when the host matches
+# the TiDB Cloud gateway pattern (gateway*.tidbcloud.com).
+
+sources:
+  tidb-source:
+    kind: tidb
+    host: ${TIDB_HOST:localhost}
+    port: ${TIDB_PORT:4000}
+    database: ${TIDB_DATABASE}
+    user: ${TIDB_USER}
+    password: ${TIDB_PASSWORD}
+tools:
+  execute_sql:
+    kind: tidb-execute-sql
+    source: tidb-source
+    description: Execute arbitrary SQL statements on TiDB. Supports SELECT, INSERT, UPDATE, DELETE, and DDL statements. Use with caution for data-modifying operations.
+  list_tables:
+    kind: tidb-list-tables
+    source: tidb-source
+    description: "Lists detailed schema information (columns, constraints, indexes, TiFlash replica count) as JSON for user-created tables. Filters by a comma-separated list of names. If names are omitted, lists all tables in user schemas. Excludes system schemas (mysql, information_schema, performance_schema, sys, METRICS_SCHEMA, INSPECTION_SCHEMA)."
+  get_query_plan:
+    kind: tidb-get-query-plan
+    source: tidb-source
+    description: "Provide information about how TiDB executes a SQL statement using EXPLAIN. Common use cases include: 1) analyze query plan to improve performance, 2) determine effectiveness of existing indexes, 3) identify if TiFlash is being used for HTAP queries. Supports 'default', 'analyze' (actual execution stats, SELECT only for safety), and 'verbose' (detailed cost info) explain types. WARNING: EXPLAIN ANALYZE actually executes the query."
+  list_active_queries:
+    kind: tidb-list-active-queries
+    source: tidb-source
+    description: Lists top N (default 10) ongoing queries from TiDB's processlist, ordered by execution time in descending order. Returns detailed information including process id, user, host, database, command, execution time, state, query text (truncated to 1000 chars), memory usage, and transaction start timestamp.
+  list_tiflash_replicas:
+    kind: tidb-list-tiflash-replicas
+    source: tidb-source
+    description: "Lists TiFlash replica status for all tables that have TiFlash replicas configured. TiFlash is TiDB's columnar storage engine that enables real-time HTAP analytics (requires TiDB 4.0+). Returns replica count, availability status, and sync progress for each table. Useful for monitoring TiFlash deployment health and identifying tables ready for analytical queries."
+toolsets:
+  data:
+    - execute_sql
+    - list_tables
+    - get_query_plan
+    - list_active_queries
+  monitor:
+    - get_query_plan
+    - list_active_queries
+    - list_tiflash_replicas
+  htap:
+    - execute_sql
+    - list_tables
+    - list_tiflash_replicas

internal/tools/tidb/tidbgetqueryplan/tidbgetqueryplan.go

@@ -0,0 +1,211 @@
+// Copyright 2026 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 tidbgetqueryplan
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"net/http"
+	"regexp"
+	"strings"
+
+	yaml "github.com/goccy/go-yaml"
+	"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
+	"github.com/googleapis/genai-toolbox/internal/sources"
+	"github.com/googleapis/genai-toolbox/internal/tools"
+	"github.com/googleapis/genai-toolbox/internal/util"
+	"github.com/googleapis/genai-toolbox/internal/util/parameters"
+)
+
+const resourceType string = "tidb-get-query-plan"
+
+// stripSQLComments removes SQL comments (both -- and /* */) and leading/trailing whitespace
+func stripSQLComments(sql string) string {
+	// Remove multi-line comments /* ... */
+	reMultiLine := regexp.MustCompile(`/\*[\s\S]*?\*/`)
+	sql = reMultiLine.ReplaceAllString(sql, "")
+	// Remove single-line comments -- ...
+	reSingleLine := regexp.MustCompile(`--[^\n]*`)
+	sql = reSingleLine.ReplaceAllString(sql, "")
+	return strings.TrimSpace(sql)
+}
+
+// isSelectOrWithStatement checks if the SQL is a SELECT or WITH (CTE) statement
+func isSelectOrWithStatement(sql string) bool {
+	normalized := strings.ToUpper(stripSQLComments(sql))
+	return strings.HasPrefix(normalized, "SELECT") || strings.HasPrefix(normalized, "WITH")
+}
+
+// containsMultipleStatements checks if SQL contains multiple statements (semicolon)
+func containsMultipleStatements(sql string) bool {
+	// Strip comments first to avoid false positives from semicolons in comments
+	sql = stripSQLComments(sql)
+	// Remove string literals to avoid false positives
+	reString := regexp.MustCompile(`'[^']*'|"[^"]*"`)
+	cleaned := reString.ReplaceAllString(sql, "")
+	return strings.Contains(cleaned, ";")
+}
+
+func init() {
+	if !tools.Register(resourceType, newConfig) {
+		panic(fmt.Sprintf("tool type %q already registered", resourceType))
+	}
+}
+
+func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
+	actual := Config{Name: name}
+	if err := decoder.DecodeContext(ctx, &actual); err != nil {
+		return nil, err
+	}
+	return actual, nil
+}
+
+type compatibleSource interface {
+	TiDBPool() *sql.DB
+	RunSQL(context.Context, string, []any) (any, error)
+}
+
+type Config struct {
+	Name         string   `yaml:"name" validate:"required"`
+	Type         string   `yaml:"type" validate:"required"`
+	Source       string   `yaml:"source" validate:"required"`
+	Description  string   `yaml:"description" validate:"required"`
+	AuthRequired []string `yaml:"authRequired"`
+}
+
+// validate interface
+var _ tools.ToolConfig = Config{}
+
+func (cfg Config) ToolConfigType() string {
+	return resourceType
+}
+
+func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
+	allParameters := parameters.Parameters{
+		parameters.NewStringParameter("sql", "The SQL query to analyze. Must be a SELECT, INSERT, UPDATE, or DELETE statement."),
+		parameters.NewStringParameterWithDefault("explain_type", "default", "Optional: The type of EXPLAIN output. Options: 'default' (basic plan), 'analyze' (actual execution stats - SELECT only), 'verbose' (detailed cost info)."),
+	}
+	paramManifest := allParameters.Manifest()
+	mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters, nil)
+
+	// finish tool setup
+	t := Tool{
+		Config:      cfg,
+		AllParams:   allParameters,
+		manifest:    tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
+		mcpManifest: mcpManifest,
+	}
+	return t, nil
+}
+
+// validate interface
+var _ tools.Tool = Tool{}
+
+type Tool struct {
+	Config
+	AllParams parameters.Parameters `yaml:"allParams"`
+
+	manifest    tools.Manifest
+	mcpManifest tools.McpManifest
+}
+
+func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
+	source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
+	if err != nil {
+		return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
+	}
+
+	paramsMap := params.AsMap()
+
+	sqlStr, ok := paramsMap["sql"].(string)
+	if !ok || strings.TrimSpace(sqlStr) == "" {
+		return nil, util.NewAgentError("'sql' parameter is required and must be a non-empty string", nil)
+	}
+
+	// Security check: reject multiple statements to prevent injection like "SELECT 1; DELETE FROM t"
+	if containsMultipleStatements(sqlStr) {
+		return nil, util.NewAgentError("multiple SQL statements are not allowed; remove any semicolons from your query", nil)
+	}
+
+	explainType, _ := paramsMap["explain_type"].(string)
+	if explainType == "" {
+		explainType = "default"
+	}
+
+	// Build the EXPLAIN statement based on the type
+	var explainStmt string
+	switch strings.ToLower(explainType) {
+	case "analyze":
+		// EXPLAIN ANALYZE actually executes the query and shows real execution stats
+		// For safety, only allow SELECT or WITH (CTE) statements
+		if !isSelectOrWithStatement(sqlStr) {
+			return nil, util.NewAgentError("EXPLAIN ANALYZE only supports SELECT statements (including WITH/CTE) for safety reasons; use 'default' or 'verbose' for other statement types", nil)
+		}
+		explainStmt = fmt.Sprintf("EXPLAIN ANALYZE %s", sqlStr)
+	case "verbose":
+		// EXPLAIN FORMAT='verbose' shows detailed cost estimation
+		explainStmt = fmt.Sprintf("EXPLAIN FORMAT='verbose' %s", sqlStr)
+	case "default":
+		explainStmt = fmt.Sprintf("EXPLAIN %s", sqlStr)
+	default:
+		return nil, util.NewAgentError(fmt.Sprintf("invalid value for explain_type: must be 'default', 'analyze', or 'verbose', but got %q", explainType), nil)
+	}
+
+	// Log the query for debugging
+	logger, err := util.LoggerFromContext(ctx)
+	if err != nil {
+		return nil, util.NewClientServerError("error getting logger", http.StatusInternalServerError, err)
+	}
+	logger.DebugContext(ctx, fmt.Sprintf("executing `%s` tool query: %s", resourceType, explainStmt))
+
+	resp, err := source.RunSQL(ctx, explainStmt, nil)
+	if err != nil {
+		return nil, util.ProcessGeneralError(err)
+	}
+	return resp, nil
+}
+
+func (t Tool) EmbedParams(ctx context.Context, paramValues parameters.ParamValues, embeddingModelsMap map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
+	return parameters.EmbedParams(ctx, t.AllParams, paramValues, embeddingModelsMap, nil)
+}
+
+func (t Tool) Manifest() tools.Manifest {
+	return t.manifest
+}
+
+func (t Tool) McpManifest() tools.McpManifest {
+	return t.mcpManifest
+}
+
+func (t Tool) Authorized(verifiedAuthServices []string) bool {
+	return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
+}
+
+func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (bool, error) {
+	return false, nil
+}
+
+func (t Tool) ToConfig() tools.ToolConfig {
+	return t.Config
+}
+
+func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
+	return "Authorization", nil
+}
+
+func (t Tool) GetParameters() parameters.Parameters {
+	return t.AllParams
+}