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

.ci/integration.cloudbuild.yaml

@@ -486,7 +486,7 @@ steps:
           "Looker" \
           looker \
           looker
-    
+
   - id: "duckdb"
     name: golang:1
     waitFor: ["compile-test-binary"]
@@ -506,7 +506,6 @@ steps:
           duckdb \
           duckdb
 
-
   - id: "alloydbwaitforoperation"
     name: golang:1
     waitFor: ["compile-test-binary"]
@@ -525,6 +524,28 @@ steps:
           "Alloydb Wait for Operation" \
           utility \
           utility/alloydbwaitforoperation
+
+  - id: "tidb"
+    name: golang:1
+    waitFor: ["compile-test-binary"]
+    entrypoint: /bin/bash
+    env:
+      - "GOPATH=/gopath"
+      - "TIDB_DATABASE=$_DATABASE_NAME"
+      - "TIDB_HOST=$_TIDB_HOST"
+      - "TIDB_PORT=$_TIDB_PORT"
+      - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
+    secretEnv: ["CLIENT_ID", "TIDB_USER", "TIDB_PASS"]
+    volumes:
+      - name: "go"
+        path: "/gopath"
+    args:
+      - -c
+      - |
+        .ci/test_with_coverage.sh \
+          "TiDB" \
+          tidb \
+          tidbsql tidbexecutesql 
 
 availableSecrets:
   secretManager:
@@ -584,7 +605,10 @@ availableSecrets:
       env: LOOKER_CLIENT_ID
     - versionName: projects/107716898620/secrets/looker_client_secret/versions/latest
       env: LOOKER_CLIENT_SECRET
-
+    - versionName: projects/107716898620/secrets/tidb_user/versions/latest
+      env: TIDB_USER
+    - versionName: projects/107716898620/secrets/tidb_pass/versions/latest
+      env: TIDB_PASS
 
 options:
   logging: CLOUD_LOGGING_ONLY
@@ -616,4 +640,6 @@ substitutions:
   _DGRAPHURL: "https://play.dgraph.io"
   _COUCHBASE_BUCKET: "couchbase-bucket"
   _COUCHBASE_SCOPE: "couchbase-scope"
-  _LOOKER_VERIFY_SSL: "true"
+  _LOOKER_VERIFY_SSL: "true"
+  _TIDB_HOST: 127.0.0.1
+  _TIDB_PORT: "4000"

cmd/root.go

@@ -95,6 +95,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"
@@ -122,6 +124,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: 1
+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 `ssl: true` for TiDB Cloud connections.
+
+- Self-Hosted TiDB
+
+    For self-hosted TiDB instances, you can optionally enable SSL by setting `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}
+            # ssl: true  # 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").                                            |
+| 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: 1
+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: 1
+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. |