This is page 8 of 33. Use http://codebase.md/googleapis/genai-toolbox?page={x} to view the full context.
# Directory Structure
```
├── .ci
│ ├── continuous.release.cloudbuild.yaml
│ ├── generate_release_table.sh
│ ├── integration.cloudbuild.yaml
│ ├── quickstart_test
│ │ ├── go.integration.cloudbuild.yaml
│ │ ├── js.integration.cloudbuild.yaml
│ │ ├── py.integration.cloudbuild.yaml
│ │ ├── run_go_tests.sh
│ │ ├── run_js_tests.sh
│ │ ├── run_py_tests.sh
│ │ └── setup_hotels_sample.sql
│ ├── test_with_coverage.sh
│ └── versioned.release.cloudbuild.yaml
├── .github
│ ├── auto-label.yaml
│ ├── blunderbuss.yml
│ ├── CODEOWNERS
│ ├── header-checker-lint.yml
│ ├── ISSUE_TEMPLATE
│ │ ├── bug_report.yml
│ │ ├── config.yml
│ │ ├── feature_request.yml
│ │ └── question.yml
│ ├── label-sync.yml
│ ├── labels.yaml
│ ├── PULL_REQUEST_TEMPLATE.md
│ ├── release-please.yml
│ ├── renovate.json5
│ ├── sync-repo-settings.yaml
│ └── workflows
│ ├── cloud_build_failure_reporter.yml
│ ├── deploy_dev_docs.yaml
│ ├── deploy_previous_version_docs.yaml
│ ├── deploy_versioned_docs.yaml
│ ├── docs_deploy.yaml
│ ├── docs_preview_clean.yaml
│ ├── docs_preview_deploy.yaml
│ ├── lint.yaml
│ ├── schedule_reporter.yml
│ ├── sync-labels.yaml
│ └── tests.yaml
├── .gitignore
├── .gitmodules
├── .golangci.yaml
├── .hugo
│ ├── archetypes
│ │ └── default.md
│ ├── assets
│ │ ├── icons
│ │ │ └── logo.svg
│ │ └── scss
│ │ ├── _styles_project.scss
│ │ └── _variables_project.scss
│ ├── go.mod
│ ├── go.sum
│ ├── hugo.toml
│ ├── layouts
│ │ ├── _default
│ │ │ └── home.releases.releases
│ │ ├── index.llms-full.txt
│ │ ├── index.llms.txt
│ │ ├── partials
│ │ │ ├── hooks
│ │ │ │ └── head-end.html
│ │ │ ├── navbar-version-selector.html
│ │ │ ├── page-meta-links.html
│ │ │ └── td
│ │ │ └── render-heading.html
│ │ ├── robot.txt
│ │ └── shortcodes
│ │ ├── include.html
│ │ ├── ipynb.html
│ │ └── regionInclude.html
│ ├── package-lock.json
│ ├── package.json
│ └── static
│ ├── favicons
│ │ ├── android-chrome-192x192.png
│ │ ├── android-chrome-512x512.png
│ │ ├── apple-touch-icon.png
│ │ ├── favicon-16x16.png
│ │ ├── favicon-32x32.png
│ │ └── favicon.ico
│ └── js
│ └── w3.js
├── CHANGELOG.md
├── cmd
│ ├── options_test.go
│ ├── options.go
│ ├── root_test.go
│ ├── root.go
│ └── version.txt
├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── DEVELOPER.md
├── Dockerfile
├── docs
│ └── en
│ ├── _index.md
│ ├── about
│ │ ├── _index.md
│ │ └── faq.md
│ ├── concepts
│ │ ├── _index.md
│ │ └── telemetry
│ │ ├── index.md
│ │ ├── telemetry_flow.png
│ │ └── telemetry_traces.png
│ ├── getting-started
│ │ ├── _index.md
│ │ ├── colab_quickstart.ipynb
│ │ ├── configure.md
│ │ ├── introduction
│ │ │ ├── _index.md
│ │ │ └── architecture.png
│ │ ├── local_quickstart_go.md
│ │ ├── local_quickstart_js.md
│ │ ├── local_quickstart.md
│ │ ├── mcp_quickstart
│ │ │ ├── _index.md
│ │ │ ├── inspector_tools.png
│ │ │ └── inspector.png
│ │ └── quickstart
│ │ ├── go
│ │ │ ├── genAI
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ ├── genkit
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ ├── langchain
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ ├── openAI
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ └── quickstart_test.go
│ │ ├── golden.txt
│ │ ├── js
│ │ │ ├── genAI
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ ├── genkit
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ ├── langchain
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ ├── llamaindex
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ └── quickstart.test.js
│ │ ├── python
│ │ │ ├── __init__.py
│ │ │ ├── adk
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ ├── core
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ ├── langchain
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ ├── llamaindex
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ └── quickstart_test.py
│ │ └── shared
│ │ ├── cloud_setup.md
│ │ ├── configure_toolbox.md
│ │ └── database_setup.md
│ ├── how-to
│ │ ├── _index.md
│ │ ├── connect_via_geminicli.md
│ │ ├── connect_via_mcp.md
│ │ ├── connect-ide
│ │ │ ├── _index.md
│ │ │ ├── alloydb_pg_admin_mcp.md
│ │ │ ├── alloydb_pg_mcp.md
│ │ │ ├── bigquery_mcp.md
│ │ │ ├── cloud_sql_mssql_admin_mcp.md
│ │ │ ├── cloud_sql_mssql_mcp.md
│ │ │ ├── cloud_sql_mysql_admin_mcp.md
│ │ │ ├── cloud_sql_mysql_mcp.md
│ │ │ ├── cloud_sql_pg_admin_mcp.md
│ │ │ ├── cloud_sql_pg_mcp.md
│ │ │ ├── firestore_mcp.md
│ │ │ ├── looker_mcp.md
│ │ │ ├── mssql_mcp.md
│ │ │ ├── mysql_mcp.md
│ │ │ ├── neo4j_mcp.md
│ │ │ ├── postgres_mcp.md
│ │ │ ├── spanner_mcp.md
│ │ │ └── sqlite_mcp.md
│ │ ├── deploy_docker.md
│ │ ├── deploy_gke.md
│ │ ├── deploy_toolbox.md
│ │ ├── export_telemetry.md
│ │ └── toolbox-ui
│ │ ├── edit-headers.gif
│ │ ├── edit-headers.png
│ │ ├── index.md
│ │ ├── optional-param-checked.png
│ │ ├── optional-param-unchecked.png
│ │ ├── run-tool.gif
│ │ ├── tools.png
│ │ └── toolsets.png
│ ├── reference
│ │ ├── _index.md
│ │ ├── cli.md
│ │ └── prebuilt-tools.md
│ ├── resources
│ │ ├── _index.md
│ │ ├── authServices
│ │ │ ├── _index.md
│ │ │ └── google.md
│ │ ├── sources
│ │ │ ├── _index.md
│ │ │ ├── alloydb-admin.md
│ │ │ ├── alloydb-pg.md
│ │ │ ├── bigquery.md
│ │ │ ├── bigtable.md
│ │ │ ├── cassandra.md
│ │ │ ├── clickhouse.md
│ │ │ ├── cloud-monitoring.md
│ │ │ ├── cloud-sql-admin.md
│ │ │ ├── cloud-sql-mssql.md
│ │ │ ├── cloud-sql-mysql.md
│ │ │ ├── cloud-sql-pg.md
│ │ │ ├── couchbase.md
│ │ │ ├── dataplex.md
│ │ │ ├── dgraph.md
│ │ │ ├── firebird.md
│ │ │ ├── firestore.md
│ │ │ ├── http.md
│ │ │ ├── looker.md
│ │ │ ├── mongodb.md
│ │ │ ├── mssql.md
│ │ │ ├── mysql.md
│ │ │ ├── neo4j.md
│ │ │ ├── oceanbase.md
│ │ │ ├── oracle.md
│ │ │ ├── postgres.md
│ │ │ ├── redis.md
│ │ │ ├── spanner.md
│ │ │ ├── sqlite.md
│ │ │ ├── tidb.md
│ │ │ ├── trino.md
│ │ │ ├── valkey.md
│ │ │ └── yugabytedb.md
│ │ └── tools
│ │ ├── _index.md
│ │ ├── alloydb
│ │ │ ├── _index.md
│ │ │ ├── alloydb-create-cluster.md
│ │ │ ├── alloydb-create-instance.md
│ │ │ ├── alloydb-create-user.md
│ │ │ ├── alloydb-get-cluster.md
│ │ │ ├── alloydb-get-instance.md
│ │ │ ├── alloydb-get-user.md
│ │ │ ├── alloydb-list-clusters.md
│ │ │ ├── alloydb-list-instances.md
│ │ │ ├── alloydb-list-users.md
│ │ │ └── alloydb-wait-for-operation.md
│ │ ├── alloydbainl
│ │ │ ├── _index.md
│ │ │ └── alloydb-ai-nl.md
│ │ ├── bigquery
│ │ │ ├── _index.md
│ │ │ ├── bigquery-analyze-contribution.md
│ │ │ ├── bigquery-conversational-analytics.md
│ │ │ ├── bigquery-execute-sql.md
│ │ │ ├── bigquery-forecast.md
│ │ │ ├── bigquery-get-dataset-info.md
│ │ │ ├── bigquery-get-table-info.md
│ │ │ ├── bigquery-list-dataset-ids.md
│ │ │ ├── bigquery-list-table-ids.md
│ │ │ ├── bigquery-search-catalog.md
│ │ │ └── bigquery-sql.md
│ │ ├── bigtable
│ │ │ ├── _index.md
│ │ │ └── bigtable-sql.md
│ │ ├── cassandra
│ │ │ ├── _index.md
│ │ │ └── cassandra-cql.md
│ │ ├── clickhouse
│ │ │ ├── _index.md
│ │ │ ├── clickhouse-execute-sql.md
│ │ │ ├── clickhouse-list-databases.md
│ │ │ ├── clickhouse-list-tables.md
│ │ │ └── clickhouse-sql.md
│ │ ├── cloudmonitoring
│ │ │ ├── _index.md
│ │ │ └── cloud-monitoring-query-prometheus.md
│ │ ├── cloudsql
│ │ │ ├── _index.md
│ │ │ ├── cloudsqlcreatedatabase.md
│ │ │ ├── cloudsqlcreateusers.md
│ │ │ ├── cloudsqlgetinstances.md
│ │ │ ├── cloudsqllistdatabases.md
│ │ │ ├── cloudsqllistinstances.md
│ │ │ ├── cloudsqlmssqlcreateinstance.md
│ │ │ ├── cloudsqlmysqlcreateinstance.md
│ │ │ ├── cloudsqlpgcreateinstances.md
│ │ │ └── cloudsqlwaitforoperation.md
│ │ ├── couchbase
│ │ │ ├── _index.md
│ │ │ └── couchbase-sql.md
│ │ ├── dataform
│ │ │ ├── _index.md
│ │ │ └── dataform-compile-local.md
│ │ ├── dataplex
│ │ │ ├── _index.md
│ │ │ ├── dataplex-lookup-entry.md
│ │ │ ├── dataplex-search-aspect-types.md
│ │ │ └── dataplex-search-entries.md
│ │ ├── dgraph
│ │ │ ├── _index.md
│ │ │ └── dgraph-dql.md
│ │ ├── firebird
│ │ │ ├── _index.md
│ │ │ ├── firebird-execute-sql.md
│ │ │ └── firebird-sql.md
│ │ ├── firestore
│ │ │ ├── _index.md
│ │ │ ├── firestore-add-documents.md
│ │ │ ├── firestore-delete-documents.md
│ │ │ ├── firestore-get-documents.md
│ │ │ ├── firestore-get-rules.md
│ │ │ ├── firestore-list-collections.md
│ │ │ ├── firestore-query-collection.md
│ │ │ ├── firestore-query.md
│ │ │ ├── firestore-update-document.md
│ │ │ └── firestore-validate-rules.md
│ │ ├── http
│ │ │ ├── _index.md
│ │ │ └── http.md
│ │ ├── looker
│ │ │ ├── _index.md
│ │ │ ├── looker-add-dashboard-element.md
│ │ │ ├── looker-conversational-analytics.md
│ │ │ ├── looker-get-dashboards.md
│ │ │ ├── looker-get-dimensions.md
│ │ │ ├── looker-get-explores.md
│ │ │ ├── looker-get-filters.md
│ │ │ ├── looker-get-looks.md
│ │ │ ├── looker-get-measures.md
│ │ │ ├── looker-get-models.md
│ │ │ ├── looker-get-parameters.md
│ │ │ ├── looker-health-analyze.md
│ │ │ ├── looker-health-pulse.md
│ │ │ ├── looker-health-vacuum.md
│ │ │ ├── looker-make-dashboard.md
│ │ │ ├── looker-make-look.md
│ │ │ ├── looker-query-sql.md
│ │ │ ├── looker-query-url.md
│ │ │ ├── looker-query.md
│ │ │ └── looker-run-look.md
│ │ ├── mongodb
│ │ │ ├── _index.md
│ │ │ ├── mongodb-aggregate.md
│ │ │ ├── mongodb-delete-many.md
│ │ │ ├── mongodb-delete-one.md
│ │ │ ├── mongodb-find-one.md
│ │ │ ├── mongodb-find.md
│ │ │ ├── mongodb-insert-many.md
│ │ │ ├── mongodb-insert-one.md
│ │ │ ├── mongodb-update-many.md
│ │ │ └── mongodb-update-one.md
│ │ ├── mssql
│ │ │ ├── _index.md
│ │ │ ├── mssql-execute-sql.md
│ │ │ ├── mssql-list-tables.md
│ │ │ └── mssql-sql.md
│ │ ├── mysql
│ │ │ ├── _index.md
│ │ │ ├── mysql-execute-sql.md
│ │ │ ├── mysql-list-active-queries.md
│ │ │ ├── mysql-list-table-fragmentation.md
│ │ │ ├── mysql-list-tables-missing-unique-indexes.md
│ │ │ ├── mysql-list-tables.md
│ │ │ └── mysql-sql.md
│ │ ├── neo4j
│ │ │ ├── _index.md
│ │ │ ├── neo4j-cypher.md
│ │ │ ├── neo4j-execute-cypher.md
│ │ │ └── neo4j-schema.md
│ │ ├── oceanbase
│ │ │ ├── _index.md
│ │ │ ├── oceanbase-execute-sql.md
│ │ │ └── oceanbase-sql.md
│ │ ├── oracle
│ │ │ ├── _index.md
│ │ │ ├── oracle-execute-sql.md
│ │ │ └── oracle-sql.md
│ │ ├── postgres
│ │ │ ├── _index.md
│ │ │ ├── postgres-execute-sql.md
│ │ │ ├── postgres-list-active-queries.md
│ │ │ ├── postgres-list-available-extensions.md
│ │ │ ├── postgres-list-installed-extensions.md
│ │ │ ├── postgres-list-tables.md
│ │ │ └── postgres-sql.md
│ │ ├── redis
│ │ │ ├── _index.md
│ │ │ └── redis.md
│ │ ├── spanner
│ │ │ ├── _index.md
│ │ │ ├── spanner-execute-sql.md
│ │ │ ├── spanner-list-tables.md
│ │ │ └── spanner-sql.md
│ │ ├── sqlite
│ │ │ ├── _index.md
│ │ │ ├── sqlite-execute-sql.md
│ │ │ └── sqlite-sql.md
│ │ ├── tidb
│ │ │ ├── _index.md
│ │ │ ├── tidb-execute-sql.md
│ │ │ └── tidb-sql.md
│ │ ├── trino
│ │ │ ├── _index.md
│ │ │ ├── trino-execute-sql.md
│ │ │ └── trino-sql.md
│ │ ├── utility
│ │ │ ├── _index.md
│ │ │ └── wait.md
│ │ ├── valkey
│ │ │ ├── _index.md
│ │ │ └── valkey.md
│ │ └── yuagbytedb
│ │ ├── _index.md
│ │ └── yugabytedb-sql.md
│ ├── samples
│ │ ├── _index.md
│ │ ├── alloydb
│ │ │ ├── _index.md
│ │ │ ├── ai-nl
│ │ │ │ ├── alloydb_ai_nl.ipynb
│ │ │ │ └── index.md
│ │ │ └── mcp_quickstart.md
│ │ ├── bigquery
│ │ │ ├── _index.md
│ │ │ ├── colab_quickstart_bigquery.ipynb
│ │ │ ├── local_quickstart.md
│ │ │ └── mcp_quickstart
│ │ │ ├── _index.md
│ │ │ ├── inspector_tools.png
│ │ │ └── inspector.png
│ │ └── looker
│ │ ├── _index.md
│ │ ├── looker_gemini_oauth
│ │ │ ├── _index.md
│ │ │ ├── authenticated.png
│ │ │ ├── authorize.png
│ │ │ └── registration.png
│ │ ├── looker_gemini.md
│ │ └── looker_mcp_inspector
│ │ ├── _index.md
│ │ ├── inspector_tools.png
│ │ └── inspector.png
│ └── sdks
│ ├── _index.md
│ ├── go-sdk.md
│ ├── js-sdk.md
│ └── python-sdk.md
├── gemini-extension.json
├── go.mod
├── go.sum
├── internal
│ ├── auth
│ │ ├── auth.go
│ │ └── google
│ │ └── google.go
│ ├── log
│ │ ├── handler.go
│ │ ├── log_test.go
│ │ ├── log.go
│ │ └── logger.go
│ ├── prebuiltconfigs
│ │ ├── prebuiltconfigs_test.go
│ │ ├── prebuiltconfigs.go
│ │ └── tools
│ │ ├── alloydb-postgres-admin.yaml
│ │ ├── alloydb-postgres-observability.yaml
│ │ ├── alloydb-postgres.yaml
│ │ ├── bigquery.yaml
│ │ ├── clickhouse.yaml
│ │ ├── cloud-sql-mssql-admin.yaml
│ │ ├── cloud-sql-mssql-observability.yaml
│ │ ├── cloud-sql-mssql.yaml
│ │ ├── cloud-sql-mysql-admin.yaml
│ │ ├── cloud-sql-mysql-observability.yaml
│ │ ├── cloud-sql-mysql.yaml
│ │ ├── cloud-sql-postgres-admin.yaml
│ │ ├── cloud-sql-postgres-observability.yaml
│ │ ├── cloud-sql-postgres.yaml
│ │ ├── dataplex.yaml
│ │ ├── firestore.yaml
│ │ ├── looker-conversational-analytics.yaml
│ │ ├── looker.yaml
│ │ ├── mssql.yaml
│ │ ├── mysql.yaml
│ │ ├── neo4j.yaml
│ │ ├── oceanbase.yaml
│ │ ├── postgres.yaml
│ │ ├── spanner-postgres.yaml
│ │ ├── spanner.yaml
│ │ └── sqlite.yaml
│ ├── server
│ │ ├── api_test.go
│ │ ├── api.go
│ │ ├── common_test.go
│ │ ├── config.go
│ │ ├── mcp
│ │ │ ├── jsonrpc
│ │ │ │ ├── jsonrpc_test.go
│ │ │ │ └── jsonrpc.go
│ │ │ ├── mcp.go
│ │ │ ├── util
│ │ │ │ └── lifecycle.go
│ │ │ ├── v20241105
│ │ │ │ ├── method.go
│ │ │ │ └── types.go
│ │ │ ├── v20250326
│ │ │ │ ├── method.go
│ │ │ │ └── types.go
│ │ │ └── v20250618
│ │ │ ├── method.go
│ │ │ └── types.go
│ │ ├── mcp_test.go
│ │ ├── mcp.go
│ │ ├── server_test.go
│ │ ├── server.go
│ │ ├── static
│ │ │ ├── assets
│ │ │ │ └── mcptoolboxlogo.png
│ │ │ ├── css
│ │ │ │ └── style.css
│ │ │ ├── index.html
│ │ │ ├── js
│ │ │ │ ├── auth.js
│ │ │ │ ├── loadTools.js
│ │ │ │ ├── mainContent.js
│ │ │ │ ├── navbar.js
│ │ │ │ ├── runTool.js
│ │ │ │ ├── toolDisplay.js
│ │ │ │ ├── tools.js
│ │ │ │ └── toolsets.js
│ │ │ ├── tools.html
│ │ │ └── toolsets.html
│ │ ├── web_test.go
│ │ └── web.go
│ ├── sources
│ │ ├── alloydbadmin
│ │ │ ├── alloydbadmin_test.go
│ │ │ └── alloydbadmin.go
│ │ ├── alloydbpg
│ │ │ ├── alloydb_pg_test.go
│ │ │ └── alloydb_pg.go
│ │ ├── bigquery
│ │ │ ├── bigquery_test.go
│ │ │ └── bigquery.go
│ │ ├── bigtable
│ │ │ ├── bigtable_test.go
│ │ │ └── bigtable.go
│ │ ├── cassandra
│ │ │ ├── cassandra_test.go
│ │ │ └── cassandra.go
│ │ ├── clickhouse
│ │ │ ├── clickhouse_test.go
│ │ │ └── clickhouse.go
│ │ ├── cloudmonitoring
│ │ │ ├── cloud_monitoring_test.go
│ │ │ └── cloud_monitoring.go
│ │ ├── cloudsqladmin
│ │ │ ├── cloud_sql_admin_test.go
│ │ │ └── cloud_sql_admin.go
│ │ ├── cloudsqlmssql
│ │ │ ├── cloud_sql_mssql_test.go
│ │ │ └── cloud_sql_mssql.go
│ │ ├── cloudsqlmysql
│ │ │ ├── cloud_sql_mysql_test.go
│ │ │ └── cloud_sql_mysql.go
│ │ ├── cloudsqlpg
│ │ │ ├── cloud_sql_pg_test.go
│ │ │ └── cloud_sql_pg.go
│ │ ├── couchbase
│ │ │ ├── couchbase_test.go
│ │ │ └── couchbase.go
│ │ ├── dataplex
│ │ │ ├── dataplex_test.go
│ │ │ └── dataplex.go
│ │ ├── dgraph
│ │ │ ├── dgraph_test.go
│ │ │ └── dgraph.go
│ │ ├── dialect.go
│ │ ├── firebird
│ │ │ ├── firebird_test.go
│ │ │ └── firebird.go
│ │ ├── firestore
│ │ │ ├── firestore_test.go
│ │ │ └── firestore.go
│ │ ├── http
│ │ │ ├── http_test.go
│ │ │ └── http.go
│ │ ├── ip_type.go
│ │ ├── looker
│ │ │ ├── looker_test.go
│ │ │ └── looker.go
│ │ ├── mongodb
│ │ │ ├── mongodb_test.go
│ │ │ └── mongodb.go
│ │ ├── mssql
│ │ │ ├── mssql_test.go
│ │ │ └── mssql.go
│ │ ├── mysql
│ │ │ ├── mysql_test.go
│ │ │ └── mysql.go
│ │ ├── neo4j
│ │ │ ├── neo4j_test.go
│ │ │ └── neo4j.go
│ │ ├── oceanbase
│ │ │ ├── oceanbase_test.go
│ │ │ └── oceanbase.go
│ │ ├── oracle
│ │ │ └── oracle.go
│ │ ├── postgres
│ │ │ ├── postgres_test.go
│ │ │ └── postgres.go
│ │ ├── redis
│ │ │ ├── redis_test.go
│ │ │ └── redis.go
│ │ ├── sources.go
│ │ ├── spanner
│ │ │ ├── spanner_test.go
│ │ │ └── spanner.go
│ │ ├── sqlite
│ │ │ ├── sqlite_test.go
│ │ │ └── sqlite.go
│ │ ├── tidb
│ │ │ ├── tidb_test.go
│ │ │ └── tidb.go
│ │ ├── trino
│ │ │ ├── trino_test.go
│ │ │ └── trino.go
│ │ ├── util.go
│ │ ├── valkey
│ │ │ ├── valkey_test.go
│ │ │ └── valkey.go
│ │ └── yugabytedb
│ │ ├── yugabytedb_test.go
│ │ └── yugabytedb.go
│ ├── telemetry
│ │ ├── instrumentation.go
│ │ └── telemetry.go
│ ├── testutils
│ │ └── testutils.go
│ ├── tools
│ │ ├── alloydb
│ │ │ ├── alloydbcreatecluster
│ │ │ │ ├── alloydbcreatecluster_test.go
│ │ │ │ └── alloydbcreatecluster.go
│ │ │ ├── alloydbcreateinstance
│ │ │ │ ├── alloydbcreateinstance_test.go
│ │ │ │ └── alloydbcreateinstance.go
│ │ │ ├── alloydbcreateuser
│ │ │ │ ├── alloydbcreateuser_test.go
│ │ │ │ └── alloydbcreateuser.go
│ │ │ ├── alloydbgetcluster
│ │ │ │ ├── alloydbgetcluster_test.go
│ │ │ │ └── alloydbgetcluster.go
│ │ │ ├── alloydbgetinstance
│ │ │ │ ├── alloydbgetinstance_test.go
│ │ │ │ └── alloydbgetinstance.go
│ │ │ ├── alloydbgetuser
│ │ │ │ ├── alloydbgetuser_test.go
│ │ │ │ └── alloydbgetuser.go
│ │ │ ├── alloydblistclusters
│ │ │ │ ├── alloydblistclusters_test.go
│ │ │ │ └── alloydblistclusters.go
│ │ │ ├── alloydblistinstances
│ │ │ │ ├── alloydblistinstances_test.go
│ │ │ │ └── alloydblistinstances.go
│ │ │ ├── alloydblistusers
│ │ │ │ ├── alloydblistusers_test.go
│ │ │ │ └── alloydblistusers.go
│ │ │ └── alloydbwaitforoperation
│ │ │ ├── alloydbwaitforoperation_test.go
│ │ │ └── alloydbwaitforoperation.go
│ │ ├── alloydbainl
│ │ │ ├── alloydbainl_test.go
│ │ │ └── alloydbainl.go
│ │ ├── bigquery
│ │ │ ├── bigqueryanalyzecontribution
│ │ │ │ ├── bigqueryanalyzecontribution_test.go
│ │ │ │ └── bigqueryanalyzecontribution.go
│ │ │ ├── bigquerycommon
│ │ │ │ ├── table_name_parser_test.go
│ │ │ │ ├── table_name_parser.go
│ │ │ │ └── util.go
│ │ │ ├── bigqueryconversationalanalytics
│ │ │ │ ├── bigqueryconversationalanalytics_test.go
│ │ │ │ └── bigqueryconversationalanalytics.go
│ │ │ ├── bigqueryexecutesql
│ │ │ │ ├── bigqueryexecutesql_test.go
│ │ │ │ └── bigqueryexecutesql.go
│ │ │ ├── bigqueryforecast
│ │ │ │ ├── bigqueryforecast_test.go
│ │ │ │ └── bigqueryforecast.go
│ │ │ ├── bigquerygetdatasetinfo
│ │ │ │ ├── bigquerygetdatasetinfo_test.go
│ │ │ │ └── bigquerygetdatasetinfo.go
│ │ │ ├── bigquerygettableinfo
│ │ │ │ ├── bigquerygettableinfo_test.go
│ │ │ │ └── bigquerygettableinfo.go
│ │ │ ├── bigquerylistdatasetids
│ │ │ │ ├── bigquerylistdatasetids_test.go
│ │ │ │ └── bigquerylistdatasetids.go
│ │ │ ├── bigquerylisttableids
│ │ │ │ ├── bigquerylisttableids_test.go
│ │ │ │ └── bigquerylisttableids.go
│ │ │ ├── bigquerysearchcatalog
│ │ │ │ ├── bigquerysearchcatalog_test.go
│ │ │ │ └── bigquerysearchcatalog.go
│ │ │ └── bigquerysql
│ │ │ ├── bigquerysql_test.go
│ │ │ └── bigquerysql.go
│ │ ├── bigtable
│ │ │ ├── bigtable_test.go
│ │ │ └── bigtable.go
│ │ ├── cassandra
│ │ │ └── cassandracql
│ │ │ ├── cassandracql_test.go
│ │ │ └── cassandracql.go
│ │ ├── clickhouse
│ │ │ ├── clickhouseexecutesql
│ │ │ │ ├── clickhouseexecutesql_test.go
│ │ │ │ └── clickhouseexecutesql.go
│ │ │ ├── clickhouselistdatabases
│ │ │ │ ├── clickhouselistdatabases_test.go
│ │ │ │ └── clickhouselistdatabases.go
│ │ │ ├── clickhouselisttables
│ │ │ │ ├── clickhouselisttables_test.go
│ │ │ │ └── clickhouselisttables.go
│ │ │ └── clickhousesql
│ │ │ ├── clickhousesql_test.go
│ │ │ └── clickhousesql.go
│ │ ├── cloudmonitoring
│ │ │ ├── cloudmonitoring_test.go
│ │ │ └── cloudmonitoring.go
│ │ ├── cloudsql
│ │ │ ├── cloudsqlcreatedatabase
│ │ │ │ ├── cloudsqlcreatedatabase_test.go
│ │ │ │ └── cloudsqlcreatedatabase.go
│ │ │ ├── cloudsqlcreateusers
│ │ │ │ ├── cloudsqlcreateusers_test.go
│ │ │ │ └── cloudsqlcreateusers.go
│ │ │ ├── cloudsqlgetinstances
│ │ │ │ ├── cloudsqlgetinstances_test.go
│ │ │ │ └── cloudsqlgetinstances.go
│ │ │ ├── cloudsqllistdatabases
│ │ │ │ ├── cloudsqllistdatabases_test.go
│ │ │ │ └── cloudsqllistdatabases.go
│ │ │ ├── cloudsqllistinstances
│ │ │ │ ├── cloudsqllistinstances_test.go
│ │ │ │ └── cloudsqllistinstances.go
│ │ │ └── cloudsqlwaitforoperation
│ │ │ ├── cloudsqlwaitforoperation_test.go
│ │ │ └── cloudsqlwaitforoperation.go
│ │ ├── cloudsqlmssql
│ │ │ └── cloudsqlmssqlcreateinstance
│ │ │ ├── cloudsqlmssqlcreateinstance_test.go
│ │ │ └── cloudsqlmssqlcreateinstance.go
│ │ ├── cloudsqlmysql
│ │ │ └── cloudsqlmysqlcreateinstance
│ │ │ ├── cloudsqlmysqlcreateinstance_test.go
│ │ │ └── cloudsqlmysqlcreateinstance.go
│ │ ├── cloudsqlpg
│ │ │ └── cloudsqlpgcreateinstances
│ │ │ ├── cloudsqlpgcreateinstances_test.go
│ │ │ └── cloudsqlpgcreateinstances.go
│ │ ├── common_test.go
│ │ ├── common.go
│ │ ├── couchbase
│ │ │ ├── couchbase_test.go
│ │ │ └── couchbase.go
│ │ ├── dataform
│ │ │ └── dataformcompilelocal
│ │ │ ├── dataformcompilelocal_test.go
│ │ │ └── dataformcompilelocal.go
│ │ ├── dataplex
│ │ │ ├── dataplexlookupentry
│ │ │ │ ├── dataplexlookupentry_test.go
│ │ │ │ └── dataplexlookupentry.go
│ │ │ ├── dataplexsearchaspecttypes
│ │ │ │ ├── dataplexsearchaspecttypes_test.go
│ │ │ │ └── dataplexsearchaspecttypes.go
│ │ │ └── dataplexsearchentries
│ │ │ ├── dataplexsearchentries_test.go
│ │ │ └── dataplexsearchentries.go
│ │ ├── dgraph
│ │ │ ├── dgraph_test.go
│ │ │ └── dgraph.go
│ │ ├── firebird
│ │ │ ├── firebirdexecutesql
│ │ │ │ ├── firebirdexecutesql_test.go
│ │ │ │ └── firebirdexecutesql.go
│ │ │ └── firebirdsql
│ │ │ ├── firebirdsql_test.go
│ │ │ └── firebirdsql.go
│ │ ├── firestore
│ │ │ ├── firestoreadddocuments
│ │ │ │ ├── firestoreadddocuments_test.go
│ │ │ │ └── firestoreadddocuments.go
│ │ │ ├── firestoredeletedocuments
│ │ │ │ ├── firestoredeletedocuments_test.go
│ │ │ │ └── firestoredeletedocuments.go
│ │ │ ├── firestoregetdocuments
│ │ │ │ ├── firestoregetdocuments_test.go
│ │ │ │ └── firestoregetdocuments.go
│ │ │ ├── firestoregetrules
│ │ │ │ ├── firestoregetrules_test.go
│ │ │ │ └── firestoregetrules.go
│ │ │ ├── firestorelistcollections
│ │ │ │ ├── firestorelistcollections_test.go
│ │ │ │ └── firestorelistcollections.go
│ │ │ ├── firestorequery
│ │ │ │ ├── firestorequery_test.go
│ │ │ │ └── firestorequery.go
│ │ │ ├── firestorequerycollection
│ │ │ │ ├── firestorequerycollection_test.go
│ │ │ │ └── firestorequerycollection.go
│ │ │ ├── firestoreupdatedocument
│ │ │ │ ├── firestoreupdatedocument_test.go
│ │ │ │ └── firestoreupdatedocument.go
│ │ │ ├── firestorevalidaterules
│ │ │ │ ├── firestorevalidaterules_test.go
│ │ │ │ └── firestorevalidaterules.go
│ │ │ └── util
│ │ │ ├── converter_test.go
│ │ │ ├── converter.go
│ │ │ ├── validator_test.go
│ │ │ └── validator.go
│ │ ├── http
│ │ │ ├── http_test.go
│ │ │ └── http.go
│ │ ├── http_method.go
│ │ ├── looker
│ │ │ ├── lookeradddashboardelement
│ │ │ │ ├── lookeradddashboardelement_test.go
│ │ │ │ └── lookeradddashboardelement.go
│ │ │ ├── lookercommon
│ │ │ │ ├── lookercommon_test.go
│ │ │ │ └── lookercommon.go
│ │ │ ├── lookerconversationalanalytics
│ │ │ │ ├── lookerconversationalanalytics_test.go
│ │ │ │ └── lookerconversationalanalytics.go
│ │ │ ├── lookergetdashboards
│ │ │ │ ├── lookergetdashboards_test.go
│ │ │ │ └── lookergetdashboards.go
│ │ │ ├── lookergetdimensions
│ │ │ │ ├── lookergetdimensions_test.go
│ │ │ │ └── lookergetdimensions.go
│ │ │ ├── lookergetexplores
│ │ │ │ ├── lookergetexplores_test.go
│ │ │ │ └── lookergetexplores.go
│ │ │ ├── lookergetfilters
│ │ │ │ ├── lookergetfilters_test.go
│ │ │ │ └── lookergetfilters.go
│ │ │ ├── lookergetlooks
│ │ │ │ ├── lookergetlooks_test.go
│ │ │ │ └── lookergetlooks.go
│ │ │ ├── lookergetmeasures
│ │ │ │ ├── lookergetmeasures_test.go
│ │ │ │ └── lookergetmeasures.go
│ │ │ ├── lookergetmodels
│ │ │ │ ├── lookergetmodels_test.go
│ │ │ │ └── lookergetmodels.go
│ │ │ ├── lookergetparameters
│ │ │ │ ├── lookergetparameters_test.go
│ │ │ │ └── lookergetparameters.go
│ │ │ ├── lookerhealthanalyze
│ │ │ │ ├── lookerhealthanalyze_test.go
│ │ │ │ └── lookerhealthanalyze.go
│ │ │ ├── lookerhealthpulse
│ │ │ │ ├── lookerhealthpulse_test.go
│ │ │ │ └── lookerhealthpulse.go
│ │ │ ├── lookerhealthvacuum
│ │ │ │ ├── lookerhealthvacuum_test.go
│ │ │ │ └── lookerhealthvacuum.go
│ │ │ ├── lookermakedashboard
│ │ │ │ ├── lookermakedashboard_test.go
│ │ │ │ └── lookermakedashboard.go
│ │ │ ├── lookermakelook
│ │ │ │ ├── lookermakelook_test.go
│ │ │ │ └── lookermakelook.go
│ │ │ ├── lookerquery
│ │ │ │ ├── lookerquery_test.go
│ │ │ │ └── lookerquery.go
│ │ │ ├── lookerquerysql
│ │ │ │ ├── lookerquerysql_test.go
│ │ │ │ └── lookerquerysql.go
│ │ │ ├── lookerqueryurl
│ │ │ │ ├── lookerqueryurl_test.go
│ │ │ │ └── lookerqueryurl.go
│ │ │ └── lookerrunlook
│ │ │ ├── lookerrunlook_test.go
│ │ │ └── lookerrunlook.go
│ │ ├── mongodb
│ │ │ ├── mongodbaggregate
│ │ │ │ ├── mongodbaggregate_test.go
│ │ │ │ └── mongodbaggregate.go
│ │ │ ├── mongodbdeletemany
│ │ │ │ ├── mongodbdeletemany_test.go
│ │ │ │ └── mongodbdeletemany.go
│ │ │ ├── mongodbdeleteone
│ │ │ │ ├── mongodbdeleteone_test.go
│ │ │ │ └── mongodbdeleteone.go
│ │ │ ├── mongodbfind
│ │ │ │ ├── mongodbfind_test.go
│ │ │ │ └── mongodbfind.go
│ │ │ ├── mongodbfindone
│ │ │ │ ├── mongodbfindone_test.go
│ │ │ │ └── mongodbfindone.go
│ │ │ ├── mongodbinsertmany
│ │ │ │ ├── mongodbinsertmany_test.go
│ │ │ │ └── mongodbinsertmany.go
│ │ │ ├── mongodbinsertone
│ │ │ │ ├── mongodbinsertone_test.go
│ │ │ │ └── mongodbinsertone.go
│ │ │ ├── mongodbupdatemany
│ │ │ │ ├── mongodbupdatemany_test.go
│ │ │ │ └── mongodbupdatemany.go
│ │ │ └── mongodbupdateone
│ │ │ ├── mongodbupdateone_test.go
│ │ │ └── mongodbupdateone.go
│ │ ├── mssql
│ │ │ ├── mssqlexecutesql
│ │ │ │ ├── mssqlexecutesql_test.go
│ │ │ │ └── mssqlexecutesql.go
│ │ │ ├── mssqllisttables
│ │ │ │ ├── mssqllisttables_test.go
│ │ │ │ └── mssqllisttables.go
│ │ │ └── mssqlsql
│ │ │ ├── mssqlsql_test.go
│ │ │ └── mssqlsql.go
│ │ ├── mysql
│ │ │ ├── mysqlcommon
│ │ │ │ └── mysqlcommon.go
│ │ │ ├── mysqlexecutesql
│ │ │ │ ├── mysqlexecutesql_test.go
│ │ │ │ └── mysqlexecutesql.go
│ │ │ ├── mysqllistactivequeries
│ │ │ │ ├── mysqllistactivequeries_test.go
│ │ │ │ └── mysqllistactivequeries.go
│ │ │ ├── mysqllisttablefragmentation
│ │ │ │ ├── mysqllisttablefragmentation_test.go
│ │ │ │ └── mysqllisttablefragmentation.go
│ │ │ ├── mysqllisttables
│ │ │ │ ├── mysqllisttables_test.go
│ │ │ │ └── mysqllisttables.go
│ │ │ ├── mysqllisttablesmissinguniqueindexes
│ │ │ │ ├── mysqllisttablesmissinguniqueindexes_test.go
│ │ │ │ └── mysqllisttablesmissinguniqueindexes.go
│ │ │ └── mysqlsql
│ │ │ ├── mysqlsql_test.go
│ │ │ └── mysqlsql.go
│ │ ├── neo4j
│ │ │ ├── neo4jcypher
│ │ │ │ ├── neo4jcypher_test.go
│ │ │ │ └── neo4jcypher.go
│ │ │ ├── neo4jexecutecypher
│ │ │ │ ├── classifier
│ │ │ │ │ ├── classifier_test.go
│ │ │ │ │ └── classifier.go
│ │ │ │ ├── neo4jexecutecypher_test.go
│ │ │ │ └── neo4jexecutecypher.go
│ │ │ └── neo4jschema
│ │ │ ├── cache
│ │ │ │ ├── cache_test.go
│ │ │ │ └── cache.go
│ │ │ ├── helpers
│ │ │ │ ├── helpers_test.go
│ │ │ │ └── helpers.go
│ │ │ ├── neo4jschema_test.go
│ │ │ ├── neo4jschema.go
│ │ │ └── types
│ │ │ └── types.go
│ │ ├── oceanbase
│ │ │ ├── oceanbaseexecutesql
│ │ │ │ ├── oceanbaseexecutesql_test.go
│ │ │ │ └── oceanbaseexecutesql.go
│ │ │ └── oceanbasesql
│ │ │ ├── oceanbasesql_test.go
│ │ │ └── oceanbasesql.go
│ │ ├── oracle
│ │ │ ├── oracleexecutesql
│ │ │ │ └── oracleexecutesql.go
│ │ │ └── oraclesql
│ │ │ └── oraclesql.go
│ │ ├── parameters_test.go
│ │ ├── parameters.go
│ │ ├── postgres
│ │ │ ├── postgresexecutesql
│ │ │ │ ├── postgresexecutesql_test.go
│ │ │ │ └── postgresexecutesql.go
│ │ │ ├── postgreslistactivequeries
│ │ │ │ ├── postgreslistactivequeries_test.go
│ │ │ │ └── postgreslistactivequeries.go
│ │ │ ├── postgreslistavailableextensions
│ │ │ │ ├── postgreslistavailableextensions_test.go
│ │ │ │ └── postgreslistavailableextensions.go
│ │ │ ├── postgreslistinstalledextensions
│ │ │ │ ├── postgreslistinstalledextensions_test.go
│ │ │ │ └── postgreslistinstalledextensions.go
│ │ │ ├── postgreslisttables
│ │ │ │ ├── postgreslisttables_test.go
│ │ │ │ └── postgreslisttables.go
│ │ │ └── postgressql
│ │ │ ├── postgressql_test.go
│ │ │ └── postgressql.go
│ │ ├── redis
│ │ │ ├── redis_test.go
│ │ │ └── redis.go
│ │ ├── spanner
│ │ │ ├── spannerexecutesql
│ │ │ │ ├── spannerexecutesql_test.go
│ │ │ │ └── spannerexecutesql.go
│ │ │ ├── spannerlisttables
│ │ │ │ ├── spannerlisttables_test.go
│ │ │ │ └── spannerlisttables.go
│ │ │ └── spannersql
│ │ │ ├── spanner_test.go
│ │ │ └── spannersql.go
│ │ ├── sqlite
│ │ │ ├── sqliteexecutesql
│ │ │ │ ├── sqliteexecutesql_test.go
│ │ │ │ └── sqliteexecutesql.go
│ │ │ └── sqlitesql
│ │ │ ├── sqlitesql_test.go
│ │ │ └── sqlitesql.go
│ │ ├── tidb
│ │ │ ├── tidbexecutesql
│ │ │ │ ├── tidbexecutesql_test.go
│ │ │ │ └── tidbexecutesql.go
│ │ │ └── tidbsql
│ │ │ ├── tidbsql_test.go
│ │ │ └── tidbsql.go
│ │ ├── tools_test.go
│ │ ├── tools.go
│ │ ├── toolsets.go
│ │ ├── trino
│ │ │ ├── trinoexecutesql
│ │ │ │ ├── trinoexecutesql_test.go
│ │ │ │ └── trinoexecutesql.go
│ │ │ └── trinosql
│ │ │ ├── trinosql_test.go
│ │ │ └── trinosql.go
│ │ ├── utility
│ │ │ └── wait
│ │ │ ├── wait_test.go
│ │ │ └── wait.go
│ │ ├── valkey
│ │ │ ├── valkey_test.go
│ │ │ └── valkey.go
│ │ └── yugabytedbsql
│ │ ├── yugabytedbsql_test.go
│ │ └── yugabytedbsql.go
│ └── util
│ └── util.go
├── LICENSE
├── logo.png
├── main.go
├── MCP-TOOLBOX-EXTENSION.md
├── README.md
└── tests
├── alloydb
│ ├── alloydb_integration_test.go
│ └── alloydb_wait_for_operation_test.go
├── alloydbainl
│ └── alloydb_ai_nl_integration_test.go
├── alloydbpg
│ └── alloydb_pg_integration_test.go
├── auth.go
├── bigquery
│ └── bigquery_integration_test.go
├── bigtable
│ └── bigtable_integration_test.go
├── cassandra
│ └── cassandra_integration_test.go
├── clickhouse
│ └── clickhouse_integration_test.go
├── cloudmonitoring
│ └── cloud_monitoring_integration_test.go
├── cloudsql
│ ├── cloud_sql_create_database_test.go
│ ├── cloud_sql_create_users_test.go
│ ├── cloud_sql_get_instances_test.go
│ ├── cloud_sql_list_databases_test.go
│ ├── cloudsql_list_instances_test.go
│ └── cloudsql_wait_for_operation_test.go
├── cloudsqlmssql
│ ├── cloud_sql_mssql_create_instance_integration_test.go
│ └── cloud_sql_mssql_integration_test.go
├── cloudsqlmysql
│ ├── cloud_sql_mysql_create_instance_integration_test.go
│ └── cloud_sql_mysql_integration_test.go
├── cloudsqlpg
│ ├── cloud_sql_pg_create_instances_test.go
│ └── cloud_sql_pg_integration_test.go
├── common.go
├── couchbase
│ └── couchbase_integration_test.go
├── dataform
│ └── dataform_integration_test.go
├── dataplex
│ └── dataplex_integration_test.go
├── dgraph
│ └── dgraph_integration_test.go
├── firebird
│ └── firebird_integration_test.go
├── firestore
│ └── firestore_integration_test.go
├── http
│ └── http_integration_test.go
├── looker
│ └── looker_integration_test.go
├── mongodb
│ └── mongodb_integration_test.go
├── mssql
│ └── mssql_integration_test.go
├── mysql
│ └── mysql_integration_test.go
├── neo4j
│ └── neo4j_integration_test.go
├── oceanbase
│ └── oceanbase_integration_test.go
├── option.go
├── oracle
│ └── oracle_integration_test.go
├── postgres
│ └── postgres_integration_test.go
├── redis
│ └── redis_test.go
├── server.go
├── source.go
├── spanner
│ └── spanner_integration_test.go
├── sqlite
│ └── sqlite_integration_test.go
├── tidb
│ └── tidb_integration_test.go
├── tool.go
├── trino
│ └── trino_integration_test.go
├── utility
│ └── wait_integration_test.go
├── valkey
│ └── valkey_test.go
└── yugabytedb
└── yugabytedb_integration_test.go
```
# Files
--------------------------------------------------------------------------------
/internal/tools/mongodb/mongodbupdateone/mongodbupdateone_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package mongodbupdateone_test
import (
"strings"
"testing"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbupdateone"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
)
func TestParseFromYamlMongoQuery(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
example_tool:
kind: mongodb-update-one
source: my-instance
description: some description
database: test_db
collection: test_coll
filterPayload: |
{ name: {{json .name}} }
filterParams:
- name: name
type: string
description: small description
updatePayload: |
{ $set : { item: {{json .item}} } }
updateParams:
- name: item
type: string
description: small description
canonical: true
upsert: true
`,
want: server.ToolConfigs{
"example_tool": mongodbupdateone.Config{
Name: "example_tool",
Kind: "mongodb-update-one",
Source: "my-instance",
AuthRequired: []string{},
Database: "test_db",
Collection: "test_coll",
Canonical: true,
FilterPayload: "{ name: {{json .name}} }\n",
FilterParams: tools.Parameters{
&tools.StringParameter{
CommonParameter: tools.CommonParameter{
Name: "name",
Type: "string",
Desc: "small description",
},
},
},
UpdatePayload: "{ $set : { item: {{json .item}} } }\n",
UpdateParams: tools.Parameters{
&tools.StringParameter{
CommonParameter: tools.CommonParameter{
Name: "item",
Type: "string",
Desc: "small description",
},
},
},
Upsert: true,
Description: "some description",
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestFailParseFromYamlMongoQuery(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
err string
}{
{
desc: "Invalid method",
in: `
tools:
example_tool:
kind: mongodb-update-one
source: my-instance
description: some description
collection: test_coll
filterPayload: |
{ name : {{json .name}} }`,
err: `unable to parse tool "example_tool" as kind "mongodb-update-one"`,
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err == nil {
t.Fatalf("expect parsing to fail")
}
errStr := err.Error()
if !strings.Contains(errStr, tc.err) {
t.Fatalf("unexpected error string: got %q, want substring %q", errStr, tc.err)
}
})
}
}
```
--------------------------------------------------------------------------------
/internal/sources/cassandra/cassandra_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package cassandra_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/sources/cassandra"
"github.com/googleapis/genai-toolbox/internal/testutils"
)
func TestParseFromYamlCassandra(t *testing.T) {
tcs := []struct {
desc string
in string
want server.SourceConfigs
}{
{
desc: "basic example (without optional fields)",
in: `
sources:
my-cassandra-instance:
kind: cassandra
hosts:
- "my-host1"
- "my-host2"
`,
want: server.SourceConfigs{
"my-cassandra-instance": cassandra.Config{
Name: "my-cassandra-instance",
Kind: cassandra.SourceKind,
Hosts: []string{"my-host1", "my-host2"},
Username: "",
Password: "",
ProtoVersion: 0,
CAPath: "",
CertPath: "",
KeyPath: "",
Keyspace: "",
EnableHostVerification: false,
},
},
},
{
desc: "with optional fields",
in: `
sources:
my-cassandra-instance:
kind: cassandra
hosts:
- "my-host1"
- "my-host2"
username: "user"
password: "pass"
keyspace: "example_keyspace"
protoVersion: 4
caPath: "path/to/ca.crt"
certPath: "path/to/cert"
keyPath: "path/to/key"
enableHostVerification: true
`,
want: server.SourceConfigs{
"my-cassandra-instance": cassandra.Config{
Name: "my-cassandra-instance",
Kind: cassandra.SourceKind,
Hosts: []string{"my-host1", "my-host2"},
Username: "user",
Password: "pass",
Keyspace: "example_keyspace",
ProtoVersion: 4,
CAPath: "path/to/ca.crt",
CertPath: "path/to/cert",
KeyPath: "path/to/key",
EnableHostVerification: true,
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Sources server.SourceConfigs `yaml:"sources"`
}{}
// Parse contents
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if !cmp.Equal(tc.want, got.Sources) {
t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources)
}
})
}
}
func TestFailParseFromYaml(t *testing.T) {
tcs := []struct {
desc string
in string
err string
}{
{
desc: "extra field",
in: `
sources:
my-cassandra-instance:
kind: cassandra
hosts:
- "my-host"
foo: bar
`,
err: "unable to parse source \"my-cassandra-instance\" as \"cassandra\": [1:1] unknown field \"foo\"\n> 1 | foo: bar\n ^\n 2 | hosts:\n 3 | - my-host\n 4 | kind: cassandra",
},
{
desc: "missing required field",
in: `
sources:
my-cassandra-instance:
kind: cassandra
`,
err: "unable to parse source \"my-cassandra-instance\" as \"cassandra\": Key: 'Config.Hosts' Error:Field validation for 'Hosts' failed on the 'required' tag",
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Sources server.SourceConfigs `yaml:"sources"`
}{}
// Parse contents
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
if err == nil {
t.Fatalf("expect parsing to fail")
}
errStr := err.Error()
if errStr != tc.err {
t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
}
})
}
}
```
--------------------------------------------------------------------------------
/internal/sources/yugabytedb/yugabytedb.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package yugabytedb
import (
"context"
"fmt"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/yugabyte/pgx/v5/pgxpool"
"go.opentelemetry.io/otel/trace"
)
const SourceKind string = "yugabytedb"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
actual := Config{Name: name}
if err := decoder.DecodeContext(ctx, &actual); err != nil {
return nil, err
}
return actual, nil
}
type Config struct {
Name string `yaml:"name" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
User string `yaml:"user" validate:"required"`
Password string `yaml:"password" validate:"required"`
Database string `yaml:"database" validate:"required"`
LoadBalance string `yaml:"loadBalance"`
TopologyKeys string `yaml:"topologyKeys"`
YBServersRefreshInterval string `yaml:"ybServersRefreshInterval"`
FallBackToTopologyKeysOnly string `yaml:"fallbackToTopologyKeysOnly"`
FailedHostReconnectDelaySeconds string `yaml:"failedHostReconnectDelaySecs"`
}
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
pool, err := initYugabyteDBConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Database, r.LoadBalance, r.TopologyKeys, r.YBServersRefreshInterval, r.FallBackToTopologyKeysOnly, r.FailedHostReconnectDelaySeconds)
if err != nil {
return nil, fmt.Errorf("unable to create pool: %w", err)
}
err = pool.Ping(ctx)
if err != nil {
return nil, fmt.Errorf("unable to connect successfully: %w", err)
}
s := &Source{
Name: r.Name,
Kind: SourceKind,
Pool: pool,
}
return s, nil
}
var _ sources.Source = &Source{}
type Source struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
Pool *pgxpool.Pool
}
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) YugabyteDBPool() *pgxpool.Pool {
return s.Pool
}
func initYugabyteDBConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, pass, dbname, loadBalance, topologyKeys, refreshInterval, explicitFallback, failedHostTTL string) (*pgxpool.Pool, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// urlExample := "postgres://username:password@localhost:5433/database_name"
i := fmt.Sprintf("postgres://%s:%s@%s:%s/%s", user, pass, host, port, dbname)
if loadBalance == "true" {
i = fmt.Sprintf("%s?load_balance=%s", i, loadBalance)
if topologyKeys != "" {
i = fmt.Sprintf("%s&topology_keys=%s", i, topologyKeys)
if explicitFallback == "true" {
i = fmt.Sprintf("%s&fallback_to_topology_keys_only=%s", i, explicitFallback)
}
}
if refreshInterval != "" {
i = fmt.Sprintf("%s&yb_servers_refresh_interval=%s", i, refreshInterval)
}
if failedHostTTL != "" {
i = fmt.Sprintf("%s&failed_host_reconnect_delay_secs=%s", i, failedHostTTL)
}
}
pool, err := pgxpool.New(ctx, i)
if err != nil {
return nil, fmt.Errorf("unable to create connection pool: %w", err)
}
return pool, nil
}
```
--------------------------------------------------------------------------------
/internal/sources/bigquery/bigquery_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquery_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/sources/bigquery"
"github.com/googleapis/genai-toolbox/internal/testutils"
)
func TestParseFromYamlBigQuery(t *testing.T) {
tcs := []struct {
desc string
in string
want server.SourceConfigs
}{
{
desc: "basic example",
in: `
sources:
my-instance:
kind: bigquery
project: my-project
`,
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "",
WriteMode: "",
},
},
},
{
desc: "all fields specified",
in: `
sources:
my-instance:
kind: bigquery
project: my-project
location: asia
writeMode: blocked
`,
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "asia",
WriteMode: "blocked",
UseClientOAuth: false,
},
},
},
{
desc: "use client auth example",
in: `
sources:
my-instance:
kind: bigquery
project: my-project
location: us
useClientOAuth: true
`,
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "us",
UseClientOAuth: true,
},
},
},
{
desc: "with allowed datasets example",
in: `
sources:
my-instance:
kind: bigquery
project: my-project
location: us
allowedDatasets:
- my_dataset
`,
want: server.SourceConfigs{
"my-instance": bigquery.Config{
Name: "my-instance",
Kind: bigquery.SourceKind,
Project: "my-project",
Location: "us",
AllowedDatasets: []string{"my_dataset"},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Sources server.SourceConfigs `yaml:"sources"`
}{}
// Parse contents
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if !cmp.Equal(tc.want, got.Sources) {
t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources)
}
})
}
}
func TestFailParseFromYaml(t *testing.T) {
tcs := []struct {
desc string
in string
err string
}{
{
desc: "extra field",
in: `
sources:
my-instance:
kind: bigquery
project: my-project
location: us
foo: bar
`,
err: "unable to parse source \"my-instance\" as \"bigquery\": [1:1] unknown field \"foo\"\n> 1 | foo: bar\n ^\n 2 | kind: bigquery\n 3 | location: us\n 4 | project: my-project",
},
{
desc: "missing required field",
in: `
sources:
my-instance:
kind: bigquery
location: us
`,
err: "unable to parse source \"my-instance\" as \"bigquery\": Key: 'Config.Project' Error:Field validation for 'Project' failed on the 'required' tag",
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Sources server.SourceConfigs `yaml:"sources"`
}{}
// Parse contents
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
if err == nil {
t.Fatalf("expect parsing to fail")
}
errStr := err.Error()
if errStr != tc.err {
t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
}
})
}
}
```
--------------------------------------------------------------------------------
/docs/en/getting-started/quickstart/shared/database_setup.md:
--------------------------------------------------------------------------------
```markdown
<!-- This file has been used in local_quickstart.md, local_quickstart_go.md & local_quickstart_js.md -->
<!-- [START database_setup] -->
In this section, we will create a database, insert some data that needs to be
accessed by our agent, and create a database user for Toolbox to connect with.
1. Connect to postgres using the `psql` command:
```bash
psql -h 127.0.0.1 -U postgres
```
Here, `postgres` denotes the default postgres superuser.
{{< notice info >}}
#### **Having trouble connecting?**
* **Password Prompt:** If you are prompted for a password for the `postgres`
user and do not know it (or a blank password doesn't work), your PostgreSQL
installation might require a password or a different authentication method.
* **`FATAL: role "postgres" does not exist`:** This error means the default
`postgres` superuser role isn't available under that name on your system.
* **`Connection refused`:** Ensure your PostgreSQL server is actually running.
You can typically check with `sudo systemctl status postgresql` and start it
with `sudo systemctl start postgresql` on Linux systems.
<br/>
#### **Common Solution**
For password issues or if the `postgres` role seems inaccessible directly, try
switching to the `postgres` operating system user first. This user often has
permission to connect without a password for local connections (this is called
peer authentication).
```bash
sudo -i -u postgres
psql -h 127.0.0.1
```
Once you are in the `psql` shell using this method, you can proceed with the
database creation steps below. Afterwards, type `\q` to exit `psql`, and then
`exit` to return to your normal user shell.
If desired, once connected to `psql` as the `postgres` OS user, you can set a
password for the `postgres` *database* user using: `ALTER USER postgres WITH
PASSWORD 'your_chosen_password';`. This would allow direct connection with `-U
postgres` and a password next time.
{{< /notice >}}
1. Create a new database and a new user:
{{< notice tip >}}
For a real application, it's best to follow the principle of least permission
and only grant the privileges your application needs.
{{< /notice >}}
```sql
CREATE USER toolbox_user WITH PASSWORD 'my-password';
CREATE DATABASE toolbox_db;
GRANT ALL PRIVILEGES ON DATABASE toolbox_db TO toolbox_user;
ALTER DATABASE toolbox_db OWNER TO toolbox_user;
```
1. End the database session:
```bash
\q
```
(If you used `sudo -i -u postgres` and then `psql`, remember you might also
need to type `exit` after `\q` to leave the `postgres` user's shell
session.)
1. Connect to your database with your new user:
```bash
psql -h 127.0.0.1 -U toolbox_user -d toolbox_db
```
1. Create a table using the following command:
```sql
CREATE TABLE hotels(
id INTEGER NOT NULL PRIMARY KEY,
name VARCHAR NOT NULL,
location VARCHAR NOT NULL,
price_tier VARCHAR NOT NULL,
checkin_date DATE NOT NULL,
checkout_date DATE NOT NULL,
booked BIT NOT NULL
);
```
1. Insert data into the table.
```sql
INSERT INTO hotels(id, name, location, price_tier, checkin_date, checkout_date, booked)
VALUES
(1, 'Hilton Basel', 'Basel', 'Luxury', '2024-04-22', '2024-04-20', B'0'),
(2, 'Marriott Zurich', 'Zurich', 'Upscale', '2024-04-14', '2024-04-21', B'0'),
(3, 'Hyatt Regency Basel', 'Basel', 'Upper Upscale', '2024-04-02', '2024-04-20', B'0'),
(4, 'Radisson Blu Lucerne', 'Lucerne', 'Midscale', '2024-04-24', '2024-04-05', B'0'),
(5, 'Best Western Bern', 'Bern', 'Upper Midscale', '2024-04-23', '2024-04-01', B'0'),
(6, 'InterContinental Geneva', 'Geneva', 'Luxury', '2024-04-23', '2024-04-28', B'0'),
(7, 'Sheraton Zurich', 'Zurich', 'Upper Upscale', '2024-04-27', '2024-04-02', B'0'),
(8, 'Holiday Inn Basel', 'Basel', 'Upper Midscale', '2024-04-24', '2024-04-09', B'0'),
(9, 'Courtyard Zurich', 'Zurich', 'Upscale', '2024-04-03', '2024-04-13', B'0'),
(10, 'Comfort Inn Bern', 'Bern', 'Midscale', '2024-04-04', '2024-04-16', B'0');
```
1. End the database session:
```bash
\q
```
<!-- [END database_setup] -->
```
--------------------------------------------------------------------------------
/tests/valkey/valkey_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package valkey
import (
"context"
"log"
"os"
"regexp"
"testing"
"time"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
"github.com/valkey-io/valkey-go"
)
var (
ValkeySourceKind = "valkey"
ValkeyToolKind = "valkey"
ValkeyAddress = os.Getenv("VALKEY_ADDRESS")
)
func getValkeyVars(t *testing.T) map[string]any {
switch "" {
case ValkeyAddress:
t.Fatal("'VALKEY_ADDRESS' not set")
}
return map[string]any{
"kind": ValkeySourceKind,
"address": []string{ValkeyAddress},
"disableCache": true,
}
}
func initValkeyClient(ctx context.Context, addr []string) (valkey.Client, error) {
// Pass in an access token getter fn for IAM auth
client, err := valkey.NewClient(valkey.ClientOption{
InitAddress: addr,
ForceSingleClient: true,
DisableCache: true,
})
if err != nil {
log.Fatalf("error creating client: %v", err)
}
// Ping the server to check connectivity (using Do)
pingCmd := client.B().Ping().Build()
_, err = client.Do(ctx, pingCmd).ToString()
if err != nil {
log.Fatalf("Failed to execute PING command: %v", err)
}
log.Println("Successfully connected to Valkey")
return client, nil
}
func TestValkeyToolEndpoints(t *testing.T) {
sourceConfig := getValkeyVars(t)
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
defer cancel()
var args []string
client, err := initValkeyClient(ctx, []string{ValkeyAddress})
if err != nil {
t.Fatalf("unable to create Valkey connection: %s", err)
}
// set up data for param tool
teardownDB := setupValkeyDB(t, ctx, client)
defer teardownDB(t)
// Write config into a file and pass it to command
toolsFile := tests.GetRedisValkeyToolsConfig(sourceConfig, ValkeyToolKind)
cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
if err != nil {
t.Fatalf("command initialization returned an error: %s", err)
}
defer cleanup()
waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
if err != nil {
t.Logf("toolbox command logs: \n%s", out)
t.Fatalf("toolbox didn't start successfully: %s", err)
}
// Get configs for tests
select1Want, mcpMyFailToolWant, invokeParamWant, invokeIdNullWant, nullWant, mcpSelect1Want, mcpInvokeParamWant := tests.GetRedisValkeyWants()
// Run tests
tests.RunToolGetTest(t)
tests.RunToolInvokeTest(t, select1Want,
tests.WithMyToolId3NameAliceWant(invokeParamWant),
tests.WithMyArrayToolWant(invokeParamWant),
tests.WithMyToolById4Want(invokeIdNullWant),
tests.WithNullWant(nullWant),
)
tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want,
tests.WithMcpMyToolId3NameAliceWant(mcpInvokeParamWant),
)
}
func setupValkeyDB(t *testing.T, ctx context.Context, client valkey.Client) func(*testing.T) {
keys := []string{"row1", "row2", "row3", "row4", "null"}
commands := [][]string{
{"HSET", keys[0], "name", "Alice", "id", "1"},
{"HSET", keys[1], "name", "Jane", "id", "2"},
{"HSET", keys[2], "name", "Sid", "id", "3"},
{"HSET", keys[3], "name", "", "id", "4"},
{"SET", keys[4], "null"},
{"HSET", tests.ServiceAccountEmail, "name", "Alice"},
}
builtCmds := make(valkey.Commands, len(commands))
for i, cmd := range commands {
builtCmds[i] = client.B().Arbitrary(cmd...).Build()
}
responses := client.DoMulti(ctx, builtCmds...)
for _, resp := range responses {
if err := resp.Error(); err != nil {
t.Fatalf("unable to insert test data: %s", err)
}
}
return func(t *testing.T) {
// tear down test
_, err := client.Do(ctx, client.B().Del().Key(keys...).Build()).AsInt64()
if err != nil {
t.Errorf("Teardown failed: %s", err)
}
}
}
```
--------------------------------------------------------------------------------
/internal/tools/bigtable/bigtable_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigtable_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/bigtable"
)
func TestParseFromYamlBigtable(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
example_tool:
kind: bigtable-sql
source: my-pg-instance
description: some description
statement: |
SELECT * FROM SQL_STATEMENT;
parameters:
- name: country
type: string
description: some description
`,
want: server.ToolConfigs{
"example_tool": bigtable.Config{
Name: "example_tool",
Kind: "bigtable-sql",
Source: "my-pg-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
AuthRequired: []string{},
Parameters: []tools.Parameter{
tools.NewStringParameter("country", "some description"),
},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestParseFromYamlWithTemplateBigtable(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
example_tool:
kind: bigtable-sql
source: my-pg-instance
description: some description
statement: |
SELECT * FROM SQL_STATEMENT;
parameters:
- name: country
type: string
description: some description
templateParameters:
- name: tableName
type: string
description: The table to select hotels from.
- name: fieldArray
type: array
description: The columns to return for the query.
items:
name: column
type: string
description: A column name that will be returned from the query.
`,
want: server.ToolConfigs{
"example_tool": bigtable.Config{
Name: "example_tool",
Kind: "bigtable-sql",
Source: "my-pg-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
AuthRequired: []string{},
Parameters: []tools.Parameter{
tools.NewStringParameter("country", "some description"),
},
TemplateParameters: []tools.Parameter{
tools.NewStringParameter("tableName", "The table to select hotels from."),
tools.NewArrayParameter("fieldArray", "The columns to return for the query.", tools.NewStringParameter("column", "A column name that will be returned from the query.")),
},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
```
--------------------------------------------------------------------------------
/internal/server/mcp/util/lifecycle.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package util
import (
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
)
const (
// SERVER_NAME is the server name used in Implementation.
SERVER_NAME = "Toolbox"
// methods that are supported
INITIALIZE = "initialize"
)
/* Initialization */
// Params to define MCP Client during initialize request.
type InitializeParams struct {
// The latest version of the Model Context Protocol that the client supports.
// The client MAY decide to support older versions as well.
ProtocolVersion string `json:"protocolVersion"`
Capabilities ClientCapabilities `json:"capabilities"`
ClientInfo Implementation `json:"clientInfo"`
}
// InitializeRequest is sent from the client to the server when it first
// connects, asking it to begin initialization.
type InitializeRequest struct {
jsonrpc.Request
Params InitializeParams `json:"params"`
}
// InitializeResult is sent after receiving an initialize request from the
// client.
type InitializeResult struct {
jsonrpc.Result
// The version of the Model Context Protocol that the server wants to use.
// This may not match the version that the client requested. If the client cannot
// support this version, it MUST disconnect.
ProtocolVersion string `json:"protocolVersion"`
Capabilities ServerCapabilities `json:"capabilities"`
ServerInfo Implementation `json:"serverInfo"`
// Instructions describing how to use the server and its features.
//
// This can be used by clients to improve the LLM's understanding of
// available tools, resources, etc. It can be thought of like a "hint" to the model.
// For example, this information MAY be added to the system prompt.
Instructions string `json:"instructions,omitempty"`
}
// InitializedNotification is sent from the client to the server after
// initialization has finished.
type InitializedNotification struct {
jsonrpc.Notification
}
// ListChange represents whether the server supports notification for changes to the capabilities.
type ListChanged struct {
ListChanged *bool `json:"listChanged,omitempty"`
}
// ClientCapabilities represents capabilities a client may support. Known
// capabilities are defined here, in this schema, but this is not a closed set: any
// client can define its own, additional capabilities.
type ClientCapabilities struct {
// Experimental, non-standard capabilities that the client supports.
Experimental map[string]interface{} `json:"experimental,omitempty"`
// Present if the client supports listing roots.
Roots *ListChanged `json:"roots,omitempty"`
// Present if the client supports sampling from an LLM.
Sampling struct{} `json:"sampling,omitempty"`
}
// ServerCapabilities represents capabilities that a server may support. Known
// capabilities are defined here, in this schema, but this is not a closed set: any
// server can define its own, additional capabilities.
type ServerCapabilities struct {
Tools *ListChanged `json:"tools,omitempty"`
}
// Base interface for metadata with name (identifier) and title (display name) properties.
type BaseMetadata struct {
// Intended for programmatic or logical use, but used as a display name in past specs
// or fallback (if title isn't present).
Name string `json:"name"`
// Intended for UI and end-user contexts — optimized to be human-readable and easily understood,
//even by those unfamiliar with domain-specific terminology.
//
// If not provided, the name should be used for display (except for Tool,
// where `annotations.title` should be given precedence over using `name`,
// if present).
Title string `json:"title,omitempty"`
}
// Implementation describes the name and version of an MCP implementation.
type Implementation struct {
BaseMetadata
Version string `json:"version"`
}
```
--------------------------------------------------------------------------------
/internal/tools/firestore/util/validator.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package util
import (
"fmt"
"regexp"
"strings"
)
// Regular expressions for validating Firestore paths
var (
// Pattern to detect absolute paths (those starting with "projects/")
absolutePathRegex = regexp.MustCompile(`^projects/[^/]+/databases/[^/]+/documents/`)
)
// PathType represents the type of Firestore path
type PathType int
const (
CollectionPath PathType = iota
DocumentPath
)
// ValidateCollectionPath validates that a path is a valid Firestore collection path.
// Collection paths must have an odd number of segments (collection/doc/collection)
func ValidateCollectionPath(path string) error {
return validatePath(path, CollectionPath)
}
// ValidateDocumentPath validates that a path is a valid Firestore document path.
// Document paths must have an even number of segments (collection/doc or collection/doc/collection/doc)
func ValidateDocumentPath(path string) error {
return validatePath(path, DocumentPath)
}
// validatePath is the common validation function for both collection and document paths
func validatePath(path string, pathType PathType) error {
pathTypeName := "document"
if pathType == CollectionPath {
pathTypeName = "collection"
}
// Check for empty path
if path == "" {
return fmt.Errorf("%s path cannot be empty", pathTypeName)
}
// Check if it's an absolute path
if absolutePathRegex.MatchString(path) {
example := "users/userId"
if pathType == CollectionPath {
example = "users"
}
return fmt.Errorf("path must be relative (e.g., '%s'), not absolute (matching pattern: ^projects/[^/]+/databases/[^/]+/documents/)", example)
}
// Split the path using strings.Split to preserve empty segments
segments := strings.Split(path, "/")
// Check for empty result
if len(segments) == 0 {
return fmt.Errorf("%s path cannot be empty or contain only slashes", pathTypeName)
}
// Check segment count based on path type
segmentCount := len(segments)
if pathType == CollectionPath && segmentCount%2 == 0 {
// Collection paths must have an odd number of segments
return fmt.Errorf("invalid collection path: must have an odd number of segments (e.g., 'collection' or 'collection/doc/subcollection'), got %d segments", segmentCount)
} else if pathType == DocumentPath && segmentCount%2 != 0 {
// Document paths must have an even number of segments
return fmt.Errorf("invalid document path: must have an even number of segments (e.g., 'collection/doc'), got %d segments", segmentCount)
}
// Validate each segment
for i, segment := range segments {
isCollectionSegment := (i % 2) == 0
if err := validateSegment(segment, isCollectionSegment); err != nil {
return fmt.Errorf("invalid segment at position %d (%s): %w", i+1, segment, err)
}
}
return nil
}
// validateSegment validates a single path segment
func validateSegment(segment string, isCollection bool) error {
segmentType := "document ID"
if isCollection {
segmentType = "collection ID"
}
// Check for empty segment
if segment == "" {
return fmt.Errorf("segment cannot be empty")
}
// Check for whitespace-only segment
if strings.TrimSpace(segment) == "" {
return fmt.Errorf("segment cannot be only whitespace")
}
// Check for single or double period
if segment == "." || segment == ".." {
return fmt.Errorf("segment cannot be '.' or '..'")
}
// Check for reserved prefix
if strings.HasPrefix(segment, "__") {
return fmt.Errorf("%s cannot start with '__' (reserved prefix)", segmentType)
}
return nil
}
// IsAbsolutePath checks if a path is an absolute Firestore path
func IsAbsolutePath(path string) bool {
return absolutePathRegex.MatchString(path)
}
// IsRelativePath checks if a path is a relative Firestore path
func IsRelativePath(path string) bool {
return path != "" && !IsAbsolutePath(path)
}
```
--------------------------------------------------------------------------------
/internal/tools/bigquery/bigquerysql/bigquerysql_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquerysql_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerysql"
)
func TestParseFromYamlBigQuery(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
example_tool:
kind: bigquery-sql
source: my-instance
description: some description
statement: |
SELECT * FROM SQL_STATEMENT;
parameters:
- name: country
type: string
description: some description
`,
want: server.ToolConfigs{
"example_tool": bigquerysql.Config{
Name: "example_tool",
Kind: "bigquery-sql",
Source: "my-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
AuthRequired: []string{},
Parameters: []tools.Parameter{
tools.NewStringParameter("country", "some description"),
},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestParseFromYamlWithTemplateBigQuery(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
example_tool:
kind: bigquery-sql
source: my-instance
description: some description
statement: |
SELECT * FROM SQL_STATEMENT;
parameters:
- name: country
type: string
description: some description
templateParameters:
- name: tableName
type: string
description: The table to select hotels from.
- name: fieldArray
type: array
description: The columns to return for the query.
items:
name: column
type: string
description: A column name that will be returned from the query.
`,
want: server.ToolConfigs{
"example_tool": bigquerysql.Config{
Name: "example_tool",
Kind: "bigquery-sql",
Source: "my-instance",
Description: "some description",
Statement: "SELECT * FROM SQL_STATEMENT;\n",
AuthRequired: []string{},
Parameters: []tools.Parameter{
tools.NewStringParameter("country", "some description"),
},
TemplateParameters: []tools.Parameter{
tools.NewStringParameter("tableName", "The table to select hotels from."),
tools.NewArrayParameter("fieldArray", "The columns to return for the query.", tools.NewStringParameter("column", "A column name that will be returned from the query.")),
},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
```
--------------------------------------------------------------------------------
/internal/sources/redis/redis.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package redis
import (
"context"
"fmt"
"time"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/redis/go-redis/v9"
"go.opentelemetry.io/otel/trace"
)
const SourceKind string = "redis"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
actual := Config{Name: name}
if err := decoder.DecodeContext(ctx, &actual); err != nil {
return nil, err
}
return actual, nil
}
type Config struct {
Name string `yaml:"name" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Address []string `yaml:"address" validate:"required"`
Username string `yaml:"username"`
Password string `yaml:"password"`
Database int `yaml:"database"`
UseGCPIAM bool `yaml:"useGCPIAM"`
ClusterEnabled bool `yaml:"clusterEnabled"`
}
func (r Config) SourceConfigKind() string {
return SourceKind
}
// RedisClient is an interface for `redis.Client` and `redis.ClusterClient
type RedisClient interface {
Do(context.Context, ...any) *redis.Cmd
}
var _ RedisClient = (*redis.Client)(nil)
var _ RedisClient = (*redis.ClusterClient)(nil)
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
client, err := initRedisClient(ctx, r)
if err != nil {
return nil, fmt.Errorf("error initializing Redis client: %s", err)
}
s := &Source{
Name: r.Name,
Kind: SourceKind,
Client: client,
}
return s, nil
}
func initRedisClient(ctx context.Context, r Config) (RedisClient, error) {
var authFn func(ctx context.Context) (username string, password string, err error)
if r.UseGCPIAM {
// Pass in an access token getter fn for IAM auth
authFn = func(ctx context.Context) (username string, password string, err error) {
token, err := sources.GetIAMAccessToken(ctx)
if err != nil {
return "", "", err
}
return "default", token, nil
}
}
var client RedisClient
var err error
if r.ClusterEnabled {
// Create a new Redis Cluster client
clusterClient := redis.NewClusterClient(&redis.ClusterOptions{
Addrs: r.Address,
// PoolSize applies per cluster node and not for the whole cluster.
PoolSize: 10,
ConnMaxIdleTime: 60 * time.Second,
MinIdleConns: 1,
CredentialsProviderContext: authFn,
Username: r.Username,
Password: r.Password,
})
err = clusterClient.ForEachShard(ctx, func(ctx context.Context, shard *redis.Client) error {
return shard.Ping(ctx).Err()
})
if err != nil {
return nil, fmt.Errorf("unable to connect to redis cluster: %s", err)
}
client = clusterClient
return client, nil
}
// Create a new Redis client
standaloneClient := redis.NewClient(&redis.Options{
Addr: r.Address[0],
PoolSize: 10,
ConnMaxIdleTime: 60 * time.Second,
MinIdleConns: 1,
DB: r.Database,
CredentialsProviderContext: authFn,
Username: r.Username,
Password: r.Password,
})
_, err = standaloneClient.Ping(ctx).Result()
if err != nil {
return nil, fmt.Errorf("unable to connect to redis: %s", err)
}
client = standaloneClient
return client, nil
}
var _ sources.Source = &Source{}
type Source struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
Client RedisClient
}
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) RedisClient() RedisClient {
return s.Client
}
```
--------------------------------------------------------------------------------
/internal/sources/cloudsqlmssql/cloud_sql_mssql.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package cloudsqlmssql
import (
"context"
"database/sql"
"fmt"
"net/url"
"slices"
"cloud.google.com/go/cloudsqlconn/sqlserver/mssql"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/util"
"go.opentelemetry.io/otel/trace"
)
const SourceKind string = "cloud-sql-mssql"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
actual := Config{Name: name, IPType: "public"} // Default IPType
if err := decoder.DecodeContext(ctx, &actual); err != nil {
return nil, err
}
return actual, nil
}
type Config struct {
// Cloud SQL MSSQL configs
Name string `yaml:"name" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Project string `yaml:"project" validate:"required"`
Region string `yaml:"region" validate:"required"`
Instance string `yaml:"instance" validate:"required"`
IPAddress string `yaml:"ipAddress" validate:"required"`
IPType sources.IPType `yaml:"ipType" validate:"required"`
User string `yaml:"user" validate:"required"`
Password string `yaml:"password" validate:"required"`
Database string `yaml:"database" validate:"required"`
}
func (r Config) SourceConfigKind() string {
// Returns Cloud SQL MSSQL source kind
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
// Initializes a Cloud SQL MSSQL source
db, err := initCloudSQLMssqlConnection(ctx, tracer, r.Name, r.Project, r.Region, r.Instance, r.IPAddress, r.IPType.String(), r.User, r.Password, r.Database)
if err != nil {
return nil, fmt.Errorf("unable to create db connection: %w", err)
}
// Verify db connection
err = db.PingContext(ctx)
if err != nil {
return nil, fmt.Errorf("unable to connect successfully: %w", err)
}
s := &Source{
Name: r.Name,
Kind: SourceKind,
Db: db,
}
return s, nil
}
var _ sources.Source = &Source{}
type Source struct {
// Cloud SQL MSSQL struct with connection pool
Name string `yaml:"name"`
Kind string `yaml:"kind"`
Db *sql.DB
}
func (s *Source) SourceKind() string {
// Returns Cloud SQL MSSQL source kind
return SourceKind
}
func (s *Source) MSSQLDB() *sql.DB {
// Returns a Cloud SQL MSSQL database connection pool
return s.Db
}
func initCloudSQLMssqlConnection(ctx context.Context, tracer trace.Tracer, name, project, region, instance, ipAddress, ipType, user, pass, dbname string) (*sql.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
userAgent, err := util.UserAgentFromContext(ctx)
if err != nil {
return nil, err
}
// Create dsn
query := url.Values{}
query.Add("app name", userAgent)
query.Add("database", dbname)
query.Add("cloudsql", fmt.Sprintf("%s:%s:%s", project, region, instance))
url := &url.URL{
Scheme: "sqlserver",
User: url.UserPassword(user, pass),
Host: ipAddress,
RawQuery: query.Encode(),
}
// Get dial options
opts, err := sources.GetCloudSQLOpts(ipType, userAgent, false)
if err != nil {
return nil, err
}
// Register sql server driver
if !slices.Contains(sql.Drivers(), "cloudsql-sqlserver-driver") {
_, err := mssql.RegisterDriver("cloudsql-sqlserver-driver", opts...)
if err != nil {
return nil, err
}
}
// Open database connection
db, err := sql.Open(
"cloudsql-sqlserver-driver",
url.String(),
)
if err != nil {
return nil, err
}
return db, nil
}
```
--------------------------------------------------------------------------------
/internal/tools/clickhouse/clickhouselistdatabases/clickhouselistdatabases.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package clickhouse
import (
"context"
"database/sql"
"fmt"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/tools"
)
type compatibleSource interface {
ClickHousePool() *sql.DB
}
var compatibleSources = []string{"clickhouse"}
const listDatabasesKind string = "clickhouse-list-databases"
func init() {
if !tools.Register(listDatabasesKind, newListDatabasesConfig) {
panic(fmt.Sprintf("tool kind %q already registered", listDatabasesKind))
}
}
func newListDatabasesConfig(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 Config struct {
Name string `yaml:"name" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Source string `yaml:"source" validate:"required"`
Description string `yaml:"description" validate:"required"`
AuthRequired []string `yaml:"authRequired"`
Parameters tools.Parameters `yaml:"parameters"`
}
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return listDatabasesKind
}
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)
}
s, ok := rawS.(compatibleSource)
if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", listDatabasesKind, compatibleSources)
}
allParameters, paramManifest, _ := tools.ProcessParameters(nil, cfg.Parameters)
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters)
t := Tool{
Name: cfg.Name,
Kind: listDatabasesKind,
Parameters: cfg.Parameters,
AllParams: allParameters,
AuthRequired: cfg.AuthRequired,
Pool: s.ClickHousePool(),
manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
mcpManifest: mcpManifest,
}
return t, nil
}
var _ tools.Tool = Tool{}
type Tool struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
AuthRequired []string `yaml:"authRequired"`
Parameters tools.Parameters `yaml:"parameters"`
AllParams tools.Parameters `yaml:"allParams"`
Pool *sql.DB
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, token tools.AccessToken) (any, error) {
// Query to list all databases
query := "SHOW DATABASES"
results, err := t.Pool.QueryContext(ctx, query)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)
}
defer results.Close()
var databases []map[string]any
for results.Next() {
var dbName string
err := results.Scan(&dbName)
if err != nil {
return nil, fmt.Errorf("unable to parse row: %w", err)
}
databases = append(databases, map[string]any{
"name": dbName,
})
}
if err := results.Err(); err != nil {
return nil, fmt.Errorf("errors encountered by results.Scan: %w", err)
}
return databases, nil
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
return tools.ParseParams(t.AllParams, data, claims)
}
func (t Tool) Manifest() tools.Manifest {
return t.manifest
}
func (t Tool) McpManifest() tools.McpManifest {
return t.mcpManifest
}
func (t Tool) Authorized(verifiedAuthServices []string) bool {
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
}
func (t Tool) RequiresClientAuthorization() bool {
return false
}
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestoreadddocuments/firestoreadddocuments_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package firestoreadddocuments_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoreadddocuments"
)
func TestParseFromYamlFirestoreAddDocuments(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
add_docs_tool:
kind: firestore-add-documents
source: my-firestore-instance
description: Add documents to Firestore collections
`,
want: server.ToolConfigs{
"add_docs_tool": firestoreadddocuments.Config{
Name: "add_docs_tool",
Kind: "firestore-add-documents",
Source: "my-firestore-instance",
Description: "Add documents to Firestore collections",
AuthRequired: []string{},
},
},
},
{
desc: "with auth requirements",
in: `
tools:
secure_add_docs:
kind: firestore-add-documents
source: prod-firestore
description: Add documents with authentication
authRequired:
- google-auth-service
- api-key-service
`,
want: server.ToolConfigs{
"secure_add_docs": firestoreadddocuments.Config{
Name: "secure_add_docs",
Kind: "firestore-add-documents",
Source: "prod-firestore",
Description: "Add documents with authentication",
AuthRequired: []string{"google-auth-service", "api-key-service"},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestParseFromYamlMultipleTools(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
in := `
tools:
add_user_docs:
kind: firestore-add-documents
source: users-firestore
description: Add user documents
authRequired:
- user-auth
add_product_docs:
kind: firestore-add-documents
source: products-firestore
description: Add product documents
add_order_docs:
kind: firestore-add-documents
source: orders-firestore
description: Add order documents
authRequired:
- user-auth
- admin-auth
`
want := server.ToolConfigs{
"add_user_docs": firestoreadddocuments.Config{
Name: "add_user_docs",
Kind: "firestore-add-documents",
Source: "users-firestore",
Description: "Add user documents",
AuthRequired: []string{"user-auth"},
},
"add_product_docs": firestoreadddocuments.Config{
Name: "add_product_docs",
Kind: "firestore-add-documents",
Source: "products-firestore",
Description: "Add product documents",
AuthRequired: []string{},
},
"add_order_docs": firestoreadddocuments.Config{
Name: "add_order_docs",
Kind: "firestore-add-documents",
Source: "orders-firestore",
Description: "Add order documents",
AuthRequired: []string{"user-auth", "admin-auth"},
},
}
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
}
```
--------------------------------------------------------------------------------
/internal/server/mcp/v20241105/types.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package v20241105
import (
"github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
"github.com/googleapis/genai-toolbox/internal/tools"
)
// SERVER_NAME is the server name used in Implementation.
const SERVER_NAME = "Toolbox"
// PROTOCOL_VERSION is the version of the MCP protocol in this package.
const PROTOCOL_VERSION = "2024-11-05"
// methods that are supported.
const (
PING = "ping"
TOOLS_LIST = "tools/list"
TOOLS_CALL = "tools/call"
)
/* Empty result */
// EmptyResult represents a response that indicates success but carries no data.
type EmptyResult jsonrpc.Result
/* Pagination */
// Cursor is an opaque token used to represent a cursor for pagination.
type Cursor string
type PaginatedRequest struct {
jsonrpc.Request
Params struct {
// An opaque token representing the current pagination position.
// If provided, the server should return results starting after this cursor.
Cursor Cursor `json:"cursor,omitempty"`
} `json:"params,omitempty"`
}
type PaginatedResult struct {
jsonrpc.Result
// An opaque token representing the pagination position after the last returned result.
// If present, there may be more results available.
NextCursor Cursor `json:"nextCursor,omitempty"`
}
/* Tools */
// Sent from the client to request a list of tools the server has.
type ListToolsRequest struct {
PaginatedRequest
}
// The server's response to a tools/list request from the client.
type ListToolsResult struct {
PaginatedResult
Tools []tools.McpManifest `json:"tools"`
}
// Used by the client to invoke a tool provided by the server.
type CallToolRequest struct {
jsonrpc.Request
Params struct {
Name string `json:"name"`
Arguments map[string]any `json:"arguments,omitempty"`
} `json:"params,omitempty"`
}
// The sender or recipient of messages and data in a conversation.
type Role string
const (
RoleUser Role = "user"
RoleAssistant Role = "assistant"
)
// Base for objects that include optional annotations for the client.
// The client can use annotations to inform how objects are used or displayed
type Annotated struct {
Annotations *struct {
// Describes who the intended customer of this object or data is.
// It can include multiple entries to indicate content useful for multiple
// audiences (e.g., `["user", "assistant"]`).
Audience []Role `json:"audience,omitempty"`
// Describes how important this data is for operating the server.
//
// A value of 1 means "most important," and indicates that the data is
// effectively required, while 0 means "least important," and indicates that
// the data is entirely optional.
//
// @TJS-type number
// @minimum 0
// @maximum 1
Priority float64 `json:"priority,omitempty"`
} `json:"annotations,omitempty"`
}
// TextContent represents text provided to or from an LLM.
type TextContent struct {
Annotated
Type string `json:"type"`
// The text content of the message.
Text string `json:"text"`
}
// The server's response to a tool call.
//
// Any errors that originate from the tool SHOULD be reported inside the result
// object, with `isError` set to true, _not_ as an MCP protocol-level error
// response. Otherwise, the LLM would not be able to see that an error occurred
// and self-correct.
//
// However, any errors in _finding_ the tool, an error indicating that the
// server does not support tool calls, or any other exceptional conditions,
// should be reported as an MCP error response.
type CallToolResult struct {
jsonrpc.Result
// Could be either a TextContent, ImageContent, or EmbeddedResources
// For Toolbox, we will only be sending TextContent
Content []TextContent `json:"content"`
// Whether the tool call ended in an error.
// If not set, this is assumed to be false (the call was successful).
IsError bool `json:"isError,omitempty"`
}
```
--------------------------------------------------------------------------------
/internal/log/handler.go:
--------------------------------------------------------------------------------
```go
// Copyright 2024 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 log
import (
"context"
"fmt"
"io"
"log/slog"
"sync"
"time"
"go.opentelemetry.io/otel/trace"
)
// ValueTextHandler is a [Handler] that writes Records to an [io.Writer] with values separated by spaces.
type ValueTextHandler struct {
h slog.Handler
mu *sync.Mutex
out io.Writer
}
// NewValueTextHandler creates a [ValueTextHandler] that writes to out, using the given options.
func NewValueTextHandler(out io.Writer, opts *slog.HandlerOptions) *ValueTextHandler {
if opts == nil {
opts = &slog.HandlerOptions{}
}
return &ValueTextHandler{
out: out,
h: slog.NewTextHandler(out, &slog.HandlerOptions{
Level: opts.Level,
AddSource: opts.AddSource,
ReplaceAttr: nil,
}),
mu: &sync.Mutex{},
}
}
func (h *ValueTextHandler) Enabled(ctx context.Context, level slog.Level) bool {
return h.h.Enabled(ctx, level)
}
func (h *ValueTextHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
return &ValueTextHandler{h: h.h.WithAttrs(attrs), out: h.out, mu: h.mu}
}
func (h *ValueTextHandler) WithGroup(name string) slog.Handler {
return &ValueTextHandler{h: h.h.WithGroup(name), out: h.out, mu: h.mu}
}
// Handle formats its argument [Record] as a single line of space-separated values.
// Example output format: 2024-11-12T15:08:11.451377-08:00 INFO "Initialized 0 sources.\n"
func (h *ValueTextHandler) Handle(ctx context.Context, r slog.Record) error {
buf := make([]byte, 0, 1024)
// time
if !r.Time.IsZero() {
buf = h.appendAttr(buf, slog.Time(slog.TimeKey, r.Time))
}
// level
buf = h.appendAttr(buf, slog.Any(slog.LevelKey, r.Level))
// message
buf = h.appendAttr(buf, slog.String(slog.MessageKey, r.Message))
r.Attrs(func(a slog.Attr) bool {
buf = h.appendAttr(buf, a)
return true
})
buf = append(buf, "\n"...)
h.mu.Lock()
defer h.mu.Unlock()
_, err := h.out.Write(buf)
return err
}
// appendAttr is responsible for formatting a single attribute
func (h *ValueTextHandler) appendAttr(buf []byte, a slog.Attr) []byte {
// Resolve the Attr's value before doing anything else.
a.Value = a.Value.Resolve()
// Ignore empty Attrs.
if a.Equal(slog.Attr{}) {
return buf
}
switch a.Value.Kind() {
case slog.KindString:
// Quote string values, to make them easy to parse.
buf = fmt.Appendf(buf, "%q ", a.Value.String())
case slog.KindTime:
// Write times in a standard way, without the monotonic time.
buf = fmt.Appendf(buf, "%s ", a.Value.Time().Format(time.RFC3339Nano))
case slog.KindGroup:
attrs := a.Value.Group()
// Ignore empty groups.
if len(attrs) == 0 {
return buf
}
for _, ga := range attrs {
buf = h.appendAttr(buf, ga)
}
default:
buf = fmt.Appendf(buf, "%s ", a.Value)
}
return buf
}
// spanContextLogHandler is an slog.Handler which adds attributes from the span
// context.
type spanContextLogHandler struct {
slog.Handler
}
// handlerWithSpanContext adds attributes from the span context.
func handlerWithSpanContext(handler slog.Handler) *spanContextLogHandler {
return &spanContextLogHandler{Handler: handler}
}
// Handle overrides slog.Handler's Handle method. This adds attributes from the
// span context to the slog.Record.
func (t *spanContextLogHandler) Handle(ctx context.Context, record slog.Record) error {
// Get the SpanContext from the golang Context.
if s := trace.SpanContextFromContext(ctx); s.IsValid() {
// Add trace context attributes following Cloud Logging structured log format described
// in https://cloud.google.com/logging/docs/structured-logging#special-payload-fields
record.AddAttrs(
slog.Any("logging.googleapis.com/trace", s.TraceID()),
)
record.AddAttrs(
slog.Any("logging.googleapis.com/spanId", s.SpanID()),
)
record.AddAttrs(
slog.Bool("logging.googleapis.com/trace_sampled", s.TraceFlags().IsSampled()),
)
}
return t.Handler.Handle(ctx, record)
}
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestoregetdocuments/firestoregetdocuments_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package firestoregetdocuments_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoregetdocuments"
)
func TestParseFromYamlFirestoreGetDocuments(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
get_docs_tool:
kind: firestore-get-documents
source: my-firestore-instance
description: Retrieve documents from Firestore by paths
`,
want: server.ToolConfigs{
"get_docs_tool": firestoregetdocuments.Config{
Name: "get_docs_tool",
Kind: "firestore-get-documents",
Source: "my-firestore-instance",
Description: "Retrieve documents from Firestore by paths",
AuthRequired: []string{},
},
},
},
{
desc: "with auth requirements",
in: `
tools:
secure_get_docs:
kind: firestore-get-documents
source: prod-firestore
description: Get documents with authentication
authRequired:
- google-auth-service
- api-key-service
`,
want: server.ToolConfigs{
"secure_get_docs": firestoregetdocuments.Config{
Name: "secure_get_docs",
Kind: "firestore-get-documents",
Source: "prod-firestore",
Description: "Get documents with authentication",
AuthRequired: []string{"google-auth-service", "api-key-service"},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestParseFromYamlMultipleTools(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
in := `
tools:
get_user_docs:
kind: firestore-get-documents
source: users-firestore
description: Get user documents
authRequired:
- user-auth
get_product_docs:
kind: firestore-get-documents
source: products-firestore
description: Get product documents
get_order_docs:
kind: firestore-get-documents
source: orders-firestore
description: Get order documents
authRequired:
- user-auth
- admin-auth
`
want := server.ToolConfigs{
"get_user_docs": firestoregetdocuments.Config{
Name: "get_user_docs",
Kind: "firestore-get-documents",
Source: "users-firestore",
Description: "Get user documents",
AuthRequired: []string{"user-auth"},
},
"get_product_docs": firestoregetdocuments.Config{
Name: "get_product_docs",
Kind: "firestore-get-documents",
Source: "products-firestore",
Description: "Get product documents",
AuthRequired: []string{},
},
"get_order_docs": firestoregetdocuments.Config{
Name: "get_order_docs",
Kind: "firestore-get-documents",
Source: "orders-firestore",
Description: "Get order documents",
AuthRequired: []string{"user-auth", "admin-auth"},
},
}
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
}
```
--------------------------------------------------------------------------------
/internal/tools/bigquery/bigquerycommon/util.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package bigquerycommon
import (
"context"
"fmt"
"sort"
"strings"
bigqueryapi "cloud.google.com/go/bigquery"
"github.com/googleapis/genai-toolbox/internal/tools"
bigqueryrestapi "google.golang.org/api/bigquery/v2"
)
// DryRunQuery performs a dry run of the SQL query to validate it and get metadata.
func DryRunQuery(ctx context.Context, restService *bigqueryrestapi.Service, projectID string, location string, sql string, params []*bigqueryrestapi.QueryParameter, connProps []*bigqueryapi.ConnectionProperty) (*bigqueryrestapi.Job, error) {
useLegacySql := false
restConnProps := make([]*bigqueryrestapi.ConnectionProperty, len(connProps))
for i, prop := range connProps {
restConnProps[i] = &bigqueryrestapi.ConnectionProperty{Key: prop.Key, Value: prop.Value}
}
jobToInsert := &bigqueryrestapi.Job{
JobReference: &bigqueryrestapi.JobReference{
ProjectId: projectID,
Location: location,
},
Configuration: &bigqueryrestapi.JobConfiguration{
DryRun: true,
Query: &bigqueryrestapi.JobConfigurationQuery{
Query: sql,
UseLegacySql: &useLegacySql,
ConnectionProperties: restConnProps,
QueryParameters: params,
},
},
}
insertResponse, err := restService.Jobs.Insert(projectID, jobToInsert).Context(ctx).Do()
if err != nil {
return nil, fmt.Errorf("failed to insert dry run job: %w", err)
}
return insertResponse, nil
}
// BQTypeStringFromToolType converts a tool parameter type string to a BigQuery standard SQL type string.
func BQTypeStringFromToolType(toolType string) (string, error) {
switch toolType {
case "string":
return "STRING", nil
case "integer":
return "INT64", nil
case "float":
return "FLOAT64", nil
case "boolean":
return "BOOL", nil
default:
return "", fmt.Errorf("unsupported tool parameter type for BigQuery: %s", toolType)
}
}
// InitializeDatasetParameters generates project and dataset tool parameters based on allowedDatasets.
func InitializeDatasetParameters(
allowedDatasets []string,
defaultProjectID string,
projectKey, datasetKey string,
projectDescription, datasetDescription string,
) (projectParam, datasetParam tools.Parameter) {
if len(allowedDatasets) > 0 {
if len(allowedDatasets) == 1 {
parts := strings.Split(allowedDatasets[0], ".")
defaultProjectID = parts[0]
datasetID := parts[1]
projectDescription += fmt.Sprintf(" Must be `%s`.", defaultProjectID)
datasetDescription += fmt.Sprintf(" Must be `%s`.", datasetID)
datasetParam = tools.NewStringParameterWithDefault(datasetKey, datasetID, datasetDescription)
} else {
datasetIDsByProject := make(map[string][]string)
for _, ds := range allowedDatasets {
parts := strings.Split(ds, ".")
project := parts[0]
dataset := parts[1]
datasetIDsByProject[project] = append(datasetIDsByProject[project], fmt.Sprintf("`%s`", dataset))
}
var datasetDescriptions, projectIDList []string
for project, datasets := range datasetIDsByProject {
sort.Strings(datasets)
projectIDList = append(projectIDList, fmt.Sprintf("`%s`", project))
datasetList := strings.Join(datasets, ", ")
datasetDescriptions = append(datasetDescriptions, fmt.Sprintf("%s from project `%s`", datasetList, project))
}
sort.Strings(projectIDList)
sort.Strings(datasetDescriptions)
projectDescription += fmt.Sprintf(" Must be one of the following: %s.", strings.Join(projectIDList, ", "))
datasetDescription += fmt.Sprintf(" Must be one of the allowed datasets: %s.", strings.Join(datasetDescriptions, "; "))
datasetParam = tools.NewStringParameter(datasetKey, datasetDescription)
}
} else {
datasetParam = tools.NewStringParameter(datasetKey, datasetDescription)
}
projectParam = tools.NewStringParameterWithDefault(projectKey, defaultProjectID, projectDescription)
return projectParam, datasetParam
}
```
--------------------------------------------------------------------------------
/docs/en/how-to/toolbox-ui/index.md:
--------------------------------------------------------------------------------
```markdown
---
title: "Toolbox UI"
type: docs
weight: 1
description: >
How to effectively use Toolbox UI.
---
Toolbox UI is a built-in web interface that allows users to visually inspect and
test out configured resources such as tools and toolsets.
## Launching Toolbox UI
To launch Toolbox's interactive UI, use the `--ui` flag.
```sh
./toolbox --ui
```
Toolbox UI will be served from the same host and port as the Toolbox Server,
with the `/ui` suffix. Once Toolbox is launched, the following INFO log with
Toolbox UI's url will be shown:
```bash
INFO "Toolbox UI is up and running at: http://localhost:5000/ui"
```
## Navigating the Tools Page
The tools page shows all tools loaded from your configuration file. This
corresponds to the default toolset (represented by an empty string). Each tool's
name on this page will exactly match its name in the configuration file.
To view details for a specific tool, click on the tool name. The main content
area will be populated with the tool name, description, and available
parameters.
### Invoking a Tool
1. Click on a Tool
1. Enter appropriate parameters in each parameter field
1. Click "Run Tool"
1. Done! Your results will appear in the response field
1. (Optional) Uncheck "Prettify JSON" to format the response as plain text
### Optional Parameters
Toolbox allows users to add [optional
parameters](../../resources/tools/#basic-parameters) with or without a default
value.
To exclude a parameter, uncheck the box to the right of an associated parameter,
and that parameter will not be included in the request body. If the parameter is
not sent, Toolbox will either use it as `nil` value or the `default` value, if
configured. If the parameter is required, Toolbox will throw an error.
When the box is checked, parameter will be sent exactly as entered in the
response field (e.g. empty string).
### Editing Headers
To edit headers, press the "Edit Headers" button to display the header modal.
Within this modal, users can make direct edits by typing into the header's text
area.
Toolbox UI validates that the headers are in correct JSON format. Other
header-related errors (e.g., incorrect header names or values required by the
tool) will be reported in the Response section after running the tool.
#### Google OAuth
Currently, Toolbox supports Google OAuth 2.0 as an AuthService, which allows
tools to utilize authorized parameters. When a tool uses an authorized
parameter, the parameter will be displayed but not editable, as it will be
populated from the authentication token.
To provide the token, add your Google OAuth ID Token to the request header using
the "Edit Headers" button and modal described above. The key should be the name
of your AuthService as defined in your tool configuration file, suffixed with
`_token`. The value should be your ID token as a string.
1. Select a tool that requires [authenticated parameters]()
1. The auth parameter's text field is greyed out. This is because it cannot be
entered manually and will be parsed from the resolved auth token
1. To update request headers with the token, select "Edit Headers"
1. (Optional) If you wish to manually edit the header, checkout the dropdown
"How to extract Google OAuth ID Token manually" for guidance on retrieving ID
token
1. To edit the header automatically, click the "Auto Setup" button that is
associated with your Auth Profile
1. Enter the Client ID defined in your tools configuration file
1. Click "Continue"
1. Click "Sign in With Google" and login with your associated google account.
This should automatically populate the header text area with your token
1. Click "Save"
1. Click "Run Tool"
```json
{
"Content-Type": "application/json",
"my-google-auth_token": "YOUR_ID_TOKEN_HERE"
}
```
## Navigating the Toolsets Page
Through the toolsets page, users can search for a specific toolset to retrieve
tools from. Simply enter the toolset name in the search bar, and press "Enter"
to retrieve the associated tools.
If the toolset name is not defined within the tools configuration file, an error
message will be displayed.
```
--------------------------------------------------------------------------------
/docs/en/resources/tools/dgraph/dgraph-dql.md:
--------------------------------------------------------------------------------
```markdown
---
title: "dgraph-dql"
type: docs
weight: 1
description: >
A "dgraph-dql" tool executes a pre-defined DQL statement against a Dgraph
database.
aliases:
- /resources/tools/dgraph-dql
---
## About
A `dgraph-dql` tool executes a pre-defined DQL statement against a Dgraph
database. It's compatible with any of the following sources:
- [dgraph](../../sources/dgraph.md)
To run a statement as a query, you need to set the config `isQuery=true`. For
upserts or mutations, set `isQuery=false`. You can also configure timeout for a
query.
> **Note:** This tool uses parameterized queries to prevent SQL injections.
> Query parameters can be used as substitutes for arbitrary expressions.
> Parameters cannot be used as substitutes for identifiers, column names, table
> names, or other parts of the query.
## Example
{{< tabpane persist="header" >}}
{{< tab header="Query" lang="yaml" >}}
tools:
search_user:
kind: dgraph-dql
source: my-dgraph-source
statement: |
query all($role: string){
users(func: has(name)) @filter(eq(role, $role) AND ge(age, 30) AND le(age, 50)) {
uid
name
email
role
age
}
}
isQuery: true
timeout: 20s
description: |
Use this tool to retrieve the details of users who are admins and are between 30 and 50 years old.
The query returns the user's name, email, role, and age.
This can be helpful when you want to fetch admin users within a specific age range.
Example: Fetch admins aged between 30 and 50:
[
{
"name": "Alice",
"role": "admin",
"age": 35
},
{
"name": "Bob",
"role": "admin",
"age": 45
}
]
parameters:
- name: $role
type: string
description: admin
{{< /tab >}}
{{< tab header="Mutation" lang="yaml" >}}
tools:
dgraph-manage-user-instance:
kind: dgraph-dql
source: my-dgraph-source
isQuery: false
statement: |
{
set {
_:user1 <name> $user1 .
_:user1 <email> $email1 .
_:user1 <role> "admin" .
_:user1 <age> "35" .
_:user2 <name> $user2 .
_:user2 <email> $email2 .
_:user2 <role> "admin" .
_:user2 <age> "45" .
}
}
description: |
Use this tool to insert or update user data into the Dgraph database.
The mutation adds or updates user details like name, email, role, and age.
Example: Add users Alice and Bob as admins with specific ages.
parameters:
- name: user1
type: string
description: Alice
- name: email1
type: string
description: [email protected]
- name: user2
type: string
description: Bob
- name: email2
type: string
description: [email protected]
{{< /tab >}}
{{< /tabpane >}}
## Reference
| **field** | **type** | **required** | **description** |
|-------------|:---------------------------------------:|:------------:|-------------------------------------------------------------------------------------------|
| kind | string | true | Must be "dgraph-dql". |
| source | string | true | Name of the source the dql query should execute on. |
| description | string | true | Description of the tool that is passed to the LLM. |
| statement | string | true | dql statement to execute |
| isQuery | boolean | false | To run statement as query set true otherwise false |
| timeout | string | false | To set timeout for query |
| parameters | [parameters](../#specifying-parameters) | false | List of [parameters](../#specifying-parameters) that will be used with the dql statement. |
```
--------------------------------------------------------------------------------
/internal/sources/trino/trino.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package trino
import (
"context"
"database/sql"
"fmt"
"net/url"
"time"
"github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
_ "github.com/trinodb/trino-go-client/trino"
"go.opentelemetry.io/otel/trace"
)
const SourceKind string = "trino"
// validate interface
var _ sources.SourceConfig = Config{}
func init() {
if !sources.Register(SourceKind, newConfig) {
panic(fmt.Sprintf("source kind %q already registered", SourceKind))
}
}
func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
actual := Config{Name: name}
if err := decoder.DecodeContext(ctx, &actual); err != nil {
return nil, err
}
return actual, nil
}
type Config struct {
Name string `yaml:"name" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Host string `yaml:"host" validate:"required"`
Port string `yaml:"port" validate:"required"`
User string `yaml:"user"`
Password string `yaml:"password"`
Catalog string `yaml:"catalog" validate:"required"`
Schema string `yaml:"schema" validate:"required"`
QueryTimeout string `yaml:"queryTimeout"`
AccessToken string `yaml:"accessToken"`
KerberosEnabled bool `yaml:"kerberosEnabled"`
SSLEnabled bool `yaml:"sslEnabled"`
}
func (r Config) SourceConfigKind() string {
return SourceKind
}
func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
pool, err := initTrinoConnectionPool(ctx, tracer, r.Name, r.Host, r.Port, r.User, r.Password, r.Catalog, r.Schema, r.QueryTimeout, r.AccessToken, r.KerberosEnabled, r.SSLEnabled)
if err != nil {
return nil, fmt.Errorf("unable to create pool: %w", err)
}
err = pool.PingContext(ctx)
if err != nil {
return nil, fmt.Errorf("unable to connect successfully: %w", err)
}
s := &Source{
Name: r.Name,
Kind: SourceKind,
Pool: pool,
}
return s, nil
}
var _ sources.Source = &Source{}
type Source struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
Pool *sql.DB
}
func (s *Source) SourceKind() string {
return SourceKind
}
func (s *Source) TrinoDB() *sql.DB {
return s.Pool
}
func initTrinoConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, password, catalog, schema, queryTimeout, accessToken string, kerberosEnabled, sslEnabled bool) (*sql.DB, error) {
//nolint:all // Reassigned ctx
ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
defer span.End()
// Build Trino DSN
dsn, err := buildTrinoDSN(host, port, user, password, catalog, schema, queryTimeout, accessToken, kerberosEnabled, sslEnabled)
if err != nil {
return nil, fmt.Errorf("failed to build DSN: %w", err)
}
db, err := sql.Open("trino", dsn)
if err != nil {
return nil, fmt.Errorf("failed to open connection: %w", err)
}
// Configure connection pool
db.SetMaxOpenConns(10)
db.SetMaxIdleConns(5)
db.SetConnMaxLifetime(time.Hour)
return db, nil
}
func buildTrinoDSN(host, port, user, password, catalog, schema, queryTimeout, accessToken string, kerberosEnabled, sslEnabled bool) (string, error) {
// Build query parameters
query := url.Values{}
query.Set("catalog", catalog)
query.Set("schema", schema)
if queryTimeout != "" {
query.Set("queryTimeout", queryTimeout)
}
if accessToken != "" {
query.Set("accessToken", accessToken)
}
if kerberosEnabled {
query.Set("KerberosEnabled", "true")
}
// Build URL
scheme := "http"
if sslEnabled {
scheme = "https"
}
u := &url.URL{
Scheme: scheme,
Host: fmt.Sprintf("%s:%s", host, port),
RawQuery: query.Encode(),
}
// Only set user and password if not empty
if user != "" && password != "" {
u.User = url.UserPassword(user, password)
} else if user != "" {
u.User = url.User(user)
}
return u.String(), nil
}
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestoregetrules/firestoregetrules_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package firestoregetrules_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoregetrules"
)
func TestParseFromYamlFirestoreGetRules(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
get_rules_tool:
kind: firestore-get-rules
source: my-firestore-instance
description: Retrieves the active Firestore security rules for the current project
`,
want: server.ToolConfigs{
"get_rules_tool": firestoregetrules.Config{
Name: "get_rules_tool",
Kind: "firestore-get-rules",
Source: "my-firestore-instance",
Description: "Retrieves the active Firestore security rules for the current project",
AuthRequired: []string{},
},
},
},
{
desc: "with auth requirements",
in: `
tools:
secure_get_rules:
kind: firestore-get-rules
source: prod-firestore
description: Get Firestore security rules with authentication
authRequired:
- google-auth-service
- admin-service
`,
want: server.ToolConfigs{
"secure_get_rules": firestoregetrules.Config{
Name: "secure_get_rules",
Kind: "firestore-get-rules",
Source: "prod-firestore",
Description: "Get Firestore security rules with authentication",
AuthRequired: []string{"google-auth-service", "admin-service"},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestParseFromYamlMultipleTools(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
in := `
tools:
get_dev_rules:
kind: firestore-get-rules
source: dev-firestore
description: Get development Firestore rules
authRequired:
- dev-auth
get_staging_rules:
kind: firestore-get-rules
source: staging-firestore
description: Get staging Firestore rules
get_prod_rules:
kind: firestore-get-rules
source: prod-firestore
description: Get production Firestore rules
authRequired:
- prod-auth
- admin-auth
`
want := server.ToolConfigs{
"get_dev_rules": firestoregetrules.Config{
Name: "get_dev_rules",
Kind: "firestore-get-rules",
Source: "dev-firestore",
Description: "Get development Firestore rules",
AuthRequired: []string{"dev-auth"},
},
"get_staging_rules": firestoregetrules.Config{
Name: "get_staging_rules",
Kind: "firestore-get-rules",
Source: "staging-firestore",
Description: "Get staging Firestore rules",
AuthRequired: []string{},
},
"get_prod_rules": firestoregetrules.Config{
Name: "get_prod_rules",
Kind: "firestore-get-rules",
Source: "prod-firestore",
Description: "Get production Firestore rules",
AuthRequired: []string{"prod-auth", "admin-auth"},
},
}
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
}
```
--------------------------------------------------------------------------------
/internal/sources/spanner/spanner_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2024 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 spanner_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/spanner"
"github.com/googleapis/genai-toolbox/internal/testutils"
)
func TestParseFromYamlSpannerDb(t *testing.T) {
tcs := []struct {
desc string
in string
want server.SourceConfigs
}{
{
desc: "basic example",
in: `
sources:
my-spanner-instance:
kind: spanner
project: my-project
instance: my-instance
database: my_db
`,
want: map[string]sources.SourceConfig{
"my-spanner-instance": spanner.Config{
Name: "my-spanner-instance",
Kind: spanner.SourceKind,
Project: "my-project",
Instance: "my-instance",
Dialect: "googlesql",
Database: "my_db",
},
},
},
{
desc: "gsql dialect",
in: `
sources:
my-spanner-instance:
kind: spanner
project: my-project
instance: my-instance
dialect: Googlesql
database: my_db
`,
want: map[string]sources.SourceConfig{
"my-spanner-instance": spanner.Config{
Name: "my-spanner-instance",
Kind: spanner.SourceKind,
Project: "my-project",
Instance: "my-instance",
Dialect: "googlesql",
Database: "my_db",
},
},
},
{
desc: "postgresql dialect",
in: `
sources:
my-spanner-instance:
kind: spanner
project: my-project
instance: my-instance
dialect: postgresql
database: my_db
`,
want: map[string]sources.SourceConfig{
"my-spanner-instance": spanner.Config{
Name: "my-spanner-instance",
Kind: spanner.SourceKind,
Project: "my-project",
Instance: "my-instance",
Dialect: "postgresql",
Database: "my_db",
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Sources server.SourceConfigs `yaml:"sources"`
}{}
// Parse contents
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if !cmp.Equal(tc.want, got.Sources) {
t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources)
}
})
}
}
func TestFailParseFromYaml(t *testing.T) {
tcs := []struct {
desc string
in string
err string
}{
{
desc: "invalid dialect",
in: `
sources:
my-spanner-instance:
kind: spanner
project: my-project
instance: my-instance
dialect: fail
database: my_db
`,
err: "unable to parse source \"my-spanner-instance\" as \"spanner\": dialect invalid: must be one of \"googlesql\", or \"postgresql\"",
},
{
desc: "extra field",
in: `
sources:
my-spanner-instance:
kind: spanner
project: my-project
instance: my-instance
database: my_db
foo: bar
`,
err: "unable to parse source \"my-spanner-instance\" as \"spanner\": [2:1] unknown field \"foo\"\n 1 | database: my_db\n> 2 | foo: bar\n ^\n 3 | instance: my-instance\n 4 | kind: spanner\n 5 | project: my-project",
},
{
desc: "missing required field",
in: `
sources:
my-spanner-instance:
kind: spanner
project: my-project
instance: my-instance
`,
err: "unable to parse source \"my-spanner-instance\" as \"spanner\": Key: 'Config.Database' Error:Field validation for 'Database' failed on the 'required' tag",
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Sources server.SourceConfigs `yaml:"sources"`
}{}
// Parse contents
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
if err == nil {
t.Fatalf("expect parsing to fail")
}
errStr := err.Error()
if errStr != tc.err {
t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
}
})
}
}
```
--------------------------------------------------------------------------------
/internal/tools/mongodb/mongodbupdatemany/mongodbupdatemany_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package mongodbupdatemany_test
import (
"strings"
"testing"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbupdatemany"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
)
func TestParseFromYamlMongoQuery(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
example_tool:
kind: mongodb-update-many
source: my-instance
description: some description
database: test_db
collection: test_coll
filterPayload: |
{ name: {{json .name}} }
filterParams:
- name: name
type: string
description: small description
canonical: true
updatePayload: |
{ $set: { name: {{json .name}} } }
updateParams:
- name: name
type: string
description: small description
`,
want: server.ToolConfigs{
"example_tool": mongodbupdatemany.Config{
Name: "example_tool",
Kind: "mongodb-update-many",
Source: "my-instance",
AuthRequired: []string{},
Database: "test_db",
Collection: "test_coll",
FilterPayload: "{ name: {{json .name}} }\n",
FilterParams: tools.Parameters{
&tools.StringParameter{
CommonParameter: tools.CommonParameter{
Name: "name",
Type: "string",
Desc: "small description",
},
},
},
UpdatePayload: "{ $set: { name: {{json .name}} } }\n",
UpdateParams: tools.Parameters{
&tools.StringParameter{
CommonParameter: tools.CommonParameter{
Name: "name",
Type: "string",
Desc: "small description",
},
},
},
Description: "some description",
Canonical: true,
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestFailParseFromYamlMongoQuery(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
err string
}{
{
desc: "Invalid method",
in: `
tools:
example_tool:
kind: mongodb-update-many
source: my-instance
description: some description
collection: test_coll
filterPayload: |
{ name : {{json .name}} }
filterParams:
- name: name
type: string
description: small description
canonical: true
updatePayload: |
{ $set: { name: {{json .name}} } }
updateParams:
- name: data
type: string
description: the content in json
`,
err: `unable to parse tool "example_tool" as kind "mongodb-update-many"`,
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err == nil {
t.Fatalf("expect parsing to fail")
}
errStr := err.Error()
if !strings.Contains(errStr, tc.err) {
t.Fatalf("unexpected error string: got %q, want substring %q", errStr, tc.err)
}
})
}
}
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestoredeletedocuments/firestoredeletedocuments_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package firestoredeletedocuments_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoredeletedocuments"
)
func TestParseFromYamlFirestoreDeleteDocuments(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
delete_docs_tool:
kind: firestore-delete-documents
source: my-firestore-instance
description: Delete documents from Firestore by paths
`,
want: server.ToolConfigs{
"delete_docs_tool": firestoredeletedocuments.Config{
Name: "delete_docs_tool",
Kind: "firestore-delete-documents",
Source: "my-firestore-instance",
Description: "Delete documents from Firestore by paths",
AuthRequired: []string{},
},
},
},
{
desc: "with auth requirements",
in: `
tools:
secure_delete_docs:
kind: firestore-delete-documents
source: prod-firestore
description: Delete documents with authentication
authRequired:
- google-auth-service
- api-key-service
`,
want: server.ToolConfigs{
"secure_delete_docs": firestoredeletedocuments.Config{
Name: "secure_delete_docs",
Kind: "firestore-delete-documents",
Source: "prod-firestore",
Description: "Delete documents with authentication",
AuthRequired: []string{"google-auth-service", "api-key-service"},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestParseFromYamlMultipleTools(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
in := `
tools:
delete_user_docs:
kind: firestore-delete-documents
source: users-firestore
description: Delete user documents
authRequired:
- user-auth
delete_product_docs:
kind: firestore-delete-documents
source: products-firestore
description: Delete product documents
delete_order_docs:
kind: firestore-delete-documents
source: orders-firestore
description: Delete order documents
authRequired:
- user-auth
- admin-auth
`
want := server.ToolConfigs{
"delete_user_docs": firestoredeletedocuments.Config{
Name: "delete_user_docs",
Kind: "firestore-delete-documents",
Source: "users-firestore",
Description: "Delete user documents",
AuthRequired: []string{"user-auth"},
},
"delete_product_docs": firestoredeletedocuments.Config{
Name: "delete_product_docs",
Kind: "firestore-delete-documents",
Source: "products-firestore",
Description: "Delete product documents",
AuthRequired: []string{},
},
"delete_order_docs": firestoredeletedocuments.Config{
Name: "delete_order_docs",
Kind: "firestore-delete-documents",
Source: "orders-firestore",
Description: "Delete order documents",
AuthRequired: []string{"user-auth", "admin-auth"},
},
}
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
}
```
--------------------------------------------------------------------------------
/internal/sources/couchbase/couchbase_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2024 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 couchbase_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/sources/couchbase"
"github.com/googleapis/genai-toolbox/internal/testutils"
)
func TestParseFromYamlCouchbase(t *testing.T) {
tcs := []struct {
desc string
in string
want server.SourceConfigs
}{
{
desc: "basic example",
in: `
sources:
my-couchbase-instance:
kind: couchbase
connectionString: localhost
username: Administrator
password: password
bucket: travel-sample
scope: inventory
`,
want: server.SourceConfigs{
"my-couchbase-instance": couchbase.Config{
Name: "my-couchbase-instance",
Kind: couchbase.SourceKind,
ConnectionString: "localhost",
Username: "Administrator",
Password: "password",
Bucket: "travel-sample",
Scope: "inventory",
},
},
},
{
desc: "with TLS configuration",
in: `
sources:
my-couchbase-instance:
kind: couchbase
connectionString: couchbases://localhost
bucket: travel-sample
scope: inventory
clientCert: /path/to/cert.pem
clientKey: /path/to/key.pem
clientCertPassword: password
clientKeyPassword: password
caCert: /path/to/ca.pem
noSslVerify: false
queryScanConsistency: 2
`,
want: server.SourceConfigs{
"my-couchbase-instance": couchbase.Config{
Name: "my-couchbase-instance",
Kind: couchbase.SourceKind,
ConnectionString: "couchbases://localhost",
Bucket: "travel-sample",
Scope: "inventory",
ClientCert: "/path/to/cert.pem",
ClientKey: "/path/to/key.pem",
ClientCertPassword: "password",
ClientKeyPassword: "password",
CACert: "/path/to/ca.pem",
NoSSLVerify: false,
QueryScanConsistency: 2,
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Sources server.SourceConfigs `yaml:"sources"`
}{}
// Parse contents
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if !cmp.Equal(tc.want, got.Sources) {
t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources)
}
})
}
}
func TestFailParseFromYaml(t *testing.T) {
tcs := []struct {
desc string
in string
err string
}{
{
desc: "extra field",
in: `
sources:
my-couchbase-instance:
kind: couchbase
connectionString: localhost
username: Administrator
password: password
bucket: travel-sample
scope: inventory
foo: bar
`,
err: "unable to parse source \"my-couchbase-instance\" as \"couchbase\": [3:1] unknown field \"foo\"\n 1 | bucket: travel-sample\n 2 | connectionString: localhost\n> 3 | foo: bar\n ^\n 4 | kind: couchbase\n 5 | password: password\n 6 | scope: inventory\n 7 | ",
},
{
desc: "missing required field",
in: `
sources:
my-couchbase-instance:
kind: couchbase
username: Administrator
password: password
bucket: travel-sample
scope: inventory
`,
err: "unable to parse source \"my-couchbase-instance\" as \"couchbase\": Key: 'Config.ConnectionString' Error:Field validation for 'ConnectionString' failed on the 'required' tag",
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Sources server.SourceConfigs `yaml:"sources"`
}{}
// Parse contents
err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
if err == nil {
t.Fatalf("expect parsing to fail")
}
errStr := err.Error()
if errStr != tc.err {
t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
}
})
}
}
```
--------------------------------------------------------------------------------
/tests/dgraph/dgraph_integration_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package dgraph
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"os"
"reflect"
"regexp"
"testing"
"time"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/tests"
)
var (
DgraphSourceKind = "dgraph"
DgraphApiKey = "api-key"
DgraphUrl = os.Getenv("DGRAPH_URL")
)
func getDgraphVars(t *testing.T) map[string]any {
if DgraphUrl == "" {
t.Fatal("'DGRAPH_URL' not set")
}
return map[string]any{
"kind": DgraphSourceKind,
"dgraphUrl": DgraphUrl,
"apiKey": DgraphApiKey,
}
}
func TestDgraphToolEndpoints(t *testing.T) {
sourceConfig := getDgraphVars(t)
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
defer cancel()
var args []string
// Write config into a file and pass it to command
toolsFile := map[string]any{
"sources": map[string]any{
"my-dgraph-instance": sourceConfig,
},
"tools": map[string]any{
"my-simple-dql-tool": map[string]any{
"kind": "dgraph-dql",
"source": "my-dgraph-instance",
"description": "Simple tool to test end to end functionality.",
"statement": "{result(func: uid(0x0)) {constant: math(1)}}",
"isQuery": true,
"timeout": "20s",
"parameters": []any{},
},
},
}
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)
}
// Test tool get endpoint
tcs := []struct {
name string
api string
want map[string]any
}{
{
name: "get my-simple-tool",
api: "http://127.0.0.1:5000/api/tool/my-simple-dql-tool/",
want: map[string]any{
"my-simple-dql-tool": map[string]any{
"description": "Simple tool to test end to end functionality.",
"parameters": []any{},
"authRequired": []any{},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.name, func(t *testing.T) {
resp, err := http.Get(tc.api)
if err != nil {
t.Fatalf("error when sending a request: %s", err)
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
t.Fatalf("response status code is not 200")
}
var body map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&body)
if err != nil {
t.Fatalf("error parsing response body")
}
got, ok := body["tools"]
if !ok {
t.Fatalf("unable to find tools in response body")
}
if !reflect.DeepEqual(got, tc.want) {
t.Fatalf("got %q, want %q", got, tc.want)
}
})
}
// Test tool invoke endpoint
invokeTcs := []struct {
name string
api string
requestBody io.Reader
want string
}{
{
name: "invoke my-simple-dql-tool",
api: "http://127.0.0.1:5000/api/tool/my-simple-dql-tool/invoke",
requestBody: bytes.NewBuffer([]byte(`{}`)),
want: "{\"result\":[{\"constant\":1}]}",
},
}
for _, tc := range invokeTcs {
t.Run(tc.name, func(t *testing.T) {
resp, err := http.Post(tc.api, "application/json", tc.requestBody)
if err != nil {
t.Fatalf("error when sending a request: %s", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
bodyBytes, _ := io.ReadAll(resp.Body)
t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
}
var body map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&body)
if err != nil {
t.Fatalf("error parsing response body")
}
got, ok := body["result"].(string)
if !ok {
t.Fatalf("unable to find result in response body")
}
if got != tc.want {
t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
}
})
}
}
```
--------------------------------------------------------------------------------
/internal/tools/looker/lookerquerysql/lookerquerysql.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package lookerquerysql
import (
"context"
"fmt"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
lookersrc "github.com/googleapis/genai-toolbox/internal/sources/looker"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/googleapis/genai-toolbox/internal/tools/looker/lookercommon"
"github.com/googleapis/genai-toolbox/internal/util"
"github.com/looker-open-source/sdk-codegen/go/rtl"
v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4"
)
const kind string = "looker-query-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 Config struct {
Name string `yaml:"name" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Source string `yaml:"source" validate:"required"`
Description string `yaml:"description" validate:"required"`
AuthRequired []string `yaml:"authRequired"`
}
// validate interface
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
// verify source exists
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
// verify the source is compatible
s, ok := rawS.(*lookersrc.Source)
if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be `looker`", kind)
}
parameters := lookercommon.GetQueryParameters()
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters)
// finish tool setup
return Tool{
Name: cfg.Name,
Kind: kind,
Parameters: parameters,
AuthRequired: cfg.AuthRequired,
UseClientOAuth: s.UseClientOAuth,
Client: s.Client,
ApiSettings: s.ApiSettings,
manifest: tools.Manifest{
Description: cfg.Description,
Parameters: parameters.Manifest(),
AuthRequired: cfg.AuthRequired,
},
mcpManifest: mcpManifest,
}, nil
}
// validate interface
var _ tools.Tool = Tool{}
type Tool struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
UseClientOAuth bool
Client *v4.LookerSDK
ApiSettings *rtl.ApiSettings
AuthRequired []string `yaml:"authRequired"`
Parameters tools.Parameters `yaml:"parameters"`
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
logger, err := util.LoggerFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
}
wq, err := lookercommon.ProcessQueryArgs(ctx, params)
if err != nil {
return nil, fmt.Errorf("error building query request: %w", err)
}
sdk, err := lookercommon.GetLookerSDK(t.UseClientOAuth, t.ApiSettings, t.Client, accessToken)
if err != nil {
return nil, fmt.Errorf("error getting sdk: %w", err)
}
resp, err := lookercommon.RunInlineQuery(ctx, sdk, wq, "sql", t.ApiSettings)
if err != nil {
return nil, fmt.Errorf("error making query request: %s", err)
}
logger.DebugContext(ctx, "resp = ", resp)
return resp, nil
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
return tools.ParseParams(t.Parameters, data, claims)
}
func (t Tool) Manifest() tools.Manifest {
return t.manifest
}
func (t Tool) McpManifest() tools.McpManifest {
return t.mcpManifest
}
func (t Tool) Authorized(verifiedAuthServices []string) bool {
return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
}
func (t Tool) RequiresClientAuthorization() bool {
return t.UseClientOAuth
}
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestorequerycollection/firestorequerycollection_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package firestorequerycollection_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools/firestore/firestorequerycollection"
)
func TestParseFromYamlFirestoreQueryCollection(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
query_users_tool:
kind: firestore-query-collection
source: my-firestore-instance
description: Query users collection with filters and ordering
`,
want: server.ToolConfigs{
"query_users_tool": firestorequerycollection.Config{
Name: "query_users_tool",
Kind: "firestore-query-collection",
Source: "my-firestore-instance",
Description: "Query users collection with filters and ordering",
AuthRequired: []string{},
},
},
},
{
desc: "with auth requirements",
in: `
tools:
secure_query_tool:
kind: firestore-query-collection
source: prod-firestore
description: Query collections with authentication
authRequired:
- google-auth-service
- api-key-service
`,
want: server.ToolConfigs{
"secure_query_tool": firestorequerycollection.Config{
Name: "secure_query_tool",
Kind: "firestore-query-collection",
Source: "prod-firestore",
Description: "Query collections with authentication",
AuthRequired: []string{"google-auth-service", "api-key-service"},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestParseFromYamlMultipleTools(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
in := `
tools:
query_users:
kind: firestore-query-collection
source: users-firestore
description: Query user documents with filtering
authRequired:
- user-auth
query_products:
kind: firestore-query-collection
source: products-firestore
description: Query product catalog
query_orders:
kind: firestore-query-collection
source: orders-firestore
description: Query customer orders with complex filters
authRequired:
- user-auth
- admin-auth
`
want := server.ToolConfigs{
"query_users": firestorequerycollection.Config{
Name: "query_users",
Kind: "firestore-query-collection",
Source: "users-firestore",
Description: "Query user documents with filtering",
AuthRequired: []string{"user-auth"},
},
"query_products": firestorequerycollection.Config{
Name: "query_products",
Kind: "firestore-query-collection",
Source: "products-firestore",
Description: "Query product catalog",
AuthRequired: []string{},
},
"query_orders": firestorequerycollection.Config{
Name: "query_orders",
Kind: "firestore-query-collection",
Source: "orders-firestore",
Description: "Query customer orders with complex filters",
AuthRequired: []string{"user-auth", "admin-auth"},
},
}
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
}
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestorevalidaterules/firestorevalidaterules_test.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package firestorevalidaterules_test
import (
"testing"
yaml "github.com/goccy/go-yaml"
"github.com/google/go-cmp/cmp"
"github.com/googleapis/genai-toolbox/internal/server"
"github.com/googleapis/genai-toolbox/internal/testutils"
"github.com/googleapis/genai-toolbox/internal/tools/firestore/firestorevalidaterules"
)
func TestParseFromYamlFirestoreValidateRules(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
tcs := []struct {
desc string
in string
want server.ToolConfigs
}{
{
desc: "basic example",
in: `
tools:
validate_rules_tool:
kind: firestore-validate-rules
source: my-firestore-instance
description: Validate Firestore security rules
`,
want: server.ToolConfigs{
"validate_rules_tool": firestorevalidaterules.Config{
Name: "validate_rules_tool",
Kind: "firestore-validate-rules",
Source: "my-firestore-instance",
Description: "Validate Firestore security rules",
AuthRequired: []string{},
},
},
},
{
desc: "with auth requirements",
in: `
tools:
secure_validate_rules:
kind: firestore-validate-rules
source: prod-firestore
description: Validate rules with authentication
authRequired:
- google-auth-service
- api-key-service
`,
want: server.ToolConfigs{
"secure_validate_rules": firestorevalidaterules.Config{
Name: "secure_validate_rules",
Kind: "firestore-validate-rules",
Source: "prod-firestore",
Description: "Validate rules with authentication",
AuthRequired: []string{"google-auth-service", "api-key-service"},
},
},
},
}
for _, tc := range tcs {
t.Run(tc.desc, func(t *testing.T) {
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
})
}
}
func TestParseFromYamlMultipleTools(t *testing.T) {
ctx, err := testutils.ContextWithNewLogger()
if err != nil {
t.Fatalf("unexpected error: %s", err)
}
in := `
tools:
validate_dev_rules:
kind: firestore-validate-rules
source: dev-firestore
description: Validate development environment rules
authRequired:
- dev-auth
validate_staging_rules:
kind: firestore-validate-rules
source: staging-firestore
description: Validate staging environment rules
validate_prod_rules:
kind: firestore-validate-rules
source: prod-firestore
description: Validate production environment rules
authRequired:
- prod-auth
- admin-auth
`
want := server.ToolConfigs{
"validate_dev_rules": firestorevalidaterules.Config{
Name: "validate_dev_rules",
Kind: "firestore-validate-rules",
Source: "dev-firestore",
Description: "Validate development environment rules",
AuthRequired: []string{"dev-auth"},
},
"validate_staging_rules": firestorevalidaterules.Config{
Name: "validate_staging_rules",
Kind: "firestore-validate-rules",
Source: "staging-firestore",
Description: "Validate staging environment rules",
AuthRequired: []string{},
},
"validate_prod_rules": firestorevalidaterules.Config{
Name: "validate_prod_rules",
Kind: "firestore-validate-rules",
Source: "prod-firestore",
Description: "Validate production environment rules",
AuthRequired: []string{"prod-auth", "admin-auth"},
},
}
got := struct {
Tools server.ToolConfigs `yaml:"tools"`
}{}
// Parse contents
err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
if err != nil {
t.Fatalf("unable to unmarshal: %s", err)
}
if diff := cmp.Diff(want, got.Tools); diff != "" {
t.Fatalf("incorrect parse: diff %v", diff)
}
}
```
--------------------------------------------------------------------------------
/internal/tools/postgres/postgreslistavailableextensions/postgreslistavailableextensions.go:
--------------------------------------------------------------------------------
```go
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package postgreslistavailableextensions
import (
"context"
"fmt"
yaml "github.com/goccy/go-yaml"
"github.com/googleapis/genai-toolbox/internal/sources"
"github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
"github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
"github.com/googleapis/genai-toolbox/internal/sources/postgres"
"github.com/googleapis/genai-toolbox/internal/tools"
"github.com/jackc/pgx/v5/pgxpool"
)
const kind string = "postgres-list-available-extensions"
const listAvailableExtensionsQuery = `
SELECT
name,
default_version,
comment as description
FROM
pg_available_extensions
ORDER BY 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 {
PostgresPool() *pgxpool.Pool
}
// validate compatible sources are still compatible
var _ compatibleSource = &alloydbpg.Source{}
var _ compatibleSource = &cloudsqlpg.Source{}
var _ compatibleSource = &postgres.Source{}
var compatibleSources = [...]string{alloydbpg.SourceKind, cloudsqlpg.SourceKind, postgres.SourceKind}
type Config struct {
Name string `yaml:"name" validate:"required"`
Kind string `yaml:"kind" validate:"required"`
Source string `yaml:"source" validate:"required"`
Description string `yaml:"description" validate:"required"`
AuthRequired []string `yaml:"authRequired"`
}
// validate interface
var _ tools.ToolConfig = Config{}
func (cfg Config) ToolConfigKind() string {
return kind
}
func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
// verify source exists
rawS, ok := srcs[cfg.Source]
if !ok {
return nil, fmt.Errorf("no source named %q configured", cfg.Source)
}
// verify the source is compatible
s, ok := rawS.(compatibleSource)
if !ok {
return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
}
parameters := tools.Parameters{}
mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters)
// finish tool setup
t := Tool{
Name: cfg.Name,
Kind: cfg.Kind,
AuthRequired: cfg.AuthRequired,
Pool: s.PostgresPool(),
manifest: tools.Manifest{
Description: cfg.Description,
Parameters: parameters.Manifest(),
AuthRequired: cfg.AuthRequired,
},
mcpManifest: mcpManifest,
}
return t, nil
}
// validate interface
var _ tools.Tool = Tool{}
type Tool struct {
Name string `yaml:"name"`
Kind string `yaml:"kind"`
AuthRequired []string `yaml:"authRequired"`
Pool *pgxpool.Pool
manifest tools.Manifest
mcpManifest tools.McpManifest
}
func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
results, err := t.Pool.Query(ctx, listAvailableExtensionsQuery)
if err != nil {
return nil, fmt.Errorf("unable to execute query: %w", err)
}
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)
}
vMap := make(map[string]any)
for i, f := range fields {
vMap[f.Name] = v[i]
}
out = append(out, vMap)
}
return out, nil
}
func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
return tools.ParamValues{}, 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() bool {
return false
}
```