Skip to content

Commit 6eaf36a

Browse files
IcemapYuan325
andauthored
feat: support tidb in data source, sql tool, and execute sql tool (#829)
This PR supports TiDB in: 1. sources - tidb: As a data source; 2. tools - tidb - tidbsql: As a prepared SQL running tool; 3. tools - tidb - tidbexecutesql: As an arbitrary SQL running tool (for development purposes). And its corresponding docs. --------- Co-authored-by: Yuan Teoh <45984206+Yuan325@users.noreply.github.com> Co-authored-by: Yuan Teoh <yuanteoh@google.com>
1 parent 051e686 commit 6eaf36a

File tree

13 files changed

+1479
-4
lines changed

13 files changed

+1479
-4
lines changed

‎.ci/integration.cloudbuild.yaml‎

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -486,7 +486,7 @@ steps:
486486
"Looker" \
487487
looker \
488488
looker
489-
489+
490490
- id: "duckdb"
491491
name: golang:1
492492
waitFor: ["compile-test-binary"]
@@ -506,7 +506,6 @@ steps:
506506
duckdb \
507507
duckdb
508508
509-
510509
- id: "alloydbwaitforoperation"
511510
name: golang:1
512511
waitFor: ["compile-test-binary"]
@@ -525,6 +524,28 @@ steps:
525524
"Alloydb Wait for Operation" \
526525
utility \
527526
utility/alloydbwaitforoperation
527+
528+
- id: "tidb"
529+
name: golang:1
530+
waitFor: ["compile-test-binary"]
531+
entrypoint: /bin/bash
532+
env:
533+
- "GOPATH=/gopath"
534+
- "TIDB_DATABASE=$_DATABASE_NAME"
535+
- "TIDB_HOST=$_TIDB_HOST"
536+
- "TIDB_PORT=$_TIDB_PORT"
537+
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
538+
secretEnv: ["CLIENT_ID", "TIDB_USER", "TIDB_PASS"]
539+
volumes:
540+
- name: "go"
541+
path: "/gopath"
542+
args:
543+
- -c
544+
- |
545+
.ci/test_with_coverage.sh \
546+
"TiDB" \
547+
tidb \
548+
tidbsql tidbexecutesql
528549
529550
availableSecrets:
530551
secretManager:
@@ -584,7 +605,10 @@ availableSecrets:
584605
env: LOOKER_CLIENT_ID
585606
- versionName: projects/107716898620/secrets/looker_client_secret/versions/latest
586607
env: LOOKER_CLIENT_SECRET
587-
608+
- versionName: projects/107716898620/secrets/tidb_user/versions/latest
609+
env: TIDB_USER
610+
- versionName: projects/107716898620/secrets/tidb_pass/versions/latest
611+
env: TIDB_PASS
588612

589613
options:
590614
logging: CLOUD_LOGGING_ONLY
@@ -616,4 +640,6 @@ substitutions:
616640
_DGRAPHURL: "https://play.dgraph.io"
617641
_COUCHBASE_BUCKET: "couchbase-bucket"
618642
_COUCHBASE_SCOPE: "couchbase-scope"
619-
_LOOKER_VERIFY_SSL: "true"
643+
_LOOKER_VERIFY_SSL: "true"
644+
_TIDB_HOST: 127.0.0.1
645+
_TIDB_PORT: "4000"

‎cmd/root.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,8 @@ import (
9595
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerexecutesql"
9696
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannersql"
9797
_ "github.com/googleapis/genai-toolbox/internal/tools/sqlitesql"
98+
_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbexecutesql"
99+
_ "github.com/googleapis/genai-toolbox/internal/tools/tidb/tidbsql"
98100
_ "github.com/googleapis/genai-toolbox/internal/tools/utility/alloydbwaitforoperation"
99101
_ "github.com/googleapis/genai-toolbox/internal/tools/utility/wait"
100102
_ "github.com/googleapis/genai-toolbox/internal/tools/valkey"
@@ -122,6 +124,7 @@ import (
122124
_ "github.com/googleapis/genai-toolbox/internal/sources/redis"
123125
_ "github.com/googleapis/genai-toolbox/internal/sources/spanner"
124126
_ "github.com/googleapis/genai-toolbox/internal/sources/sqlite"
127+
_ "github.com/googleapis/genai-toolbox/internal/sources/tidb"
125128
_ "github.com/googleapis/genai-toolbox/internal/sources/valkey"
126129
)
127130

‎docs/en/resources/sources/tidb.md‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
---
2+
title: "TiDB"
3+
type: docs
4+
weight: 1
5+
description: >
6+
TiDB is a distributed SQL database that combines the best of traditional RDBMS and NoSQL databases.
7+
8+
---
9+
10+
## About
11+
12+
[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.
13+
14+
[tidb-docs]: https://docs.pingcap.com/tidb/stable
15+
16+
## Requirements
17+
18+
### Database User
19+
20+
This source uses standard MySQL protocol authentication. You will need to [create a TiDB user][tidb-users] to login to the database with.
21+
22+
For TiDB Cloud users, you can create database users through the TiDB Cloud console.
23+
24+
[tidb-users]: https://docs.pingcap.com/tidb/stable/user-account-management
25+
26+
## SSL Configuration
27+
28+
- TiDB Cloud
29+
30+
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.
31+
32+
- Self-Hosted TiDB
33+
34+
For self-hosted TiDB instances, you can optionally enable SSL by setting `ssl: true` in your configuration.
35+
36+
## Example
37+
38+
- TiDB Cloud
39+
40+
```yaml
41+
sources:
42+
my-tidb-cloud-source:
43+
kind: tidb
44+
host: gateway01.us-west-2.prod.aws.tidbcloud.com
45+
port: 4000
46+
database: my_db
47+
user: ${TIDB_USERNAME}
48+
password: ${TIDB_PASSWORD}
49+
# SSL is automatically enabled for TiDB Cloud
50+
```
51+
52+
- Self-Hosted TiDB
53+
54+
```yaml
55+
sources:
56+
my-tidb-source:
57+
kind: tidb
58+
host: 127.0.0.1
59+
port: 4000
60+
database: my_db
61+
user: ${TIDB_USERNAME}
62+
password: ${TIDB_PASSWORD}
63+
# ssl: true # Optional: enable SSL for secure connections
64+
```
65+
66+
{{< notice tip >}}
67+
Use environment variable replacement with the format ${ENV_NAME}
68+
instead of hardcoding your secrets into the configuration file.
69+
{{< /notice >}}
70+
71+
## Reference
72+
73+
| **field** | **type** | **required** | **description** |
74+
|-----------|:--------:|:------------:|--------------------------------------------------------------------------------------------|
75+
| kind | string | true | Must be "tidb". |
76+
| host | string | true | IP address or hostname to connect to (e.g. "127.0.0.1" or "gateway01.*.tidbcloud.com"). |
77+
| port | string | true | Port to connect to (typically "4000" for TiDB). |
78+
| database | string | true | Name of the TiDB database to connect to (e.g. "my_db"). |
79+
| user | string | true | Name of the TiDB user to connect as (e.g. "my-tidb-user"). |
80+
| password | string | true | Password of the TiDB user (e.g. "my-password"). |
81+
| ssl | boolean | false | Whether to use SSL/TLS encryption. Automatically enabled for TiDB Cloud instances. |
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
---
2+
title: "TiDB"
3+
type: docs
4+
weight: 1
5+
description: >
6+
Tools that work with TiDB Sources, such as TiDB Cloud and self-hosted TiDB.
7+
---
Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
---
2+
title: "tidb-execute-sql"
3+
type: docs
4+
weight: 1
5+
description: >
6+
A "tidb-execute-sql" tool executes a SQL statement against a TiDB
7+
database.
8+
aliases:
9+
- /resources/tools/tidb-execute-sql
10+
---
11+
12+
## About
13+
14+
A `tidb-execute-sql` tool executes a SQL statement against a TiDB
15+
database. It's compatible with the following source:
16+
17+
- [tidb](../sources/tidb.md)
18+
19+
`tidb-execute-sql` takes one input parameter `sql` and run the sql
20+
statement against the `source`.
21+
22+
> **Note:** This tool is intended for developer assistant workflows with
23+
> human-in-the-loop and shouldn't be used for production agents.
24+
25+
## Example
26+
27+
```yaml
28+
tools:
29+
execute_sql_tool:
30+
kind: tidb-execute-sql
31+
source: my-tidb-instance
32+
description: Use this tool to execute sql statement.
33+
```
34+
35+
## Reference
36+
37+
| **field** | **type** | **required** | **description** |
38+
|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------|
39+
| kind | string | true | Must be "tidb-execute-sql". |
40+
| source | string | true | Name of the source the SQL should execute on. |
41+
| description | string | true | Description of the tool that is passed to the LLM. |
Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
---
2+
title: "tidb-sql"
3+
type: docs
4+
weight: 1
5+
description: >
6+
A "tidb-sql" tool executes a pre-defined SQL statement against a TiDB
7+
database.
8+
aliases:
9+
- /resources/tools/tidb-sql
10+
---
11+
12+
## About
13+
14+
A `tidb-sql` tool executes a pre-defined SQL statement against a TiDB
15+
database. It's compatible with the following source:
16+
17+
- [tidb](../sources/tidb.md)
18+
19+
The specified SQL statement is executed as a [prepared statement][tidb-prepare],
20+
and expects parameters in the SQL query to be in the form of placeholders `?`.
21+
22+
[tidb-prepare]: https://docs.pingcap.com/tidb/stable/sql-prepared-plan-cache
23+
24+
## Example
25+
26+
> **Note:** This tool uses parameterized queries to prevent SQL injections.
27+
> Query parameters can be used as substitutes for arbitrary expressions.
28+
> Parameters cannot be used as substitutes for identifiers, column names, table
29+
> names, or other parts of the query.
30+
31+
```yaml
32+
tools:
33+
search_flights_by_number:
34+
kind: tidb-sql
35+
source: my-tidb-instance
36+
statement: |
37+
SELECT * FROM flights
38+
WHERE airline = ?
39+
AND flight_number = ?
40+
LIMIT 10
41+
description: |
42+
Use this tool to get information for a specific flight.
43+
Takes an airline code and flight number and returns info on the flight.
44+
Do NOT use this tool with a flight id. Do NOT guess an airline code or flight number.
45+
A airline code is a code for an airline service consisting of two-character
46+
airline designator and followed by flight number, which is 1 to 4 digit number.
47+
For example, if given CY 0123, the airline is "CY", and flight_number is "123".
48+
Another example for this is DL 1234, the airline is "DL", and flight_number is "1234".
49+
If the tool returns more than one option choose the date closes to today.
50+
Example:
51+
{{
52+
"airline": "CY",
53+
"flight_number": "888",
54+
}}
55+
Example:
56+
{{
57+
"airline": "DL",
58+
"flight_number": "1234",
59+
}}
60+
parameters:
61+
- name: airline
62+
type: string
63+
description: Airline unique 2 letter identifier
64+
- name: flight_number
65+
type: string
66+
description: 1 to 4 digit number
67+
```
68+
69+
### Example with Template Parameters
70+
71+
> **Note:** This tool allows direct modifications to the SQL statement,
72+
> including identifiers, column names, and table names. **This makes it more
73+
> vulnerable to SQL injections**. Using basic parameters only (see above) is
74+
> recommended for performance and safety reasons. For more details, please check
75+
> [templateParameters](_index#template-parameters).
76+
77+
```yaml
78+
tools:
79+
list_table:
80+
kind: tidb-sql
81+
source: my-tidb-instance
82+
statement: |
83+
SELECT * FROM {{.tableName}};
84+
description: |
85+
Use this tool to list all information from a specific table.
86+
Example:
87+
{{
88+
"tableName": "flights",
89+
}}
90+
templateParameters:
91+
- name: tableName
92+
type: string
93+
description: Table to select from
94+
```
95+
96+
## Reference
97+
98+
| **field** | **type** | **required** | **description** |
99+
|--------------------|:------------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------------------------------------------------|
100+
| kind | string | true | Must be "tidb-sql". |
101+
| source | string | true | Name of the source the SQL should execute on. |
102+
| description | string | true | Description of the tool that is passed to the LLM. |
103+
| statement | string | true | SQL statement to execute on. |
104+
| parameters | [parameters](_index#specifying-parameters) | false | List of [parameters](_index#specifying-parameters) that will be inserted into the SQL statement. |
105+
| templateParameters | [templateParameters](_index#template-parameters) | false | List of [templateParameters](_index#template-parameters) that will be inserted into the SQL statement before executing prepared statement. |

0 commit comments

Comments
 (0)