feat: support tidb in data source, sql tool, and execute sql tool

cmd/root.go

@@ -92,6 +92,8 @@ import (
 	_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerexecutesql"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannersql"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/sqlitesql"
+	_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbexecutesql"
+	_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbsql"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/utility/alloydbwaitforoperation"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/utility/wait"
 	_ "github.com/googleapis/genai-toolbox/internal/tools/valkey"
@@ -118,6 +120,7 @@ import (
 	_ "github.com/googleapis/genai-toolbox/internal/sources/redis"
 	_ "github.com/googleapis/genai-toolbox/internal/sources/spanner"
 	_ "github.com/googleapis/genai-toolbox/internal/sources/sqlite"
+	_ "github.com/googleapis/genai-toolbox/internal/sources/tidb"
 	_ "github.com/googleapis/genai-toolbox/internal/sources/valkey"
 )
 

docs/en/resources/sources/tidb.md

@@ -0,0 +1,81 @@
+---
+title: "TiDB"
+type: docs
+weight: 2
+description: >
+  TiDB is a distributed SQL database that combines the best of traditional RDBMS and NoSQL databases.
+
+---
+
+## About
+
+[TiDB][tidb-docs] is an open-source distributed SQL database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads. It is MySQL-compatible and features horizontal scalability, strong consistency, and high availability.
+
+[tidb-docs]: https://docs.pingcap.com/tidb/stable
+
+## Requirements
+
+### Database User
+
+This source uses standard MySQL protocol authentication. You will need to [create a TiDB user][tidb-users] to login to the database with.
+
+For TiDB Cloud users, you can create database users through the TiDB Cloud console.
+
+[tidb-users]: https://docs.pingcap.com/tidb/stable/user-account-management
+
+## SSL Configuration
+
+- TiDB Cloud
+
+    For TiDB Cloud instances, SSL is automatically enabled when the hostname matches the TiDB Cloud pattern (`gateway*.*.*.tidbcloud.com`). You don't need to explicitly set `use_ssl: true` for TiDB Cloud connections.
+
+- Self-Hosted TiDB
+
+    For self-hosted TiDB instances, you can optionally enable SSL by setting `use_ssl: true` in your configuration.
+
+## Example
+
+- TiDB Cloud
+
+    ```yaml
+    sources:
+        my-tidb-cloud-source:
+            kind: tidb
+            host: gateway01.us-west-2.prod.aws.tidbcloud.com
+            port: 4000
+            database: my_db
+            user: ${TIDB_USERNAME}
+            password: ${TIDB_PASSWORD}
+            # SSL is automatically enabled for TiDB Cloud
+    ```
+
+- Self-Hosted TiDB
+
+    ```yaml
+    sources:
+        my-tidb-source:
+            kind: tidb
+            host: 127.0.0.1
+            port: 4000
+            database: my_db
+            user: ${TIDB_USERNAME}
+            password: ${TIDB_PASSWORD}
+            use_ssl: false  # Optional: enable SSL for secure connections
+    ```
+
+{{< notice tip >}}
+Use environment variable replacement with the format ${ENV_NAME}
+instead of hardcoding your secrets into the configuration file.
+{{< /notice >}}
+
+## Reference
+
+| **field** | **type** | **required** | **description**                                                                             |
+|-----------|:--------:|:------------:|---------------------------------------------------------------------------------------------|
+| kind      |  string  |     true     | Must be "tidb".                                                                             |
+| host      |  string  |     true     | IP address or hostname to connect to (e.g. "127.0.0.1" or "gateway01.*.tidbcloud.com").    |
+| port      |  string  |     true     | Port to connect to (typically "4000" for TiDB).                                            |
+| database  |  string  |     true     | Name of the TiDB database to connect to (e.g. "my_db").                                    |
+| user      |  string  |     true     | Name of the TiDB user to connect as (e.g. "my-tidb-user").                                 |
+| password  |  string  |     true     | Password of the TiDB user (e.g. "my-password").                                            |
+| use_ssl   |  boolean |    false     | Whether to use SSL/TLS encryption. Automatically enabled for TiDB Cloud instances.         | 

docs/en/resources/tools/tidb/_index.md

@@ -0,0 +1,7 @@
+---
+title: "TiDB"
+type: docs
+weight: 2
+description: > 
+  Tools that work with TiDB Sources, such as TiDB Cloud and self-hosted TiDB.
+---

docs/en/resources/tools/tidb/tidb-execute-sql.md

@@ -0,0 +1,41 @@
+---
+title: "tidb-execute-sql"
+type: docs
+weight: 2
+description: > 
+  A "tidb-execute-sql" tool executes a SQL statement against a TiDB
+  database.
+aliases:
+- /resources/tools/tidb-execute-sql
+---
+
+## About
+
+A `tidb-execute-sql` tool executes a SQL statement against a TiDB
+database. It's compatible with the following source:
+
+- [tidb](../sources/tidb.md)
+
+`tidb-execute-sql` takes one input parameter `sql` and run the sql
+statement against the `source`.
+
+> **Note:** This tool is intended for developer assistant workflows with
+> human-in-the-loop and shouldn't be used for production agents.
+
+## Example
+
+```yaml
+tools:
+ execute_sql_tool:
+    kind: tidb-execute-sql
+    source: my-tidb-instance
+    description: Use this tool to execute sql statement.
+```
+
+## Reference
+
+| **field**   |                  **type**                  | **required** | **description**                                                                                  |
+|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------|
+| kind        |                   string                   |     true     | Must be "tidb-execute-sql".                                                                     |
+| source      |                   string                   |     true     | Name of the source the SQL should execute on.                                                    |
+| description |                   string                   |     true     | Description of the tool that is passed to the LLM.                                               |

docs/en/resources/tools/tidb/tidb-sql.md

@@ -0,0 +1,105 @@
+---
+title: "tidb-sql"
+type: docs
+weight: 1
+description: > 
+  A "tidb-sql" tool executes a pre-defined SQL statement against a TiDB
+  database.
+aliases:
+- /resources/tools/tidb-sql
+---
+
+## About
+
+A `tidb-sql` tool executes a pre-defined SQL statement against a TiDB
+database. It's compatible with the following source:
+
+- [tidb](../sources/tidb.md)
+
+The specified SQL statement is executed as a [prepared statement][tidb-prepare],
+and expects parameters in the SQL query to be in the form of placeholders `?`.
+
+[tidb-prepare]: https://docs.pingcap.com/tidb/stable/sql-prepared-plan-cache
+
+## Example
+
+> **Note:** This tool uses parameterized queries to prevent SQL injections.
+> Query parameters can be used as substitutes for arbitrary expressions.
+> Parameters cannot be used as substitutes for identifiers, column names, table
+> names, or other parts of the query.
+
+```yaml
+tools:
+ search_flights_by_number:
+    kind: tidb-sql
+    source: my-tidb-instance
+    statement: |
+      SELECT * FROM flights
+      WHERE airline = ?
+      AND flight_number = ?
+      LIMIT 10
+    description: |
+      Use this tool to get information for a specific flight.
+      Takes an airline code and flight number and returns info on the flight.
+      Do NOT use this tool with a flight id. Do NOT guess an airline code or flight number.
+      A airline code is a code for an airline service consisting of two-character
+      airline designator and followed by flight number, which is 1 to 4 digit number.
+      For example, if given CY 0123, the airline is "CY", and flight_number is "123".
+      Another example for this is DL 1234, the airline is "DL", and flight_number is "1234".
+      If the tool returns more than one option choose the date closes to today.
+      Example:
+      {{
+          "airline": "CY",
+          "flight_number": "888",
+      }}
+      Example:
+      {{
+          "airline": "DL",
+          "flight_number": "1234",
+      }}
+    parameters:
+      - name: airline
+        type: string
+        description: Airline unique 2 letter identifier
+      - name: flight_number
+        type: string
+        description: 1 to 4 digit number
+```
+
+### Example with Template Parameters
+
+> **Note:** This tool allows direct modifications to the SQL statement,
+> including identifiers, column names, and table names. **This makes it more
+> vulnerable to SQL injections**. Using basic parameters only (see above) is
+> recommended for performance and safety reasons. For more details, please check
+> [templateParameters](_index#template-parameters).
+
+```yaml
+tools:
+ list_table:
+    kind: tidb-sql
+    source: my-tidb-instance
+    statement: |
+      SELECT * FROM {{.tableName}};
+    description: |
+      Use this tool to list all information from a specific table.
+      Example:
+      {{
+          "tableName": "flights",
+      }}
+    templateParameters:
+      - name: tableName
+        type: string
+        description: Table to select from
+```
+
+## Reference
+
+| **field**          |                  **type**                        | **required** | **description**                                                                                                                            |
+|--------------------|:------------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------------------------------------------------|
+| kind               |                   string                         |     true     | Must be "tidb-sql".                                                                                                                       |
+| source             |                   string                         |     true     | Name of the source the SQL should execute on.                                                                                              |
+| description        |                   string                         |     true     | Description of the tool that is passed to the LLM.                                                                                         |
+| statement          |                   string                         |     true     | SQL statement to execute on.                                                                                                               |
+| parameters         | [parameters](_index#specifying-parameters)       |    false     | List of [parameters](_index#specifying-parameters) that will be inserted into the SQL statement.                                           |
+| templateParameters | [templateParameters](_index#template-parameters) |    false     | List of [templateParameters](_index#template-parameters) that will be inserted into the SQL statement before executing prepared statement. |

internal/sources/tidb/tidb.go

@@ -0,0 +1,133 @@
+// 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 tidb
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"regexp"
+
+	_ "github.com/go-sql-driver/mysql"
+	"github.com/goccy/go-yaml"
+	"github.com/googleapis/genai-toolbox/internal/sources"
+	"go.opentelemetry.io/otel/trace"
+)
+
+const SourceKind string = "tidb"
+const TiDBCloudHostPattern string = `gateway\d{2}\.(.+)\.(prod|dev|staging)\.(.+)\.tidbcloud\.com`
+
+// validate interface
+var _ sources.SourceConfig = Config{}
+
+func init() {
+	if !sources.Register(SourceKind, newConfig) {
+		panic(fmt.Sprintf("source kind %q already registered", SourceKind))
+	}
+}
+
+func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
+	actual := Config{Name: name}
+	if err := decoder.DecodeContext(ctx, &actual); err != nil {
+		return nil, err
+	}
+
+	// If the host is a TiDB Cloud instance, force to use SSL
+	if isTiDBCloudHost(actual.Host) {
+		actual.UseSSL = true
+	}
+
+	return actual, nil
+}
+
+type Config struct {
+	Name     string `yaml:"name" validate:"required"`
+	Kind     string `yaml:"kind" validate:"required"`
+	Host     string `yaml:"host" validate:"required"`
+	Port     string `yaml:"port" validate:"required"`
+	User     string `yaml:"user" validate:"required"`
+	Password string `yaml:"password" validate:"required"`
+	Database string `yaml:"database" validate:"required"`
+	UseSSL   bool   `yaml:"use_ssl"`
+}
+
+func (r Config) SourceConfigKind() string {
+	return SourceKind
+}
+
+func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
+	pool, err := initTiDBConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database, r.UseSSL)
+	if err != nil {
+		return nil, fmt.Errorf("unable to create pool: %w", err)
+	}
+
+	err = pool.PingContext(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("unable to connect successfully: %w", err)
+	}
+
+	s := &Source{
+		Name: r.Name,
+		Kind: SourceKind,
+		Pool: pool,
+	}
+	return s, nil
+}
+
+var _ sources.Source = &Source{}
+
+type Source struct {
+	Name string `yaml:"name"`
+	Kind string `yaml:"kind"`
+	Pool *sql.DB
+}
+
+func (s *Source) SourceKind() string {
+	return SourceKind
+}
+
+func (s *Source) TiDBPool() *sql.DB {
+	return s.Pool
+}
+
+func isTiDBCloudHost(host string) bool {
+	pattern := `gateway\d{2}\.(.+)\.(prod|dev|staging)\.(.+)\.tidbcloud\.com`
+	match, err := regexp.MatchString(pattern, host)
+	if err != nil {
+		return false
+	}
+	return match
+}
+
+func initTiDBConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string, useSSL bool) (*sql.DB, error) {
+	//nolint:all // Reassigned ctx
+	ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
+	defer span.End()
+
+	// If the host is a TiDB Cloud instance, force to use SSL
+	if isTiDBCloudHost(host) {
+		useSSL = true
+	}
+
+	// Configure the driver to connect to the database
+	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?parseTime=true&charset=utf8mb4&tls=%t", user, pass, host, port, dbname, useSSL)
+
+	// Interact with the driver directly as you normally would
+	pool, err := sql.Open("mysql", dsn)
+	if err != nil {
+		return nil, fmt.Errorf("sql.Open: %w", err)
+	}
+	return pool, nil
+}