Compare commits

...

2 Commits

Author SHA1 Message Date
dependabot[bot]
f1a5451beb chore(deps): bump axios
Bumps [axios](https://github.com/axios/axios) from 1.12.2 to 1.13.5.
- [Release notes](https://github.com/axios/axios/releases)
- [Changelog](https://github.com/axios/axios/blob/v1.x/CHANGELOG.md)
- [Commits](https://github.com/axios/axios/compare/v1.12.2...v1.13.5)

---
updated-dependencies:
- dependency-name: axios
  dependency-version: 1.13.5
  dependency-type: indirect
...

Signed-off-by: dependabot[bot] <support@github.com>
2026-02-11 22:05:58 +00:00
Virag Tripathi
1fdd99a9b6 feat(cockroachdb): add CockroachDB integration with cockroach-go (#2006)
Add support for CockroachDB v25.4.0+ using the official cockroach-go/v2
library for automatic transaction retry.

- Add CockroachDB source with ExecuteTxWithRetry using crdbpgx.ExecuteTx
- Implement 4 tools: execute-sql, sql, list-tables, list-schemas
- Use UUID primary keys (CockroachDB best practice)
- Add unit tests for source and all tools
- Add integration tests with retry verification
- Update Cloud Build configuration for CI

Fixes #2005

## Description

> Should include a concise description of the changes (bug or feature),
it's
> impact, along with a summary of the solution

## 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:

- [ ] Make sure you reviewed

[CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md)
- [ ] Make sure to open an issue as a

[bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose)
  before writing your code! That way we can discuss the change, evaluate
  designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)
- [ ] Make sure to add `!` if this involve a breaking change

🛠️ Fixes #<issue_number_goes_here>

---------

Co-authored-by: duwenxin99 <duwenxin@google.com>
Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>
2026-02-11 17:05:03 -05:00
23 changed files with 4052 additions and 9 deletions

View File

@@ -354,6 +354,30 @@ steps:
postgressql \
postgresexecutesql
- id: "cockroachdb"
name: golang:1
waitFor: ["compile-test-binary"]
entrypoint: /bin/bash
env:
- "GOPATH=/gopath"
- "COCKROACHDB_DATABASE=$_DATABASE_NAME"
- "COCKROACHDB_PORT=$_COCKROACHDB_PORT"
- "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
secretEnv: ["COCKROACHDB_USER", "COCKROACHDB_HOST","CLIENT_ID"]
volumes:
- name: "go"
path: "/gopath"
args:
- -c
- |
.ci/test_with_coverage.sh \
"CockroachDB" \
cockroachdb \
cockroachdbsql \
cockroachdbexecutesql \
cockroachdblisttables \
cockroachdblistschemas
- id: "spanner"
name: golang:1
waitFor: ["compile-test-binary"]
@@ -1129,6 +1153,11 @@ availableSecrets:
env: MARIADB_HOST
- versionName: projects/$PROJECT_ID/secrets/mongodb_uri/versions/latest
env: MONGODB_URI
- versionName: projects/$PROJECT_ID/secrets/cockroachdb_user/versions/latest
env: COCKROACHDB_USER
- versionName: projects/$PROJECT_ID/secrets/cockroachdb_host/versions/latest
env: COCKROACHDB_HOST
options:
logging: CLOUD_LOGGING_ONLY
@@ -1189,6 +1218,9 @@ substitutions:
_SINGLESTORE_PORT: "3308"
_SINGLESTORE_DATABASE: "singlestore"
_SINGLESTORE_USER: "root"
_COCKROACHDB_HOST: 127.0.0.1
_COCKROACHDB_PORT: "26257"
_COCKROACHDB_USER: "root"
_MARIADB_PORT: "3307"
_MARIADB_DATABASE: test_database
_SNOWFLAKE_DATABASE: "test"

View File

@@ -110,6 +110,10 @@ import (
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlmysql/cloudsqlmysqlcreateinstance"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlpg/cloudsqlpgcreateinstances"
_ "github.com/googleapis/genai-toolbox/internal/tools/cloudsqlpg/cloudsqlpgupgradeprecheck"
_ "github.com/googleapis/genai-toolbox/internal/tools/cockroachdb/cockroachdbexecutesql"
_ "github.com/googleapis/genai-toolbox/internal/tools/cockroachdb/cockroachdblistschemas"
_ "github.com/googleapis/genai-toolbox/internal/tools/cockroachdb/cockroachdblisttables"
_ "github.com/googleapis/genai-toolbox/internal/tools/cockroachdb/cockroachdbsql"
_ "github.com/googleapis/genai-toolbox/internal/tools/couchbase"
_ "github.com/googleapis/genai-toolbox/internal/tools/dataform/dataformcompilelocal"
_ "github.com/googleapis/genai-toolbox/internal/tools/dataplex/dataplexlookupentry"
@@ -256,6 +260,7 @@ import (
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmssql"
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql"
_ "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
_ "github.com/googleapis/genai-toolbox/internal/sources/cockroachdb"
_ "github.com/googleapis/genai-toolbox/internal/sources/couchbase"
_ "github.com/googleapis/genai-toolbox/internal/sources/dataplex"
_ "github.com/googleapis/genai-toolbox/internal/sources/dgraph"

View File

@@ -342,16 +342,17 @@
"node_modules/asynckit": {
"version": "0.4.0",
"resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz",
"integrity": "sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q=="
"integrity": "sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q==",
"license": "MIT"
},
"node_modules/axios": {
"version": "1.12.2",
"resolved": "https://registry.npmjs.org/axios/-/axios-1.12.2.tgz",
"integrity": "sha512-vMJzPewAlRyOgxV2dU0Cuz2O8zzzx9VYtbJOaBgXFeLc4IV/Eg50n4LowmehOOR61S8ZMpc2K5Sa7g6A4jfkUw==",
"version": "1.13.5",
"resolved": "https://registry.npmjs.org/axios/-/axios-1.13.5.tgz",
"integrity": "sha512-cz4ur7Vb0xS4/KUN0tPWe44eqxrIu31me+fbang3ijiNscE129POzipJJA6zniq2C/Z6sJCjMimjS8Lc/GAs8Q==",
"license": "MIT",
"dependencies": {
"follow-redirects": "^1.15.6",
"form-data": "^4.0.4",
"follow-redirects": "^1.15.11",
"form-data": "^4.0.5",
"proxy-from-env": "^1.1.0"
}
},
@@ -391,6 +392,7 @@
"version": "1.0.2",
"resolved": "https://registry.npmjs.org/call-bind-apply-helpers/-/call-bind-apply-helpers-1.0.2.tgz",
"integrity": "sha512-Sp1ablJ0ivDkSzjcaJdxEunN5/XvksFJ2sMBFfq6x0ryhQV/2b/KwFe21cMpmHtPOSij8K99/wSfoEuTObmuMQ==",
"license": "MIT",
"dependencies": {
"es-errors": "^1.3.0",
"function-bind": "^1.1.2"
@@ -403,6 +405,7 @@
"version": "1.0.8",
"resolved": "https://registry.npmjs.org/combined-stream/-/combined-stream-1.0.8.tgz",
"integrity": "sha512-FQN4MRfuJeHf7cBbBMJFXhKSDq+2kAArBlmRBvcvFE5BB1HZKXtSFASDhdlz9zOYwxh8lDdnvmMOe/+5cdoEdg==",
"license": "MIT",
"dependencies": {
"delayed-stream": "~1.0.0"
},
@@ -446,6 +449,7 @@
"version": "1.0.0",
"resolved": "https://registry.npmjs.org/delayed-stream/-/delayed-stream-1.0.0.tgz",
"integrity": "sha512-ZySD7Nf91aLB0RxL4KGrKHBXl7Eds1DAmEdcoVawXnLD7SDhpNgtuII2aAkg7a7QS41jxPSZ17p4VdGnMHk3MQ==",
"license": "MIT",
"engines": {
"node": ">=0.4.0"
}
@@ -505,6 +509,7 @@
"version": "1.0.1",
"resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz",
"integrity": "sha512-KIN/nDJBQRcXw0MLVhZE9iQHmG68qAVIBg9CqmUYjmQIhgij9U5MFvrqkUL5FbtyyzZuOeOt0zdeRe4UY7ct+A==",
"license": "MIT",
"dependencies": {
"call-bind-apply-helpers": "^1.0.1",
"es-errors": "^1.3.0",
@@ -537,6 +542,7 @@
"version": "1.0.1",
"resolved": "https://registry.npmjs.org/es-define-property/-/es-define-property-1.0.1.tgz",
"integrity": "sha512-e3nRfgfUZ4rNGL232gUgX06QNyyez04KdjFrF+LTRoOXmrOgFKDg4BCdsjW8EnT69eqdYGmRpJwiPVYNrCaW3g==",
"license": "MIT",
"engines": {
"node": ">= 0.4"
}
@@ -545,6 +551,7 @@
"version": "1.3.0",
"resolved": "https://registry.npmjs.org/es-errors/-/es-errors-1.3.0.tgz",
"integrity": "sha512-Zf5H2Kxt2xjTvbJvP2ZWLEICxA6j+hAmMzIlypy4xcBg1vKVnx89Wy0GbS+kf5cwCVFFzdCFh2XSCFNULS6csw==",
"license": "MIT",
"engines": {
"node": ">= 0.4"
}
@@ -553,6 +560,7 @@
"version": "1.1.1",
"resolved": "https://registry.npmjs.org/es-object-atoms/-/es-object-atoms-1.1.1.tgz",
"integrity": "sha512-FGgH2h8zKNim9ljj7dankFPcICIK9Cp5bm+c2gQSYePhpaG5+esrLODihIorn+Pe6FGJzWhXQotPv73jTaldXA==",
"license": "MIT",
"dependencies": {
"es-errors": "^1.3.0"
},
@@ -564,6 +572,7 @@
"version": "2.1.0",
"resolved": "https://registry.npmjs.org/es-set-tostringtag/-/es-set-tostringtag-2.1.0.tgz",
"integrity": "sha512-j6vWzfrGVfyXxge+O0x5sh6cvxAog0a/4Rdd2K36zCMV5eJ+/+tOAngRO8cODMNWbVRdVlmGZQL2YS3yR8bIUA==",
"license": "MIT",
"dependencies": {
"es-errors": "^1.3.0",
"get-intrinsic": "^1.2.6",
@@ -621,9 +630,10 @@
}
},
"node_modules/form-data": {
"version": "4.0.4",
"resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.4.tgz",
"integrity": "sha512-KrGhL9Q4zjj0kiUt5OO4Mr/A/jlI2jDYs5eHBpYHPcBEVSiipAvn2Ko2HnPe20rmcuuvMHNdZFp+4IlGTMF0Ow==",
"version": "4.0.5",
"resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.5.tgz",
"integrity": "sha512-8RipRLol37bNs2bhoV67fiTEvdTrbMUYcFTiy3+wuuOnUog2QBHCZWXDRijWQfAkhBj2Uf5UnVaiWwA5vdd82w==",
"license": "MIT",
"dependencies": {
"asynckit": "^0.4.0",
"combined-stream": "^1.0.8",
@@ -650,6 +660,7 @@
"version": "1.1.2",
"resolved": "https://registry.npmjs.org/function-bind/-/function-bind-1.1.2.tgz",
"integrity": "sha512-7XHNxH7qX9xG5mIwxkhumTox/MIRNcOgDrxWsMt2pAr23WHp6MrRlN7FBSFpCpr+oVO0F744iUgR82nJMfG2SA==",
"license": "MIT",
"funding": {
"url": "https://github.com/sponsors/ljharb"
}
@@ -684,6 +695,7 @@
"version": "1.3.0",
"resolved": "https://registry.npmjs.org/get-intrinsic/-/get-intrinsic-1.3.0.tgz",
"integrity": "sha512-9fSjSaos/fRIVIp+xSJlE6lfwhES7LNtKaCBIamHsjr2na1BiABJPo0mOjjz8GJDURarmCPGqaiVg5mfjb98CQ==",
"license": "MIT",
"dependencies": {
"call-bind-apply-helpers": "^1.0.2",
"es-define-property": "^1.0.1",
@@ -707,6 +719,7 @@
"version": "1.0.1",
"resolved": "https://registry.npmjs.org/get-proto/-/get-proto-1.0.1.tgz",
"integrity": "sha512-sTSfBjoXBp89JvIKIefqw7U2CCebsc74kiY6awiGogKtoSGbgjYE/G/+l9sF3MWFPNc9IcoOC4ODfKHfxFmp0g==",
"license": "MIT",
"dependencies": {
"dunder-proto": "^1.0.1",
"es-object-atoms": "^1.0.0"
@@ -744,6 +757,7 @@
"version": "1.2.0",
"resolved": "https://registry.npmjs.org/gopd/-/gopd-1.2.0.tgz",
"integrity": "sha512-ZUKRh6/kUFoAiTAtTYPZJ3hw9wNxx+BIBOijnlG9PnrJsCcSjs1wyyD6vJpaYtgnzDrKYRSqf3OO6Rfa93xsRg==",
"license": "MIT",
"engines": {
"node": ">= 0.4"
},
@@ -767,6 +781,7 @@
"version": "1.1.0",
"resolved": "https://registry.npmjs.org/has-symbols/-/has-symbols-1.1.0.tgz",
"integrity": "sha512-1cDNdwJ2Jaohmb3sg4OmKaMBwuC48sYni5HUw2DvsC8LjGTLK9h+eb1X6RyuOHe4hT0ULCW68iomhjUoKUqlPQ==",
"license": "MIT",
"engines": {
"node": ">= 0.4"
},
@@ -778,6 +793,7 @@
"version": "1.0.2",
"resolved": "https://registry.npmjs.org/has-tostringtag/-/has-tostringtag-1.0.2.tgz",
"integrity": "sha512-NqADB8VjPFLM2V0VvHUewwwsw0ZWBaIdgo+ieHtK3hasLz4qeCRjYcqfB6AQrBggRKppKF8L52/VqdVsO47Dlw==",
"license": "MIT",
"dependencies": {
"has-symbols": "^1.0.3"
},
@@ -792,6 +808,7 @@
"version": "2.0.2",
"resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.2.tgz",
"integrity": "sha512-0hJU9SCPvmMzIBdZFqNPXWa6dqh7WdH0cII9y+CyS8rG3nL48Bclra9HmKhVVUHyPWNH5Y7xDwAB7bfgSjkUMQ==",
"license": "MIT",
"dependencies": {
"function-bind": "^1.1.2"
},
@@ -989,6 +1006,7 @@
"version": "1.1.0",
"resolved": "https://registry.npmjs.org/math-intrinsics/-/math-intrinsics-1.1.0.tgz",
"integrity": "sha512-/IXtbwEk5HTPyEwyKX6hGkYXxM9nbj64B+ilVJnC/R6B0pH5G4V3b0pVbL7DBj4tkhBAppbQUlf6F6Xl9LHu1g==",
"license": "MIT",
"engines": {
"node": ">= 0.4"
}
@@ -997,6 +1015,7 @@
"version": "1.52.0",
"resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.52.0.tgz",
"integrity": "sha512-sPU4uV7dYlvtWJxwwxHD0PuihVNiE7TyAbQ5SWxDCB9mUYvOgroQOwYQQOKPJ8CIbE+1ETVlOoK1UC2nU3gYvg==",
"license": "MIT",
"engines": {
"node": ">= 0.6"
}
@@ -1005,6 +1024,7 @@
"version": "2.1.35",
"resolved": "https://registry.npmjs.org/mime-types/-/mime-types-2.1.35.tgz",
"integrity": "sha512-ZDY+bPm5zTTF+YpCrAU9nK0UgICYPT0QtT1NZWFv4s++TNkcgVaT0g6+4R2uI4MjQjzysHB1zxuWL50hzaeXiw==",
"license": "MIT",
"dependencies": {
"mime-db": "1.52.0"
},

View File

@@ -0,0 +1,242 @@
---
title: "CockroachDB"
type: docs
weight: 1
description: >
CockroachDB is a distributed SQL database built for cloud applications.
---
## About
[CockroachDB][crdb-docs] is a distributed SQL database designed for cloud-native applications. It provides strong consistency, horizontal scalability, and built-in resilience with automatic failover and recovery. CockroachDB uses the PostgreSQL wire protocol, making it compatible with many PostgreSQL tools and drivers while providing unique features like multi-region deployments and distributed transactions.
**Minimum Version:** CockroachDB v25.1 or later is recommended for full tool compatibility.
[crdb-docs]: https://www.cockroachlabs.com/docs/
## Available Tools
- [`cockroachdb-sql`](../tools/cockroachdb/cockroachdb-sql.md)
Execute SQL queries as prepared statements in CockroachDB (alias for execute-sql).
- [`cockroachdb-execute-sql`](../tools/cockroachdb/cockroachdb-execute-sql.md)
Run parameterized SQL statements in CockroachDB.
- [`cockroachdb-list-schemas`](../tools/cockroachdb/cockroachdb-list-schemas.md)
List schemas in a CockroachDB database.
- [`cockroachdb-list-tables`](../tools/cockroachdb/cockroachdb-list-tables.md)
List tables in a CockroachDB database.
## Requirements
### Database User
This source uses standard authentication. You will need to [create a CockroachDB user][crdb-users] to login to the database with. For CockroachDB Cloud deployments, SSL/TLS is required.
[crdb-users]: https://www.cockroachlabs.com/docs/stable/create-user.html
### SSL/TLS Configuration
CockroachDB Cloud clusters require SSL/TLS connections. Use the `queryParams` section to configure SSL settings:
- **For CockroachDB Cloud**: Use `sslmode: require` at minimum
- **For self-hosted with certificates**: Use `sslmode: verify-full` with certificate paths
- **For local development only**: Use `sslmode: disable` (not recommended for production)
## Example
```yaml
sources:
my_cockroachdb:
type: cockroachdb
host: your-cluster.cockroachlabs.cloud
port: "26257"
user: myuser
password: mypassword
database: defaultdb
maxRetries: 5
retryBaseDelay: 500ms
queryParams:
sslmode: require
application_name: my-app
# MCP Security Settings (recommended for production)
readOnlyMode: true # Read-only by default (MCP best practice)
enableWriteMode: false # Set to true to allow write operations
maxRowLimit: 1000 # Limit query results
queryTimeoutSec: 30 # Prevent long-running queries
enableTelemetry: true # Enable observability
telemetryVerbose: false # Set true for detailed logs
clusterID: "my-cluster" # Optional identifier
tools:
list_expenses:
type: cockroachdb-sql
source: my_cockroachdb
description: List all expenses
statement: SELECT id, description, amount, category FROM expenses WHERE user_id = $1
parameters:
- name: user_id
type: string
description: The user's ID
describe_expenses:
type: cockroachdb-describe-table
source: my_cockroachdb
description: Describe the expenses table schema
list_expenses_indexes:
type: cockroachdb-list-indexes
source: my_cockroachdb
description: List indexes on the expenses table
```
## Configuration Parameters
### Required Parameters
| Parameter | Type | Description |
|-----------|------|-------------|
| `type` | string | Must be `cockroachdb` |
| `host` | string | The hostname or IP address of the CockroachDB cluster |
| `port` | string | The port number (typically "26257") |
| `user` | string | The database user name |
| `database` | string | The database name to connect to |
### Optional Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `password` | string | "" | The database password (can be empty for certificate-based auth) |
| `maxRetries` | integer | 5 | Maximum number of connection retry attempts |
| `retryBaseDelay` | string | "500ms" | Base delay between retry attempts (exponential backoff) |
| `queryParams` | map | {} | Additional connection parameters (e.g., SSL configuration) |
### MCP Security Parameters
CockroachDB integration includes security features following the [Model Context Protocol (MCP)](https://modelcontextprotocol.io/) specification:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `readOnlyMode` | boolean | true | Enables read-only mode by default (MCP requirement) |
| `enableWriteMode` | boolean | false | Explicitly enable write operations (INSERT/UPDATE/DELETE/CREATE/DROP) |
| `maxRowLimit` | integer | 1000 | Maximum rows returned per SELECT query (auto-adds LIMIT clause) |
| `queryTimeoutSec` | integer | 30 | Query timeout in seconds to prevent long-running queries |
| `enableTelemetry` | boolean | true | Enable structured logging of tool invocations |
| `telemetryVerbose` | boolean | false | Enable detailed JSON telemetry output |
| `clusterID` | string | "" | Optional cluster identifier for telemetry |
### Query Parameters
Common query parameters for CockroachDB connections:
| Parameter | Values | Description |
|-----------|--------|-------------|
| `sslmode` | `disable`, `require`, `verify-ca`, `verify-full` | SSL/TLS mode (CockroachDB Cloud requires `require` or higher) |
| `sslrootcert` | file path | Path to root certificate for SSL verification |
| `sslcert` | file path | Path to client certificate |
| `sslkey` | file path | Path to client key |
| `application_name` | string | Application name for connection tracking |
## Best Practices
### Security and MCP Compliance
**Read-Only by Default**: The integration follows MCP best practices by defaulting to read-only mode. This prevents accidental data modifications:
```yaml
sources:
my_cockroachdb:
readOnlyMode: true # Default behavior
enableWriteMode: false # Explicit write opt-in required
```
To enable write operations:
```yaml
sources:
my_cockroachdb:
readOnlyMode: false # Disable read-only protection
enableWriteMode: true # Explicitly allow writes
```
**Query Limits**: Automatic row limits prevent excessive data retrieval:
- SELECT queries automatically get `LIMIT 1000` appended (configurable via `maxRowLimit`)
- Queries are terminated after 30 seconds (configurable via `queryTimeoutSec`)
**Observability**: Structured telemetry provides visibility into tool usage:
- Tool invocations are logged with status, latency, and row counts
- SQL queries are redacted to protect sensitive values
- Set `telemetryVerbose: true` for detailed JSON logs
### Use UUID Primary Keys
CockroachDB performs best with UUID primary keys rather than sequential integers to avoid transaction hotspots:
```sql
CREATE TABLE expenses (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
description TEXT,
amount DECIMAL(10,2)
);
```
### Automatic Transaction Retry
This source uses the official `cockroach-go/v2` library which provides automatic transaction retry for serialization conflicts. For write operations requiring explicit transaction control, tools can use the `ExecuteTxWithRetry` method.
### Multi-Region Deployments
CockroachDB supports multi-region deployments with automatic data distribution. Configure your cluster's regions and survival goals separately from the Toolbox configuration. The source will connect to any node in the cluster.
### Connection Pooling
The source maintains a connection pool to the CockroachDB cluster. The pool automatically handles:
- Load balancing across cluster nodes
- Connection retry with exponential backoff
- Health checking of connections
## Troubleshooting
### SSL/TLS Errors
If you encounter "server requires encryption" errors:
1. For CockroachDB Cloud, ensure `sslmode` is set to `require` or higher:
```yaml
queryParams:
sslmode: require
```
2. For certificate verification, download your cluster's root certificate and configure:
```yaml
queryParams:
sslmode: verify-full
sslrootcert: /path/to/ca.crt
```
### Connection Timeouts
If experiencing connection timeouts:
1. Check network connectivity to the CockroachDB cluster
2. Verify firewall rules allow connections on port 26257
3. For CockroachDB Cloud, ensure IP allowlisting is configured
4. Increase `maxRetries` or `retryBaseDelay` if needed
### Transaction Retry Errors
CockroachDB may encounter serializable transaction conflicts. The integration automatically handles these retries using the cockroach-go library. If you see retry-related errors, check:
1. Database load and contention
2. Query patterns that might cause conflicts
3. Consider using `SELECT FOR UPDATE` for explicit locking
## Additional Resources
- [CockroachDB Documentation](https://www.cockroachlabs.com/docs/)
- [CockroachDB Best Practices](https://www.cockroachlabs.com/docs/stable/performance-best-practices-overview.html)
- [Multi-Region Capabilities](https://www.cockroachlabs.com/docs/stable/multiregion-overview.html)
- [Connection Parameters](https://www.cockroachlabs.com/docs/stable/connection-parameters.html)

View File

@@ -0,0 +1,273 @@
---
title: "cockroachdb-execute-sql"
type: docs
weight: 1
description: >
Execute ad-hoc SQL statements against a CockroachDB database.
---
## About
A `cockroachdb-execute-sql` tool executes ad-hoc SQL statements against a CockroachDB database. This tool is designed for interactive workflows where the SQL query is provided dynamically at runtime, making it ideal for developer assistance and exploratory data analysis.
The tool takes a single `sql` parameter containing the SQL statement to execute and returns the query results.
> **Note:** This tool is intended for developer assistant workflows with human-in-the-loop and shouldn't be used for production agents. For production use cases with predefined queries, use [cockroachdb-sql](./cockroachdb-sql.md) instead.
## Example
```yaml
sources:
my_cockroachdb:
type: cockroachdb
host: your-cluster.cockroachlabs.cloud
port: "26257"
user: myuser
password: mypassword
database: defaultdb
queryParams:
sslmode: require
tools:
execute_sql:
type: cockroachdb-execute-sql
source: my_cockroachdb
description: Execute any SQL statement against the CockroachDB database
```
## Usage Examples
### Simple SELECT Query
```json
{
"sql": "SELECT * FROM users LIMIT 10"
}
```
### Query with Aggregations
```json
{
"sql": "SELECT category, COUNT(*) as count, SUM(amount) as total FROM expenses GROUP BY category ORDER BY total DESC"
}
```
### Database Introspection
```json
{
"sql": "SHOW TABLES"
}
```
```json
{
"sql": "SHOW COLUMNS FROM expenses"
}
```
### Multi-Region Information
```json
{
"sql": "SHOW REGIONS FROM DATABASE defaultdb"
}
```
```json
{
"sql": "SHOW ZONE CONFIGURATIONS"
}
```
## CockroachDB-Specific Features
### Check Cluster Version
```json
{
"sql": "SELECT version()"
}
```
### View Node Status
```json
{
"sql": "SELECT node_id, address, locality, is_live FROM crdb_internal.gossip_nodes"
}
```
### Check Replication Status
```json
{
"sql": "SELECT range_id, start_key, end_key, replicas, lease_holder FROM crdb_internal.ranges LIMIT 10"
}
```
### View Table Regions
```json
{
"sql": "SHOW REGIONS FROM TABLE expenses"
}
```
## Configuration
### Required Fields
| Field | Type | Description |
|-------|------|-------------|
| `type` | string | Must be `cockroachdb-execute-sql` |
| `source` | string | Name of the CockroachDB source to use |
| `description` | string | Human-readable description for the LLM |
### Optional Fields
| Field | Type | Description |
|-------|------|-------------|
| `authRequired` | array | List of authentication services required |
## Parameters
The tool accepts a single runtime parameter:
| Parameter | Type | Description |
|-----------|------|-------------|
| `sql` | string | The SQL statement to execute |
## Best Practices
### Use for Exploration, Not Production
This tool is ideal for:
- Interactive database exploration
- Ad-hoc analysis and reporting
- Debugging and troubleshooting
- Schema inspection
For production use cases, use [cockroachdb-sql](./cockroachdb-sql.md) with parameterized queries.
### Be Cautious with Data Modification
While this tool can execute any SQL statement, be careful with:
- `INSERT`, `UPDATE`, `DELETE` statements
- `DROP` or `ALTER` statements
- Schema changes in production
### Use LIMIT for Large Results
Always use `LIMIT` clauses when exploring data:
```sql
SELECT * FROM large_table LIMIT 100
```
### Leverage CockroachDB's SQL Extensions
CockroachDB supports PostgreSQL syntax plus extensions:
```sql
-- Show database survival goal
SHOW SURVIVAL GOAL FROM DATABASE defaultdb;
-- View zone configurations
SHOW ZONE CONFIGURATION FOR TABLE expenses;
-- Check table localities
SHOW CREATE TABLE expenses;
```
## Error Handling
The tool will return descriptive errors for:
- **Syntax errors**: Invalid SQL syntax
- **Permission errors**: Insufficient user privileges
- **Connection errors**: Network or authentication issues
- **Runtime errors**: Constraint violations, type mismatches, etc.
## Security Considerations
### SQL Injection Risk
Since this tool executes arbitrary SQL, it should only be used with:
- Trusted users in interactive sessions
- Human-in-the-loop workflows
- Development and testing environments
Never expose this tool directly to end users without proper authorization controls.
### Use Authentication
Configure the `authRequired` field to restrict access:
```yaml
tools:
execute_sql:
type: cockroachdb-execute-sql
source: my_cockroachdb
description: Execute SQL statements
authRequired:
- my-auth-service
```
### Read-Only Users
For safer exploration, create read-only database users:
```sql
CREATE USER readonly_user;
GRANT SELECT ON DATABASE defaultdb TO readonly_user;
```
## Common Use Cases
### Database Administration
```sql
-- View database size
SELECT
table_name,
pg_size_pretty(pg_total_relation_size(table_name::regclass)) AS size
FROM information_schema.tables
WHERE table_schema = 'public'
ORDER BY pg_total_relation_size(table_name::regclass) DESC;
```
### Performance Analysis
```sql
-- Find slow queries
SELECT query, count, mean_latency
FROM crdb_internal.statement_statistics
WHERE mean_latency > INTERVAL '1 second'
ORDER BY mean_latency DESC
LIMIT 10;
```
### Data Quality Checks
```sql
-- Find NULL values
SELECT COUNT(*) as null_count
FROM expenses
WHERE description IS NULL OR amount IS NULL;
-- Find duplicates
SELECT user_id, email, COUNT(*) as count
FROM users
GROUP BY user_id, email
HAVING COUNT(*) > 1;
```
## See Also
- [cockroachdb-sql](./cockroachdb-sql.md) - For parameterized, production-ready queries
- [cockroachdb-list-tables](./cockroachdb-list-tables.md) - List tables in the database
- [cockroachdb-list-schemas](./cockroachdb-list-schemas.md) - List database schemas
- [CockroachDB Source](../../sources/cockroachdb.md) - Source configuration reference
- [CockroachDB SQL Reference](https://www.cockroachlabs.com/docs/stable/sql-statements.html) - Official SQL documentation

View File

@@ -0,0 +1,305 @@
---
title: "cockroachdb-list-schemas"
type: docs
weight: 1
description: >
List schemas in a CockroachDB database.
---
## About
The `cockroachdb-list-schemas` tool retrieves a list of schemas (namespaces) in a CockroachDB database. Schemas are used to organize database objects such as tables, views, and functions into logical groups.
This tool is useful for:
- Understanding database organization
- Discovering available schemas
- Multi-tenant application analysis
- Schema-level access control planning
## Example
```yaml
sources:
my_cockroachdb:
type: cockroachdb
host: your-cluster.cockroachlabs.cloud
port: "26257"
user: myuser
password: mypassword
database: defaultdb
queryParams:
sslmode: require
tools:
list_schemas:
type: cockroachdb-list-schemas
source: my_cockroachdb
description: List all schemas in the database
```
## Configuration
### Required Fields
| Field | Type | Description |
|-------|------|-------------|
| `type` | string | Must be `cockroachdb-list-schemas` |
| `source` | string | Name of the CockroachDB source to use |
| `description` | string | Human-readable description for the LLM |
### Optional Fields
| Field | Type | Description |
|-------|------|-------------|
| `authRequired` | array | List of authentication services required |
## Output Structure
The tool returns a list of schemas with the following information:
```json
[
{
"catalog_name": "defaultdb",
"schema_name": "public",
"is_user_defined": true
},
{
"catalog_name": "defaultdb",
"schema_name": "analytics",
"is_user_defined": true
}
]
```
### Fields
| Field | Type | Description |
|-------|------|-------------|
| `catalog_name` | string | The database (catalog) name |
| `schema_name` | string | The schema name |
| `is_user_defined` | boolean | Whether this is a user-created schema (excludes system schemas) |
## Usage Example
```json
{}
```
No parameters are required. The tool automatically lists all user-defined schemas.
## Default Schemas
CockroachDB includes several standard schemas:
- **`public`**: The default schema for user objects
- **`pg_catalog`**: PostgreSQL system catalog (excluded from results)
- **`information_schema`**: SQL standard metadata views (excluded from results)
- **`crdb_internal`**: CockroachDB internal metadata (excluded from results)
- **`pg_extension`**: PostgreSQL extension objects (excluded from results)
The tool filters out system schemas and only returns user-defined schemas.
## Schema Management in CockroachDB
### Creating Schemas
```sql
CREATE SCHEMA analytics;
```
### Using Schemas
```sql
-- Create table in specific schema
CREATE TABLE analytics.revenue (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
amount DECIMAL(10,2),
date DATE
);
-- Query from specific schema
SELECT * FROM analytics.revenue;
```
### Schema Search Path
The search path determines which schemas are searched for unqualified object names:
```sql
-- Show current search path
SHOW search_path;
-- Set search path
SET search_path = analytics, public;
```
## Multi-Tenant Applications
Schemas are commonly used for multi-tenant applications:
```sql
-- Create schema per tenant
CREATE SCHEMA tenant_acme;
CREATE SCHEMA tenant_globex;
-- Create same table structure in each schema
CREATE TABLE tenant_acme.orders (...);
CREATE TABLE tenant_globex.orders (...);
```
The `cockroachdb-list-schemas` tool helps discover all tenant schemas:
```yaml
tools:
list_tenants:
type: cockroachdb-list-schemas
source: my_cockroachdb
description: |
List all tenant schemas in the database.
Each schema represents a separate tenant's data namespace.
```
## Best Practices
### Use Schemas for Organization
Group related tables into schemas:
```sql
CREATE SCHEMA sales;
CREATE SCHEMA inventory;
CREATE SCHEMA hr;
CREATE TABLE sales.orders (...);
CREATE TABLE inventory.products (...);
CREATE TABLE hr.employees (...);
```
### Schema Naming Conventions
Use clear, descriptive schema names:
- Lowercase names
- Use underscores for multi-word names
- Avoid reserved keywords
- Use prefixes for grouped schemas (e.g., `tenant_`, `app_`)
### Schema-Level Permissions
Schemas enable fine-grained access control:
```sql
-- Grant access to specific schema
GRANT USAGE ON SCHEMA analytics TO analyst_role;
GRANT SELECT ON ALL TABLES IN SCHEMA analytics TO analyst_role;
-- Revoke access
REVOKE ALL ON SCHEMA hr FROM public;
```
## Integration with Other Tools
### Combined with List Tables
```yaml
tools:
list_schemas:
type: cockroachdb-list-schemas
source: my_cockroachdb
description: List all schemas first
list_tables:
type: cockroachdb-list-tables
source: my_cockroachdb
description: |
List tables in the database.
Use list_schemas first to understand schema organization.
```
### Schema Discovery Workflow
1. Call `cockroachdb-list-schemas` to discover schemas
2. Call `cockroachdb-list-tables` to see tables in each schema
3. Generate queries using fully qualified names: `schema.table`
## Common Use Cases
### Discover Database Structure
```yaml
tools:
discover_schemas:
type: cockroachdb-list-schemas
source: my_cockroachdb
description: |
Discover how the database is organized into schemas.
Use this to understand the logical grouping of tables.
```
### Multi-Tenant Analysis
```yaml
tools:
list_tenant_schemas:
type: cockroachdb-list-schemas
source: my_cockroachdb
description: |
List all tenant schemas (each tenant has their own schema).
Schema names follow the pattern: tenant_<company_name>
```
### Schema Migration Planning
```yaml
tools:
audit_schemas:
type: cockroachdb-list-schemas
source: my_cockroachdb
description: |
Audit existing schemas before migration.
Identifies all schemas that need to be migrated.
```
## Error Handling
The tool handles common errors:
- **Connection errors**: Returns connection failure details
- **Permission errors**: Returns error if user lacks USAGE privilege
- **Empty results**: Returns empty array if no user schemas exist
## Permissions Required
To list schemas, the user needs:
- `CONNECT` privilege on the database
- No specific schema privileges required for listing
To query objects within schemas, the user needs:
- `USAGE` privilege on the schema
- Appropriate object privileges (SELECT, INSERT, etc.)
## CockroachDB-Specific Features
### System Schemas
CockroachDB includes PostgreSQL-compatible system schemas plus CockroachDB-specific ones:
- `crdb_internal.*`: CockroachDB internal metadata and statistics
- `pg_catalog.*`: PostgreSQL system catalog
- `information_schema.*`: SQL standard information schema
These are automatically filtered from the results.
### User-Defined Flag
The `is_user_defined` field helps distinguish:
- `true`: User-created schemas
- `false`: System schemas (already filtered out)
## See Also
- [cockroachdb-sql](./cockroachdb-sql.md) - Execute parameterized queries
- [cockroachdb-execute-sql](./cockroachdb-execute-sql.md) - Execute ad-hoc SQL
- [cockroachdb-list-tables](./cockroachdb-list-tables.md) - List tables in the database
- [CockroachDB Source](../../sources/cockroachdb.md) - Source configuration reference
- [CockroachDB Schema Design](https://www.cockroachlabs.com/docs/stable/schema-design-overview.html) - Official documentation

View File

@@ -0,0 +1,344 @@
---
title: "cockroachdb-list-tables"
type: docs
weight: 1
description: >
List tables in a CockroachDB database with schema details.
---
## About
The `cockroachdb-list-tables` tool retrieves a list of tables from a CockroachDB database. It provides detailed information about table structure, including columns, constraints, indexes, and foreign key relationships.
This tool is useful for:
- Database schema discovery
- Understanding table relationships
- Generating context for AI-powered database queries
- Documentation and analysis
## Example
```yaml
sources:
my_cockroachdb:
type: cockroachdb
host: your-cluster.cockroachlabs.cloud
port: "26257"
user: myuser
password: mypassword
database: defaultdb
queryParams:
sslmode: require
tools:
list_all_tables:
type: cockroachdb-list-tables
source: my_cockroachdb
description: List all user tables in the database with their structure
```
## Configuration
### Required Fields
| Field | Type | Description |
|-------|------|-------------|
| `type` | string | Must be `cockroachdb-list-tables` |
| `source` | string | Name of the CockroachDB source to use |
| `description` | string | Human-readable description for the LLM |
### Optional Fields
| Field | Type | Description |
|-------|------|-------------|
| `authRequired` | array | List of authentication services required |
## Parameters
The tool accepts optional runtime parameters:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `table_names` | array | all tables | List of specific table names to retrieve |
| `output_format` | string | "detailed" | Output format: "simple" or "detailed" |
## Output Formats
### Simple Format
Returns basic table information:
- Table name
- Row count estimate
- Size information
```json
{
"table_names": ["users"],
"output_format": "simple"
}
```
### Detailed Format (Default)
Returns comprehensive table information:
- Table name and schema
- All columns with types and constraints
- Primary keys
- Foreign keys and relationships
- Indexes
- Check constraints
- Table size and row counts
```json
{
"table_names": ["users", "orders"],
"output_format": "detailed"
}
```
## Usage Examples
### List All Tables
```json
{}
```
### List Specific Tables
```json
{
"table_names": ["users", "orders", "expenses"]
}
```
### Simple Output
```json
{
"output_format": "simple"
}
```
## Output Structure
### Simple Format Output
```json
{
"table_name": "users",
"estimated_rows": 1000,
"size": "128 KB"
}
```
### Detailed Format Output
```json
{
"table_name": "users",
"schema": "public",
"columns": [
{
"name": "id",
"type": "UUID",
"nullable": false,
"default": "gen_random_uuid()"
},
{
"name": "email",
"type": "STRING",
"nullable": false,
"default": null
},
{
"name": "created_at",
"type": "TIMESTAMP",
"nullable": false,
"default": "now()"
}
],
"primary_key": ["id"],
"indexes": [
{
"name": "users_pkey",
"columns": ["id"],
"unique": true,
"primary": true
},
{
"name": "users_email_idx",
"columns": ["email"],
"unique": true,
"primary": false
}
],
"foreign_keys": [],
"constraints": [
{
"name": "users_email_check",
"type": "CHECK",
"definition": "email ~* '^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\\.[A-Z|a-z]{2,}$'"
}
]
}
```
## CockroachDB-Specific Information
### UUID Primary Keys
The tool recognizes CockroachDB's recommended UUID primary key pattern:
```sql
CREATE TABLE users (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
...
);
```
### Multi-Region Tables
For multi-region tables, the output includes locality information:
```json
{
"table_name": "users",
"locality": "REGIONAL BY ROW",
"regions": ["us-east-1", "us-west-2", "eu-west-1"]
}
```
### Interleaved Tables
The tool shows parent-child relationships for interleaved tables (legacy feature):
```json
{
"table_name": "order_items",
"interleaved_in": "orders"
}
```
## Best Practices
### Use for Schema Discovery
The tool is ideal for helping AI assistants understand your database structure:
```yaml
tools:
discover_schema:
type: cockroachdb-list-tables
source: my_cockroachdb
description: |
Use this tool first to understand the database schema before generating queries.
It shows all tables, their columns, data types, and relationships.
```
### Filter Large Schemas
For databases with many tables, specify relevant tables:
```json
{
"table_names": ["users", "orders", "products"],
"output_format": "detailed"
}
```
### Use Simple Format for Overviews
When you need just table names and sizes:
```json
{
"output_format": "simple"
}
```
## Excluded Tables
The tool automatically excludes system tables and schemas:
- `pg_catalog.*` - PostgreSQL system catalog
- `information_schema.*` - SQL standard information schema
- `crdb_internal.*` - CockroachDB internal tables
- `pg_extension.*` - PostgreSQL extension tables
Only user-created tables in the public schema (and other user schemas) are returned.
## Error Handling
The tool handles common errors:
- **Table not found**: Returns empty result for non-existent tables
- **Permission errors**: Returns error if user lacks SELECT privileges
- **Connection errors**: Returns connection failure details
## Integration with AI Assistants
### Prompt Example
```yaml
tools:
list_tables:
type: cockroachdb-list-tables
source: my_cockroachdb
description: |
Lists all tables in the database with detailed schema information.
Use this tool to understand:
- What tables exist
- What columns each table has
- Data types and constraints
- Relationships between tables (foreign keys)
- Available indexes
Always call this tool before generating SQL queries to ensure
you use correct table and column names.
```
## Common Use Cases
### Generate Context for Queries
```json
{}
```
This provides comprehensive schema information that helps AI assistants generate accurate SQL queries.
### Analyze Table Structure
```json
{
"table_names": ["users"],
"output_format": "detailed"
}
```
Perfect for understanding a specific table's structure, constraints, and relationships.
### Quick Schema Overview
```json
{
"output_format": "simple"
}
```
Gets a quick list of tables with basic statistics.
## Performance Considerations
- **Simple format** is faster for large databases
- **Detailed format** queries system tables extensively
- Specifying `table_names` reduces query time
- Results are fetched in a single query for efficiency
## See Also
- [cockroachdb-sql](./cockroachdb-sql.md) - Execute parameterized queries
- [cockroachdb-execute-sql](./cockroachdb-execute-sql.md) - Execute ad-hoc SQL
- [cockroachdb-list-schemas](./cockroachdb-list-schemas.md) - List database schemas
- [CockroachDB Source](../../sources/cockroachdb.md) - Source configuration reference
- [CockroachDB Schema Design](https://www.cockroachlabs.com/docs/stable/schema-design-overview.html) - Best practices

View File

@@ -0,0 +1,291 @@
---
title: "cockroachdb-sql"
type: docs
weight: 1
description: >
Execute parameterized SQL queries in CockroachDB.
---
## About
The `cockroachdb-sql` tool allows you to execute parameterized SQL queries against a CockroachDB database. This tool supports prepared statements with parameter binding, template parameters for dynamic query construction, and automatic transaction retry for resilience against serialization conflicts.
## Example
```yaml
sources:
my_cockroachdb:
type: cockroachdb
host: your-cluster.cockroachlabs.cloud
port: "26257"
user: myuser
password: mypassword
database: defaultdb
queryParams:
sslmode: require
tools:
get_user_orders:
type: cockroachdb-sql
source: my_cockroachdb
description: Get all orders for a specific user
statement: |
SELECT o.id, o.order_date, o.total_amount, o.status
FROM orders o
WHERE o.user_id = $1
ORDER BY o.order_date DESC
parameters:
- name: user_id
type: string
description: The UUID of the user
```
## Configuration
### Required Fields
| Field | Type | Description |
|-------|------|-------------|
| `type` | string | Must be `cockroachdb-sql` |
| `source` | string | Name of the CockroachDB source to use |
| `description` | string | Human-readable description of what the tool does |
| `statement` | string | The SQL query to execute |
### Optional Fields
| Field | Type | Description |
|-------|------|-------------|
| `parameters` | array | List of parameter definitions for the query |
| `templateParameters` | array | List of template parameters for dynamic query construction |
| `authRequired` | array | List of authentication services required |
## Parameters
Parameters allow you to safely pass values into your SQL queries using prepared statements. CockroachDB uses PostgreSQL-style parameter placeholders: `$1`, `$2`, etc.
### Parameter Types
- `string`: Text values
- `number`: Numeric values (integers or decimals)
- `boolean`: True/false values
- `array`: Array of values
### Example with Multiple Parameters
```yaml
tools:
filter_expenses:
type: cockroachdb-sql
source: my_cockroachdb
description: Filter expenses by category and date range
statement: |
SELECT id, description, amount, category, expense_date
FROM expenses
WHERE user_id = $1
AND category = $2
AND expense_date >= $3
AND expense_date <= $4
ORDER BY expense_date DESC
parameters:
- name: user_id
type: string
description: The user's UUID
- name: category
type: string
description: Expense category (e.g., "Food", "Transport")
- name: start_date
type: string
description: Start date in YYYY-MM-DD format
- name: end_date
type: string
description: End date in YYYY-MM-DD format
```
## Template Parameters
Template parameters enable dynamic query construction by replacing placeholders in the SQL statement before parameter binding. This is useful for dynamic table names, column names, or query structure.
### Example with Template Parameters
```yaml
tools:
get_column_data:
type: cockroachdb-sql
source: my_cockroachdb
description: Get data from a specific column
statement: |
SELECT {{column_name}}
FROM {{table_name}}
WHERE user_id = $1
LIMIT 100
templateParameters:
- name: table_name
type: string
description: The table to query
- name: column_name
type: string
description: The column to retrieve
parameters:
- name: user_id
type: string
description: The user's UUID
```
## Best Practices
### Use UUID Primary Keys
CockroachDB performs best with UUID primary keys to avoid transaction hotspots:
```sql
CREATE TABLE orders (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
user_id UUID NOT NULL,
order_date TIMESTAMP DEFAULT now(),
total_amount DECIMAL(10,2)
);
```
### Use Indexes for Performance
Create indexes on frequently queried columns:
```sql
CREATE INDEX idx_orders_user_id ON orders(user_id);
CREATE INDEX idx_orders_date ON orders(order_date DESC);
```
### Use JOINs Efficiently
CockroachDB supports standard SQL JOINs. Keep joins efficient by:
- Adding appropriate indexes
- Using UUIDs for foreign keys
- Limiting result sets with WHERE clauses
```yaml
tools:
get_user_with_orders:
type: cockroachdb-sql
source: my_cockroachdb
description: Get user details with their recent orders
statement: |
SELECT u.name, u.email, o.id as order_id, o.order_date, o.total_amount
FROM users u
LEFT JOIN orders o ON u.id = o.user_id
WHERE u.id = $1
ORDER BY o.order_date DESC
LIMIT 10
parameters:
- name: user_id
type: string
description: The user's UUID
```
### Handle NULL Values
Use COALESCE or NULL checks when dealing with nullable columns:
```sql
SELECT id, description, COALESCE(notes, 'No notes') as notes
FROM expenses
WHERE user_id = $1
```
## Error Handling
The tool automatically handles:
- **Connection errors**: Retried with exponential backoff
- **Serialization conflicts**: Automatically retried using cockroach-go library
- **Invalid parameters**: Returns descriptive error messages
- **SQL syntax errors**: Returns database error details
## Advanced Usage
### Aggregations
```yaml
tools:
expense_summary:
type: cockroachdb-sql
source: my_cockroachdb
description: Get expense summary by category for a user
statement: |
SELECT
category,
COUNT(*) as count,
SUM(amount) as total_amount,
AVG(amount) as avg_amount
FROM expenses
WHERE user_id = $1
AND expense_date >= $2
GROUP BY category
ORDER BY total_amount DESC
parameters:
- name: user_id
type: string
description: The user's UUID
- name: start_date
type: string
description: Start date in YYYY-MM-DD format
```
### Window Functions
```yaml
tools:
running_total:
type: cockroachdb-sql
source: my_cockroachdb
description: Get running total of expenses
statement: |
SELECT
expense_date,
amount,
SUM(amount) OVER (ORDER BY expense_date) as running_total
FROM expenses
WHERE user_id = $1
ORDER BY expense_date
parameters:
- name: user_id
type: string
description: The user's UUID
```
### Common Table Expressions (CTEs)
```yaml
tools:
top_spenders:
type: cockroachdb-sql
source: my_cockroachdb
description: Find top spending users
statement: |
WITH user_totals AS (
SELECT
user_id,
SUM(amount) as total_spent
FROM expenses
WHERE expense_date >= $1
GROUP BY user_id
)
SELECT
u.name,
u.email,
ut.total_spent
FROM user_totals ut
JOIN users u ON ut.user_id = u.id
ORDER BY ut.total_spent DESC
LIMIT 10
parameters:
- name: start_date
type: string
description: Start date in YYYY-MM-DD format
```
## See Also
- [cockroachdb-execute-sql](./cockroachdb-execute-sql.md) - For ad-hoc SQL execution
- [cockroachdb-list-tables](./cockroachdb-list-tables.md) - List tables in the database
- [cockroachdb-list-schemas](./cockroachdb-list-schemas.md) - List database schemas
- [CockroachDB Source](../../sources/cockroachdb.md) - Source configuration reference

1
go.mod
View File

@@ -21,6 +21,7 @@ require (
github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace v1.30.0
github.com/apache/cassandra-gocql-driver/v2 v2.0.0
github.com/cenkalti/backoff/v5 v5.0.3
github.com/cockroachdb/cockroach-go/v2 v2.4.2
github.com/couchbase/gocb/v2 v2.11.1
github.com/couchbase/tools-common/http v1.0.9
github.com/elastic/elastic-transport-go/v8 v8.8.0

4
go.sum
View File

@@ -800,6 +800,8 @@ github.com/cncf/xds/go v0.0.0-20230105202645-06c439db220b/go.mod h1:eXthEFrGJvWH
github.com/cncf/xds/go v0.0.0-20230607035331-e9ce68804cb4/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 h1:aQ3y1lwWyqYPiWZThqv1aFbZMiM9vblcSArJRf2Irls=
github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8=
github.com/cockroachdb/cockroach-go/v2 v2.4.2 h1:QB0ozDWQUUJ0GP8Zw63X/qHefPTCpLvtfCs6TLrPgyE=
github.com/cockroachdb/cockroach-go/v2 v2.4.2/go.mod h1:9U179XbCx4qFWtNhc7BiWLPfuyMVQ7qdAhfrwLz1vH0=
github.com/containerd/continuity v0.4.5 h1:ZRoN1sXq9u7V6QoHMcVWGhOwDFqZ4B9i5H6un1Wh0x4=
github.com/containerd/continuity v0.4.5/go.mod h1:/lNJvtJKUQStBzpVQ1+rasXO1LAWtUQssk28EZvJ3nE=
github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI=
@@ -960,6 +962,8 @@ github.com/godror/godror v0.49.6 h1:ts4ZGw8uLJ42e1D7aXmVuSrld0/lzUzmIUjuUuQOgGM=
github.com/godror/godror v0.49.6/go.mod h1:kTMcxZzRw73RT5kn9v3JkBK4kHI6dqowHotqV72ebU8=
github.com/godror/knownpb v0.3.0 h1:+caUdy8hTtl7X05aPl3tdL540TvCcaQA6woZQroLZMw=
github.com/godror/knownpb v0.3.0/go.mod h1:PpTyfJwiOEAzQl7NtVCM8kdPCnp3uhxsZYIzZ5PV4zU=
github.com/gofrs/flock v0.12.1 h1:MTLVXXHf8ekldpJk3AKicLij9MdwOWkZ+a/jHHZby9E=
github.com/gofrs/flock v0.12.1/go.mod h1:9zxTsyu5xtJ9DK+1tFZyibEV7y3uwDxPPfbxeeHCoD0=
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
github.com/golang-jwt/jwt/v5 v5.3.0 h1:pv4AsKCKKZuqlgs5sUmn4x8UlGa0kEVt/puTpKx9vvo=
github.com/golang-jwt/jwt/v5 v5.3.0/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=

View File

@@ -0,0 +1,430 @@
// Copyright 2026 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 cockroachdb
import (
"context"
"encoding/json"
"fmt"
"log/slog"
"math"
"net/url"
"regexp"
"strings"
"time"
crdbpgx "github.com/cockroachdb/cockroach-go/v2/crdb/crdbpgxv5"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"go.opentelemetry.io/otel/trace"
)
const SourceKind string = "cockroachdb"
const SourceType string = "cockroachdb"
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) {
// MCP compliance: Read-only by default, require explicit opt-in for writes
actual := Config{
Name: name,
MaxRetries: 5,
RetryBaseDelay: "500ms",
ReadOnlyMode: true, // MCP requirement: read-only by default
EnableWriteMode: false, // Must be explicitly enabled
MaxRowLimit: 1000, // MCP requirement: limit query results
QueryTimeoutSec: 30, // MCP requirement: prevent long-running queries
EnableTelemetry: true, // MCP requirement: observability
TelemetryVerbose: false,
}
if err := decoder.DecodeContext(ctx, &actual); err != nil {
return nil, err
}
// Security validation: If EnableWriteMode is true, ReadOnlyMode should be false
if actual.EnableWriteMode {
actual.ReadOnlyMode = false
}
return actual, nil
}
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"type" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
User string `yaml:"user" validate:"required"`
Password string `yaml:"password"`
Database string `yaml:"database" validate:"required"`
QueryParams map[string]string `yaml:"queryParams"`
MaxRetries int `yaml:"maxRetries"`
RetryBaseDelay string `yaml:"retryBaseDelay"`
// MCP Security Features
ReadOnlyMode bool `yaml:"readOnlyMode"` // Default: true (enforced in Initialize)
EnableWriteMode bool `yaml:"enableWriteMode"` // Explicit opt-in for write operations
MaxRowLimit int `yaml:"maxRowLimit"` // Default: 1000
QueryTimeoutSec int `yaml:"queryTimeoutSec"` // Default: 30
// Observability
EnableTelemetry bool `yaml:"enableTelemetry"` // Default: true
TelemetryVerbose bool `yaml:"telemetryVerbose"` // Default: false
ClusterID string `yaml:"clusterID"` // Optional cluster identifier for telemetry
}
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) SourceConfigType() string {
return SourceType
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
retryBaseDelay, err := time.ParseDuration(r.RetryBaseDelay)
if err != nil {
return nil, fmt.Errorf("invalid retryBaseDelay: %w", err)
}
pool, err := initCockroachDBConnectionPoolWithRetry(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database, r.QueryParams, r.MaxRetries, retryBaseDelay)
if err != nil {
return nil, fmt.Errorf("unable to create pool: %w", err)
}
s := &Source{
Config: r,
Pool: pool,
}
return s, nil
}
var _ sources.Source = &Source{}
type Source struct {
Config
Pool *pgxpool.Pool
}
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) SourceType() string {
return SourceType
}
func (s *Source) ToConfig() sources.SourceConfig {
return s.Config
}
func (s *Source) CockroachDBPool() *pgxpool.Pool {
return s.Pool
}
func (s *Source) PostgresPool() *pgxpool.Pool {
return s.Pool
}
// ExecuteTxWithRetry executes a function within a transaction with automatic retry logic
// using the official CockroachDB retry mechanism from cockroach-go/v2
func (s *Source) ExecuteTxWithRetry(ctx context.Context, fn func(pgx.Tx) error) error {
return crdbpgx.ExecuteTx(ctx, s.Pool, pgx.TxOptions{}, fn)
}
// Query executes a query using the connection pool with MCP security enforcement.
// For read-only queries, connection-level retry is sufficient.
// For write operations requiring transaction retry, use ExecuteTxWithRetry directly.
// Note: Callers should manage context timeouts as needed.
func (s *Source) Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error) {
// MCP Security Check 1: Enforce write operation restrictions
if err := s.CanExecuteWrite(sql); err != nil {
return nil, err
}
// MCP Security Check 2: Apply query limits (row limit)
modifiedSQL, err := s.ApplyQueryLimits(sql)
if err != nil {
return nil, err
}
return s.Pool.Query(ctx, modifiedSQL, args...)
}
// ============================================================================
// MCP Security & Observability Features
// ============================================================================
// TelemetryEvent represents a structured telemetry event for MCP tool calls
type TelemetryEvent struct {
Timestamp time.Time `json:"timestamp"`
ToolName string `json:"tool_name"`
ClusterID string `json:"cluster_id"`
Database string `json:"database"`
User string `json:"user"`
SQLRedacted string `json:"sql_redacted"` // Query with values redacted
Status string `json:"status"` // "success" | "failure"
ErrorCode string `json:"error_code,omitempty"`
ErrorMsg string `json:"error_msg,omitempty"`
LatencyMs int64 `json:"latency_ms"`
RowsAffected int64 `json:"rows_affected,omitempty"`
Metadata map[string]string `json:"metadata,omitempty"`
}
// StructuredError represents an MCP-compliant error with error codes
type StructuredError struct {
Code string `json:"error_code"`
Message string `json:"message"`
Details map[string]any `json:"details,omitempty"`
}
func (e *StructuredError) Error() string {
return fmt.Sprintf("[%s] %s", e.Code, e.Message)
}
// MCP Error Codes
const (
ErrCodeUnauthorized = "CRDB_UNAUTHORIZED"
ErrCodeReadOnlyViolation = "CRDB_READONLY_VIOLATION"
ErrCodeQueryTimeout = "CRDB_QUERY_TIMEOUT"
ErrCodeRowLimitExceeded = "CRDB_ROW_LIMIT_EXCEEDED"
ErrCodeInvalidSQL = "CRDB_INVALID_SQL"
ErrCodeConnectionFailed = "CRDB_CONNECTION_FAILED"
ErrCodeWriteModeRequired = "CRDB_WRITE_MODE_REQUIRED"
ErrCodeQueryExecutionFailed = "CRDB_QUERY_EXECUTION_FAILED"
)
// SQLStatementType represents the type of SQL statement
type SQLStatementType int
const (
SQLTypeUnknown SQLStatementType = iota
SQLTypeSelect
SQLTypeInsert
SQLTypeUpdate
SQLTypeDelete
SQLTypeDDL // CREATE, ALTER, DROP
SQLTypeTruncate
SQLTypeExplain
SQLTypeShow
SQLTypeSet
)
// ClassifySQL analyzes a SQL statement and returns its type
func ClassifySQL(sql string) SQLStatementType {
// Normalize: trim and convert to uppercase for analysis
normalized := strings.TrimSpace(strings.ToUpper(sql))
if normalized == "" {
return SQLTypeUnknown
}
// Remove comments
normalized = regexp.MustCompile(`--.*`).ReplaceAllString(normalized, "")
normalized = regexp.MustCompile(`/\*.*?\*/`).ReplaceAllString(normalized, "")
normalized = strings.TrimSpace(normalized)
// Check statement type
switch {
case strings.HasPrefix(normalized, "SELECT"):
return SQLTypeSelect
case strings.HasPrefix(normalized, "INSERT"):
return SQLTypeInsert
case strings.HasPrefix(normalized, "UPDATE"):
return SQLTypeUpdate
case strings.HasPrefix(normalized, "DELETE"):
return SQLTypeDelete
case strings.HasPrefix(normalized, "TRUNCATE"):
return SQLTypeTruncate
case strings.HasPrefix(normalized, "CREATE"):
return SQLTypeDDL
case strings.HasPrefix(normalized, "ALTER"):
return SQLTypeDDL
case strings.HasPrefix(normalized, "DROP"):
return SQLTypeDDL
case strings.HasPrefix(normalized, "EXPLAIN"):
return SQLTypeExplain
case strings.HasPrefix(normalized, "SHOW"):
return SQLTypeShow
case strings.HasPrefix(normalized, "SET"):
return SQLTypeSet
default:
return SQLTypeUnknown
}
}
// IsWriteOperation returns true if the SQL statement modifies data
func IsWriteOperation(sqlType SQLStatementType) bool {
switch sqlType {
case SQLTypeInsert, SQLTypeUpdate, SQLTypeDelete, SQLTypeTruncate, SQLTypeDDL:
return true
default:
return false
}
}
// IsReadOnlyMode returns whether the source is in read-only mode
func (s *Source) IsReadOnlyMode() bool {
return s.ReadOnlyMode && !s.EnableWriteMode
}
// CanExecuteWrite checks if a write operation is allowed
func (s *Source) CanExecuteWrite(sql string) error {
sqlType := ClassifySQL(sql)
if IsWriteOperation(sqlType) && s.IsReadOnlyMode() {
return &StructuredError{
Code: ErrCodeReadOnlyViolation,
Message: "Write operations are not allowed in read-only mode. Set enableWriteMode: true to allow writes.",
Details: map[string]any{
"sql_type": sqlType,
"read_only_mode": s.ReadOnlyMode,
"enable_write_mode": s.EnableWriteMode,
},
}
}
return nil
}
// ApplyQueryLimits applies row limits to a SQL query for MCP security compliance.
// Context timeout management is the responsibility of the caller (following Go best practices).
// Returns potentially modified SQL with LIMIT clause for SELECT queries.
func (s *Source) ApplyQueryLimits(sql string) (string, error) {
sqlType := ClassifySQL(sql)
// Apply row limit only to SELECT queries
if sqlType == SQLTypeSelect && s.MaxRowLimit > 0 {
// Check if query already has LIMIT clause
normalized := strings.ToUpper(sql)
if !strings.Contains(normalized, " LIMIT ") {
// Add LIMIT clause - trim trailing whitespace and semicolon
sql = strings.TrimSpace(sql)
sql = strings.TrimSuffix(sql, ";")
sql = fmt.Sprintf("%s LIMIT %d", sql, s.MaxRowLimit)
}
}
return sql, nil
}
// RedactSQL redacts sensitive values from SQL for telemetry
func RedactSQL(sql string) string {
// Redact string literals
sql = regexp.MustCompile(`'[^']*'`).ReplaceAllString(sql, "'***'")
// Redact numbers that might be sensitive
sql = regexp.MustCompile(`\b\d{10,}\b`).ReplaceAllString(sql, "***")
return sql
}
// EmitTelemetry logs a telemetry event in structured JSON format
func (s *Source) EmitTelemetry(ctx context.Context, event TelemetryEvent) {
if !s.EnableTelemetry {
return
}
// Set cluster ID if not already set
if event.ClusterID == "" {
event.ClusterID = s.ClusterID
if event.ClusterID == "" {
event.ClusterID = s.Database // Fallback to database name
}
}
// Set database and user
if event.Database == "" {
event.Database = s.Database
}
if event.User == "" {
event.User = s.User
}
// Log as structured JSON
if s.TelemetryVerbose {
jsonBytes, _ := json.Marshal(event)
slog.Info("CockroachDB MCP Telemetry", "event", string(jsonBytes))
} else {
// Minimal logging
slog.Info("CockroachDB MCP",
"tool", event.ToolName,
"status", event.Status,
"latency_ms", event.LatencyMs,
"error_code", event.ErrorCode,
)
}
}
func initCockroachDBConnectionPoolWithRetry(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname string, queryParams map[string]string, maxRetries int, baseDelay time.Duration) (*pgxpool.Pool, error) {
//nolint:all
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)
if err != nil {
userAgent = "genai-toolbox"
}
if queryParams == nil {
queryParams = make(map[string]string)
}
if _, ok := queryParams["application_name"]; !ok {
queryParams["application_name"] = userAgent
}
connURL := &url.URL{
Scheme: "postgres",
User: url.UserPassword(user, pass),
Host: fmt.Sprintf("%s:%s", host, port),
Path: dbname,
RawQuery: ConvertParamMapToRawQuery(queryParams),
}
var pool *pgxpool.Pool
for attempt := 0; attempt <= maxRetries; attempt++ {
pool, err = pgxpool.New(ctx, connURL.String())
if err == nil {
err = pool.Ping(ctx)
}
if err == nil {
return pool, nil
}
if attempt < maxRetries {
backoff := baseDelay * time.Duration(math.Pow(2, float64(attempt)))
time.Sleep(backoff)
}
}
return nil, fmt.Errorf("failed to connect to CockroachDB after %d retries: %w", maxRetries, err)
}
func ConvertParamMapToRawQuery(queryParams map[string]string) string {
values := url.Values{}
for k, v := range queryParams {
values.Add(k, v)
}
return values.Encode()
}

View File

@@ -0,0 +1,224 @@
// Copyright 2026 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 cockroachdb
import (
"context"
"strings"
"testing"
"github.com/goccy/go-yaml"
)
func TestCockroachDBSourceConfig(t *testing.T) {
tests := []struct {
name string
yaml string
}{
{
name: "valid config",
yaml: `
name: test-cockroachdb
type: cockroachdb
host: localhost
port: "26257"
user: root
password: ""
database: defaultdb
maxRetries: 5
retryBaseDelay: 500ms
queryParams:
sslmode: disable
`,
},
{
name: "with optional queryParams",
yaml: `
name: test-cockroachdb
type: cockroachdb
host: localhost
port: "26257"
user: root
password: testpass
database: testdb
queryParams:
sslmode: require
sslcert: /path/to/cert
`,
},
{
name: "with custom retry settings",
yaml: `
name: test-cockroachdb
type: cockroachdb
host: localhost
port: "26257"
user: root
password: ""
database: defaultdb
maxRetries: 10
retryBaseDelay: 1s
`,
},
{
name: "without password (insecure mode)",
yaml: `
name: test-cockroachdb
type: cockroachdb
host: localhost
port: "26257"
user: root
database: defaultdb
`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
decoder := yaml.NewDecoder(strings.NewReader(tt.yaml))
cfg, err := newConfig(context.Background(), "test", decoder)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if cfg == nil {
t.Fatal("expected config but got nil")
}
// Verify it's the right type
cockroachCfg, ok := cfg.(Config)
if !ok {
t.Fatalf("expected Config type, got %T", cfg)
}
// Verify SourceConfigType
if cockroachCfg.SourceConfigType() != SourceType {
t.Errorf("expected SourceConfigType %q, got %q", SourceType, cockroachCfg.SourceConfigType())
}
t.Logf("✅ Config parsed successfully: %+v", cockroachCfg)
})
}
}
func TestCockroachDBSourceType(t *testing.T) {
yamlContent := `
name: test-cockroachdb
type: cockroachdb
host: localhost
port: "26257"
user: root
password: ""
database: defaultdb
`
decoder := yaml.NewDecoder(strings.NewReader(yamlContent))
cfg, err := newConfig(context.Background(), "test", decoder)
if err != nil {
t.Fatalf("failed to create config: %v", err)
}
if cfg.SourceConfigType() != "cockroachdb" {
t.Errorf("expected SourceConfigType 'cockroachdb', got %q", cfg.SourceConfigType())
}
}
func TestCockroachDBDefaultValues(t *testing.T) {
yamlContent := `
name: test-cockroachdb
type: cockroachdb
host: localhost
port: "26257"
user: root
password: ""
database: defaultdb
`
decoder := yaml.NewDecoder(strings.NewReader(yamlContent))
cfg, err := newConfig(context.Background(), "test", decoder)
if err != nil {
t.Fatalf("failed to create config: %v", err)
}
cockroachCfg, ok := cfg.(Config)
if !ok {
t.Fatalf("expected Config type")
}
// Check default values
if cockroachCfg.MaxRetries != 5 {
t.Errorf("expected default MaxRetries 5, got %d", cockroachCfg.MaxRetries)
}
if cockroachCfg.RetryBaseDelay != "500ms" {
t.Errorf("expected default RetryBaseDelay '500ms', got %q", cockroachCfg.RetryBaseDelay)
}
t.Logf("✅ Default values set correctly")
}
func TestConvertParamMapToRawQuery(t *testing.T) {
tests := []struct {
name string
params map[string]string
want []string // Expected substrings in any order
}{
{
name: "empty params",
params: map[string]string{},
want: []string{},
},
{
name: "single param",
params: map[string]string{
"sslmode": "disable",
},
want: []string{"sslmode=disable"},
},
{
name: "multiple params",
params: map[string]string{
"sslmode": "require",
"application_name": "test-app",
},
want: []string{"sslmode=require", "application_name=test-app"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := ConvertParamMapToRawQuery(tt.params)
if len(tt.want) == 0 {
if result != "" {
t.Errorf("expected empty string, got %q", result)
}
return
}
// Check that all expected substrings are in the result
for _, want := range tt.want {
if !contains(result, want) {
t.Errorf("expected result to contain %q, got %q", want, result)
}
}
t.Logf("✅ Query string: %s", result)
})
}
}
func contains(s, substr string) bool {
return strings.Contains(s, substr)
}

View File

@@ -0,0 +1,455 @@
// Copyright 2026 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 cockroachdb
import (
"context"
"strings"
"testing"
"time"
yaml "github.com/goccy/go-yaml"
)
// TestClassifySQL tests SQL statement classification
func TestClassifySQL(t *testing.T) {
tests := []struct {
name string
sql string
expected SQLStatementType
}{
{"SELECT", "SELECT * FROM users", SQLTypeSelect},
{"SELECT with spaces", " SELECT * FROM users ", SQLTypeSelect},
{"SELECT with comment", "-- comment\nSELECT * FROM users", SQLTypeSelect},
{"INSERT", "INSERT INTO users (name) VALUES ('alice')", SQLTypeInsert},
{"UPDATE", "UPDATE users SET name='bob' WHERE id=1", SQLTypeUpdate},
{"DELETE", "DELETE FROM users WHERE id=1", SQLTypeDelete},
{"CREATE TABLE", "CREATE TABLE users (id UUID PRIMARY KEY)", SQLTypeDDL},
{"ALTER TABLE", "ALTER TABLE users ADD COLUMN email STRING", SQLTypeDDL},
{"DROP TABLE", "DROP TABLE users", SQLTypeDDL},
{"TRUNCATE", "TRUNCATE TABLE users", SQLTypeTruncate},
{"EXPLAIN", "EXPLAIN SELECT * FROM users", SQLTypeExplain},
{"SHOW", "SHOW TABLES", SQLTypeShow},
{"SET", "SET application_name = 'myapp'", SQLTypeSet},
{"Empty", "", SQLTypeUnknown},
{"Lowercase select", "select * from users", SQLTypeSelect},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := ClassifySQL(tt.sql)
if result != tt.expected {
t.Errorf("ClassifySQL(%q) = %v, want %v", tt.sql, result, tt.expected)
}
})
}
}
// TestIsWriteOperation tests write operation detection
func TestIsWriteOperation(t *testing.T) {
tests := []struct {
sqlType SQLStatementType
expected bool
}{
{SQLTypeSelect, false},
{SQLTypeInsert, true},
{SQLTypeUpdate, true},
{SQLTypeDelete, true},
{SQLTypeTruncate, true},
{SQLTypeDDL, true},
{SQLTypeExplain, false},
{SQLTypeShow, false},
{SQLTypeSet, false},
{SQLTypeUnknown, false},
}
for _, tt := range tests {
t.Run(tt.sqlType.String(), func(t *testing.T) {
result := IsWriteOperation(tt.sqlType)
if result != tt.expected {
t.Errorf("IsWriteOperation(%v) = %v, want %v", tt.sqlType, result, tt.expected)
}
})
}
}
// Helper for SQLStatementType to string
func (s SQLStatementType) String() string {
switch s {
case SQLTypeSelect:
return "SELECT"
case SQLTypeInsert:
return "INSERT"
case SQLTypeUpdate:
return "UPDATE"
case SQLTypeDelete:
return "DELETE"
case SQLTypeDDL:
return "DDL"
case SQLTypeTruncate:
return "TRUNCATE"
case SQLTypeExplain:
return "EXPLAIN"
case SQLTypeShow:
return "SHOW"
case SQLTypeSet:
return "SET"
default:
return "UNKNOWN"
}
}
// TestCanExecuteWrite tests write operation enforcement
func TestCanExecuteWrite(t *testing.T) {
tests := []struct {
name string
readOnlyMode bool
enableWriteMode bool
sql string
expectError bool
errorCode string
}{
{
name: "SELECT in read-only mode",
readOnlyMode: true,
enableWriteMode: false,
sql: "SELECT * FROM users",
expectError: false,
},
{
name: "INSERT in read-only mode",
readOnlyMode: true,
enableWriteMode: false,
sql: "INSERT INTO users (name) VALUES ('alice')",
expectError: true,
errorCode: ErrCodeReadOnlyViolation,
},
{
name: "INSERT with write mode enabled",
readOnlyMode: false,
enableWriteMode: true,
sql: "INSERT INTO users (name) VALUES ('alice')",
expectError: false,
},
{
name: "CREATE TABLE in read-only mode",
readOnlyMode: true,
enableWriteMode: false,
sql: "CREATE TABLE test (id UUID PRIMARY KEY)",
expectError: true,
errorCode: ErrCodeReadOnlyViolation,
},
{
name: "CREATE TABLE with write mode enabled",
readOnlyMode: false,
enableWriteMode: true,
sql: "CREATE TABLE test (id UUID PRIMARY KEY)",
expectError: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
source := &Source{
Config: Config{
ReadOnlyMode: tt.readOnlyMode,
EnableWriteMode: tt.enableWriteMode,
},
}
err := source.CanExecuteWrite(tt.sql)
if tt.expectError {
if err == nil {
t.Errorf("Expected error but got nil")
return
}
structErr, ok := err.(*StructuredError)
if !ok {
t.Errorf("Expected StructuredError but got %T", err)
return
}
if structErr.Code != tt.errorCode {
t.Errorf("Expected error code %s but got %s", tt.errorCode, structErr.Code)
}
} else {
if err != nil {
t.Errorf("Expected no error but got: %v", err)
}
}
})
}
}
// TestApplyQueryLimits tests query limit application
func TestApplyQueryLimits(t *testing.T) {
tests := []struct {
name string
sql string
maxRowLimit int
expectedSQL string
shouldAddLimit bool
}{
{
name: "SELECT without LIMIT",
sql: "SELECT * FROM users",
maxRowLimit: 100,
expectedSQL: "SELECT * FROM users LIMIT 100",
shouldAddLimit: true,
},
{
name: "SELECT with existing LIMIT",
sql: "SELECT * FROM users LIMIT 50",
maxRowLimit: 100,
expectedSQL: "SELECT * FROM users LIMIT 50",
shouldAddLimit: false,
},
{
name: "SELECT without LIMIT and semicolon",
sql: "SELECT * FROM users;",
maxRowLimit: 100,
expectedSQL: "SELECT * FROM users LIMIT 100",
shouldAddLimit: true,
},
{
name: "SELECT with trailing newline and semicolon",
sql: "SELECT * FROM users;\n",
maxRowLimit: 100,
expectedSQL: "SELECT * FROM users LIMIT 100",
shouldAddLimit: true,
},
{
name: "SELECT with multiline and semicolon",
sql: "\n\tSELECT *\n\tFROM users\n\tORDER BY id;\n",
maxRowLimit: 100,
expectedSQL: "SELECT *\n\tFROM users\n\tORDER BY id LIMIT 100",
shouldAddLimit: true,
},
{
name: "INSERT should not have LIMIT added",
sql: "INSERT INTO users (name) VALUES ('alice')",
maxRowLimit: 100,
expectedSQL: "INSERT INTO users (name) VALUES ('alice')",
shouldAddLimit: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
source := &Source{
Config: Config{
MaxRowLimit: tt.maxRowLimit,
QueryTimeoutSec: 0, // Timeout now managed by caller
},
}
modifiedSQL, err := source.ApplyQueryLimits(tt.sql)
if err != nil {
t.Errorf("Unexpected error: %v", err)
return
}
if modifiedSQL != tt.expectedSQL {
t.Errorf("Expected SQL:\n%s\nGot:\n%s", tt.expectedSQL, modifiedSQL)
}
})
}
}
// TestApplyQueryTimeout tests that timeout is managed by caller (not source)
func TestApplyQueryTimeout(t *testing.T) {
source := &Source{
Config: Config{
QueryTimeoutSec: 5, // Documented recommended timeout
MaxRowLimit: 0, // Don't add LIMIT
},
}
// Caller creates timeout context (following Go best practices)
ctx := context.Background()
ctx, cancel := context.WithTimeout(ctx, time.Duration(source.QueryTimeoutSec)*time.Second)
defer cancel()
// Apply query limits (doesn't modify context anymore)
modifiedSQL, err := source.ApplyQueryLimits("SELECT * FROM users")
if err != nil {
t.Errorf("Unexpected error: %v", err)
return
}
// Verify context has deadline (managed by caller)
deadline, ok := ctx.Deadline()
if !ok {
t.Error("Expected deadline to be set but it wasn't")
return
}
// Verify deadline is approximately 5 seconds from now
expectedDeadline := time.Now().Add(5 * time.Second)
diff := deadline.Sub(expectedDeadline)
if diff < 0 {
diff = -diff
}
// Allow 1 second tolerance
if diff > time.Second {
t.Errorf("Deadline diff too large: %v", diff)
}
// Verify SQL is unchanged (LIMIT not added since MaxRowLimit=0)
if modifiedSQL != "SELECT * FROM users" {
t.Errorf("Expected SQL unchanged, got: %s", modifiedSQL)
}
}
// TestRedactSQL tests SQL redaction for telemetry
func TestRedactSQL(t *testing.T) {
tests := []struct {
name string
sql string
expected string
}{
{
name: "String literal redaction",
sql: "SELECT * FROM users WHERE name='alice' AND email='alice@example.com'",
expected: "SELECT * FROM users WHERE name='***' AND email='***'",
},
{
name: "Long number redaction",
sql: "SELECT * FROM users WHERE ssn=1234567890123",
expected: "SELECT * FROM users WHERE ssn=***",
},
{
name: "Short numbers not redacted",
sql: "SELECT * FROM users WHERE age=25",
expected: "SELECT * FROM users WHERE age=25",
},
{
name: "Multiple sensitive values",
sql: "INSERT INTO users (name, email, phone) VALUES ('bob', 'bob@example.com', '5551234567')",
expected: "INSERT INTO users (name, email, phone) VALUES ('***', '***', '***')",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := RedactSQL(tt.sql)
if result != tt.expected {
t.Errorf("RedactSQL:\nGot: %s\nExpected: %s", result, tt.expected)
}
})
}
}
// TestIsReadOnlyMode tests read-only mode detection
func TestIsReadOnlyMode(t *testing.T) {
tests := []struct {
name string
readOnlyMode bool
enableWriteMode bool
expected bool
}{
{"Read-only by default", true, false, true},
{"Write mode enabled", false, true, false},
{"Both false", false, false, false},
{"Read-only overridden by write mode", true, true, false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
source := &Source{
Config: Config{
ReadOnlyMode: tt.readOnlyMode,
EnableWriteMode: tt.enableWriteMode,
},
}
result := source.IsReadOnlyMode()
if result != tt.expected {
t.Errorf("IsReadOnlyMode() = %v, want %v", result, tt.expected)
}
})
}
}
// TestStructuredError tests error formatting
func TestStructuredError(t *testing.T) {
err := &StructuredError{
Code: ErrCodeReadOnlyViolation,
Message: "Write operations not allowed",
Details: map[string]any{
"sql_type": "INSERT",
},
}
errorStr := err.Error()
if !strings.Contains(errorStr, ErrCodeReadOnlyViolation) {
t.Errorf("Error string should contain error code: %s", errorStr)
}
if !strings.Contains(errorStr, "Write operations not allowed") {
t.Errorf("Error string should contain message: %s", errorStr)
}
}
// TestDefaultSecuritySettings tests that security defaults are correct
func TestDefaultSecuritySettings(t *testing.T) {
ctx := context.Background()
// Create a minimal YAML config
yamlData := `name: test
type: cockroachdb
host: localhost
port: "26257"
user: root
database: defaultdb
`
var cfg Config
if err := yaml.Unmarshal([]byte(yamlData), &cfg); err != nil {
t.Fatalf("Failed to unmarshal YAML: %v", err)
}
// Apply defaults through newConfig logic manually
cfg.MaxRetries = 5
cfg.RetryBaseDelay = "500ms"
cfg.ReadOnlyMode = true
cfg.EnableWriteMode = false
cfg.MaxRowLimit = 1000
cfg.QueryTimeoutSec = 30
cfg.EnableTelemetry = true
cfg.TelemetryVerbose = false
_ = ctx // prevent unused
// Verify MCP security defaults
if !cfg.ReadOnlyMode {
t.Error("ReadOnlyMode should be true by default")
}
if cfg.EnableWriteMode {
t.Error("EnableWriteMode should be false by default")
}
if cfg.MaxRowLimit != 1000 {
t.Errorf("MaxRowLimit should be 1000, got %d", cfg.MaxRowLimit)
}
if cfg.QueryTimeoutSec != 30 {
t.Errorf("QueryTimeoutSec should be 30, got %d", cfg.QueryTimeoutSec)
}
if !cfg.EnableTelemetry {
t.Error("EnableTelemetry should be true by default")
}
}

View File

@@ -0,0 +1,186 @@
// Copyright 2026 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 cockroachdbexecutesql
import (
"context"
"fmt"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/cockroachdb"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"github.com/jackc/pgx/v5"
)
const kind string = "cockroachdb-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 {
Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error)
}
var compatibleSources = [...]string{cockroachdb.SourceKind}
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"type" validate:"required"`
Source string `yaml:"source" validate:"required"`
Description string `yaml:"description" validate:"required"`
AuthRequired []string `yaml:"authRequired"`
}
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string {
return kind
}
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
_, ok = rawS.(compatibleSource)
if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
}
sqlParameter := parameters.NewStringParameter("sql", "The sql to execute.")
params := parameters.Parameters{sqlParameter}
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, params, nil)
t := Tool{
Config: cfg,
Parameters: params,
manifest: tools.Manifest{Description: cfg.Description, Parameters: params.Manifest(), AuthRequired: cfg.AuthRequired},
mcpManifest: mcpManifest,
}
return t, nil
}
var _ tools.Tool = Tool{}
type Tool struct {
Config
Parameters parameters.Parameters `yaml:"parameters"`
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
}
paramsMap := params.AsMap()
sql, ok := paramsMap["sql"].(string)
if !ok {
return nil, fmt.Errorf("unable to get cast %s", paramsMap["sql"])
}
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("error getting logger: %s", err)
}
logger.DebugContext(ctx, fmt.Sprintf("executing `%s` tool query: %s", kind, sql))
results, err := source.Query(ctx, sql)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)
}
defer results.Close()
fields := results.FieldDescriptions()
var out []any
for results.Next() {
v, err := results.Values()
if err != nil {
return nil, fmt.Errorf("unable to parse row: %w", err)
}
row := orderedmap.Row{}
for i, f := range fields {
row.Add(f.Name, v[i])
}
out = append(out, row)
}
if err := results.Err(); err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)
}
return out, nil
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
return parameters.ParseParams(t.Parameters, data, claims)
}
func (t Tool) EmbedParams(ctx context.Context, params parameters.ParamValues, models map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
return params, nil
}
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(_ tools.SourceProvider) (bool, error) {
return false, nil
}
func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
return "Authorization", nil
}
func (t Tool) GetParameters() parameters.Parameters {
return t.Parameters
}

View File

@@ -0,0 +1,81 @@
// Copyright 2026 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 cockroachdbexecutesql_test
import (
"testing"
"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/cockroachdb/cockroachdbexecutesql"
)
func TestParseFromYamlCockroachDBExecuteSQL(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: `
kind: tools
name: execute_sql_tool
type: cockroachdb-execute-sql
source: my-crdb-instance
description: Execute SQL on CockroachDB
`,
want: server.ToolConfigs{
"execute_sql_tool": cockroachdbexecutesql.Config{
Name: "execute_sql_tool",
Type: "cockroachdb-execute-sql",
Source: "my-crdb-instance",
Description: "Execute SQL on CockroachDB",
AuthRequired: []string{},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
// Parse contents
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestCockroachDBExecuteSQLToolConfigType(t *testing.T) {
cfg := cockroachdbexecutesql.Config{
Name: "test-tool",
Type: "cockroachdb-execute-sql",
Source: "test-source",
Description: "test description",
}
if cfg.ToolConfigType() != "cockroachdb-execute-sql" {
t.Errorf("expected ToolConfigType 'cockroachdb-execute-sql', got %q", cfg.ToolConfigType())
}
}

View File

@@ -0,0 +1,187 @@
// Copyright 2026 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 cockroachdblistschemas
import (
"context"
"fmt"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/cockroachdb"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"github.com/jackc/pgx/v5"
)
const kind string = "cockroachdb-list-schemas"
const listSchemasStatement = `
SELECT
catalog_name,
schema_name,
crdb_is_user_defined
FROM information_schema.schemata
WHERE schema_name NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'crdb_internal', 'pg_extension')
AND schema_name NOT LIKE 'pg_temp_%'
AND schema_name NOT LIKE 'pg_toast_temp_%'
ORDER BY schema_name;
`
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 {
Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error)
}
var _ compatibleSource = &cockroachdb.Source{}
var compatibleSources = [...]string{cockroachdb.SourceKind}
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"type" validate:"required"`
Source string `yaml:"source" validate:"required"`
Description string `yaml:"description" validate:"required"`
AuthRequired []string `yaml:"authRequired"`
}
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string {
return kind
}
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
_, ok = rawS.(compatibleSource)
if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
}
allParameters := parameters.Parameters{}
paramManifest := allParameters.Manifest()
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters, nil)
t := Tool{
Config: cfg,
AllParams: allParameters,
manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
mcpManifest: mcpManifest,
}
return t, nil
}
var _ tools.Tool = Tool{}
type Tool struct {
Config
AllParams parameters.Parameters `yaml:"allParams"`
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
}
results, err := source.Query(ctx, listSchemasStatement)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)
}
defer results.Close()
fields := results.FieldDescriptions()
var out []map[string]any
for results.Next() {
values, err := results.Values()
if err != nil {
return nil, fmt.Errorf("unable to parse row: %w", err)
}
rowMap := make(map[string]any)
for i, field := range fields {
rowMap[string(field.Name)] = values[i]
}
out = append(out, rowMap)
}
if err := results.Err(); err != nil {
return nil, fmt.Errorf("error reading query results: %w", err)
}
return out, nil
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
return parameters.ParseParams(t.AllParams, data, claims)
}
func (t Tool) EmbedParams(ctx context.Context, params parameters.ParamValues, models map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
return params, nil
}
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(_ tools.SourceProvider) (bool, error) {
return false, nil
}
func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
return "Authorization", nil
}
func (t Tool) GetParameters() parameters.Parameters {
return t.AllParams
}

View File

@@ -0,0 +1,81 @@
// Copyright 2026 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 cockroachdblistschemas_test
import (
"testing"
"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/cockroachdb/cockroachdblistschemas"
)
func TestParseFromYamlCockroachDBListSchemas(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: `
kind: tools
name: list_schemas_tool
type: cockroachdb-list-schemas
source: my-crdb-instance
description: List schemas in CockroachDB
`,
want: server.ToolConfigs{
"list_schemas_tool": cockroachdblistschemas.Config{
Name: "list_schemas_tool",
Type: "cockroachdb-list-schemas",
Source: "my-crdb-instance",
Description: "List schemas in CockroachDB",
AuthRequired: []string{},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
// Parse contents
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestCockroachDBListSchemasToolConfigType(t *testing.T) {
cfg := cockroachdblistschemas.Config{
Name: "test-tool",
Type: "cockroachdb-list-schemas",
Source: "test-source",
Description: "test description",
}
if cfg.ToolConfigType() != "cockroachdb-list-schemas" {
t.Errorf("expected ToolConfigType 'cockroachdb-list-schemas', got %q", cfg.ToolConfigType())
}
}

View File

@@ -0,0 +1,261 @@
// Copyright 2026 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 cockroachdblisttables
import (
"context"
"fmt"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/cockroachdb"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"github.com/jackc/pgx/v5"
)
const kind string = "cockroachdb-list-tables"
const listTablesStatement = `
WITH desired_relkinds AS (
SELECT ARRAY['r', 'p']::char[] AS kinds
),
table_info AS (
SELECT
t.oid AS table_oid,
ns.nspname AS schema_name,
t.relname AS table_name,
pg_get_userbyid(t.relowner) AS table_owner,
obj_description(t.oid, 'pg_class') AS table_comment,
t.relkind AS object_kind
FROM
pg_class t
JOIN
pg_namespace ns ON ns.oid = t.relnamespace
CROSS JOIN desired_relkinds dk
WHERE
t.relkind = ANY(dk.kinds)
AND (NULLIF(TRIM($1), '') IS NULL OR t.relname = ANY(string_to_array($1,',')))
AND ns.nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'crdb_internal', 'pg_extension')
AND ns.nspname NOT LIKE 'pg_temp_%' AND ns.nspname NOT LIKE 'pg_toast_temp_%'
),
columns_info AS (
SELECT
att.attrelid AS table_oid, att.attname AS column_name, format_type(att.atttypid, att.atttypmod) AS data_type,
att.attnum AS column_ordinal_position, att.attnotnull AS is_not_nullable,
pg_get_expr(ad.adbin, ad.adrelid) AS column_default, col_description(att.attrelid, att.attnum) AS column_comment
FROM pg_attribute att LEFT JOIN pg_attrdef ad ON att.attrelid = ad.adrelid AND att.attnum = ad.adnum
JOIN table_info ti ON att.attrelid = ti.table_oid WHERE att.attnum > 0 AND NOT att.attisdropped
),
constraints_info AS (
SELECT
con.conrelid AS table_oid, con.conname AS constraint_name, pg_get_constraintdef(con.oid) AS constraint_definition,
CASE con.contype WHEN 'p' THEN 'PRIMARY KEY' WHEN 'f' THEN 'FOREIGN KEY' WHEN 'u' THEN 'UNIQUE' WHEN 'c' THEN 'CHECK' ELSE con.contype::text END AS constraint_type,
(SELECT array_agg(att.attname ORDER BY u.attposition) FROM unnest(con.conkey) WITH ORDINALITY AS u(attnum, attposition) JOIN pg_attribute att ON att.attrelid = con.conrelid AND att.attnum = u.attnum) AS constraint_columns,
NULLIF(con.confrelid, 0)::regclass AS foreign_key_referenced_table,
(SELECT array_agg(att.attname ORDER BY u.attposition) FROM unnest(con.confkey) WITH ORDINALITY AS u(attnum, attposition) JOIN pg_attribute att ON att.attrelid = con.confrelid AND att.attnum = u.attnum WHERE con.contype = 'f') AS foreign_key_referenced_columns
FROM pg_constraint con JOIN table_info ti ON con.conrelid = ti.table_oid
),
indexes_info AS (
SELECT
idx.indrelid AS table_oid, ic.relname AS index_name, pg_get_indexdef(idx.indexrelid) AS index_definition,
idx.indisunique AS is_unique, idx.indisprimary AS is_primary, am.amname AS index_method,
(SELECT array_agg(att.attname ORDER BY u.ord) FROM unnest(idx.indkey::int[]) WITH ORDINALITY AS u(colidx, ord) LEFT JOIN pg_attribute att ON att.attrelid = idx.indrelid AND att.attnum = u.colidx WHERE u.colidx <> 0) AS index_columns
FROM pg_index idx JOIN pg_class ic ON ic.oid = idx.indexrelid JOIN pg_am am ON am.oid = ic.relam JOIN table_info ti ON idx.indrelid = ti.table_oid
)
SELECT
ti.schema_name,
ti.table_name AS object_name,
CASE
WHEN $2 = 'simple' THEN
json_build_object('name', ti.table_name)
ELSE
json_build_object(
'schema_name', ti.schema_name,
'object_name', ti.table_name,
'object_type', CASE ti.object_kind
WHEN 'r' THEN 'TABLE'
WHEN 'p' THEN 'PARTITIONED TABLE'
ELSE ti.object_kind::text
END,
'owner', ti.table_owner,
'comment', ti.table_comment,
'columns', COALESCE((SELECT json_agg(json_build_object('column_name',ci.column_name,'data_type',ci.data_type,'ordinal_position',ci.column_ordinal_position,'is_not_nullable',ci.is_not_nullable,'column_default',ci.column_default,'column_comment',ci.column_comment) ORDER BY ci.column_ordinal_position) FROM columns_info ci WHERE ci.table_oid = ti.table_oid), '[]'::json),
'constraints', COALESCE((SELECT json_agg(json_build_object('constraint_name',cons.constraint_name,'constraint_type',cons.constraint_type,'constraint_definition',cons.constraint_definition,'constraint_columns',cons.constraint_columns,'foreign_key_referenced_table',cons.foreign_key_referenced_table,'foreign_key_referenced_columns',cons.foreign_key_referenced_columns)) FROM constraints_info cons WHERE cons.table_oid = ti.table_oid), '[]'::json),
'indexes', COALESCE((SELECT json_agg(json_build_object('index_name',ii.index_name,'index_definition',ii.index_definition,'is_unique',ii.is_unique,'is_primary',ii.is_primary,'index_method',ii.index_method,'index_columns',ii.index_columns)) FROM indexes_info ii WHERE ii.table_oid = ti.table_oid), '[]'::json)
)
END AS object_details
FROM table_info ti ORDER BY ti.schema_name, ti.table_name;
`
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 {
Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error)
}
var _ compatibleSource = &cockroachdb.Source{}
var compatibleSources = [...]string{cockroachdb.SourceKind}
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"type" validate:"required"`
Source string `yaml:"source" validate:"required"`
Description string `yaml:"description" validate:"required"`
AuthRequired []string `yaml:"authRequired"`
}
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string {
return kind
}
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
_, ok = rawS.(compatibleSource)
if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
}
allParameters := parameters.Parameters{
parameters.NewStringParameterWithDefault("table_names", "", "Optional: A comma-separated list of table names. If empty, details for all tables will be listed."),
parameters.NewStringParameterWithDefault("output_format", "detailed", "Optional: Use 'simple' for names only or 'detailed' for full info."),
}
paramManifest := allParameters.Manifest()
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters, nil)
t := Tool{
Config: cfg,
AllParams: allParameters,
manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
mcpManifest: mcpManifest,
}
return t, nil
}
var _ tools.Tool = Tool{}
type Tool struct {
Config
AllParams parameters.Parameters `yaml:"allParams"`
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
}
paramsMap := params.AsMap()
tableNames, ok := paramsMap["table_names"].(string)
if !ok {
return nil, fmt.Errorf("invalid 'table_names' parameter; expected a string")
}
outputFormat, _ := paramsMap["output_format"].(string)
if outputFormat != "simple" && outputFormat != "detailed" {
return nil, fmt.Errorf("invalid value for output_format: must be 'simple' or 'detailed', but got %q", outputFormat)
}
results, err := source.Query(ctx, listTablesStatement, tableNames, outputFormat)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)
}
defer results.Close()
fields := results.FieldDescriptions()
var out []map[string]any
for results.Next() {
values, err := results.Values()
if err != nil {
return nil, fmt.Errorf("unable to parse row: %w", err)
}
rowMap := make(map[string]any)
for i, field := range fields {
rowMap[string(field.Name)] = values[i]
}
out = append(out, rowMap)
}
if err := results.Err(); err != nil {
return nil, fmt.Errorf("error reading query results: %w", err)
}
return out, nil
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
return parameters.ParseParams(t.AllParams, data, claims)
}
func (t Tool) EmbedParams(ctx context.Context, params parameters.ParamValues, models map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
return params, nil
}
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(_ tools.SourceProvider) (bool, error) {
return false, nil
}
func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
return "Authorization", nil
}
func (t Tool) GetParameters() parameters.Parameters {
return t.AllParams
}

View File

@@ -0,0 +1,81 @@
// Copyright 2026 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 cockroachdblisttables_test
import (
"testing"
"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/cockroachdb/cockroachdblisttables"
)
func TestParseFromYamlCockroachDBListTables(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: `
kind: tools
name: list_tables_tool
type: cockroachdb-list-tables
source: my-crdb-instance
description: List tables in CockroachDB
`,
want: server.ToolConfigs{
"list_tables_tool": cockroachdblisttables.Config{
Name: "list_tables_tool",
Type: "cockroachdb-list-tables",
Source: "my-crdb-instance",
Description: "List tables in CockroachDB",
AuthRequired: []string{},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
// Parse contents
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestCockroachDBListTablesToolConfigType(t *testing.T) {
cfg := cockroachdblisttables.Config{
Name: "test-tool",
Type: "cockroachdb-list-tables",
Source: "test-source",
Description: "test description",
}
if cfg.ToolConfigType() != "cockroachdb-list-tables" {
t.Errorf("expected ToolConfigType 'cockroachdb-list-tables', got %q", cfg.ToolConfigType())
}
}

View File

@@ -0,0 +1,192 @@
// Copyright 2026 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 cockroachdbsql
import (
"context"
"fmt"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/embeddingmodels"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/cockroachdb"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/util/orderedmap"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
"github.com/jackc/pgx/v5"
)
const kind string = "cockroachdb-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 {
Query(ctx context.Context, sql string, args ...interface{}) (pgx.Rows, error)
}
var _ compatibleSource = &cockroachdb.Source{}
var compatibleSources = [...]string{cockroachdb.SourceKind}
type Config struct {
Name string `yaml:"name" validate:"required"`
Type string `yaml:"type" 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 parameters.Parameters `yaml:"parameters"`
TemplateParameters parameters.Parameters `yaml:"templateParameters"`
}
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) ToolConfigType() string {
return kind
}
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
_, 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 := parameters.ProcessParameters(cfg.TemplateParameters, cfg.Parameters)
if err != nil {
return nil, err
}
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters, nil)
t := Tool{
Config: cfg,
AllParams: allParameters,
manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
mcpManifest: mcpManifest,
}
return t, nil
}
var _ tools.Tool = Tool{}
type Tool struct {
Config
AllParams parameters.Parameters `yaml:"allParams"`
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) {
source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Source, t.Name, t.Type)
if err != nil {
return nil, err
}
paramsMap := params.AsMap()
newStatement, err := parameters.ResolveTemplateParams(t.TemplateParameters, t.Statement, paramsMap)
if err != nil {
return nil, fmt.Errorf("unable to extract template params %w", err)
}
newParams, err := parameters.GetParams(t.Parameters, paramsMap)
if err != nil {
return nil, fmt.Errorf("unable to extract standard params %w", err)
}
sliceParams := newParams.AsSlice()
results, err := source.Query(ctx, newStatement, sliceParams...)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)
}
defer results.Close()
fields := results.FieldDescriptions()
var out []any
for results.Next() {
v, err := results.Values()
if err != nil {
return nil, fmt.Errorf("unable to parse row: %w", err)
}
row := orderedmap.Row{}
for i, f := range fields {
row.Add(f.Name, v[i])
}
out = append(out, row)
}
if err := results.Err(); err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)
}
return out, nil
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) {
return parameters.ParseParams(t.AllParams, data, claims)
}
func (t Tool) EmbedParams(ctx context.Context, params parameters.ParamValues, models map[string]embeddingmodels.EmbeddingModel) (parameters.ParamValues, error) {
return params, nil
}
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(_ tools.SourceProvider) (bool, error) {
return false, nil
}
func (t Tool) ToConfig() tools.ToolConfig {
return t.Config
}
func (t Tool) GetAuthTokenHeaderName(resourceMgr tools.SourceProvider) (string, error) {
return "Authorization", nil
}
func (t Tool) GetParameters() parameters.Parameters {
return t.AllParams
}

View File

@@ -0,0 +1,93 @@
// Copyright 2026 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 cockroachdbsql_test
import (
"testing"
"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/cockroachdb/cockroachdbsql"
"github.com/googleapis/genai-toolbox/internal/util/parameters"
)
func TestParseFromYamlCockroachDB(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: `
kind: tools
name: example_tool
type: cockroachdb-sql
source: my-crdb-instance
description: some description
statement: |
SELECT * FROM SQL_STATEMENT;
parameters:
- name: user_id
type: string
description: user id parameter
`,
want: server.ToolConfigs{
"example_tool": cockroachdbsql.Config{
Name: "example_tool",
Type: "cockroachdb-sql",
Source: "my-crdb-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
AuthRequired: []string{},
Parameters: []parameters.Parameter{
parameters.NewStringParameter("user_id", "user id parameter"),
},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
// Parse contents
_, _, _, got, _, _, err := server.UnmarshalResourceConfig(ctx, testutils.FormatYaml(tc.in))
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestCockroachDBSQLToolConfigType(t *testing.T) {
cfg := cockroachdbsql.Config{
Name: "test-tool",
Type: "cockroachdb-sql",
Source: "test-source",
Description: "test description",
Statement: "SELECT 1",
}
if cfg.ToolConfigType() != "cockroachdb-sql" {
t.Errorf("expected ToolConfigType 'cockroachdb-sql', got %q", cfg.ToolConfigType())
}
}

View File

@@ -0,0 +1,220 @@
// Copyright 2026 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 cockroachdb
import (
"context"
"fmt"
"net/url"
"os"
"regexp"
"strings"
"testing"
"time"
"github.com/google/uuid"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
"github.com/jackc/pgx/v5/pgxpool"
)
var (
CockroachDBSourceKind = "cockroachdb"
CockroachDBToolKind = "cockroachdb-sql"
CockroachDBDatabase = getEnvOrDefault("COCKROACHDB_DATABASE", "defaultdb")
CockroachDBHost = getEnvOrDefault("COCKROACHDB_HOST", "localhost")
CockroachDBPort = getEnvOrDefault("COCKROACHDB_PORT", "26257")
CockroachDBUser = getEnvOrDefault("COCKROACHDB_USER", "root")
CockroachDBPass = getEnvOrDefault("COCKROACHDB_PASS", "")
)
func getEnvOrDefault(key, defaultValue string) string {
if value := os.Getenv(key); value != "" {
return value
}
return defaultValue
}
func getCockroachDBVars(t *testing.T) map[string]any {
if CockroachDBHost == "" {
t.Skip("COCKROACHDB_HOST not set, skipping CockroachDB integration test")
}
return map[string]any{
"type": CockroachDBSourceKind,
"host": CockroachDBHost,
"port": CockroachDBPort,
"database": CockroachDBDatabase,
"user": CockroachDBUser,
"password": CockroachDBPass,
"maxRetries": 5,
"retryBaseDelay": "500ms",
"queryParams": map[string]string{
"sslmode": "disable",
},
}
}
func initCockroachDBConnectionPool(host, port, user, pass, dbname string) (*pgxpool.Pool, error) {
connURL := &url.URL{
Scheme: "postgres",
User: url.UserPassword(user, pass),
Host: fmt.Sprintf("%s:%s", host, port),
Path: dbname,
RawQuery: "sslmode=disable&application_name=cockroachdb-integration-test",
}
pool, err := pgxpool.New(context.Background(), connURL.String())
if err != nil {
return nil, fmt.Errorf("unable to create connection pool: %w", err)
}
return pool, nil
}
func TestCockroachDB(t *testing.T) {
sourceConfig := getCockroachDBVars(t)
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
var args []string
pool, err := initCockroachDBConnectionPool(CockroachDBHost, CockroachDBPort, CockroachDBUser, CockroachDBPass, CockroachDBDatabase)
if err != nil {
t.Fatalf("unable to create cockroachdb connection pool: %s", err)
}
// Note: Don't defer pool.Close() here - the pool is only used for test setup/teardown.
// Closing it explicitly can cause hangs if the server's pool is still active.
// The pool will be cleaned up when the test exits.
// Verify CockroachDB version
var version string
err = pool.QueryRow(ctx, "SELECT version()").Scan(&version)
if err != nil {
t.Fatalf("failed to query version: %s", err)
}
if !strings.Contains(version, "CockroachDB") {
t.Fatalf("not connected to CockroachDB, got: %s", version)
}
t.Logf("✅ Connected to: %s", version)
// cleanup test environment
tests.CleanupPostgresTables(t, ctx, pool)
// create table names with UUID suffix
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 (using CockroachDB explicit INT primary keys)
createParamTableStmt, insertParamTableStmt, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, paramTestParams := tests.GetCockroachDBParamToolInfo(tableNameParam)
teardownTable1 := tests.SetupPostgresSQLTable(t, ctx, pool, createParamTableStmt, insertParamTableStmt, tableNameParam, paramTestParams)
defer teardownTable1(t)
// set up data for auth tool
createAuthTableStmt, insertAuthTableStmt, authToolStmt, authTestParams := tests.GetCockroachDBAuthToolInfo(tableNameAuth)
teardownTable2 := tests.SetupPostgresSQLTable(t, ctx, pool, createAuthTableStmt, insertAuthTableStmt, tableNameAuth, authTestParams)
defer teardownTable2(t)
// Write config into a file and pass it to command
toolsFile := tests.GetToolsConfig(sourceConfig, CockroachDBToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
// Add execute-sql tool with write-enabled source (CockroachDB MCP security requires explicit opt-in)
toolsFile = addCockroachDBExecuteSqlConfig(t, toolsFile, sourceConfig)
tmplSelectCombined, tmplSelectFilterCombined := tests.GetPostgresSQLTmplToolStatement()
toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CockroachDBToolKind, 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 (use CockroachDB-specific expectations)
select1Want, mcpMyFailToolWant, createTableStatement, mcpSelect1Want := tests.GetCockroachDBWants()
// Run required integration test suites (per CONTRIBUTING.md)
t.Run("ToolGetTest", func(t *testing.T) {
tests.RunToolGetTest(t)
})
t.Run("ToolInvokeTest", func(t *testing.T) {
tests.RunToolInvokeTest(t, select1Want)
})
t.Run("MCPToolCallMethod", func(t *testing.T) {
tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want)
})
t.Run("ExecuteSqlToolInvokeTest", func(t *testing.T) {
tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want)
})
t.Run("ToolInvokeWithTemplateParameters", func(t *testing.T) {
tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam)
})
t.Logf("✅✅✅ All CockroachDB integration tests passed!")
}
// addCockroachDBExecuteSqlConfig adds execute-sql tool with write-enabled source
// CockroachDB has MCP security enabled by default, so execute-sql needs a separate source with enableWriteMode
func addCockroachDBExecuteSqlConfig(t *testing.T, config map[string]any, baseSourceConfig map[string]any) map[string]any {
// Add write-enabled source for execute-sql tool
sources, ok := config["sources"].(map[string]any)
if !ok {
t.Fatalf("unable to get sources from config")
}
// Create a copy of the base source config with write mode enabled
writeEnabledSource := make(map[string]any)
for k, v := range baseSourceConfig {
writeEnabledSource[k] = v
}
writeEnabledSource["enableWriteMode"] = true
writeEnabledSource["readOnlyMode"] = false
sources["my-write-instance"] = writeEnabledSource
// Add tools using the write-enabled source
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{
"type": "cockroachdb-execute-sql",
"source": "my-write-instance",
"description": "Tool to execute sql",
}
tools["my-auth-exec-sql-tool"] = map[string]any{
"type": "cockroachdb-execute-sql",
"source": "my-write-instance",
"description": "Tool to execute sql",
"authRequired": []string{
"my-google-auth",
},
}
return config
}

View File

@@ -528,6 +528,41 @@ func GetPostgresSQLTmplToolStatement() (string, string) {
return tmplSelectCombined, tmplSelectFilterCombined
}
// GetCockroachDBParamToolInfo returns statements and param for my-tool cockroachdb-sql type
// Uses explicit INT PRIMARY KEY instead of SERIAL to ensure deterministic IDs
func GetCockroachDBParamToolInfo(tableName string) (string, string, string, string, string, string, []any) {
createStatement := fmt.Sprintf("CREATE TABLE %s (id INT 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 * FROM %s WHERE id = $1 OR name = $2 ORDER BY id;", tableName)
idParamStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = $1;", tableName)
nameParamStatement := fmt.Sprintf("SELECT * FROM %s WHERE name = $1;", tableName)
arrayToolStatement := fmt.Sprintf("SELECT * FROM %s WHERE id = ANY($1) AND name = ANY($2) ORDER BY id;", tableName)
params := []any{"Alice", "Jane", "Sid", nil}
return createStatement, insertStatement, toolStatement, idParamStatement, nameParamStatement, arrayToolStatement, params
}
// GetCockroachDBAuthToolInfo returns statements and param of my-auth-tool for cockroachdb-sql type
// Uses explicit INT PRIMARY KEY instead of SERIAL to ensure deterministic IDs
func GetCockroachDBAuthToolInfo(tableName string) (string, string, string, []any) {
createStatement := fmt.Sprintf("CREATE TABLE %s (id INT 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", ServiceAccountEmail, "Jane", "janedoe@gmail.com"}
return createStatement, insertStatement, toolStatement, params
}
// GetCockroachDBWants return the expected wants for cockroachdb
func GetCockroachDBWants() (string, string, string, string) {
select1Want := "[{\"?column?\":1}]"
// CockroachDB formats syntax errors differently than PostgreSQL:
// - Uses lowercase for SQL keywords in error messages
// - Uses format: 'at or near "token": syntax error' instead of 'syntax error at or near "TOKEN"'
mcpMyFailToolWant := `{"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"unable to execute query: ERROR: at or near \"selec\": syntax error (SQLSTATE 42601)"}],"isError":true}}`
createTableStatement := `"CREATE TABLE t (id INT PRIMARY KEY, name TEXT)"`
mcpSelect1Want := `{"jsonrpc":"2.0","id":"invoke my-auth-required-tool","result":{"content":[{"type":"text","text":"{\"?column?\":1}"}]}}`
return select1Want, mcpMyFailToolWant, createTableStatement, mcpSelect1Want
}
// GetMSSQLParamToolInfo returns statements and param for my-tool mssql-sql type
func GetMSSQLParamToolInfo(tableName string) (string, string, string, string, string, string, []any) {
createStatement := fmt.Sprintf("CREATE TABLE %s (id INT IDENTITY(1,1) PRIMARY KEY, name VARCHAR(255));", tableName)