mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-08 15:14:00 -05:00
feat(singlestore): Add SingleStore Source and Tools (#1333)
## Description --- - This PR adds SingleStore database source and tools. The code is mostly based on MySQL source and tools, and it uses the same go-mysql driver. - https://github.com/singlestore-labs/singlestoredb-dev-image can be used to deploy a test SingleStore instance. In this PR the default port is set to 3308 so the command would be ```docker run \ -d --name singlestoredb-dev \ -e ROOT_PASSWORD="YOUR SINGLESTORE ROOT PASSWORD" \ -p 3308:3306 ghcr.io/singlestore-labs/singlestoredb-dev:latest ``` ## PR Checklist --- > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [x] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [x] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/langchain-google-alloydb-pg-python/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [x] Make sure to add `!` if this involve a breaking change 🛠️ Fixes https://github.com/googleapis/genai-toolbox/issues/1348 --------- Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>
This commit is contained in:
Pavlo Mishchenko
@@ -784,6 +784,28 @@ steps:
|
||||
"Serverless Spark" \
|
||||
serverlessspark
|
||||
|
||||
- id: "singlestore"
|
||||
name: golang:1
|
||||
waitFor: ["compile-test-binary"]
|
||||
entrypoint: /bin/bash
|
||||
env:
|
||||
- "GOPATH=/gopath"
|
||||
- "SINGLESTORE_PORT=$_SINGLESTORE_PORT"
|
||||
- "SINGLESTORE_USER=$_SINGLESTORE_USER"
|
||||
- "SINGLESTORE_DATABASE=$_SINGLESTORE_DATABASE"
|
||||
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
|
||||
secretEnv: ["SINGLESTORE_PASSWORD", "SINGLESTORE_HOST", "CLIENT_ID"]
|
||||
volumes:
|
||||
- name: "go"
|
||||
path: "/gopath"
|
||||
args:
|
||||
- -c
|
||||
- |
|
||||
.ci/test_with_coverage.sh \
|
||||
"SingleStore" \
|
||||
singlestore \
|
||||
singlestore
|
||||
|
||||
availableSecrets:
|
||||
secretManager:
|
||||
- versionName: projects/$PROJECT_ID/secrets/cloud_sql_pg_user/versions/latest
|
||||
@@ -892,6 +914,10 @@ availableSecrets:
|
||||
env: ORACLE_PASS
|
||||
- versionName: projects/$PROJECT_ID/secrets/oracle_host/versions/latest
|
||||
env: ORACLE_HOST
|
||||
- versionName: projects/$PROJECT_ID/secrets/singlestore_pass/versions/latest
|
||||
env: SINGLESTORE_PASSWORD
|
||||
- versionName: projects/$PROJECT_ID/secrets/singlestore_host/versions/latest
|
||||
env: SINGLESTORE_HOST
|
||||
|
||||
options:
|
||||
logging: CLOUD_LOGGING_ONLY
|
||||
@@ -947,3 +973,7 @@ substitutions:
|
||||
_YUGABYTEDB_PORT: "5433"
|
||||
_YUGABYTEDB_LOADBALANCE: "false"
|
||||
_ORACLE_SERVER_NAME: "FREEPDB1"
|
||||
_SINGLESTORE_HOST: 127.0.0.1
|
||||
_SINGLESTORE_PORT: "3308"
|
||||
_SINGLESTORE_DATABASE: "singlestore"
|
||||
_SINGLESTORE_USER: "root"
|
||||
|
||||
@@ -168,6 +168,8 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparkcancelbatch"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparkgetbatch"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparklistbatches"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/singlestore/singlestoreexecutesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/singlestore/singlestoresql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerexecutesql"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerlisttables"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannersql"
|
||||
@@ -212,6 +214,7 @@ import (
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/postgres"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/redis"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/serverlessspark"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/singlestore"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/spanner"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/sqlite"
|
||||
_ "github.com/googleapis/genai-toolbox/internal/sources/tidb"
|
||||
|
||||
63
docs/en/resources/sources/singlestore.md
Normal file
63
docs/en/resources/sources/singlestore.md
Normal file
@@ -0,0 +1,63 @@
|
||||
---
|
||||
title: "SingleStore"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
SingleStore is the cloud-native database built with speed and scale to power data-intensive applications.
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
[SingleStore][singlestore-docs] is a distributed SQL database built to power intelligent applications. It is both relational and multi-model, enabling developers to easily build and scale applications and workloads.
|
||||
|
||||
SingleStore is built around Universal Storage which combines in-memory rowstore and on-disk columnstore data formats to deliver a single table type that is optimized to handle both transactional and analytical workloads.
|
||||
|
||||
[singlestore-docs]: https://docs.singlestore.com/
|
||||
|
||||
## Available Tools
|
||||
|
||||
- [`singlestore-sql`](../tools/singlestore/singlestore-sql.md)
|
||||
Execute pre-defined prepared SQL queries in SingleStore.
|
||||
|
||||
- [`singlestore-execute-sql`](../tools/singlestore/singlestore-execute-sql.md)
|
||||
Run parameterized SQL queries in SingleStore.
|
||||
|
||||
## Requirements
|
||||
|
||||
### Database User
|
||||
|
||||
This source only uses standard authentication. You will need to [create a
|
||||
database user][singlestore-user] to login to the database with.
|
||||
|
||||
[singlestore-user]: https://docs.singlestore.com/cloud/reference/sql-reference/security-management-commands/create-user/
|
||||
|
||||
## Example
|
||||
|
||||
```yaml
|
||||
sources:
|
||||
my-singlestore-source:
|
||||
kind: singlestore
|
||||
host: 127.0.0.1
|
||||
port: 3306
|
||||
database: my_db
|
||||
user: ${USER_NAME}
|
||||
password: ${PASSWORD}
|
||||
queryTimeout: 30s # Optional: query timeout duration
|
||||
```
|
||||
|
||||
{{< 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 "singlestore". |
|
||||
| host | string | true | IP address to connect to (e.g. "127.0.0.1"). |
|
||||
| port | string | true | Port to connect to (e.g. "3306"). |
|
||||
| database | string | true | Name of the SingleStore database to connect to (e.g. "my_db"). |
|
||||
| user | string | true | Name of the SingleStore database user to connect as (e.g. "admin"). |
|
||||
| password | string | true | Password of the SingleStore database user. |
|
||||
| queryTimeout | string | false | Maximum time to wait for query execution (e.g. "30s", "2m"). By default, no timeout is applied. |
|
||||
7
docs/en/resources/tools/singlestore/_index.md
Normal file
7
docs/en/resources/tools/singlestore/_index.md
Normal file
@@ -0,0 +1,7 @@
|
||||
---
|
||||
title: "SingleStore"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
Tools that work with SingleStore Sources
|
||||
---
|
||||
@@ -0,0 +1,41 @@
|
||||
---
|
||||
title: "singlestore-execute-sql"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
A "singlestore-execute-sql" tool executes a SQL statement against a SingleStore
|
||||
database.
|
||||
aliases:
|
||||
- /resources/tools/singlestore-execute-sql
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
A `singlestore-execute-sql` tool executes a SQL statement against a SingleStore
|
||||
database. It's compatible with the following sources:
|
||||
|
||||
- [singlestore](../../sources/singlestore.md)
|
||||
|
||||
`singlestore-execute-sql` takes one input parameter `sql` and runs 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: singlestore-execute-sql
|
||||
source: my-s2-instance
|
||||
description: Use this tool to execute sql statement
|
||||
```
|
||||
|
||||
## Reference
|
||||
|
||||
| **field** | **type** | **required** | **description** |
|
||||
|-------------|:------------------------------------------:|:------------:|--------------------------------------------------------------------------------------------------|
|
||||
| kind | string | true | Must be "singlestore-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. |
|
||||
102
docs/en/resources/tools/singlestore/singlestore-sql.md
Normal file
102
docs/en/resources/tools/singlestore/singlestore-sql.md
Normal file
@@ -0,0 +1,102 @@
|
||||
---
|
||||
title: "singlestore-sql"
|
||||
type: docs
|
||||
weight: 1
|
||||
description: >
|
||||
A "singlestore-sql" tool executes a pre-defined SQL statement against a SingleStore
|
||||
database.
|
||||
aliases:
|
||||
- /resources/tools/singlestore-sql
|
||||
---
|
||||
|
||||
## About
|
||||
|
||||
A `singlestore-execute-sql` tool executes a SQL statement against a SingleStore
|
||||
database. It's compatible with the following sources:
|
||||
|
||||
- [singlestore](../../sources/singlestore.md)
|
||||
|
||||
The specified SQL statement expects parameters in the SQL query to be in the form of placeholders `?`.
|
||||
|
||||
## 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: singlestore-sql
|
||||
source: my-s2-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](..#template-parameters).
|
||||
|
||||
```yaml
|
||||
tools:
|
||||
list_table:
|
||||
kind: singlestore-sql
|
||||
source: my-s2-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 "singlestore-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](../#specifying-parameters) | false | List of [parameters](../#specifying-parameters) that will be inserted into the SQL statement. |
|
||||
| templateParameters | [templateParameters](..#template-parameters) | false | List of [templateParameters](..#template-parameters) that will be inserted into the SQL statement before executing prepared statement. |
|
||||
@@ -47,6 +47,7 @@ var expectedToolSources = []string{
|
||||
"oceanbase",
|
||||
"postgres",
|
||||
"serverless-spark",
|
||||
"singlestore",
|
||||
"spanner-postgres",
|
||||
"spanner",
|
||||
"sqlite",
|
||||
@@ -118,6 +119,7 @@ func TestGetPrebuiltTool(t *testing.T) {
|
||||
mssql_config, _ := Get("mssql")
|
||||
oceanbase_config, _ := Get("oceanbase")
|
||||
postgresconfig, _ := Get("postgres")
|
||||
singlestore_config, _ := Get("singlestore")
|
||||
spanner_config, _ := Get("spanner")
|
||||
spannerpg_config, _ := Get("spanner-postgres")
|
||||
mindsdb_config, _ := Get("mindsdb")
|
||||
@@ -190,6 +192,9 @@ func TestGetPrebuiltTool(t *testing.T) {
|
||||
if len(postgresconfig) <= 0 {
|
||||
t.Fatalf("unexpected error: could not fetch postgres prebuilt tools yaml")
|
||||
}
|
||||
if len(singlestore_config) <= 0 {
|
||||
t.Fatalf("unexpected error: could not fetch singlestore prebuilt tools yaml")
|
||||
}
|
||||
if len(spanner_config) <= 0 {
|
||||
t.Fatalf("unexpected error: could not fetch spanner prebuilt tools yaml")
|
||||
}
|
||||
|
||||
193
internal/prebuiltconfigs/tools/singlestore.yaml
Normal file
193
internal/prebuiltconfigs/tools/singlestore.yaml
Normal file
@@ -0,0 +1,193 @@
|
||||
# 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.
|
||||
|
||||
sources:
|
||||
singlestore-source:
|
||||
kind: singlestore
|
||||
host: ${SINGLESTORE_HOST}
|
||||
port: ${SINGLESTORE_PORT}
|
||||
database: ${SINGLESTORE_DATABASE}
|
||||
user: ${SINGLESTORE_USER}
|
||||
password: ${SINGLESTORE_PASSWORD}
|
||||
queryTimeout: 30s # Optional
|
||||
tools:
|
||||
execute_sql:
|
||||
kind: singlestore-execute-sql
|
||||
source: singlestore-source
|
||||
description: Use this tool to execute SQL.
|
||||
list_tables:
|
||||
kind: singlestore-sql
|
||||
source: singlestore-source
|
||||
description: "Lists detailed schema information (object type, columns, constraints, indexes, triggers, comment) as JSON for user-created tables (ordinary or partitioned). Filters by a comma-separated list of names. If names are omitted, lists all tables in user schemas."
|
||||
statement: |
|
||||
WITH constraint_columns_cte AS (
|
||||
SELECT
|
||||
KCU.CONSTRAINT_SCHEMA,
|
||||
KCU.CONSTRAINT_NAME,
|
||||
KCU.TABLE_NAME,
|
||||
JSON_AGG(KCU.COLUMN_NAME ORDER BY KCU.ORDINAL_POSITION) AS constraint_columns
|
||||
FROM
|
||||
INFORMATION_SCHEMA.KEY_COLUMN_USAGE KCU
|
||||
GROUP BY
|
||||
KCU.CONSTRAINT_SCHEMA, KCU.CONSTRAINT_NAME, KCU.TABLE_NAME
|
||||
),
|
||||
foreign_key_columns_cte AS (
|
||||
SELECT
|
||||
FKCU.CONSTRAINT_SCHEMA,
|
||||
FKCU.CONSTRAINT_NAME,
|
||||
FKCU.TABLE_NAME,
|
||||
JSON_AGG(FKCU.REFERENCED_COLUMN_NAME ORDER BY FKCU.ORDINAL_POSITION) AS foreign_key_referenced_columns
|
||||
FROM
|
||||
INFORMATION_SCHEMA.KEY_COLUMN_USAGE FKCU
|
||||
WHERE
|
||||
FKCU.REFERENCED_TABLE_NAME IS NOT NULL
|
||||
GROUP BY
|
||||
FKCU.CONSTRAINT_SCHEMA, FKCU.CONSTRAINT_NAME, FKCU.TABLE_NAME
|
||||
),
|
||||
table_owners AS (
|
||||
SELECT DISTINCT
|
||||
U.TABLE_SCHEMA,
|
||||
FIRST_VALUE(IFNULL(U.GRANTEE, 'N/A')) OVER (PARTITION BY U.TABLE_SCHEMA ORDER BY U.GRANTEE) AS owner
|
||||
FROM
|
||||
INFORMATION_SCHEMA.SCHEMA_PRIVILEGES U
|
||||
),
|
||||
table_columns AS (
|
||||
SELECT
|
||||
C.TABLE_SCHEMA,
|
||||
C.TABLE_NAME,
|
||||
JSON_AGG(
|
||||
JSON_BUILD_OBJECT(
|
||||
'column_name', C.COLUMN_NAME,
|
||||
'data_type', C.COLUMN_TYPE,
|
||||
'ordinal_position', C.ORDINAL_POSITION,
|
||||
'is_not_nullable', IF(C.IS_NULLABLE = 'NO', TRUE, FALSE),
|
||||
'column_default', C.COLUMN_DEFAULT,
|
||||
'column_comment', IFNULL(C.COLUMN_COMMENT, '')
|
||||
) ORDER BY C.ORDINAL_POSITION
|
||||
) AS columns_json
|
||||
FROM
|
||||
INFORMATION_SCHEMA.COLUMNS C
|
||||
GROUP BY
|
||||
C.TABLE_SCHEMA, C.TABLE_NAME
|
||||
),
|
||||
table_indexes AS (
|
||||
SELECT
|
||||
S.TABLE_SCHEMA,
|
||||
S.TABLE_NAME,
|
||||
JSON_AGG(
|
||||
JSON_BUILD_OBJECT(
|
||||
'index_name', S.INDEX_NAME,
|
||||
'is_unique', IF(S.NON_UNIQUE = 0, TRUE, FALSE),
|
||||
'is_primary', IF(S.INDEX_NAME = 'PRIMARY', TRUE, FALSE),
|
||||
'index_columns', S.INDEX_COLUMNS_ARRAY
|
||||
)
|
||||
) AS indexes_json
|
||||
FROM (
|
||||
SELECT
|
||||
S.TABLE_SCHEMA,
|
||||
S.TABLE_NAME,
|
||||
S.INDEX_NAME,
|
||||
MIN(S.NON_UNIQUE) AS NON_UNIQUE,
|
||||
JSON_AGG(S.COLUMN_NAME ORDER BY S.SEQ_IN_INDEX) AS INDEX_COLUMNS_ARRAY
|
||||
FROM
|
||||
INFORMATION_SCHEMA.STATISTICS S
|
||||
GROUP BY
|
||||
S.TABLE_SCHEMA, S.TABLE_NAME, S.INDEX_NAME
|
||||
) S
|
||||
GROUP BY
|
||||
S.TABLE_SCHEMA, S.TABLE_NAME
|
||||
),
|
||||
table_constraints AS (
|
||||
SELECT
|
||||
TC.TABLE_SCHEMA,
|
||||
TC.TABLE_NAME,
|
||||
JSON_AGG(
|
||||
JSON_BUILD_OBJECT(
|
||||
'constraint_name', TC.CONSTRAINT_NAME,
|
||||
'constraint_type',
|
||||
CASE TC.CONSTRAINT_TYPE
|
||||
WHEN 'PRIMARY KEY' THEN 'PRIMARY KEY'
|
||||
WHEN 'FOREIGN KEY' THEN 'FOREIGN KEY'
|
||||
WHEN 'UNIQUE' THEN 'UNIQUE'
|
||||
ELSE TC.CONSTRAINT_TYPE
|
||||
END,
|
||||
'constraint_definition', '',
|
||||
'constraint_columns', IFNULL(CC.constraint_columns, JSON_BUILD_ARRAY()),
|
||||
'foreign_key_referenced_table', IF(TC.CONSTRAINT_TYPE = 'FOREIGN KEY', RC.REFERENCED_TABLE_NAME, NULL),
|
||||
'foreign_key_referenced_columns', IF(TC.CONSTRAINT_TYPE = 'FOREIGN KEY', IFNULL(FKC.foreign_key_referenced_columns, JSON_BUILD_ARRAY()), NULL)
|
||||
)
|
||||
) AS constraints_json
|
||||
FROM
|
||||
INFORMATION_SCHEMA.TABLE_CONSTRAINTS TC
|
||||
LEFT JOIN INFORMATION_SCHEMA.REFERENTIAL_CONSTRAINTS RC
|
||||
ON TC.CONSTRAINT_SCHEMA = RC.CONSTRAINT_SCHEMA
|
||||
AND TC.CONSTRAINT_NAME = RC.CONSTRAINT_NAME
|
||||
AND TC.TABLE_NAME = RC.TABLE_NAME
|
||||
LEFT JOIN constraint_columns_cte CC
|
||||
ON TC.CONSTRAINT_SCHEMA = CC.CONSTRAINT_SCHEMA
|
||||
AND TC.CONSTRAINT_NAME = CC.CONSTRAINT_NAME
|
||||
AND TC.TABLE_NAME = CC.TABLE_NAME
|
||||
LEFT JOIN foreign_key_columns_cte FKC
|
||||
ON TC.CONSTRAINT_SCHEMA = FKC.CONSTRAINT_SCHEMA
|
||||
AND TC.CONSTRAINT_NAME = FKC.CONSTRAINT_NAME
|
||||
AND TC.TABLE_NAME = FKC.TABLE_NAME
|
||||
GROUP BY
|
||||
TC.TABLE_SCHEMA, TC.TABLE_NAME
|
||||
)
|
||||
SELECT
|
||||
T.TABLE_SCHEMA AS schema_name,
|
||||
T.TABLE_NAME AS object_name,
|
||||
JSON_BUILD_OBJECT(
|
||||
'schema_name', T.TABLE_SCHEMA,
|
||||
'object_name', T.TABLE_NAME,
|
||||
'object_type', 'TABLE',
|
||||
'owner', IFNULL(TOW.owner, 'N/A'),
|
||||
'comment', IFNULL(T.TABLE_COMMENT, ''),
|
||||
'columns', IFNULL(TC.columns_json, JSON_BUILD_ARRAY()),
|
||||
'indexes', IFNULL(TI.indexes_json, JSON_BUILD_ARRAY()),
|
||||
'constraints', IFNULL(TCN.constraints_json, JSON_BUILD_ARRAY()),
|
||||
'triggers', JSON_BUILD_ARRAY()
|
||||
) AS object_details
|
||||
FROM
|
||||
INFORMATION_SCHEMA.TABLES T
|
||||
CROSS JOIN (SELECT ? AS table_names_param) AS variables
|
||||
LEFT JOIN table_owners TOW
|
||||
ON T.TABLE_SCHEMA = TOW.TABLE_SCHEMA
|
||||
LEFT JOIN table_columns TC
|
||||
ON T.TABLE_SCHEMA = TC.TABLE_SCHEMA
|
||||
AND T.TABLE_NAME = TC.TABLE_NAME
|
||||
LEFT JOIN table_indexes TI
|
||||
ON T.TABLE_SCHEMA = TI.TABLE_SCHEMA
|
||||
AND T.TABLE_NAME = TI.TABLE_NAME
|
||||
LEFT JOIN table_constraints TCN
|
||||
ON T.TABLE_SCHEMA = TCN.TABLE_SCHEMA
|
||||
AND T.TABLE_NAME = TCN.TABLE_NAME
|
||||
WHERE
|
||||
T.TABLE_SCHEMA NOT IN ('cluster', 'information_schema', 'memsql')
|
||||
AND T.TABLE_TYPE = 'BASE TABLE'
|
||||
AND (NULLIF(TRIM(variables.table_names_param), '') IS NULL OR
|
||||
CONCAT(',', variables.table_names_param, ',') LIKE CONCAT('%,', T.TABLE_NAME, ',%'))
|
||||
ORDER BY
|
||||
T.TABLE_SCHEMA, T.TABLE_NAME
|
||||
|
||||
parameters:
|
||||
- name: table_names
|
||||
type: string
|
||||
description: "Optional: A comma-separated list of table names. If empty, details for all tables in user-accessible schemas will be listed."
|
||||
default: ""
|
||||
|
||||
toolsets:
|
||||
singlestore-database-tools:
|
||||
- execute_sql
|
||||
- list_tables
|
||||
140
internal/sources/singlestore/singlestore.go
Normal file
140
internal/sources/singlestore/singlestore.go
Normal file
@@ -0,0 +1,140 @@
|
||||
// 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 singlestore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"net/url"
|
||||
"fmt"
|
||||
"time"
|
||||
"strings"
|
||||
|
||||
_ "github.com/go-sql-driver/mysql"
|
||||
"github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
)
|
||||
|
||||
// SourceKind for SingleStore source
|
||||
const SourceKind string = "singlestore"
|
||||
|
||||
// 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
|
||||
}
|
||||
return actual, nil
|
||||
}
|
||||
|
||||
// Config holds the configuration parameters for connecting to a SingleStore database.
|
||||
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"`
|
||||
QueryTimeout string `yaml:"queryTimeout"`
|
||||
}
|
||||
|
||||
// SourceConfigKind returns the kind of the source configuration.
|
||||
func (r Config) SourceConfigKind() string {
|
||||
return SourceKind
|
||||
}
|
||||
|
||||
// Initialize sets up the SingleStore connection pool and returns a Source.
|
||||
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
|
||||
pool, err := initSingleStoreConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database, r.QueryTimeout)
|
||||
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{}
|
||||
|
||||
// Source represents a SingleStore database source and holds its connection pool.
|
||||
type Source struct {
|
||||
Name string `yaml:"name"`
|
||||
Kind string `yaml:"kind"`
|
||||
Pool *sql.DB
|
||||
}
|
||||
|
||||
// SourceKind returns the kind of the source configuration.
|
||||
func (s *Source) SourceKind() string {
|
||||
return SourceKind
|
||||
}
|
||||
|
||||
// SingleStorePool returns the underlying *sql.DB connection pool for SingleStore.
|
||||
func (s *Source) SingleStorePool() *sql.DB {
|
||||
return s.Pool
|
||||
}
|
||||
|
||||
func initSingleStoreConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, queryTimeout string) (*sql.DB, error) {
|
||||
//nolint:all // Reassigned ctx
|
||||
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
|
||||
defer span.End()
|
||||
|
||||
// Configure the driver to connect to the database
|
||||
dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?parseTime=true&vector_type_project_format=JSON", user, pass, host, port, dbname)
|
||||
|
||||
// Add connection attributes to DSN
|
||||
customAttrs := []string{"_connector_name"}
|
||||
customAttrValues := []string{"MCP toolbox for Databases"}
|
||||
|
||||
customAttrStrs := make([]string, len(customAttrs))
|
||||
for i := range customAttrs {
|
||||
customAttrStrs[i] = fmt.Sprintf("%s:%s", customAttrs[i], customAttrValues[i])
|
||||
}
|
||||
dsn += "&connectionAttributes=" + url.QueryEscape(strings.Join(customAttrStrs, ","))
|
||||
|
||||
// Add query timeout to DSN if specified
|
||||
if queryTimeout != "" {
|
||||
timeout, err := time.ParseDuration(queryTimeout)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid queryTimeout %q: %w", queryTimeout, err)
|
||||
}
|
||||
dsn += "&readTimeout=" + timeout.String()
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
153
internal/sources/singlestore/singlestore_test.go
Normal file
153
internal/sources/singlestore/singlestore_test.go
Normal file
@@ -0,0 +1,153 @@
|
||||
package singlestore_test
|
||||
|
||||
// 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.
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/singlestore"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
)
|
||||
|
||||
func TestParseFromYaml(t *testing.T) {
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.SourceConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
sources:
|
||||
my-s2-instance:
|
||||
kind: singlestore
|
||||
host: 0.0.0.0
|
||||
port: my-port
|
||||
database: my_db
|
||||
user: my_user
|
||||
password: my_pass
|
||||
`,
|
||||
want: server.SourceConfigs{
|
||||
"my-s2-instance": singlestore.Config{
|
||||
Name: "my-s2-instance",
|
||||
Kind: singlestore.SourceKind,
|
||||
Host: "0.0.0.0",
|
||||
Port: "my-port",
|
||||
Database: "my_db",
|
||||
User: "my_user",
|
||||
Password: "my_pass",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
desc: "with query timeout",
|
||||
in: `
|
||||
sources:
|
||||
my-s2-instance:
|
||||
kind: singlestore
|
||||
host: 0.0.0.0
|
||||
port: my-port
|
||||
database: my_db
|
||||
user: my_user
|
||||
password: my_pass
|
||||
queryTimeout: 45s
|
||||
`,
|
||||
want: server.SourceConfigs{
|
||||
"my-s2-instance": singlestore.Config{
|
||||
Name: "my-s2-instance",
|
||||
Kind: singlestore.SourceKind,
|
||||
Host: "0.0.0.0",
|
||||
Port: "my-port",
|
||||
Database: "my_db",
|
||||
User: "my_user",
|
||||
Password: "my_pass",
|
||||
QueryTimeout: "45s",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Sources server.SourceConfigs `yaml:"sources"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if !cmp.Equal(tc.want, got.Sources) {
|
||||
t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func TestFailParseFromYaml(t *testing.T) {
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
err string
|
||||
}{
|
||||
{
|
||||
desc: "extra field",
|
||||
in: `
|
||||
sources:
|
||||
my-s2-instance:
|
||||
kind: singlestore
|
||||
host: 0.0.0.0
|
||||
port: my-port
|
||||
database: my_db
|
||||
user: my_user
|
||||
password: my_pass
|
||||
foo: bar
|
||||
`,
|
||||
err: "unable to parse source \"my-s2-instance\" as \"singlestore\": [2:1] unknown field \"foo\"\n 1 | database: my_db\n> 2 | foo: bar\n ^\n 3 | host: 0.0.0.0\n 4 | kind: singlestore\n 5 | password: my_pass\n 6 | ",
|
||||
},
|
||||
{
|
||||
desc: "missing required field",
|
||||
in: `
|
||||
sources:
|
||||
my-s2-instance:
|
||||
kind: singlestore
|
||||
port: my-port
|
||||
database: my_db
|
||||
user: my_user
|
||||
password: my_pass
|
||||
`,
|
||||
err: "unable to parse source \"my-s2-instance\" as \"singlestore\": Key: 'Config.Host' Error:Field validation for 'Host' failed on the 'required' tag",
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Sources server.SourceConfigs `yaml:"sources"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
|
||||
if err == nil {
|
||||
t.Fatalf("expect parsing to fail")
|
||||
}
|
||||
errStr := err.Error()
|
||||
if errStr != tc.err {
|
||||
t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
// 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 singlestoreexecutesql
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/singlestore"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlcommon"
|
||||
"github.com/googleapis/genai-toolbox/internal/util"
|
||||
)
|
||||
|
||||
const kind string = "singlestore-execute-sql"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
SingleStorePool() *sql.DB
|
||||
}
|
||||
|
||||
// validate compatible sources are still compatible
|
||||
var _ compatibleSource = &singlestore.Source{}
|
||||
|
||||
var compatibleSources = [...]string{singlestore.SourceKind}
|
||||
|
||||
// Config represents the configuration for the singlestore-execute-sql tool.
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Description string `yaml:"description" validate:"required"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
// ToolConfigKind returns the kind of the tool configuration.
|
||||
func (cfg Config) ToolConfigKind() string {
|
||||
return kind
|
||||
}
|
||||
|
||||
// Initialize sets up the Tool using the provided sources map.
|
||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||
// verify source exists
|
||||
rawS, ok := srcs[cfg.Source]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
|
||||
}
|
||||
|
||||
// verify the source is compatible
|
||||
s, ok := rawS.(compatibleSource)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
sqlParameter := tools.NewStringParameter("sql", "The sql to execute.")
|
||||
parameters := tools.Parameters{sqlParameter}
|
||||
|
||||
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters)
|
||||
|
||||
// finish tool setup
|
||||
t := Tool{
|
||||
Name: cfg.Name,
|
||||
Kind: kind,
|
||||
Parameters: parameters,
|
||||
AuthRequired: cfg.AuthRequired,
|
||||
Pool: s.SingleStorePool(),
|
||||
manifest: tools.Manifest{Description: cfg.Description, Parameters: parameters.Manifest(), AuthRequired: cfg.AuthRequired},
|
||||
mcpManifest: mcpManifest,
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
// Tool represents a tool for executing SQL queries on a SingleStore database.
|
||||
type Tool struct {
|
||||
Name string `yaml:"name"`
|
||||
Kind string `yaml:"kind"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
Parameters tools.Parameters `yaml:"parameters"`
|
||||
|
||||
Pool *sql.DB
|
||||
manifest tools.Manifest
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
// Invoke executes the provided SQL query using the tool's database connection and returns the results.
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
|
||||
paramsMap := params.AsMap()
|
||||
sql, ok := paramsMap["sql"].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unable to get cast %s", paramsMap["sql"])
|
||||
}
|
||||
|
||||
// Log the query executed for debugging.
|
||||
logger, err := util.LoggerFromContext(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error getting logger: %s", err)
|
||||
}
|
||||
logger.DebugContext(ctx, "executing `%s` tool query: %s", kind, sql)
|
||||
|
||||
results, err := t.Pool.QueryContext(ctx, sql)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
cols, err := results.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
|
||||
}
|
||||
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
|
||||
colTypes, err := results.ColumnTypes()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get column types: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
err := results.Scan(values...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
vMap[name] = nil
|
||||
continue
|
||||
}
|
||||
|
||||
vMap[name], err = mysqlcommon.ConvertToType(colTypes[i], val)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("errors encountered when converting values: %w", err)
|
||||
}
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
return tools.ParseParams(t.Parameters, data, claims)
|
||||
}
|
||||
|
||||
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() bool {
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
// 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 singlestoreexecutesql_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/singlestore/singlestoreexecutesql"
|
||||
)
|
||||
|
||||
func TestParseFromYamlExecuteSql(t *testing.T) {
|
||||
ctx, err := testutils.ContextWithNewLogger()
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.ToolConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
tools:
|
||||
example_tool:
|
||||
kind: singlestore-execute-sql
|
||||
source: my-instance
|
||||
description: some description
|
||||
authRequired:
|
||||
- my-google-auth-service
|
||||
- other-auth-service
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": singlestoreexecutesql.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "singlestore-execute-sql",
|
||||
Source: "my-instance",
|
||||
Description: "some description",
|
||||
AuthRequired: []string{"my-google-auth-service", "other-auth-service"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Tools server.ToolConfigs `yaml:"tools"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
}
|
||||
234
internal/tools/singlestore/singlestoresql/singlestoresql.go
Normal file
234
internal/tools/singlestore/singlestoresql/singlestoresql.go
Normal file
@@ -0,0 +1,234 @@
|
||||
// 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 singlestoresql
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources"
|
||||
"github.com/googleapis/genai-toolbox/internal/sources/singlestore"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlcommon"
|
||||
)
|
||||
|
||||
const kind string = "singlestore-sql"
|
||||
|
||||
func init() {
|
||||
if !tools.Register(kind, newConfig) {
|
||||
panic(fmt.Sprintf("tool kind %q already registered", kind))
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
SingleStorePool() *sql.DB
|
||||
}
|
||||
|
||||
// validate compatible sources are still compatible
|
||||
var _ compatibleSource = &singlestore.Source{}
|
||||
|
||||
var compatibleSources = [...]string{singlestore.SourceKind}
|
||||
|
||||
// Config defines the configuration for a SingleStore SQL tool.
|
||||
type Config struct {
|
||||
Name string `yaml:"name" validate:"required"`
|
||||
Kind string `yaml:"kind" validate:"required"`
|
||||
Source string `yaml:"source" validate:"required"`
|
||||
Description string `yaml:"description" validate:"required"`
|
||||
Statement string `yaml:"statement" validate:"required"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
Parameters tools.Parameters `yaml:"parameters"`
|
||||
TemplateParameters tools.Parameters `yaml:"templateParameters"`
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.ToolConfig = Config{}
|
||||
|
||||
// ToolConfigKind returns the kind of the tool configuration.
|
||||
func (cfg Config) ToolConfigKind() string {
|
||||
return kind
|
||||
}
|
||||
|
||||
// Initialize sets up and returns a new Tool instance based on the provided configuration and available sources.
|
||||
// It verifies that the specified source exists and is compatible, processes tool parameters, and constructs
|
||||
// the necessary manifests for tool operation. Returns an initialized Tool or an error if setup fails.
|
||||
//
|
||||
// Parameters:
|
||||
// srcs - a map of available sources, keyed by source name.
|
||||
//
|
||||
// Returns:
|
||||
// tools.Tool - the initialized tool instance.
|
||||
// error - an error if the source is missing, incompatible, or setup fails.
|
||||
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
|
||||
// verify source exists
|
||||
rawS, ok := srcs[cfg.Source]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
|
||||
}
|
||||
|
||||
// verify the source is compatible
|
||||
s, ok := rawS.(compatibleSource)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
|
||||
}
|
||||
|
||||
allParameters, paramManifest, err := tools.ProcessParameters(cfg.TemplateParameters, cfg.Parameters)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters)
|
||||
|
||||
// finish tool setup
|
||||
t := Tool{
|
||||
Name: cfg.Name,
|
||||
Kind: kind,
|
||||
Parameters: cfg.Parameters,
|
||||
TemplateParameters: cfg.TemplateParameters,
|
||||
AllParams: allParameters,
|
||||
Statement: cfg.Statement,
|
||||
AuthRequired: cfg.AuthRequired,
|
||||
Pool: s.SingleStorePool(),
|
||||
manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
|
||||
mcpManifest: mcpManifest,
|
||||
}
|
||||
return t, nil
|
||||
}
|
||||
|
||||
// validate interface
|
||||
var _ tools.Tool = Tool{}
|
||||
|
||||
// Tool represents a SingleStore SQL tool instance with its configuration, parameters, and database connection.
|
||||
type Tool struct {
|
||||
Name string `yaml:"name"`
|
||||
Kind string `yaml:"kind"`
|
||||
AuthRequired []string `yaml:"authRequired"`
|
||||
Parameters tools.Parameters `yaml:"parameters"`
|
||||
TemplateParameters tools.Parameters `yaml:"templateParameters"`
|
||||
AllParams tools.Parameters `yaml:"allParams"`
|
||||
|
||||
Pool *sql.DB
|
||||
Statement string
|
||||
manifest tools.Manifest
|
||||
mcpManifest tools.McpManifest
|
||||
}
|
||||
|
||||
// Invoke executes the SQL statement defined in the Tool using the provided context and parameter values.
|
||||
// It resolves template parameters and standard parameters, executes the query, and processes the result rows.
|
||||
// Each row is returned as a map with column names as keys and their corresponding values, handling special
|
||||
// cases for JSON and string types. Returns a slice of maps representing the result set, or an error if any
|
||||
// step fails.
|
||||
//
|
||||
// Parameters:
|
||||
// ctx - The context for controlling cancellation and timeouts.
|
||||
// params - The parameter values to be used for the SQL statement.
|
||||
//
|
||||
// Returns:
|
||||
// - A slice of maps, where each map represents a row with column names as keys.
|
||||
// - An error if template resolution, parameter extraction, query execution, or result processing fails.
|
||||
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
|
||||
paramsMap := params.AsMap()
|
||||
newStatement, err := tools.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to extract template params %w", err)
|
||||
}
|
||||
|
||||
newParams, err := tools.GetParams(t.Parameters, paramsMap)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to extract standard params %w", err)
|
||||
}
|
||||
|
||||
sliceParams := newParams.AsSlice()
|
||||
results, err := t.Pool.QueryContext(ctx, newStatement, sliceParams...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to execute query: %w", err)
|
||||
}
|
||||
|
||||
cols, err := results.Columns()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
|
||||
}
|
||||
|
||||
// create an array of values for each column, which can be re-used to scan each row
|
||||
rawValues := make([]any, len(cols))
|
||||
values := make([]any, len(cols))
|
||||
for i := range rawValues {
|
||||
values[i] = &rawValues[i]
|
||||
}
|
||||
defer results.Close()
|
||||
|
||||
colTypes, err := results.ColumnTypes()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to get column types: %w", err)
|
||||
}
|
||||
|
||||
var out []any
|
||||
for results.Next() {
|
||||
err := results.Scan(values...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to parse row: %w", err)
|
||||
}
|
||||
vMap := make(map[string]any)
|
||||
for i, name := range cols {
|
||||
val := rawValues[i]
|
||||
if val == nil {
|
||||
vMap[name] = nil
|
||||
continue
|
||||
}
|
||||
|
||||
vMap[name], err = mysqlcommon.ConvertToType(colTypes[i], val)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("errors encountered when converting values: %w", err)
|
||||
}
|
||||
}
|
||||
out = append(out, vMap)
|
||||
}
|
||||
|
||||
if err := results.Err(); err != nil {
|
||||
return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
|
||||
return tools.ParseParams(t.AllParams, data, claims)
|
||||
}
|
||||
|
||||
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() bool {
|
||||
return false
|
||||
}
|
||||
175
internal/tools/singlestore/singlestoresql/singlestoresql_test.go
Normal file
175
internal/tools/singlestore/singlestoresql/singlestoresql_test.go
Normal file
@@ -0,0 +1,175 @@
|
||||
// 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 singlestoresql_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
yaml "github.com/goccy/go-yaml"
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/googleapis/genai-toolbox/internal/server"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools"
|
||||
"github.com/googleapis/genai-toolbox/internal/tools/singlestore/singlestoresql"
|
||||
)
|
||||
|
||||
func TestParseFromYamlSingleStore(t *testing.T) {
|
||||
ctx, err := testutils.ContextWithNewLogger()
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.ToolConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
tools:
|
||||
example_tool:
|
||||
kind: singlestore-sql
|
||||
source: my-singlestore-instance
|
||||
description: some description
|
||||
statement: |
|
||||
SELECT * FROM SQL_STATEMENT;
|
||||
authRequired:
|
||||
- my-google-auth-service
|
||||
- other-auth-service
|
||||
parameters:
|
||||
- name: country
|
||||
type: string
|
||||
description: some description
|
||||
authServices:
|
||||
- name: my-google-auth-service
|
||||
field: user_id
|
||||
- name: other-auth-service
|
||||
field: user_id
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": singlestoresql.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "singlestore-sql",
|
||||
Source: "my-singlestore-instance",
|
||||
Description: "some description",
|
||||
Statement: "SELECT * FROM SQL_STATEMENT;\n",
|
||||
AuthRequired: []string{"my-google-auth-service", "other-auth-service"},
|
||||
Parameters: []tools.Parameter{
|
||||
tools.NewStringParameterWithAuth("country", "some description",
|
||||
[]tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_id"},
|
||||
{Name: "other-auth-service", Field: "user_id"}}),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Tools server.ToolConfigs `yaml:"tools"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseFromYamlWithTemplateParamsSingleStore(t *testing.T) {
|
||||
ctx, err := testutils.ContextWithNewLogger()
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %s", err)
|
||||
}
|
||||
tcs := []struct {
|
||||
desc string
|
||||
in string
|
||||
want server.ToolConfigs
|
||||
}{
|
||||
{
|
||||
desc: "basic example",
|
||||
in: `
|
||||
tools:
|
||||
example_tool:
|
||||
kind: singlestore-sql
|
||||
source: my-singlestore-instance
|
||||
description: some description
|
||||
statement: |
|
||||
SELECT * FROM SQL_STATEMENT;
|
||||
authRequired:
|
||||
- my-google-auth-service
|
||||
- other-auth-service
|
||||
parameters:
|
||||
- name: country
|
||||
type: string
|
||||
description: some description
|
||||
authServices:
|
||||
- name: my-google-auth-service
|
||||
field: user_id
|
||||
- name: other-auth-service
|
||||
field: user_id
|
||||
templateParameters:
|
||||
- name: tableName
|
||||
type: string
|
||||
description: The table to select hotels from.
|
||||
- name: fieldArray
|
||||
type: array
|
||||
description: The columns to return for the query.
|
||||
items:
|
||||
name: column
|
||||
type: string
|
||||
description: A column name that will be returned from the query.
|
||||
`,
|
||||
want: server.ToolConfigs{
|
||||
"example_tool": singlestoresql.Config{
|
||||
Name: "example_tool",
|
||||
Kind: "singlestore-sql",
|
||||
Source: "my-singlestore-instance",
|
||||
Description: "some description",
|
||||
Statement: "SELECT * FROM SQL_STATEMENT;\n",
|
||||
AuthRequired: []string{"my-google-auth-service", "other-auth-service"},
|
||||
Parameters: []tools.Parameter{
|
||||
tools.NewStringParameterWithAuth("country", "some description",
|
||||
[]tools.ParamAuthService{{Name: "my-google-auth-service", Field: "user_id"},
|
||||
{Name: "other-auth-service", Field: "user_id"}}),
|
||||
},
|
||||
TemplateParameters: []tools.Parameter{
|
||||
tools.NewStringParameter("tableName", "The table to select hotels from."),
|
||||
tools.NewArrayParameter("fieldArray", "The columns to return for the query.", tools.NewStringParameter("column", "A column name that will be returned from the query.")),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, tc := range tcs {
|
||||
t.Run(tc.desc, func(t *testing.T) {
|
||||
got := struct {
|
||||
Tools server.ToolConfigs `yaml:"tools"`
|
||||
}{}
|
||||
// Parse contents
|
||||
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to unmarshal: %s", err)
|
||||
}
|
||||
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
|
||||
t.Fatalf("incorrect parse: diff %v", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
238
tests/singlestore/singlestore_integration_test.go
Normal file
238
tests/singlestore/singlestore_integration_test.go
Normal file
@@ -0,0 +1,238 @@
|
||||
// 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 singlestore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/googleapis/genai-toolbox/internal/testutils"
|
||||
"github.com/googleapis/genai-toolbox/tests"
|
||||
)
|
||||
|
||||
var (
|
||||
SingleStoreSourceKind = "singlestore"
|
||||
SingleStoreToolKind = "singlestore-sql"
|
||||
SingleStoreDatabase = os.Getenv("SINGLESTORE_DATABASE")
|
||||
SingleStoreHost = os.Getenv("SINGLESTORE_HOST")
|
||||
SingleStorePort = os.Getenv("SINGLESTORE_PORT")
|
||||
SingleStoreUser = os.Getenv("SINGLESTORE_USER")
|
||||
SingleStorePass = os.Getenv("SINGLESTORE_PASSWORD")
|
||||
)
|
||||
|
||||
func getSingleStoreVars(t *testing.T) map[string]any {
|
||||
switch "" {
|
||||
case SingleStoreDatabase:
|
||||
t.Fatal("'SINGLESTORE_DATABASE' not set")
|
||||
case SingleStoreHost:
|
||||
t.Fatal("'SINGLESTORE_HOST' not set")
|
||||
case SingleStorePort:
|
||||
t.Fatal("'SINGLESTORE_PORT' not set")
|
||||
case SingleStoreUser:
|
||||
t.Fatal("'SINGLESTORE_USER' not set")
|
||||
case SingleStorePass:
|
||||
t.Fatal("'SINGLESTORE_PASSWORD' not set")
|
||||
}
|
||||
|
||||
return map[string]any{
|
||||
"kind": SingleStoreSourceKind,
|
||||
"host": SingleStoreHost,
|
||||
"port": SingleStorePort,
|
||||
"database": SingleStoreDatabase,
|
||||
"user": SingleStoreUser,
|
||||
"password": SingleStorePass,
|
||||
}
|
||||
}
|
||||
|
||||
// getSingleStoreParamToolInfo returns statements and params for my-tool
|
||||
func getSingleStoreParamToolInfo(tableName string) (string, string, string, string, string, string, []any) {
|
||||
createStatement := fmt.Sprintf("CREATE TABLE %s (id BIGINT NOT NULL PRIMARY KEY, name VARCHAR(255));", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (id, name) VALUES (?, ?), (?, ?), (?, ?), (?, ?);", tableName)
|
||||
toolStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = ? OR name = ? ORDER BY id;", tableName)
|
||||
idParamStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = ? ORDER BY id;", tableName)
|
||||
nameParamStatement := fmt.Sprintf("SELECT * FROM %s WHERE name = ? ORDER BY id;", tableName)
|
||||
// SingleStore doesn't support array parameters in IN clause unlike some other databases
|
||||
arrayToolStmt := ""
|
||||
insertParams := []any{1, "Alice", 2, "Jane", 3, "Sid", 4, nil}
|
||||
return createStatement, insertStatement, toolStatement, idParamStatement, nameParamStatement, arrayToolStmt, insertParams
|
||||
}
|
||||
|
||||
// getSingleStoreAuthToolInfo returns statements and param of my-auth-tool
|
||||
func getSingleStoreAuthToolInfo(tableName string) (string, string, string, []any) {
|
||||
createStatement := fmt.Sprintf("CREATE TABLE %s (id BIGINT NOT NULL AUTO_INCREMENT PRIMARY KEY, name VARCHAR(255), email VARCHAR(255));", tableName)
|
||||
insertStatement := fmt.Sprintf("INSERT INTO %s (name, email) VALUES (?, ?), (?, ?)", tableName)
|
||||
toolStatement := fmt.Sprintf("SELECT name FROM %s WHERE email = ?;", tableName)
|
||||
params := []any{"Alice", tests.ServiceAccountEmail, "Jane", "janedoe@gmail.com"}
|
||||
return createStatement, insertStatement, toolStatement, params
|
||||
}
|
||||
|
||||
// getSingleStoreTmplToolStatement returns statements and param for template parameter test cases for singlestore-sql kind
|
||||
func getSingleStoreTmplToolStatement() (string, string) {
|
||||
tmplSelectCombined := "SELECT * FROM {{.tableName}} WHERE id = ?"
|
||||
tmplSelectFilterCombined := "SELECT * FROM {{.tableName}} WHERE {{.columnFilter}} = ?"
|
||||
return tmplSelectCombined, tmplSelectFilterCombined
|
||||
}
|
||||
|
||||
// getSingleStoreWants return the expected wants for singlestore
|
||||
func getSingleStoreWants() (string, string, string, string) {
|
||||
select1Want := "[{\"1\":1}]"
|
||||
mcpMyFailToolWant := `{"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"unable to execute query: Error 1064 (42000): You have an error in your SQL syntax; check the manual that corresponds to your MySQL server version for the right syntax to use near 'SELEC 1' at line 1"}],"isError":true}}`
|
||||
createTableStatement := `"CREATE TABLE t (id BIGINT PRIMARY KEY, name TEXT)"`
|
||||
mcpSelect1Want := `{"jsonrpc":"2.0","id":"invoke my-auth-required-tool","result":{"content":[{"type":"text","text":"{\"1\":1}"}]}}`
|
||||
return select1Want, mcpMyFailToolWant, createTableStatement, mcpSelect1Want
|
||||
}
|
||||
|
||||
// setupSingleStoreTable creates and inserts data into a table of tool
|
||||
// compatible with singlestore-sql tool
|
||||
func setupSingleStoreTable(t *testing.T, ctx context.Context, pool *sql.DB, createStatement, insertStatement, tableName string, params []any) func(*testing.T) {
|
||||
err := pool.PingContext(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to connect to test database: %s", err)
|
||||
}
|
||||
|
||||
// Create table
|
||||
_, err = pool.QueryContext(ctx, createStatement)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create test table %s: %s", tableName, err)
|
||||
}
|
||||
|
||||
// Insert test data
|
||||
_, err = pool.QueryContext(ctx, insertStatement, params...)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to insert test data: %s", err)
|
||||
}
|
||||
|
||||
return func(t *testing.T) {
|
||||
// tear down test
|
||||
_, err = pool.ExecContext(ctx, fmt.Sprintf("DROP TABLE %s;", tableName))
|
||||
if err != nil {
|
||||
t.Errorf("Teardown failed: %s", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func getSingleStoreToolsConfig(sourceConfig map[string]any, toolKind, paramToolStatement, idParamToolStmt, nameParamToolStmt, arrayToolStatement, authToolStatement string) map[string]any {
|
||||
toolsFile := tests.GetToolsConfig(sourceConfig, toolKind, paramToolStatement, idParamToolStmt, nameParamToolStmt, arrayToolStatement, authToolStatement)
|
||||
|
||||
toolsMap, ok := toolsFile["tools"].(map[string]any)
|
||||
if !ok {
|
||||
return toolsFile
|
||||
}
|
||||
// Remove tools that are not supported
|
||||
delete(toolsMap, "my-array-tool")
|
||||
|
||||
toolsFile["tools"] = toolsMap
|
||||
return toolsFile
|
||||
}
|
||||
|
||||
// addSingleStoreExecuteSQLConfig gets the tools config for `singlestore-execute-sql`
|
||||
func addSingleStoreExecuteSQLConfig(t *testing.T, config map[string]any) map[string]any {
|
||||
tools, ok := config["tools"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("unable to get tools from config")
|
||||
}
|
||||
tools["my-exec-sql-tool"] = map[string]any{
|
||||
"kind": "singlestore-execute-sql",
|
||||
"source": "my-instance",
|
||||
"description": "Tool to execute sql",
|
||||
}
|
||||
tools["my-auth-exec-sql-tool"] = map[string]any{
|
||||
"kind": "singlestore-execute-sql",
|
||||
"source": "my-instance",
|
||||
"description": "Tool to execute sql",
|
||||
"authRequired": []string{
|
||||
"my-google-auth",
|
||||
},
|
||||
}
|
||||
config["tools"] = tools
|
||||
return config
|
||||
}
|
||||
|
||||
// Copied over from singlestore.go
|
||||
func initSingleStoreConnectionPool(host, port, user, pass, dbname string) (*sql.DB, error) {
|
||||
dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?parseTime=true", user, pass, host, port, dbname)
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
func TestSingleStoreToolEndpoints(t *testing.T) {
|
||||
sourceConfig := getSingleStoreVars(t)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
|
||||
defer cancel()
|
||||
|
||||
var args []string
|
||||
|
||||
pool, err := initSingleStoreConnectionPool(SingleStoreHost, SingleStorePort, SingleStoreUser, SingleStorePass, SingleStoreDatabase)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create SingleStore connection pool: %s", err)
|
||||
}
|
||||
|
||||
// create table name with UUID
|
||||
tableNameParam := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
|
||||
// set up data for param tool
|
||||
createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := getSingleStoreParamToolInfo(tableNameParam)
|
||||
teardownTable1 := setupSingleStoreTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
|
||||
defer teardownTable1(t)
|
||||
|
||||
// set up data for auth tool
|
||||
createAuthTableStmt, insertAuthTableStmt, authToolStmt, authTestParams := getSingleStoreAuthToolInfo(tableNameAuth)
|
||||
teardownTable2 := setupSingleStoreTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
|
||||
defer teardownTable2(t)
|
||||
|
||||
// Write config into a file and pass it to command
|
||||
toolsFile := getSingleStoreToolsConfig(sourceConfig, SingleStoreToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
|
||||
toolsFile = addSingleStoreExecuteSQLConfig(t, toolsFile)
|
||||
tmplSelectCombined, tmplSelectFilterCombined := getSingleStoreTmplToolStatement()
|
||||
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, SingleStoreToolKind, tmplSelectCombined, tmplSelectFilterCombined, "")
|
||||
|
||||
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
|
||||
if err != nil {
|
||||
t.Fatalf("command initialization returned an error: %s", err)
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
|
||||
if err != nil {
|
||||
t.Logf("toolbox command logs: \n%s", out)
|
||||
t.Fatalf("toolbox didn't start successfully: %s", err)
|
||||
}
|
||||
|
||||
// Get configs for tests
|
||||
select1Want, mcpMyFailToolWant, createTableStatement, mcpSelect1Want := getSingleStoreWants()
|
||||
|
||||
// Run tests
|
||||
tests.RunToolGetTest(t)
|
||||
tests.RunToolInvokeTest(t, select1Want, tests.DisableArrayTest())
|
||||
tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want)
|
||||
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
|
||||
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam)
|
||||
}
|
||||
Reference in New Issue
Block a user