Compare commits

...

2 Commits

Author SHA1 Message Date
Yuan Teoh
fcca884f63 ci: update build release to include c and c++ compiler 2025-08-06 16:33:14 -07:00
Averi Kitsch
3e7c3ab6f0 Revert "revert: 'feat: add support for DuckDB' (#1079)"
This reverts commit f79cdd6144.
2025-08-05 14:00:47 -07:00
14 changed files with 933 additions and 27 deletions

View File

@@ -20,9 +20,22 @@ steps:
docker buildx create --name container-builder --driver docker-container --bootstrap --use
docker buildx build --platform linux/amd64,linux/arm64 --build-arg COMMIT_SHA=$(git rev-parse HEAD) -t ${_DOCKER_URI}:$REF_NAME --push .
- id: "install-packages"
name: ubuntu
waitFor: ['-']
env:
- 'GOPATH=/gopath'
volumes:
- name: 'go'
path: '/gopath'
script: |
apt-get update
apt-get install -y build-essential gcc-aarch64-linux-gnu g++-aarch64-linux-gnu libc6-dev-arm64-cross
- id: "install-dependencies"
name: golang:1
waitFor: ['-']
waitFor:
- "install-packages"
env:
- 'GOPATH=/gopath'
volumes:
@@ -42,7 +55,7 @@ steps:
path: '/gopath'
script: |
#!/usr/bin/env bash
CGO_ENABLED=0 GOOS=linux GOARCH=amd64 \
CGO_ENABLED=1 GOOS=linux GOARCH=amd64 \
go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.commitSha=$(git rev-parse HEAD)" -o toolbox.linux.amd64
- id: "store-linux-amd64"
@@ -59,13 +72,17 @@ steps:
- "install-dependencies"
env:
- 'GOPATH=/gopath'
- 'CGO_ENABLED=1'
- 'GOOS=darwin'
- 'GOARCH=arm64'
- 'CC=aarch64-linux-gnu-gcc'
- 'CXX=aarch64-linux-gnu-g++'
volumes:
- name: 'go'
path: '/gopath'
script: |
#!/usr/bin/env bash
CGO_ENABLED=0 GOOS=darwin GOARCH=arm64 \
go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.commitSha=$(git rev-parse HEAD)" -o toolbox.darwin.arm64
go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.commitSha=$(git rev-parse HEAD)" -o toolbox.darwin.arm64
- id: "store-darwin-arm64"
name: "gcr.io/cloud-builders/gcloud:latest"
@@ -86,7 +103,7 @@ steps:
path: '/gopath'
script: |
#!/usr/bin/env bash
CGO_ENABLED=0 GOOS=darwin GOARCH=amd64 \
CGO_ENABLED=1 GOOS=darwin GOARCH=amd64 \
go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.commitSha=$(git rev-parse HEAD)" -o toolbox.darwin.amd64
- id: "store-darwin-amd64"
@@ -108,7 +125,7 @@ steps:
path: '/gopath'
script: |
#!/usr/bin/env bash
CGO_ENABLED=0 GOOS=windows GOARCH=amd64 \
CGO_ENABLED=1 GOOS=windows GOARCH=amd64 \
go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.commitSha=$(git rev-parse HEAD)" -o toolbox.windows.amd64
- id: "store-windows-amd64"

View File

@@ -62,6 +62,7 @@ steps:
postgressql \
postgresexecutesql
- id: "alloydb-pg"
name: golang:1
waitFor: ["compile-test-binary"]
@@ -120,7 +121,8 @@ steps:
- "BIGTABLE_PROJECT=$PROJECT_ID"
- "BIGTABLE_INSTANCE=$_BIGTABLE_INSTANCE"
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
secretEnv: ["CLIENT_ID"]
secretEnv:
["CLIENT_ID"]
volumes:
- name: "go"
path: "/gopath"
@@ -151,7 +153,7 @@ steps:
"BigQuery" \
bigquery \
bigquery
- id: "dataplex"
name: golang:1
waitFor: ["compile-test-binary"]
@@ -272,7 +274,8 @@ steps:
- "CLOUD_SQL_MYSQL_DATABASE=$_DATABASE_NAME"
- "CLOUD_SQL_MYSQL_REGION=$_REGION"
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
secretEnv: ["CLOUD_SQL_MYSQL_USER", "CLOUD_SQL_MYSQL_PASS", "CLIENT_ID"]
secretEnv:
["CLOUD_SQL_MYSQL_USER", "CLOUD_SQL_MYSQL_PASS", "CLIENT_ID"]
volumes:
- name: "go"
path: "/gopath"
@@ -384,7 +387,7 @@ steps:
sqlite
- id: "couchbase"
name: golang:1
name : golang:1
waitFor: ["compile-test-binary"]
entrypoint: /bin/bash
env:
@@ -392,8 +395,7 @@ steps:
- "COUCHBASE_SCOPE=$_COUCHBASE_SCOPE"
- "COUCHBASE_BUCKET=$_COUCHBASE_BUCKET"
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
secretEnv:
["COUCHBASE_CONNECTION", "COUCHBASE_USER", "COUCHBASE_PASS", "CLIENT_ID"]
secretEnv: ["COUCHBASE_CONNECTION", "COUCHBASE_USER", "COUCHBASE_PASS", "CLIENT_ID"]
volumes:
- name: "go"
path: "/gopath"
@@ -406,7 +408,7 @@ steps:
couchbase
- id: "redis"
name: golang:1
name : golang:1
waitFor: ["compile-test-binary"]
entrypoint: /bin/bash
env:
@@ -423,9 +425,9 @@ steps:
"Redis" \
redis \
redis
- id: "valkey"
name: golang:1
name : golang:1
waitFor: ["compile-test-binary"]
entrypoint: /bin/bash
env:
@@ -473,13 +475,7 @@ steps:
- "FIRESTORE_PROJECT=$PROJECT_ID"
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
- "LOOKER_VERIFY_SSL=$_LOOKER_VERIFY_SSL"
secretEnv:
[
"CLIENT_ID",
"LOOKER_BASE_URL",
"LOOKER_CLIENT_ID",
"LOOKER_CLIENT_SECRET",
]
secretEnv: ["CLIENT_ID", "LOOKER_BASE_URL", "LOOKER_CLIENT_ID", "LOOKER_CLIENT_SECRET"]
volumes:
- name: "go"
path: "/gopath"
@@ -491,6 +487,25 @@ steps:
looker \
looker
- id: "duckdb"
name: golang:1
waitFor: ["compile-test-binary"]
entrypoint: /bin/bash
env:
- "GOPATH=/gopath"
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
volumes:
- name: "go"
path: "/gopath"
secretEnv: ["CLIENT_ID"]
args:
- -c
- |
.ci/test_with_coverage.sh \
"DuckDB" \
duckdb \
duckdb
- id: "alloydbwaitforoperation"
name: golang:1
waitFor: ["compile-test-binary"]
@@ -530,8 +545,8 @@ steps:
.ci/test_with_coverage.sh \
"TiDB" \
tidb \
tidbsql tidbexecutesql
tidbsql tidbexecutesql
availableSecrets:
secretManager:
- versionName: projects/$PROJECT_ID/secrets/cloud_sql_pg_user/versions/latest

View File

@@ -23,9 +23,21 @@ ARG TARGETARCH
ARG BUILD_TYPE="container.dev"
ARG COMMIT_SHA=""
RUN apt-get update && \
apt-get install -y --no-install-recommends \
build-essential \
gcc-aarch64-linux-gnu \
g++-aarch64-linux-gnu \
libc6-dev-arm64-cross \
&& rm -rf /var/lib/apt/lists/*
RUN go get ./...
RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} \
go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.buildType=container.${BUILD_TYPE} -X github.com/googleapis/genai-toolbox/cmd.commitSha=${COMMIT_SHA}"
ENV CGO_ENABLED=1
ENV GOOS=$TARGETOS
ENV GOARCH=$TARGETARCH
RUN if [ "$TARGETARCH" = "arm64" ]; then export CC=aarch64-linux-gnu-gcc CXX=aarch64-linux-gnu-g++; fi \
&& go build -ldflags "-X github.com/googleapis/genai-toolbox/cmd.buildType=container.${BUILD_TYPE} -X github.com/googleapis/genai-toolbox/cmd.commitSha=${COMMIT_SHA}"
# Final Stage
FROM gcr.io/distroless/static:nonroot

View File

@@ -54,6 +54,7 @@ import (
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexlookupentry"
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexsearchentries"
_ "github.com/googleapis/genai-toolbox/internal/tools/dgraph"
_ "github.com/googleapis/genai-toolbox/internal/tools/duckdbsql"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoredeletedocuments"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoregetdocuments"
_ "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoregetrules"
@@ -111,6 +112,7 @@ import (
_ "github.com/googleapis/genai-toolbox/internal/sources/couchbase"
_ "github.com/googleapis/genai-toolbox/internal/sources/dataplex"
_ "github.com/googleapis/genai-toolbox/internal/sources/dgraph"
_ "github.com/googleapis/genai-toolbox/internal/sources/duckdb"
_ "github.com/googleapis/genai-toolbox/internal/sources/firestore"
_ "github.com/googleapis/genai-toolbox/internal/sources/http"
_ "github.com/googleapis/genai-toolbox/internal/sources/looker"

View File

@@ -0,0 +1,73 @@
---
title: DuckDB
linkTitle: DuckDB
type: docs
weight: 1
description: >
DuckDB is an in-process SQL OLAP database management system designed for analytical query processing.
---
## About
[DuckDB](https://duckdb.org/) is an embedded analytical database management system that runs in-process with the client application. It is optimized for analytical workloads, providing high performance for complex queries with minimal setup.
DuckDB has the following notable characteristics:
- In-process, serverless database engine
- Supports complex SQL queries for analytical processing
- Can operate on in-memory or persistent storage
- Zero-configuration - no external dependencies or server setup required
- Highly optimized for columnar data storage and query execution
For more details, refer to the [DuckDB Documentation](https://duckdb.org/).
## Available Tools
- [`duckdb-sql`](../tools/duckdb/duckdb-sql.md)
Execute pre-defined prepared SQL queries in DuckDB.
## Requirements
### Database File
To use DuckDB, you can either:
- Specify a file path for a persistent database stored on the filesystem
- Omit the file path to use an in-memory database
## Example
For a persistent DuckDB database:
```yaml
sources:
my-duckdb:
kind: "duckdb"
dbFilePath: "/path/to/database.db"
configuration:
memory_limit: "2GB"
threads: "4"
```
For an in-memory DuckDB database:
```yaml
sources:
my-duckdb-memory:
name: "my-duckdb-memory"
kind: "duckdb"
```
## Reference
### Configuration Fields
| **field** | **type** | **required** | **description** |
|-------------------|:-----------------:|:------------:|---------------------------------------------------------------------------------|
| kind | string | true | Must be "duckdb". |
| dbFilePath | string | false | Path to the DuckDB database file. Omit for an in-memory database. |
| configuration | map[string]string | false | Additional DuckDB configuration options (e.g., `memory_limit`, `threads`). |
For a complete list of available configuration options, refer to the [DuckDB Configuration Documentation](https://duckdb.org/docs/stable/configuration/overview.html#local-configuration-options).
For more details on the Go implementation, see the [go-duckdb package documentation](https://pkg.go.dev/github.com/scottlepp/go-duckdb#section-readme).

View File

@@ -0,0 +1,7 @@
---
title: "DuckDB"
type: docs
weight: 1
description: >
Tools that work with DuckDB Sources.
---

View File

@@ -0,0 +1,80 @@
---
title: "duckdb-sql"
type: docs
weight: 1
description: >
Execute SQL statements against a DuckDB database using the DuckDB SQL tools configuration.
aliases:
- /resources/tools/duckdb-sql
---
## About
A `duckdb-sql` tool executes a pre-defined SQL statement against a [DuckDB](https://duckdb.org/) database. It is compatible with any DuckDB source configuration as defined in the [DuckDB source documentation](../../sources/duckdb.md).
The specified SQL statement is executed as a prepared statement, and parameters are inserted according to their position: e.g., `$1` is the first parameter, `$2` is the second, and so on. If template parameters are included, they are resolved before execution of the prepared statement.
DuckDB's SQL dialect closely follows the conventions of the PostgreSQL dialect, with a few exceptions listed in the [DuckDB PostgreSQL Compatibility documentation](https://duckdb.org/docs/stable/sql/dialect/postgresql_compatibility.html). For an introduction to DuckDB's SQL dialect, refer to the [DuckDB SQL Introduction](https://duckdb.org/docs/stable/sql/introduction).
### Concepts
DuckDB is a relational database management system (RDBMS). Data is stored in relations (tables), where each table is a named collection of rows. Each row in a table has the same set of named columns, each with a specific data type. Tables are stored within schemas, and a collection of schemas constitutes the entire database.
For more details, see the [DuckDB SQL Introduction](https://duckdb.org/docs/stable/sql/introduction).
## Example
> **Note:** This tool uses parameterized queries to prevent SQL injections. Query parameters can be used as substitutes for arbitrary expressions but cannot be used for identifiers, column names, table names, or other parts of the query.
```yaml
tools:
search-users:
kind: duckdb-sql
source: my-duckdb
description: Search users by name and age
statement: SELECT * FROM users WHERE name LIKE $1 AND age >= $2
parameters:
- name: name
type: string
description: The name to search for
- name: min_age
type: integer
description: Minimum age
```
## Example with Template Parameters
> **Note:** Template parameters allow direct modifications to the SQL statement, including identifiers, column names, and table names, which makes them more vulnerable to SQL injections. Using basic parameters (see above) is recommended for performance and safety. For more details, see the [templateParameters](../#template-parameters) section.
```yaml
tools:
list_table:
kind: duckdb-sql
source: my-duckdb
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
### Configuration Fields
| **field** | **type** | **required** | **description** |
|--------------------|:-------------------------------:|:------------:|--------------------------------------------------------------------------------------------------------------------------------------------|
| kind | string | true | Must be "duckdb-sql". |
| source | string | true | Name of the DuckDB source configuration (see [DuckDB source documentation](../../sources/duckdb.md)). |
| description | string | true | Description of the tool that is passed to the LLM. |
| statement | string | true | The SQL statement to execute. |
| authRequired | []string | false | List of authentication requirements for the tool (if any). |
| parameters | [parameters](../#specifying-parameters) | false | List of parameters that will be inserted into the SQL statement |
| templateParameters | [templateParameters](../#template-parameters) | false | List of template parameters that will be inserted into the SQL statement before executing the prepared statement. |

12
go.mod
View File

@@ -51,8 +51,15 @@ require (
require (
github.com/andybalholm/cascadia v1.3.3 // indirect
github.com/duckdb/duckdb-go-bindings v0.1.17 // indirect
github.com/duckdb/duckdb-go-bindings/darwin-amd64 v0.1.12 // indirect
github.com/duckdb/duckdb-go-bindings/darwin-arm64 v0.1.12 // indirect
github.com/duckdb/duckdb-go-bindings/linux-amd64 v0.1.12 // indirect
github.com/duckdb/duckdb-go-bindings/linux-arm64 v0.1.12 // indirect
github.com/duckdb/duckdb-go-bindings/windows-amd64 v0.1.12 // indirect
github.com/marcboeker/go-duckdb/arrowmapping v0.0.10 // indirect
github.com/marcboeker/go-duckdb/mapping v0.0.11 // indirect
golang.org/x/exp v0.0.0-20250620022241-b7579e27df2b // indirect
gonum.org/v1/gonum v0.16.0 // indirect
)
require (
@@ -72,6 +79,7 @@ require (
github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.53.0 // indirect
github.com/PuerkitoBio/goquery v1.10.3 // indirect
github.com/ajg/form v1.5.1 // indirect
github.com/apache/arrow-go/v18 v18.4.0 // indirect
github.com/apache/arrow/go/v15 v15.0.2 // indirect
github.com/cenkalti/backoff/v5 v5.0.2 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
@@ -92,6 +100,7 @@ require (
github.com/go-logr/stdr v1.2.2 // indirect
github.com/go-playground/locales v0.14.1 // indirect
github.com/go-playground/universal-translator v0.18.1 // indirect
github.com/go-viper/mapstructure/v2 v2.3.0 // indirect
github.com/goccy/go-json v0.10.5 // indirect
github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 // indirect
github.com/golang-sql/sqlexp v0.1.0 // indirect
@@ -111,6 +120,7 @@ require (
github.com/klauspost/compress v1.18.0 // indirect
github.com/klauspost/cpuid/v2 v2.2.11 // indirect
github.com/leodido/go-urn v1.4.0 // indirect
github.com/marcboeker/go-duckdb/v2 v2.3.4
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect

29
go.sum
View File

@@ -674,14 +674,20 @@ github.com/ajstarks/deck/generate v0.0.0-20210309230005-c3f852c02e19/go.mod h1:T
github.com/ajstarks/svgo v0.0.0-20180226025133-644b8db467af/go.mod h1:K08gAheRH3/J6wwsYMMT4xOr94bZjxIelGM0+d/wbFw=
github.com/ajstarks/svgo v0.0.0-20211024235047-1546f124cd8b/go.mod h1:1KcenG0jGWcpt8ov532z81sp/kMMUG485J2InIOyADM=
github.com/andybalholm/brotli v1.0.4/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ=
github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
github.com/andybalholm/cascadia v1.3.3 h1:AG2YHrzJIm4BZ19iwJ/DAua6Btl3IwJX+VI4kktS1LM=
github.com/andybalholm/cascadia v1.3.3/go.mod h1:xNd9bqTn98Ln4DwST8/nG+H0yuB8Hmgu1YHNnWw0GeA=
github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
github.com/apache/arrow-go/v18 v18.4.0 h1:/RvkGqH517iY8bZKc4FD5/kkdwXJGjxf28JIXbJ/oB0=
github.com/apache/arrow-go/v18 v18.4.0/go.mod h1:Aawvwhj8x2jURIzD9Moy72cF0FyJXOpkYpdmGRHcw14=
github.com/apache/arrow/go/v10 v10.0.1/go.mod h1:YvhnlEePVnBS4+0z3fhPfUy7W1Ikj0Ih0vcRo/gZ1M0=
github.com/apache/arrow/go/v11 v11.0.0/go.mod h1:Eg5OsL5H+e299f7u5ssuXsuHQVEGC4xei5aX110hRiI=
github.com/apache/arrow/go/v15 v15.0.2 h1:60IliRbiyTWCWjERBCkO1W4Qun9svcYoZrSLcyOsMLE=
github.com/apache/arrow/go/v15 v15.0.2/go.mod h1:DGXsR3ajT524njufqf95822i+KTh+yea1jass9YXgjA=
github.com/apache/thrift v0.16.0/go.mod h1:PHK3hniurgQaNMZYaCLEqXKsYK8upmhPbmdP2FXSqgU=
github.com/apache/thrift v0.22.0 h1:r7mTJdj51TMDe6RtcmNdQxgn9XcyfGDOzegMDRg47uc=
github.com/apache/thrift v0.22.0/go.mod h1:1e7J/O1Ae6ZQMTYdy9xa3w9k+XHWPfRvdPyJeynQ+/g=
github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
github.com/boombuler/barcode v1.0.0/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
github.com/boombuler/barcode v1.0.1/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
@@ -743,6 +749,18 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8Yc
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3ebgob9U8Nd0kOddGdZWjyMGR8Wziv+TBNwSE=
github.com/duckdb/duckdb-go-bindings v0.1.17 h1:SjpRwrJ7v0vqnIvLeVFHlhuS72+Lp8xxQ5jIER2LZP4=
github.com/duckdb/duckdb-go-bindings v0.1.17/go.mod h1:pBnfviMzANT/9hi4bg+zW4ykRZZPCXlVuvBWEcZofkc=
github.com/duckdb/duckdb-go-bindings/darwin-amd64 v0.1.12 h1:8CLBnsq9YDhi2Gmt3sjSUeXxMzyMQAKefjqUy9zVPFk=
github.com/duckdb/duckdb-go-bindings/darwin-amd64 v0.1.12/go.mod h1:Ezo7IbAfB8NP7CqPIN8XEHKUg5xdRRQhcPPlCXImXYA=
github.com/duckdb/duckdb-go-bindings/darwin-arm64 v0.1.12 h1:wjO4I0GhMh2xIpiUgRpzuyOT4KxXLoUS/rjU7UUVvCE=
github.com/duckdb/duckdb-go-bindings/darwin-arm64 v0.1.12/go.mod h1:eS7m/mLnPQgVF4za1+xTyorKRBuK0/BA44Oy6DgrGXI=
github.com/duckdb/duckdb-go-bindings/linux-amd64 v0.1.12 h1:HzKQi2C+1jzmwANsPuYH6x9Sfw62SQTjNAEq3OySKFI=
github.com/duckdb/duckdb-go-bindings/linux-amd64 v0.1.12/go.mod h1:1GOuk1PixiESxLaCGFhag+oFi7aP+9W8byymRAvunBk=
github.com/duckdb/duckdb-go-bindings/linux-arm64 v0.1.12 h1:YGSR7AFLw2gJ7IbgLE6DkKYmgKv1LaRSd/ZKF1yh2oE=
github.com/duckdb/duckdb-go-bindings/linux-arm64 v0.1.12/go.mod h1:o7crKMpT2eOIi5/FY6HPqaXcvieeLSqdXXaXbruGX7w=
github.com/duckdb/duckdb-go-bindings/windows-amd64 v0.1.12 h1:2aduW6fnFnT2Q45PlIgHbatsPOxV9WSZ5B2HzFfxaxA=
github.com/duckdb/duckdb-go-bindings/windows-amd64 v0.1.12/go.mod h1:IlOhJdVKUJCAPj3QsDszUo8DVdvp1nBFp4TUJVdw99s=
github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
@@ -818,6 +836,8 @@ github.com/go-playground/validator/v10 v10.27.0/go.mod h1:I5QpIEbmr8On7W0TktmJAu
github.com/go-sql-driver/mysql v1.9.3 h1:U/N249h2WzJ3Ukj8SowVFjdtZKfu9vlLZxjPXV1aweo=
github.com/go-sql-driver/mysql v1.9.3/go.mod h1:qn46aNg1333BRMNU69Lq93t8du/dwxI64Gl8i5p1WMU=
github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
github.com/go-viper/mapstructure/v2 v2.3.0 h1:27XbWsHIqhbdR5TIC911OfYvgSaW93HM+dX7970Q7jk=
github.com/go-viper/mapstructure/v2 v2.3.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
github.com/goccy/go-json v0.9.11/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4=
github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
@@ -1001,6 +1021,7 @@ github.com/jung-kurt/gofpdf v1.0.3-0.20190309125859-24315acbbda5/go.mod h1:7Id9E
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8=
github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
github.com/klauspost/asmfmt v1.3.2 h1:4Ri7ox3EwapiOjCki+hw14RyKk201CN4rzyCJRFLpK4=
github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE=
github.com/klauspost/compress v1.15.9/go.mod h1:PhcZ0MbTNciWF3rruxRgKxI5NkcHHrHUDtV4Yw2GlzU=
github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
@@ -1025,6 +1046,12 @@ github.com/looker-open-source/sdk-codegen/go v0.25.10/go.mod h1:YM/IYSsTPk7I54j4
github.com/lyft/protoc-gen-star v0.6.0/go.mod h1:TGAoBVkt8w7MPG72TrKIu85MIdXwDuzJYeZuUPFPNwA=
github.com/lyft/protoc-gen-star v0.6.1/go.mod h1:TGAoBVkt8w7MPG72TrKIu85MIdXwDuzJYeZuUPFPNwA=
github.com/lyft/protoc-gen-star/v2 v2.0.1/go.mod h1:RcCdONR2ScXaYnQC5tUzxzlpA3WVYF7/opLeUgcQs/o=
github.com/marcboeker/go-duckdb/arrowmapping v0.0.10 h1:G1W+GVnUefR8uy7jHdNO+CRMsmFG5mFPIHVAespfFCA=
github.com/marcboeker/go-duckdb/arrowmapping v0.0.10/go.mod h1:jccUb8TYD0p5TsEEeN4SXuslNJHo23QaKOqKD+U6uFU=
github.com/marcboeker/go-duckdb/mapping v0.0.11 h1:fusN1b1l7Myxafifp596I6dNLNhN5Uv/rw31qAqBwqw=
github.com/marcboeker/go-duckdb/mapping v0.0.11/go.mod h1:aYBjFLgfKO0aJIbDtXPiaL5/avRQISveX/j9tMf9JhU=
github.com/marcboeker/go-duckdb/v2 v2.3.4 h1:o98wrefPbH0IdJRix4pF0+jZiXoFQ+FSR8InMsCUZD0=
github.com/marcboeker/go-duckdb/v2 v2.3.4/go.mod h1:8adNrftF4Ye29XMrpIl5NYNosTVsZu1mz3C82WdVvrk=
github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
@@ -1032,7 +1059,9 @@ github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D
github.com/mattn/go-sqlite3 v1.14.14/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU=
github.com/microsoft/go-mssqldb v1.9.2 h1:nY8TmFMQOHpm2qVWo6y4I2mAmVdZqlGiMGAYt64Ibbs=
github.com/microsoft/go-mssqldb v1.9.2/go.mod h1:GBbW9ASTiDC+mpgWDGKdm3FnFLTUsLYN3iFL90lQ+PA=
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY=
github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI=
github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3/go.mod h1:RagcQ7I8IeTMnF8JTXieKnO4Z6JCsikNEzj0DwauVzE=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=

View File

@@ -0,0 +1,128 @@
// 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 duckdb
import (
"context"
"database/sql"
"fmt"
"net/url"
"strings"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
_ "github.com/marcboeker/go-duckdb/v2"
"go.opentelemetry.io/otel/trace"
)
const SourceKind string = "duckdb"
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
}
type Source struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
Db *sql.DB
}
// SourceKind implements sources.Source.
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) DuckDb() *sql.DB {
return s.Db
}
// validate Source
var _ sources.Source = &Source{}
type Config struct {
Name string `yaml:"name" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
DatabaseFile string `yaml:"dbFilePath,omitempty"`
Configuration map[string]string `yaml:"configuration,omitempty"`
}
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
db, err := initDuckDbConnection(ctx, tracer, r.Name, r.DatabaseFile, r.Configuration)
if err != nil {
return nil, fmt.Errorf("unable to create db connection: %w", err)
}
err = db.PingContext(context.Background())
if err != nil {
return nil, fmt.Errorf("unable to connect sucessfully: %w", err)
}
s := &Source{
Name: r.Name,
Kind: r.Kind,
Db: db,
}
return s, nil
}
// validate interface
var _ sources.SourceConfig = Config{}
func initDuckDbConnection(ctx context.Context, tracer trace.Tracer, name string, dbFilePath string, duckdbConfiguration map[string]string) (*sql.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
var configStr = getDuckDbConfiguration(dbFilePath, duckdbConfiguration)
//Open database connection
db, err := sql.Open("duckdb", configStr)
if err != nil {
return nil, fmt.Errorf("unable to open duckdb connection: %w", err)
}
return db, nil
}
func getDuckDbConfiguration(dbFilePath string, duckdbConfiguration map[string]string) string {
if len(duckdbConfiguration) == 0 {
return dbFilePath
}
params := url.Values{}
for key, value := range duckdbConfiguration {
params.Set(key, value)
}
var configStr strings.Builder
configStr.WriteString(dbFilePath)
configStr.WriteString("?")
configStr.WriteString(params.Encode())
return configStr.String()
}

View File

@@ -0,0 +1,84 @@
// 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 duckdb_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/sources/duckdb"
"github.com/googleapis/genai-toolbox/internal/testutils"
)
func TestParserFromYamlDuckDb(t *testing.T) {
config := make(map[string]string)
config["access_mode"] = "READ_ONLY"
config["threads"] = "4"
tcs := []struct {
desc string
in string
want server.SourceConfigs
}{
{
desc: "basic example",
in: `
sources:
my-duckdb:
kind: duckdb
`,
want: server.SourceConfigs{
"my-duckdb": duckdb.Config{
Name: "my-duckdb",
Kind: duckdb.SourceKind,
},
},
},
{
desc: "with custom configuration",
in: `
sources:
my-duckdb:
kind: duckdb
configuration:
access_mode: READ_ONLY
threads: 4
`,
want: server.SourceConfigs{
"my-duckdb": duckdb.Config{
Name: "my-duckdb",
Kind: duckdb.SourceKind,
Configuration: config,
},
},
},
}
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)
}
})
}
}

View File

@@ -0,0 +1,210 @@
// 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 duckdbsql
import (
"context"
"database/sql"
"fmt"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/duckdb"
"github.com/googleapis/genai-toolbox/internal/tools"
)
const kind string = "duckdb-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 {
DuckDb() *sql.DB
}
// validate compatible sources are still compatible
var _ compatibleSource = &duckdb.Source{}
var compatibleSources = [...]string{duckdb.SourceKind}
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"`
}
// Initialize implements tools.ToolConfig.
func (c Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
// verify source exists
rawS, ok := srcs[c.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", c.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, paramMcpManifest := tools.ProcessParameters(c.TemplateParameters, c.Parameters)
mcpManifest := tools.McpManifest{
Name: c.Name,
Description: c.Description,
InputSchema: paramMcpManifest,
}
// finish tool setup
t := Tool{
Name: c.Name,
Kind: kind,
Parameters: c.Parameters,
TemplateParameters: c.TemplateParameters,
AllParams: allParameters,
Statement: c.Statement,
AuthRequired: c.AuthRequired,
Db: s.DuckDb(),
manifest: tools.Manifest{Description: c.Description, Parameters: paramManifest, AuthRequired: c.AuthRequired},
mcpManifest: mcpManifest,
}
return t, nil
}
// ToolConfigKind implements tools.ToolConfig.
func (c Config) ToolConfigKind() string {
return kind
}
var _ tools.ToolConfig = Config{}
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"`
Db *sql.DB
Statement string `yaml:"statement"`
manifest tools.Manifest
mcpManifest tools.McpManifest
}
// Authorized implements tools.Tool.
func (t Tool) Authorized(verifiedAuthSources []string) bool {
return tools.IsAuthorized(t.AuthRequired, verifiedAuthSources)
}
// Invoke implements tools.Tool.
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues) (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()
// Execute the SQL query with parameters
rows, err := t.Db.QueryContext(ctx, newStatement, sliceParams...)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)
}
defer rows.Close()
// Get column names
cols, err := rows.Columns()
if err != nil {
return nil, fmt.Errorf("unable to get column names: %w", err)
}
values := make([]any, len(cols))
valuePtrs := make([]any, len(cols))
for i := range values {
valuePtrs[i] = &values[i]
}
// Prepare the result slice
var result []any
// Iterate through the rows
for rows.Next() {
// Scan the row into the value pointers
if err := rows.Scan(valuePtrs...); err != nil {
return nil, fmt.Errorf("unable to scan row: %w", err)
}
// Create a map for this row
rowMap := make(map[string]interface{})
for i, col := range cols {
val := values[i]
// Handle nil values
if val == nil {
rowMap[col] = nil
continue
}
// Store the value in the map
rowMap[col] = val
}
result = append(result, rowMap)
}
if err = rows.Close(); err != nil {
return nil, fmt.Errorf("unable to close rows: %w", err)
}
if err = rows.Err(); err != nil {
return nil, fmt.Errorf("error iterating rows: %w", err)
}
return result, nil
}
// Manifest implements tools.Tool.
func (t Tool) Manifest() tools.Manifest {
return t.manifest
}
// McpManifest implements tools.Tool.
func (t Tool) McpManifest() tools.McpManifest {
return t.mcpManifest
}
// ParseParams implements tools.Tool.
func (t Tool) ParseParams(data map[string]any, claimsMap map[string]map[string]any) (tools.ParamValues, error) {
return tools.ParseParams(t.AllParams, data, claimsMap)
}
var _ tools.Tool = Tool{}

View File

@@ -0,0 +1,87 @@
// 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 duckdbsql_test
import (
"testing"
"github.com/googleapis/genai-toolbox/internal/tools/duckdbsql"
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"
)
func TestParseFromYamlDuckDb(t *testing.T) {
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
example_tool:
kind: duckdb-sql
source: my-duckdb-instance
description: some description
statement: |
select * from hotel WHERE name = $hotel;
parameters:
- name: hotel
type: string
description: hotel parameter description
`,
want: server.ToolConfigs{
"example_tool": duckdbsql.Config{
Name: "example_tool",
Kind: "duckdb-sql",
Source: "my-duckdb-instance",
Description: "some description",
Statement: "select * from hotel WHERE name = $hotel;\n",
AuthRequired: []string{},
Parameters: []tools.Parameter{
tools.NewStringParameter("hotel", "hotel parameter description"),
},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Create a context with a logger
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unable to create context with logger: %s", err)
}
// Parse contents with context
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)
}
})
}
}

View File

@@ -0,0 +1,152 @@
// 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 duckdb
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 (
DuckDbKind = "duckdb-sql"
dbPath = "/tmp/users.db"
)
func getDuckDbVars() map[string]any {
return map[string]any{
"kind": "duckdb",
"dbFilePath": dbPath,
"configuration": map[string]any{
"access_mode": "READ_WRITE",
},
}
}
func setupDuckDb(t *testing.T, createParamStmt, insertParamStmt, createAuthStmt, insertAuthStmt string, params []any, authparams []any) {
// Remove any existing database file to ensure a clean state
os.Remove(dbPath)
// Open a connection to DuckDB
db, err := sql.Open("duckdb", dbPath)
if err != nil {
t.Fatalf("Failed to open DuckDB connection: %v", err)
}
defer db.Close()
_, err = db.Exec(createParamStmt, params...)
if err != nil {
t.Fatalf("Failed to create table: %v", err)
}
_, err = db.Exec(createAuthStmt, authparams...)
if err != nil {
t.Fatalf("Failed to create table: %v", err)
}
_, err = db.Exec(insertParamStmt, params...)
if err != nil {
t.Fatalf("Failed to insert initial data: %v", err)
}
_, err = db.Exec(insertAuthStmt, authparams...)
if err != nil {
t.Fatalf("Failed to create table: %v", err)
}
}
func TestDuckDb(t *testing.T) {
sourceConfig := getDuckDbVars()
var args []string
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
defer cancel()
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(), "-", "")
createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, paramToolStmt2, arrayToolStmt, paramTestParams := GetDuckDbParamToolInfo(tableNameParam)
createAuthTableStmt, insertAuthTableStmt, authToolStmt, authTestParams := GetDuckDbAuthToolInfo(tableNameAuth)
setupDuckDb(t, createParamTableStmt, insertParamTableStmt, createAuthTableStmt, insertAuthTableStmt, paramTestParams, authTestParams)
toolsFile := tests.GetToolsConfig(sourceConfig, DuckDbKind, paramToolStmt, idParamToolStmt, paramToolStmt2, arrayToolStmt, authToolStmt)
tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, DuckDbKind, tmplSelectCombined, tmplSelectFilterCombined, "")
defer os.Remove(dbPath)
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)
}
tests.RunToolGetTest(t)
select1Want, failInvocationWant, _ := GetDuckDbWants()
_, invokeParamWantNull, nullWant, _ := tests.GetNonSpannerInvokeParamWant()
invokeParamWant := "[{\"name\":\"Alice\"},{\"name\":\"Sid\"}]"
mcpInvokeParamWant := "{\"jsonrpc\":\"2.0\",\"id\":\"my-tool\",\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"{\\\"name\\\":\\\"Alice\\\"}\"},{\"type\":\"text\",\"text\":\"{\\\"name\\\":\\\"Sid\\\"}\"}]}}"
tests.RunToolInvokeTest(t, select1Want, invokeParamWant, invokeParamWantNull, nullWant, true, true)
tests.RunMCPToolCallMethod(t, mcpInvokeParamWant, failInvocationWant)
templateParamTestConfig := tests.NewTemplateParameterTestConfig(
tests.WithSelectAllWant("[{\"age\":21,\"id\":1,\"name\":\"Alex\"},{\"age\":100,\"id\":2,\"name\":\"Alice\"}]"),
tests.WithSelect1Want("[{\"age\":21,\"id\":1,\"name\":\"Alex\"}]"),
tests.WithReplaceNameFieldArray(`["name"]`),
tests.WithReplaceNameColFilter("name"),
tests.WithCreateColArray(`["id INT","name VARCHAR(20)","age INT"]`),
tests.WithInsert1Want("[{\"Count\":1}]"),
)
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, templateParamTestConfig)
}
func GetDuckDbParamToolInfo(tableName string) (string, string, string, string, string, string, []any) {
createStatement := fmt.Sprintf("CREATE TABLE %s (id INTEGER PRIMARY KEY, name TEXT);", tableName)
insertStatement := fmt.Sprintf("INSERT INTO %s (id, name) VALUES (1, $1), (2, $2), (3, $3), (4, $4);", tableName)
toolStatement := fmt.Sprintf("SELECT * EXCLUDE (id) FROM %s WHERE id = $1 OR name = $2 order by id;", tableName)
idParamStatement := fmt.Sprintf("SELECT * FROM %s WHERE id IN (SELECT unnest(list_value($1)) AS id);", tableName)
toolStatement2 := fmt.Sprintf("SELECT name FROM %s WHERE id = list_extract(list_value($1), 1);", tableName)
arrayToolStatement := fmt.Sprintf("SELECT name FROM %s WHERE id = ANY($1) AND name = ANY($2) order by name;", tableName)
params := []any{"Alice", "Jane", "Sid", nil}
return createStatement, insertStatement, toolStatement, idParamStatement, toolStatement2, arrayToolStatement, params
}
// GetDuckDbAuthToolInfo returns statements and param of my-auth-tool for duckdb-sql kind
func GetDuckDbAuthToolInfo(tableName string) (string, string, string, []any) {
createStatement := fmt.Sprintf("CREATE TABLE %s (id INTEGER PRIMARY KEY, name TEXT, email TEXT);", tableName)
insertStatement := fmt.Sprintf("INSERT INTO %s (id, name, email) VALUES (1, $1, $2), (2, $3, $4)", tableName)
toolStatement := fmt.Sprintf("SELECT name FROM %s WHERE email = $1;", tableName)
params := []any{"Alice", tests.ServiceAccountEmail, "Jane", "janedoe@gmail.com"}
return createStatement, insertStatement, toolStatement, params
}
func GetDuckDbWants() (string, string, string) {
select1Want := "[{\"1\":1}]"
failInvocationWant := `{"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"unable to execute query: Parser Error: syntax error at or near \"SELEC\""}],"isError":true}}`
createTableStatement := `"CREATE TABLE t (id SERIAL PRIMARY KEY, name TEXT)"`
return select1Want, failInvocationWant, createTableStatement
}