This is page 11 of 45. Use http://codebase.md/googleapis/genai-toolbox?lines=true&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/log/handler.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2024 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package log
16 |
17 | import (
18 | "context"
19 | "fmt"
20 | "io"
21 | "log/slog"
22 | "sync"
23 | "time"
24 |
25 | "go.opentelemetry.io/otel/trace"
26 | )
27 |
28 | // ValueTextHandler is a [Handler] that writes Records to an [io.Writer] with values separated by spaces.
29 | type ValueTextHandler struct {
30 | h slog.Handler
31 | mu *sync.Mutex
32 | out io.Writer
33 | }
34 |
35 | // NewValueTextHandler creates a [ValueTextHandler] that writes to out, using the given options.
36 | func NewValueTextHandler(out io.Writer, opts *slog.HandlerOptions) *ValueTextHandler {
37 | if opts == nil {
38 | opts = &slog.HandlerOptions{}
39 | }
40 | return &ValueTextHandler{
41 | out: out,
42 | h: slog.NewTextHandler(out, &slog.HandlerOptions{
43 | Level: opts.Level,
44 | AddSource: opts.AddSource,
45 | ReplaceAttr: nil,
46 | }),
47 | mu: &sync.Mutex{},
48 | }
49 | }
50 |
51 | func (h *ValueTextHandler) Enabled(ctx context.Context, level slog.Level) bool {
52 | return h.h.Enabled(ctx, level)
53 | }
54 |
55 | func (h *ValueTextHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
56 | return &ValueTextHandler{h: h.h.WithAttrs(attrs), out: h.out, mu: h.mu}
57 | }
58 |
59 | func (h *ValueTextHandler) WithGroup(name string) slog.Handler {
60 | return &ValueTextHandler{h: h.h.WithGroup(name), out: h.out, mu: h.mu}
61 | }
62 |
63 | // Handle formats its argument [Record] as a single line of space-separated values.
64 | // Example output format: 2024-11-12T15:08:11.451377-08:00 INFO "Initialized 0 sources.\n"
65 | func (h *ValueTextHandler) Handle(ctx context.Context, r slog.Record) error {
66 | buf := make([]byte, 0, 1024)
67 |
68 | // time
69 | if !r.Time.IsZero() {
70 | buf = h.appendAttr(buf, slog.Time(slog.TimeKey, r.Time))
71 | }
72 | // level
73 | buf = h.appendAttr(buf, slog.Any(slog.LevelKey, r.Level))
74 | // message
75 | buf = h.appendAttr(buf, slog.String(slog.MessageKey, r.Message))
76 |
77 | r.Attrs(func(a slog.Attr) bool {
78 | buf = h.appendAttr(buf, a)
79 | return true
80 | })
81 | buf = append(buf, "\n"...)
82 |
83 | h.mu.Lock()
84 | defer h.mu.Unlock()
85 | _, err := h.out.Write(buf)
86 | return err
87 | }
88 |
89 | // appendAttr is responsible for formatting a single attribute
90 | func (h *ValueTextHandler) appendAttr(buf []byte, a slog.Attr) []byte {
91 | // Resolve the Attr's value before doing anything else.
92 | a.Value = a.Value.Resolve()
93 | // Ignore empty Attrs.
94 | if a.Equal(slog.Attr{}) {
95 | return buf
96 | }
97 | switch a.Value.Kind() {
98 | case slog.KindString:
99 | // Quote string values, to make them easy to parse.
100 | buf = fmt.Appendf(buf, "%q ", a.Value.String())
101 | case slog.KindTime:
102 | // Write times in a standard way, without the monotonic time.
103 | buf = fmt.Appendf(buf, "%s ", a.Value.Time().Format(time.RFC3339Nano))
104 | case slog.KindGroup:
105 | attrs := a.Value.Group()
106 | // Ignore empty groups.
107 | if len(attrs) == 0 {
108 | return buf
109 | }
110 | for _, ga := range attrs {
111 | buf = h.appendAttr(buf, ga)
112 | }
113 | default:
114 | buf = fmt.Appendf(buf, "%s ", a.Value)
115 | }
116 |
117 | return buf
118 | }
119 |
120 | // spanContextLogHandler is an slog.Handler which adds attributes from the span
121 | // context.
122 | type spanContextLogHandler struct {
123 | slog.Handler
124 | }
125 |
126 | // handlerWithSpanContext adds attributes from the span context.
127 | func handlerWithSpanContext(handler slog.Handler) *spanContextLogHandler {
128 | return &spanContextLogHandler{Handler: handler}
129 | }
130 |
131 | // Handle overrides slog.Handler's Handle method. This adds attributes from the
132 | // span context to the slog.Record.
133 | func (t *spanContextLogHandler) Handle(ctx context.Context, record slog.Record) error {
134 | // Get the SpanContext from the golang Context.
135 | if s := trace.SpanContextFromContext(ctx); s.IsValid() {
136 | // Add trace context attributes following Cloud Logging structured log format described
137 | // in https://cloud.google.com/logging/docs/structured-logging#special-payload-fields
138 | record.AddAttrs(
139 | slog.Any("logging.googleapis.com/trace", s.TraceID()),
140 | )
141 | record.AddAttrs(
142 | slog.Any("logging.googleapis.com/spanId", s.SpanID()),
143 | )
144 | record.AddAttrs(
145 | slog.Bool("logging.googleapis.com/trace_sampled", s.TraceFlags().IsSampled()),
146 | )
147 | }
148 | return t.Handler.Handle(ctx, record)
149 | }
150 |
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestoregetdocuments/firestoregetdocuments_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package firestoregetdocuments_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/testutils"
24 | "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoregetdocuments"
25 | )
26 |
27 | func TestParseFromYamlFirestoreGetDocuments(t *testing.T) {
28 | ctx, err := testutils.ContextWithNewLogger()
29 | if err != nil {
30 | t.Fatalf("unexpected error: %s", err)
31 | }
32 | tcs := []struct {
33 | desc string
34 | in string
35 | want server.ToolConfigs
36 | }{
37 | {
38 | desc: "basic example",
39 | in: `
40 | tools:
41 | get_docs_tool:
42 | kind: firestore-get-documents
43 | source: my-firestore-instance
44 | description: Retrieve documents from Firestore by paths
45 | `,
46 | want: server.ToolConfigs{
47 | "get_docs_tool": firestoregetdocuments.Config{
48 | Name: "get_docs_tool",
49 | Kind: "firestore-get-documents",
50 | Source: "my-firestore-instance",
51 | Description: "Retrieve documents from Firestore by paths",
52 | AuthRequired: []string{},
53 | },
54 | },
55 | },
56 | {
57 | desc: "with auth requirements",
58 | in: `
59 | tools:
60 | secure_get_docs:
61 | kind: firestore-get-documents
62 | source: prod-firestore
63 | description: Get documents with authentication
64 | authRequired:
65 | - google-auth-service
66 | - api-key-service
67 | `,
68 | want: server.ToolConfigs{
69 | "secure_get_docs": firestoregetdocuments.Config{
70 | Name: "secure_get_docs",
71 | Kind: "firestore-get-documents",
72 | Source: "prod-firestore",
73 | Description: "Get documents with authentication",
74 | AuthRequired: []string{"google-auth-service", "api-key-service"},
75 | },
76 | },
77 | },
78 | }
79 | for _, tc := range tcs {
80 | t.Run(tc.desc, func(t *testing.T) {
81 | got := struct {
82 | Tools server.ToolConfigs `yaml:"tools"`
83 | }{}
84 | // Parse contents
85 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
86 | if err != nil {
87 | t.Fatalf("unable to unmarshal: %s", err)
88 | }
89 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
90 | t.Fatalf("incorrect parse: diff %v", diff)
91 | }
92 | })
93 | }
94 | }
95 |
96 | func TestParseFromYamlMultipleTools(t *testing.T) {
97 | ctx, err := testutils.ContextWithNewLogger()
98 | if err != nil {
99 | t.Fatalf("unexpected error: %s", err)
100 | }
101 | in := `
102 | tools:
103 | get_user_docs:
104 | kind: firestore-get-documents
105 | source: users-firestore
106 | description: Get user documents
107 | authRequired:
108 | - user-auth
109 | get_product_docs:
110 | kind: firestore-get-documents
111 | source: products-firestore
112 | description: Get product documents
113 | get_order_docs:
114 | kind: firestore-get-documents
115 | source: orders-firestore
116 | description: Get order documents
117 | authRequired:
118 | - user-auth
119 | - admin-auth
120 | `
121 | want := server.ToolConfigs{
122 | "get_user_docs": firestoregetdocuments.Config{
123 | Name: "get_user_docs",
124 | Kind: "firestore-get-documents",
125 | Source: "users-firestore",
126 | Description: "Get user documents",
127 | AuthRequired: []string{"user-auth"},
128 | },
129 | "get_product_docs": firestoregetdocuments.Config{
130 | Name: "get_product_docs",
131 | Kind: "firestore-get-documents",
132 | Source: "products-firestore",
133 | Description: "Get product documents",
134 | AuthRequired: []string{},
135 | },
136 | "get_order_docs": firestoregetdocuments.Config{
137 | Name: "get_order_docs",
138 | Kind: "firestore-get-documents",
139 | Source: "orders-firestore",
140 | Description: "Get order documents",
141 | AuthRequired: []string{"user-auth", "admin-auth"},
142 | },
143 | }
144 |
145 | got := struct {
146 | Tools server.ToolConfigs `yaml:"tools"`
147 | }{}
148 | // Parse contents
149 | err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
150 | if err != nil {
151 | t.Fatalf("unable to unmarshal: %s", err)
152 | }
153 | if diff := cmp.Diff(want, got.Tools); diff != "" {
154 | t.Fatalf("incorrect parse: diff %v", diff)
155 | }
156 | }
157 |
```
--------------------------------------------------------------------------------
/internal/tools/bigquery/bigquerycommon/util.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package bigquerycommon
16 |
17 | import (
18 | "context"
19 | "fmt"
20 | "sort"
21 | "strings"
22 |
23 | bigqueryapi "cloud.google.com/go/bigquery"
24 | "github.com/googleapis/genai-toolbox/internal/tools"
25 | bigqueryrestapi "google.golang.org/api/bigquery/v2"
26 | )
27 |
28 | // DryRunQuery performs a dry run of the SQL query to validate it and get metadata.
29 | func DryRunQuery(ctx context.Context, restService *bigqueryrestapi.Service, projectID string, location string, sql string, params []*bigqueryrestapi.QueryParameter, connProps []*bigqueryapi.ConnectionProperty) (*bigqueryrestapi.Job, error) {
30 | useLegacySql := false
31 |
32 | restConnProps := make([]*bigqueryrestapi.ConnectionProperty, len(connProps))
33 | for i, prop := range connProps {
34 | restConnProps[i] = &bigqueryrestapi.ConnectionProperty{Key: prop.Key, Value: prop.Value}
35 | }
36 |
37 | jobToInsert := &bigqueryrestapi.Job{
38 | JobReference: &bigqueryrestapi.JobReference{
39 | ProjectId: projectID,
40 | Location: location,
41 | },
42 | Configuration: &bigqueryrestapi.JobConfiguration{
43 | DryRun: true,
44 | Query: &bigqueryrestapi.JobConfigurationQuery{
45 | Query: sql,
46 | UseLegacySql: &useLegacySql,
47 | ConnectionProperties: restConnProps,
48 | QueryParameters: params,
49 | },
50 | },
51 | }
52 |
53 | insertResponse, err := restService.Jobs.Insert(projectID, jobToInsert).Context(ctx).Do()
54 | if err != nil {
55 | return nil, fmt.Errorf("failed to insert dry run job: %w", err)
56 | }
57 | return insertResponse, nil
58 | }
59 |
60 | // BQTypeStringFromToolType converts a tool parameter type string to a BigQuery standard SQL type string.
61 | func BQTypeStringFromToolType(toolType string) (string, error) {
62 | switch toolType {
63 | case "string":
64 | return "STRING", nil
65 | case "integer":
66 | return "INT64", nil
67 | case "float":
68 | return "FLOAT64", nil
69 | case "boolean":
70 | return "BOOL", nil
71 | default:
72 | return "", fmt.Errorf("unsupported tool parameter type for BigQuery: %s", toolType)
73 | }
74 | }
75 |
76 | // InitializeDatasetParameters generates project and dataset tool parameters based on allowedDatasets.
77 | func InitializeDatasetParameters(
78 | allowedDatasets []string,
79 | defaultProjectID string,
80 | projectKey, datasetKey string,
81 | projectDescription, datasetDescription string,
82 | ) (projectParam, datasetParam tools.Parameter) {
83 | if len(allowedDatasets) > 0 {
84 | if len(allowedDatasets) == 1 {
85 | parts := strings.Split(allowedDatasets[0], ".")
86 | defaultProjectID = parts[0]
87 | datasetID := parts[1]
88 | projectDescription += fmt.Sprintf(" Must be `%s`.", defaultProjectID)
89 | datasetDescription += fmt.Sprintf(" Must be `%s`.", datasetID)
90 | datasetParam = tools.NewStringParameterWithDefault(datasetKey, datasetID, datasetDescription)
91 | } else {
92 | datasetIDsByProject := make(map[string][]string)
93 | for _, ds := range allowedDatasets {
94 | parts := strings.Split(ds, ".")
95 | project := parts[0]
96 | dataset := parts[1]
97 | datasetIDsByProject[project] = append(datasetIDsByProject[project], fmt.Sprintf("`%s`", dataset))
98 | }
99 |
100 | var datasetDescriptions, projectIDList []string
101 | for project, datasets := range datasetIDsByProject {
102 | sort.Strings(datasets)
103 | projectIDList = append(projectIDList, fmt.Sprintf("`%s`", project))
104 | datasetList := strings.Join(datasets, ", ")
105 | datasetDescriptions = append(datasetDescriptions, fmt.Sprintf("%s from project `%s`", datasetList, project))
106 | }
107 | sort.Strings(projectIDList)
108 | sort.Strings(datasetDescriptions)
109 | projectDescription += fmt.Sprintf(" Must be one of the following: %s.", strings.Join(projectIDList, ", "))
110 | datasetDescription += fmt.Sprintf(" Must be one of the allowed datasets: %s.", strings.Join(datasetDescriptions, "; "))
111 | datasetParam = tools.NewStringParameter(datasetKey, datasetDescription)
112 | }
113 | } else {
114 | datasetParam = tools.NewStringParameter(datasetKey, datasetDescription)
115 | }
116 |
117 | projectParam = tools.NewStringParameterWithDefault(projectKey, defaultProjectID, projectDescription)
118 |
119 | return projectParam, datasetParam
120 | }
121 |
```
--------------------------------------------------------------------------------
/docs/en/how-to/toolbox-ui/index.md:
--------------------------------------------------------------------------------
```markdown
1 | ---
2 | title: "Toolbox UI"
3 | type: docs
4 | weight: 1
5 | description: >
6 | How to effectively use Toolbox UI.
7 | ---
8 |
9 | Toolbox UI is a built-in web interface that allows users to visually inspect and
10 | test out configured resources such as tools and toolsets.
11 |
12 | ## Launching Toolbox UI
13 |
14 | To launch Toolbox's interactive UI, use the `--ui` flag.
15 |
16 | ```sh
17 | ./toolbox --ui
18 | ```
19 |
20 | Toolbox UI will be served from the same host and port as the Toolbox Server,
21 | with the `/ui` suffix. Once Toolbox is launched, the following INFO log with
22 | Toolbox UI's url will be shown:
23 |
24 | ```bash
25 | INFO "Toolbox UI is up and running at: http://localhost:5000/ui"
26 | ```
27 |
28 | ## Navigating the Tools Page
29 |
30 | The tools page shows all tools loaded from your configuration file. This
31 | corresponds to the default toolset (represented by an empty string). Each tool's
32 | name on this page will exactly match its name in the configuration file.
33 |
34 | To view details for a specific tool, click on the tool name. The main content
35 | area will be populated with the tool name, description, and available
36 | parameters.
37 |
38 | 
39 |
40 | ### Invoking a Tool
41 |
42 | 1. Click on a Tool
43 | 1. Enter appropriate parameters in each parameter field
44 | 1. Click "Run Tool"
45 | 1. Done! Your results will appear in the response field
46 | 1. (Optional) Uncheck "Prettify JSON" to format the response as plain text
47 |
48 | 
49 |
50 | ### Optional Parameters
51 |
52 | Toolbox allows users to add [optional
53 | parameters](../../resources/tools/#basic-parameters) with or without a default
54 | value.
55 |
56 | To exclude a parameter, uncheck the box to the right of an associated parameter,
57 | and that parameter will not be included in the request body. If the parameter is
58 | not sent, Toolbox will either use it as `nil` value or the `default` value, if
59 | configured. If the parameter is required, Toolbox will throw an error.
60 |
61 | When the box is checked, parameter will be sent exactly as entered in the
62 | response field (e.g. empty string).
63 |
64 | 
65 |
66 | 
67 |
68 | ### Editing Headers
69 |
70 | To edit headers, press the "Edit Headers" button to display the header modal.
71 | Within this modal, users can make direct edits by typing into the header's text
72 | area.
73 |
74 | Toolbox UI validates that the headers are in correct JSON format. Other
75 | header-related errors (e.g., incorrect header names or values required by the
76 | tool) will be reported in the Response section after running the tool.
77 |
78 | 
79 |
80 | #### Google OAuth
81 |
82 | Currently, Toolbox supports Google OAuth 2.0 as an AuthService, which allows
83 | tools to utilize authorized parameters. When a tool uses an authorized
84 | parameter, the parameter will be displayed but not editable, as it will be
85 | populated from the authentication token.
86 |
87 | To provide the token, add your Google OAuth ID Token to the request header using
88 | the "Edit Headers" button and modal described above. The key should be the name
89 | of your AuthService as defined in your tool configuration file, suffixed with
90 | `_token`. The value should be your ID token as a string.
91 |
92 | 1. Select a tool that requires [authenticated parameters]()
93 | 1. The auth parameter's text field is greyed out. This is because it cannot be
94 | entered manually and will be parsed from the resolved auth token
95 | 1. To update request headers with the token, select "Edit Headers"
96 | 1. (Optional) If you wish to manually edit the header, checkout the dropdown
97 | "How to extract Google OAuth ID Token manually" for guidance on retrieving ID
98 | token
99 | 1. To edit the header automatically, click the "Auto Setup" button that is
100 | associated with your Auth Profile
101 | 1. Enter the Client ID defined in your tools configuration file
102 | 1. Click "Continue"
103 | 1. Click "Sign in With Google" and login with your associated google account.
104 | This should automatically populate the header text area with your token
105 | 1. Click "Save"
106 | 1. Click "Run Tool"
107 |
108 | ```json
109 | {
110 | "Content-Type": "application/json",
111 | "my-google-auth_token": "YOUR_ID_TOKEN_HERE"
112 | }
113 | ```
114 |
115 | 
116 |
117 | ## Navigating the Toolsets Page
118 |
119 | Through the toolsets page, users can search for a specific toolset to retrieve
120 | tools from. Simply enter the toolset name in the search bar, and press "Enter"
121 | to retrieve the associated tools.
122 |
123 | If the toolset name is not defined within the tools configuration file, an error
124 | message will be displayed.
125 |
126 | 
127 |
```
--------------------------------------------------------------------------------
/docs/en/resources/tools/dgraph/dgraph-dql.md:
--------------------------------------------------------------------------------
```markdown
1 | ---
2 | title: "dgraph-dql"
3 | type: docs
4 | weight: 1
5 | description: >
6 | A "dgraph-dql" tool executes a pre-defined DQL statement against a Dgraph
7 | database.
8 | aliases:
9 | - /resources/tools/dgraph-dql
10 | ---
11 |
12 | ## About
13 |
14 | A `dgraph-dql` tool executes a pre-defined DQL statement against a Dgraph
15 | database. It's compatible with any of the following sources:
16 |
17 | - [dgraph](../../sources/dgraph.md)
18 |
19 | To run a statement as a query, you need to set the config `isQuery=true`. For
20 | upserts or mutations, set `isQuery=false`. You can also configure timeout for a
21 | query.
22 |
23 | > **Note:** This tool uses parameterized queries to prevent SQL injections.
24 | > Query parameters can be used as substitutes for arbitrary expressions.
25 | > Parameters cannot be used as substitutes for identifiers, column names, table
26 | > names, or other parts of the query.
27 |
28 | ## Example
29 |
30 | {{< tabpane persist="header" >}}
31 | {{< tab header="Query" lang="yaml" >}}
32 |
33 | tools:
34 | search_user:
35 | kind: dgraph-dql
36 | source: my-dgraph-source
37 | statement: |
38 | query all($role: string){
39 | users(func: has(name)) @filter(eq(role, $role) AND ge(age, 30) AND le(age, 50)) {
40 | uid
41 | name
42 | email
43 | role
44 | age
45 | }
46 | }
47 | isQuery: true
48 | timeout: 20s
49 | description: |
50 | Use this tool to retrieve the details of users who are admins and are between 30 and 50 years old.
51 | The query returns the user's name, email, role, and age.
52 | This can be helpful when you want to fetch admin users within a specific age range.
53 | Example: Fetch admins aged between 30 and 50:
54 | [
55 | {
56 | "name": "Alice",
57 | "role": "admin",
58 | "age": 35
59 | },
60 | {
61 | "name": "Bob",
62 | "role": "admin",
63 | "age": 45
64 | }
65 | ]
66 | parameters:
67 | - name: $role
68 | type: string
69 | description: admin
70 |
71 | {{< /tab >}}
72 | {{< tab header="Mutation" lang="yaml" >}}
73 |
74 | tools:
75 | dgraph-manage-user-instance:
76 | kind: dgraph-dql
77 | source: my-dgraph-source
78 | isQuery: false
79 | statement: |
80 | {
81 | set {
82 | _:user1 <name> $user1 .
83 | _:user1 <email> $email1 .
84 | _:user1 <role> "admin" .
85 | _:user1 <age> "35" .
86 |
87 | _:user2 <name> $user2 .
88 | _:user2 <email> $email2 .
89 | _:user2 <role> "admin" .
90 | _:user2 <age> "45" .
91 | }
92 | }
93 | description: |
94 | Use this tool to insert or update user data into the Dgraph database.
95 | The mutation adds or updates user details like name, email, role, and age.
96 | Example: Add users Alice and Bob as admins with specific ages.
97 | parameters:
98 | - name: user1
99 | type: string
100 | description: Alice
101 | - name: email1
102 | type: string
103 | description: [email protected]
104 | - name: user2
105 | type: string
106 | description: Bob
107 | - name: email2
108 | type: string
109 | description: [email protected]
110 |
111 | {{< /tab >}}
112 | {{< /tabpane >}}
113 |
114 | ## Reference
115 |
116 | | **field** | **type** | **required** | **description** |
117 | |-------------|:---------------------------------------:|:------------:|-------------------------------------------------------------------------------------------|
118 | | kind | string | true | Must be "dgraph-dql". |
119 | | source | string | true | Name of the source the dql query should execute on. |
120 | | description | string | true | Description of the tool that is passed to the LLM. |
121 | | statement | string | true | dql statement to execute |
122 | | isQuery | boolean | false | To run statement as query set true otherwise false |
123 | | timeout | string | false | To set timeout for query |
124 | | parameters | [parameters](../#specifying-parameters) | false | List of [parameters](../#specifying-parameters) that will be used with the dql statement. |
125 |
```
--------------------------------------------------------------------------------
/internal/sources/trino/trino.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package trino
16 |
17 | import (
18 | "context"
19 | "database/sql"
20 | "fmt"
21 | "net/url"
22 | "time"
23 |
24 | "github.com/goccy/go-yaml"
25 | "github.com/googleapis/genai-toolbox/internal/sources"
26 | _ "github.com/trinodb/trino-go-client/trino"
27 | "go.opentelemetry.io/otel/trace"
28 | )
29 |
30 | const SourceKind string = "trino"
31 |
32 | // validate interface
33 | var _ sources.SourceConfig = Config{}
34 |
35 | func init() {
36 | if !sources.Register(SourceKind, newConfig) {
37 | panic(fmt.Sprintf("source kind %q already registered", SourceKind))
38 | }
39 | }
40 |
41 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
42 | actual := Config{Name: name}
43 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
44 | return nil, err
45 | }
46 | return actual, nil
47 | }
48 |
49 | type Config struct {
50 | Name string `yaml:"name" validate:"required"`
51 | Kind string `yaml:"kind" validate:"required"`
52 | Host string `yaml:"host" validate:"required"`
53 | Port string `yaml:"port" validate:"required"`
54 | User string `yaml:"user"`
55 | Password string `yaml:"password"`
56 | Catalog string `yaml:"catalog" validate:"required"`
57 | Schema string `yaml:"schema" validate:"required"`
58 | QueryTimeout string `yaml:"queryTimeout"`
59 | AccessToken string `yaml:"accessToken"`
60 | KerberosEnabled bool `yaml:"kerberosEnabled"`
61 | SSLEnabled bool `yaml:"sslEnabled"`
62 | }
63 |
64 | func (r Config) SourceConfigKind() string {
65 | return SourceKind
66 | }
67 |
68 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
69 | 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)
70 | if err != nil {
71 | return nil, fmt.Errorf("unable to create pool: %w", err)
72 | }
73 |
74 | err = pool.PingContext(ctx)
75 | if err != nil {
76 | return nil, fmt.Errorf("unable to connect successfully: %w", err)
77 | }
78 |
79 | s := &Source{
80 | Name: r.Name,
81 | Kind: SourceKind,
82 | Pool: pool,
83 | }
84 | return s, nil
85 | }
86 |
87 | var _ sources.Source = &Source{}
88 |
89 | type Source struct {
90 | Name string `yaml:"name"`
91 | Kind string `yaml:"kind"`
92 | Pool *sql.DB
93 | }
94 |
95 | func (s *Source) SourceKind() string {
96 | return SourceKind
97 | }
98 |
99 | func (s *Source) TrinoDB() *sql.DB {
100 | return s.Pool
101 | }
102 |
103 | func initTrinoConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, password, catalog, schema, queryTimeout, accessToken string, kerberosEnabled, sslEnabled bool) (*sql.DB, error) {
104 | //nolint:all // Reassigned ctx
105 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
106 | defer span.End()
107 |
108 | // Build Trino DSN
109 | dsn, err := buildTrinoDSN(host, port, user, password, catalog, schema, queryTimeout, accessToken, kerberosEnabled, sslEnabled)
110 | if err != nil {
111 | return nil, fmt.Errorf("failed to build DSN: %w", err)
112 | }
113 |
114 | db, err := sql.Open("trino", dsn)
115 | if err != nil {
116 | return nil, fmt.Errorf("failed to open connection: %w", err)
117 | }
118 |
119 | // Configure connection pool
120 | db.SetMaxOpenConns(10)
121 | db.SetMaxIdleConns(5)
122 | db.SetConnMaxLifetime(time.Hour)
123 |
124 | return db, nil
125 | }
126 |
127 | func buildTrinoDSN(host, port, user, password, catalog, schema, queryTimeout, accessToken string, kerberosEnabled, sslEnabled bool) (string, error) {
128 | // Build query parameters
129 | query := url.Values{}
130 | query.Set("catalog", catalog)
131 | query.Set("schema", schema)
132 | if queryTimeout != "" {
133 | query.Set("queryTimeout", queryTimeout)
134 | }
135 | if accessToken != "" {
136 | query.Set("accessToken", accessToken)
137 | }
138 | if kerberosEnabled {
139 | query.Set("KerberosEnabled", "true")
140 | }
141 |
142 | // Build URL
143 | scheme := "http"
144 | if sslEnabled {
145 | scheme = "https"
146 | }
147 |
148 | u := &url.URL{
149 | Scheme: scheme,
150 | Host: fmt.Sprintf("%s:%s", host, port),
151 | RawQuery: query.Encode(),
152 | }
153 |
154 | // Only set user and password if not empty
155 | if user != "" && password != "" {
156 | u.User = url.UserPassword(user, password)
157 | } else if user != "" {
158 | u.User = url.User(user)
159 | }
160 |
161 | return u.String(), nil
162 | }
163 |
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestoregetrules/firestoregetrules_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package firestoregetrules_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/testutils"
24 | "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoregetrules"
25 | )
26 |
27 | func TestParseFromYamlFirestoreGetRules(t *testing.T) {
28 | ctx, err := testutils.ContextWithNewLogger()
29 | if err != nil {
30 | t.Fatalf("unexpected error: %s", err)
31 | }
32 | tcs := []struct {
33 | desc string
34 | in string
35 | want server.ToolConfigs
36 | }{
37 | {
38 | desc: "basic example",
39 | in: `
40 | tools:
41 | get_rules_tool:
42 | kind: firestore-get-rules
43 | source: my-firestore-instance
44 | description: Retrieves the active Firestore security rules for the current project
45 | `,
46 | want: server.ToolConfigs{
47 | "get_rules_tool": firestoregetrules.Config{
48 | Name: "get_rules_tool",
49 | Kind: "firestore-get-rules",
50 | Source: "my-firestore-instance",
51 | Description: "Retrieves the active Firestore security rules for the current project",
52 | AuthRequired: []string{},
53 | },
54 | },
55 | },
56 | {
57 | desc: "with auth requirements",
58 | in: `
59 | tools:
60 | secure_get_rules:
61 | kind: firestore-get-rules
62 | source: prod-firestore
63 | description: Get Firestore security rules with authentication
64 | authRequired:
65 | - google-auth-service
66 | - admin-service
67 | `,
68 | want: server.ToolConfigs{
69 | "secure_get_rules": firestoregetrules.Config{
70 | Name: "secure_get_rules",
71 | Kind: "firestore-get-rules",
72 | Source: "prod-firestore",
73 | Description: "Get Firestore security rules with authentication",
74 | AuthRequired: []string{"google-auth-service", "admin-service"},
75 | },
76 | },
77 | },
78 | }
79 | for _, tc := range tcs {
80 | t.Run(tc.desc, func(t *testing.T) {
81 | got := struct {
82 | Tools server.ToolConfigs `yaml:"tools"`
83 | }{}
84 | // Parse contents
85 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
86 | if err != nil {
87 | t.Fatalf("unable to unmarshal: %s", err)
88 | }
89 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
90 | t.Fatalf("incorrect parse: diff %v", diff)
91 | }
92 | })
93 | }
94 | }
95 |
96 | func TestParseFromYamlMultipleTools(t *testing.T) {
97 | ctx, err := testutils.ContextWithNewLogger()
98 | if err != nil {
99 | t.Fatalf("unexpected error: %s", err)
100 | }
101 | in := `
102 | tools:
103 | get_dev_rules:
104 | kind: firestore-get-rules
105 | source: dev-firestore
106 | description: Get development Firestore rules
107 | authRequired:
108 | - dev-auth
109 | get_staging_rules:
110 | kind: firestore-get-rules
111 | source: staging-firestore
112 | description: Get staging Firestore rules
113 | get_prod_rules:
114 | kind: firestore-get-rules
115 | source: prod-firestore
116 | description: Get production Firestore rules
117 | authRequired:
118 | - prod-auth
119 | - admin-auth
120 | `
121 | want := server.ToolConfigs{
122 | "get_dev_rules": firestoregetrules.Config{
123 | Name: "get_dev_rules",
124 | Kind: "firestore-get-rules",
125 | Source: "dev-firestore",
126 | Description: "Get development Firestore rules",
127 | AuthRequired: []string{"dev-auth"},
128 | },
129 | "get_staging_rules": firestoregetrules.Config{
130 | Name: "get_staging_rules",
131 | Kind: "firestore-get-rules",
132 | Source: "staging-firestore",
133 | Description: "Get staging Firestore rules",
134 | AuthRequired: []string{},
135 | },
136 | "get_prod_rules": firestoregetrules.Config{
137 | Name: "get_prod_rules",
138 | Kind: "firestore-get-rules",
139 | Source: "prod-firestore",
140 | Description: "Get production Firestore rules",
141 | AuthRequired: []string{"prod-auth", "admin-auth"},
142 | },
143 | }
144 |
145 | got := struct {
146 | Tools server.ToolConfigs `yaml:"tools"`
147 | }{}
148 | // Parse contents
149 | err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
150 | if err != nil {
151 | t.Fatalf("unable to unmarshal: %s", err)
152 | }
153 | if diff := cmp.Diff(want, got.Tools); diff != "" {
154 | t.Fatalf("incorrect parse: diff %v", diff)
155 | }
156 | }
157 |
```
--------------------------------------------------------------------------------
/internal/sources/spanner/spanner_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2024 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package spanner_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/sources"
24 | "github.com/googleapis/genai-toolbox/internal/sources/spanner"
25 | "github.com/googleapis/genai-toolbox/internal/testutils"
26 | )
27 |
28 | func TestParseFromYamlSpannerDb(t *testing.T) {
29 | tcs := []struct {
30 | desc string
31 | in string
32 | want server.SourceConfigs
33 | }{
34 | {
35 | desc: "basic example",
36 | in: `
37 | sources:
38 | my-spanner-instance:
39 | kind: spanner
40 | project: my-project
41 | instance: my-instance
42 | database: my_db
43 | `,
44 | want: map[string]sources.SourceConfig{
45 | "my-spanner-instance": spanner.Config{
46 | Name: "my-spanner-instance",
47 | Kind: spanner.SourceKind,
48 | Project: "my-project",
49 | Instance: "my-instance",
50 | Dialect: "googlesql",
51 | Database: "my_db",
52 | },
53 | },
54 | },
55 | {
56 | desc: "gsql dialect",
57 | in: `
58 | sources:
59 | my-spanner-instance:
60 | kind: spanner
61 | project: my-project
62 | instance: my-instance
63 | dialect: Googlesql
64 | database: my_db
65 | `,
66 | want: map[string]sources.SourceConfig{
67 | "my-spanner-instance": spanner.Config{
68 | Name: "my-spanner-instance",
69 | Kind: spanner.SourceKind,
70 | Project: "my-project",
71 | Instance: "my-instance",
72 | Dialect: "googlesql",
73 | Database: "my_db",
74 | },
75 | },
76 | },
77 | {
78 | desc: "postgresql dialect",
79 | in: `
80 | sources:
81 | my-spanner-instance:
82 | kind: spanner
83 | project: my-project
84 | instance: my-instance
85 | dialect: postgresql
86 | database: my_db
87 | `,
88 | want: map[string]sources.SourceConfig{
89 | "my-spanner-instance": spanner.Config{
90 | Name: "my-spanner-instance",
91 | Kind: spanner.SourceKind,
92 | Project: "my-project",
93 | Instance: "my-instance",
94 | Dialect: "postgresql",
95 | Database: "my_db",
96 | },
97 | },
98 | },
99 | }
100 | for _, tc := range tcs {
101 | t.Run(tc.desc, func(t *testing.T) {
102 | got := struct {
103 | Sources server.SourceConfigs `yaml:"sources"`
104 | }{}
105 | // Parse contents
106 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
107 | if err != nil {
108 | t.Fatalf("unable to unmarshal: %s", err)
109 | }
110 | if !cmp.Equal(tc.want, got.Sources) {
111 | t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources)
112 | }
113 | })
114 | }
115 |
116 | }
117 |
118 | func TestFailParseFromYaml(t *testing.T) {
119 | tcs := []struct {
120 | desc string
121 | in string
122 | err string
123 | }{
124 | {
125 | desc: "invalid dialect",
126 | in: `
127 | sources:
128 | my-spanner-instance:
129 | kind: spanner
130 | project: my-project
131 | instance: my-instance
132 | dialect: fail
133 | database: my_db
134 | `,
135 | err: "unable to parse source \"my-spanner-instance\" as \"spanner\": dialect invalid: must be one of \"googlesql\", or \"postgresql\"",
136 | },
137 | {
138 | desc: "extra field",
139 | in: `
140 | sources:
141 | my-spanner-instance:
142 | kind: spanner
143 | project: my-project
144 | instance: my-instance
145 | database: my_db
146 | foo: bar
147 | `,
148 | 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",
149 | },
150 | {
151 | desc: "missing required field",
152 | in: `
153 | sources:
154 | my-spanner-instance:
155 | kind: spanner
156 | project: my-project
157 | instance: my-instance
158 | `,
159 | err: "unable to parse source \"my-spanner-instance\" as \"spanner\": Key: 'Config.Database' Error:Field validation for 'Database' failed on the 'required' tag",
160 | },
161 | }
162 | for _, tc := range tcs {
163 | t.Run(tc.desc, func(t *testing.T) {
164 | got := struct {
165 | Sources server.SourceConfigs `yaml:"sources"`
166 | }{}
167 | // Parse contents
168 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
169 | if err == nil {
170 | t.Fatalf("expect parsing to fail")
171 | }
172 | errStr := err.Error()
173 | if errStr != tc.err {
174 | t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
175 | }
176 | })
177 | }
178 | }
179 |
```
--------------------------------------------------------------------------------
/internal/tools/mongodb/mongodbupdatemany/mongodbupdatemany_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package mongodbupdatemany_test
16 |
17 | import (
18 | "strings"
19 | "testing"
20 |
21 | "github.com/googleapis/genai-toolbox/internal/tools"
22 | "github.com/googleapis/genai-toolbox/internal/tools/mongodb/mongodbupdatemany"
23 |
24 | yaml "github.com/goccy/go-yaml"
25 | "github.com/google/go-cmp/cmp"
26 | "github.com/googleapis/genai-toolbox/internal/server"
27 | "github.com/googleapis/genai-toolbox/internal/testutils"
28 | )
29 |
30 | func TestParseFromYamlMongoQuery(t *testing.T) {
31 | ctx, err := testutils.ContextWithNewLogger()
32 | if err != nil {
33 | t.Fatalf("unexpected error: %s", err)
34 | }
35 | tcs := []struct {
36 | desc string
37 | in string
38 | want server.ToolConfigs
39 | }{
40 | {
41 | desc: "basic example",
42 | in: `
43 | tools:
44 | example_tool:
45 | kind: mongodb-update-many
46 | source: my-instance
47 | description: some description
48 | database: test_db
49 | collection: test_coll
50 | filterPayload: |
51 | { name: {{json .name}} }
52 | filterParams:
53 | - name: name
54 | type: string
55 | description: small description
56 | canonical: true
57 | updatePayload: |
58 | { $set: { name: {{json .name}} } }
59 | updateParams:
60 | - name: name
61 | type: string
62 | description: small description
63 | `,
64 | want: server.ToolConfigs{
65 | "example_tool": mongodbupdatemany.Config{
66 | Name: "example_tool",
67 | Kind: "mongodb-update-many",
68 | Source: "my-instance",
69 | AuthRequired: []string{},
70 | Database: "test_db",
71 | Collection: "test_coll",
72 | FilterPayload: "{ name: {{json .name}} }\n",
73 | FilterParams: tools.Parameters{
74 | &tools.StringParameter{
75 | CommonParameter: tools.CommonParameter{
76 | Name: "name",
77 | Type: "string",
78 | Desc: "small description",
79 | },
80 | },
81 | },
82 | UpdatePayload: "{ $set: { name: {{json .name}} } }\n",
83 | UpdateParams: tools.Parameters{
84 | &tools.StringParameter{
85 | CommonParameter: tools.CommonParameter{
86 | Name: "name",
87 | Type: "string",
88 | Desc: "small description",
89 | },
90 | },
91 | },
92 | Description: "some description",
93 | Canonical: true,
94 | },
95 | },
96 | },
97 | }
98 | for _, tc := range tcs {
99 | t.Run(tc.desc, func(t *testing.T) {
100 | got := struct {
101 | Tools server.ToolConfigs `yaml:"tools"`
102 | }{}
103 | // Parse contents
104 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
105 | if err != nil {
106 | t.Fatalf("unable to unmarshal: %s", err)
107 | }
108 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
109 | t.Fatalf("incorrect parse: diff %v", diff)
110 | }
111 | })
112 | }
113 |
114 | }
115 |
116 | func TestFailParseFromYamlMongoQuery(t *testing.T) {
117 | ctx, err := testutils.ContextWithNewLogger()
118 | if err != nil {
119 | t.Fatalf("unexpected error: %s", err)
120 | }
121 | tcs := []struct {
122 | desc string
123 | in string
124 | err string
125 | }{
126 | {
127 | desc: "Invalid method",
128 | in: `
129 | tools:
130 | example_tool:
131 | kind: mongodb-update-many
132 | source: my-instance
133 | description: some description
134 | collection: test_coll
135 | filterPayload: |
136 | { name : {{json .name}} }
137 | filterParams:
138 | - name: name
139 | type: string
140 | description: small description
141 | canonical: true
142 | updatePayload: |
143 | { $set: { name: {{json .name}} } }
144 | updateParams:
145 | - name: data
146 | type: string
147 | description: the content in json
148 | `,
149 | err: `unable to parse tool "example_tool" as kind "mongodb-update-many"`,
150 | },
151 | }
152 | for _, tc := range tcs {
153 | t.Run(tc.desc, func(t *testing.T) {
154 | got := struct {
155 | Tools server.ToolConfigs `yaml:"tools"`
156 | }{}
157 | // Parse contents
158 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
159 | if err == nil {
160 | t.Fatalf("expect parsing to fail")
161 | }
162 | errStr := err.Error()
163 | if !strings.Contains(errStr, tc.err) {
164 | t.Fatalf("unexpected error string: got %q, want substring %q", errStr, tc.err)
165 | }
166 | })
167 | }
168 |
169 | }
170 |
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestoredeletedocuments/firestoredeletedocuments_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package firestoredeletedocuments_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/testutils"
24 | "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestoredeletedocuments"
25 | )
26 |
27 | func TestParseFromYamlFirestoreDeleteDocuments(t *testing.T) {
28 | ctx, err := testutils.ContextWithNewLogger()
29 | if err != nil {
30 | t.Fatalf("unexpected error: %s", err)
31 | }
32 | tcs := []struct {
33 | desc string
34 | in string
35 | want server.ToolConfigs
36 | }{
37 | {
38 | desc: "basic example",
39 | in: `
40 | tools:
41 | delete_docs_tool:
42 | kind: firestore-delete-documents
43 | source: my-firestore-instance
44 | description: Delete documents from Firestore by paths
45 | `,
46 | want: server.ToolConfigs{
47 | "delete_docs_tool": firestoredeletedocuments.Config{
48 | Name: "delete_docs_tool",
49 | Kind: "firestore-delete-documents",
50 | Source: "my-firestore-instance",
51 | Description: "Delete documents from Firestore by paths",
52 | AuthRequired: []string{},
53 | },
54 | },
55 | },
56 | {
57 | desc: "with auth requirements",
58 | in: `
59 | tools:
60 | secure_delete_docs:
61 | kind: firestore-delete-documents
62 | source: prod-firestore
63 | description: Delete documents with authentication
64 | authRequired:
65 | - google-auth-service
66 | - api-key-service
67 | `,
68 | want: server.ToolConfigs{
69 | "secure_delete_docs": firestoredeletedocuments.Config{
70 | Name: "secure_delete_docs",
71 | Kind: "firestore-delete-documents",
72 | Source: "prod-firestore",
73 | Description: "Delete documents with authentication",
74 | AuthRequired: []string{"google-auth-service", "api-key-service"},
75 | },
76 | },
77 | },
78 | }
79 | for _, tc := range tcs {
80 | t.Run(tc.desc, func(t *testing.T) {
81 | got := struct {
82 | Tools server.ToolConfigs `yaml:"tools"`
83 | }{}
84 | // Parse contents
85 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
86 | if err != nil {
87 | t.Fatalf("unable to unmarshal: %s", err)
88 | }
89 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
90 | t.Fatalf("incorrect parse: diff %v", diff)
91 | }
92 | })
93 | }
94 | }
95 |
96 | func TestParseFromYamlMultipleTools(t *testing.T) {
97 | ctx, err := testutils.ContextWithNewLogger()
98 | if err != nil {
99 | t.Fatalf("unexpected error: %s", err)
100 | }
101 | in := `
102 | tools:
103 | delete_user_docs:
104 | kind: firestore-delete-documents
105 | source: users-firestore
106 | description: Delete user documents
107 | authRequired:
108 | - user-auth
109 | delete_product_docs:
110 | kind: firestore-delete-documents
111 | source: products-firestore
112 | description: Delete product documents
113 | delete_order_docs:
114 | kind: firestore-delete-documents
115 | source: orders-firestore
116 | description: Delete order documents
117 | authRequired:
118 | - user-auth
119 | - admin-auth
120 | `
121 | want := server.ToolConfigs{
122 | "delete_user_docs": firestoredeletedocuments.Config{
123 | Name: "delete_user_docs",
124 | Kind: "firestore-delete-documents",
125 | Source: "users-firestore",
126 | Description: "Delete user documents",
127 | AuthRequired: []string{"user-auth"},
128 | },
129 | "delete_product_docs": firestoredeletedocuments.Config{
130 | Name: "delete_product_docs",
131 | Kind: "firestore-delete-documents",
132 | Source: "products-firestore",
133 | Description: "Delete product documents",
134 | AuthRequired: []string{},
135 | },
136 | "delete_order_docs": firestoredeletedocuments.Config{
137 | Name: "delete_order_docs",
138 | Kind: "firestore-delete-documents",
139 | Source: "orders-firestore",
140 | Description: "Delete order documents",
141 | AuthRequired: []string{"user-auth", "admin-auth"},
142 | },
143 | }
144 |
145 | got := struct {
146 | Tools server.ToolConfigs `yaml:"tools"`
147 | }{}
148 | // Parse contents
149 | err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
150 | if err != nil {
151 | t.Fatalf("unable to unmarshal: %s", err)
152 | }
153 | if diff := cmp.Diff(want, got.Tools); diff != "" {
154 | t.Fatalf("incorrect parse: diff %v", diff)
155 | }
156 | }
157 |
```
--------------------------------------------------------------------------------
/internal/sources/couchbase/couchbase_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2024 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package couchbase_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/sources/couchbase"
24 | "github.com/googleapis/genai-toolbox/internal/testutils"
25 | )
26 |
27 | func TestParseFromYamlCouchbase(t *testing.T) {
28 | tcs := []struct {
29 | desc string
30 | in string
31 | want server.SourceConfigs
32 | }{
33 | {
34 | desc: "basic example",
35 | in: `
36 | sources:
37 | my-couchbase-instance:
38 | kind: couchbase
39 | connectionString: localhost
40 | username: Administrator
41 | password: password
42 | bucket: travel-sample
43 | scope: inventory
44 | `,
45 | want: server.SourceConfigs{
46 | "my-couchbase-instance": couchbase.Config{
47 | Name: "my-couchbase-instance",
48 | Kind: couchbase.SourceKind,
49 | ConnectionString: "localhost",
50 | Username: "Administrator",
51 | Password: "password",
52 | Bucket: "travel-sample",
53 | Scope: "inventory",
54 | },
55 | },
56 | },
57 | {
58 | desc: "with TLS configuration",
59 | in: `
60 | sources:
61 | my-couchbase-instance:
62 | kind: couchbase
63 | connectionString: couchbases://localhost
64 | bucket: travel-sample
65 | scope: inventory
66 | clientCert: /path/to/cert.pem
67 | clientKey: /path/to/key.pem
68 | clientCertPassword: password
69 | clientKeyPassword: password
70 | caCert: /path/to/ca.pem
71 | noSslVerify: false
72 | queryScanConsistency: 2
73 | `,
74 | want: server.SourceConfigs{
75 | "my-couchbase-instance": couchbase.Config{
76 | Name: "my-couchbase-instance",
77 | Kind: couchbase.SourceKind,
78 | ConnectionString: "couchbases://localhost",
79 | Bucket: "travel-sample",
80 | Scope: "inventory",
81 | ClientCert: "/path/to/cert.pem",
82 | ClientKey: "/path/to/key.pem",
83 | ClientCertPassword: "password",
84 | ClientKeyPassword: "password",
85 | CACert: "/path/to/ca.pem",
86 | NoSSLVerify: false,
87 | QueryScanConsistency: 2,
88 | },
89 | },
90 | },
91 | }
92 | for _, tc := range tcs {
93 | t.Run(tc.desc, func(t *testing.T) {
94 | got := struct {
95 | Sources server.SourceConfigs `yaml:"sources"`
96 | }{}
97 | // Parse contents
98 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
99 | if err != nil {
100 | t.Fatalf("unable to unmarshal: %s", err)
101 | }
102 | if !cmp.Equal(tc.want, got.Sources) {
103 | t.Fatalf("incorrect parse: want %v, got %v", tc.want, got.Sources)
104 | }
105 | })
106 | }
107 | }
108 |
109 | func TestFailParseFromYaml(t *testing.T) {
110 | tcs := []struct {
111 | desc string
112 | in string
113 | err string
114 | }{
115 | {
116 | desc: "extra field",
117 | in: `
118 | sources:
119 | my-couchbase-instance:
120 | kind: couchbase
121 | connectionString: localhost
122 | username: Administrator
123 | password: password
124 | bucket: travel-sample
125 | scope: inventory
126 | foo: bar
127 | `,
128 | 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 | ",
129 | },
130 | {
131 | desc: "missing required field",
132 | in: `
133 | sources:
134 | my-couchbase-instance:
135 | kind: couchbase
136 | username: Administrator
137 | password: password
138 | bucket: travel-sample
139 | scope: inventory
140 | `,
141 | err: "unable to parse source \"my-couchbase-instance\" as \"couchbase\": Key: 'Config.ConnectionString' Error:Field validation for 'ConnectionString' failed on the 'required' tag",
142 | },
143 | }
144 | for _, tc := range tcs {
145 | t.Run(tc.desc, func(t *testing.T) {
146 | got := struct {
147 | Sources server.SourceConfigs `yaml:"sources"`
148 | }{}
149 | // Parse contents
150 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
151 | if err == nil {
152 | t.Fatalf("expect parsing to fail")
153 | }
154 | errStr := err.Error()
155 | if errStr != tc.err {
156 | t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
157 | }
158 | })
159 | }
160 | }
161 |
```
--------------------------------------------------------------------------------
/tests/dgraph/dgraph_integration_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package dgraph
16 |
17 | import (
18 | "bytes"
19 | "context"
20 | "encoding/json"
21 | "io"
22 | "net/http"
23 | "os"
24 | "reflect"
25 | "regexp"
26 | "testing"
27 | "time"
28 |
29 | "github.com/googleapis/genai-toolbox/internal/testutils"
30 | "github.com/googleapis/genai-toolbox/tests"
31 | )
32 |
33 | var (
34 | DgraphSourceKind = "dgraph"
35 | DgraphApiKey = "api-key"
36 | DgraphUrl = os.Getenv("DGRAPH_URL")
37 | )
38 |
39 | func getDgraphVars(t *testing.T) map[string]any {
40 | if DgraphUrl == "" {
41 | t.Fatal("'DGRAPH_URL' not set")
42 | }
43 | return map[string]any{
44 | "kind": DgraphSourceKind,
45 | "dgraphUrl": DgraphUrl,
46 | "apiKey": DgraphApiKey,
47 | }
48 | }
49 |
50 | func TestDgraphToolEndpoints(t *testing.T) {
51 | sourceConfig := getDgraphVars(t)
52 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
53 | defer cancel()
54 |
55 | var args []string
56 |
57 | // Write config into a file and pass it to command
58 | toolsFile := map[string]any{
59 | "sources": map[string]any{
60 | "my-dgraph-instance": sourceConfig,
61 | },
62 | "tools": map[string]any{
63 | "my-simple-dql-tool": map[string]any{
64 | "kind": "dgraph-dql",
65 | "source": "my-dgraph-instance",
66 | "description": "Simple tool to test end to end functionality.",
67 | "statement": "{result(func: uid(0x0)) {constant: math(1)}}",
68 | "isQuery": true,
69 | "timeout": "20s",
70 | "parameters": []any{},
71 | },
72 | },
73 | }
74 | cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
75 | if err != nil {
76 | t.Fatalf("command initialization returned an error: %s", err)
77 | }
78 | defer cleanup()
79 |
80 | waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
81 | defer cancel()
82 | out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
83 | if err != nil {
84 | t.Logf("toolbox command logs: \n%s", out)
85 | t.Fatalf("toolbox didn't start successfully: %s", err)
86 | }
87 |
88 | // Test tool get endpoint
89 | tcs := []struct {
90 | name string
91 | api string
92 | want map[string]any
93 | }{
94 | {
95 | name: "get my-simple-tool",
96 | api: "http://127.0.0.1:5000/api/tool/my-simple-dql-tool/",
97 | want: map[string]any{
98 | "my-simple-dql-tool": map[string]any{
99 | "description": "Simple tool to test end to end functionality.",
100 | "parameters": []any{},
101 | "authRequired": []any{},
102 | },
103 | },
104 | },
105 | }
106 | for _, tc := range tcs {
107 | t.Run(tc.name, func(t *testing.T) {
108 | resp, err := http.Get(tc.api)
109 | if err != nil {
110 | t.Fatalf("error when sending a request: %s", err)
111 | }
112 | defer resp.Body.Close()
113 | if resp.StatusCode != 200 {
114 | t.Fatalf("response status code is not 200")
115 | }
116 |
117 | var body map[string]interface{}
118 | err = json.NewDecoder(resp.Body).Decode(&body)
119 | if err != nil {
120 | t.Fatalf("error parsing response body")
121 | }
122 |
123 | got, ok := body["tools"]
124 | if !ok {
125 | t.Fatalf("unable to find tools in response body")
126 | }
127 | if !reflect.DeepEqual(got, tc.want) {
128 | t.Fatalf("got %q, want %q", got, tc.want)
129 | }
130 | })
131 | }
132 |
133 | // Test tool invoke endpoint
134 | invokeTcs := []struct {
135 | name string
136 | api string
137 | requestBody io.Reader
138 | want string
139 | }{
140 | {
141 | name: "invoke my-simple-dql-tool",
142 | api: "http://127.0.0.1:5000/api/tool/my-simple-dql-tool/invoke",
143 | requestBody: bytes.NewBuffer([]byte(`{}`)),
144 | want: "{\"result\":[{\"constant\":1}]}",
145 | },
146 | }
147 | for _, tc := range invokeTcs {
148 | t.Run(tc.name, func(t *testing.T) {
149 | resp, err := http.Post(tc.api, "application/json", tc.requestBody)
150 | if err != nil {
151 | t.Fatalf("error when sending a request: %s", err)
152 | }
153 | defer resp.Body.Close()
154 | if resp.StatusCode != http.StatusOK {
155 | bodyBytes, _ := io.ReadAll(resp.Body)
156 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
157 | }
158 |
159 | var body map[string]interface{}
160 | err = json.NewDecoder(resp.Body).Decode(&body)
161 | if err != nil {
162 | t.Fatalf("error parsing response body")
163 | }
164 | got, ok := body["result"].(string)
165 |
166 | if !ok {
167 | t.Fatalf("unable to find result in response body")
168 | }
169 |
170 | if got != tc.want {
171 | t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
172 | }
173 | })
174 | }
175 | }
176 |
```
--------------------------------------------------------------------------------
/internal/tools/looker/lookerquerysql/lookerquerysql.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 | package lookerquerysql
15 |
16 | import (
17 | "context"
18 | "fmt"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/googleapis/genai-toolbox/internal/sources"
22 | lookersrc "github.com/googleapis/genai-toolbox/internal/sources/looker"
23 | "github.com/googleapis/genai-toolbox/internal/tools"
24 | "github.com/googleapis/genai-toolbox/internal/tools/looker/lookercommon"
25 | "github.com/googleapis/genai-toolbox/internal/util"
26 |
27 | "github.com/looker-open-source/sdk-codegen/go/rtl"
28 | v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4"
29 | )
30 |
31 | const kind string = "looker-query-sql"
32 |
33 | func init() {
34 | if !tools.Register(kind, newConfig) {
35 | panic(fmt.Sprintf("tool kind %q already registered", kind))
36 | }
37 | }
38 |
39 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
40 | actual := Config{Name: name}
41 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
42 | return nil, err
43 | }
44 | return actual, nil
45 | }
46 |
47 | type Config struct {
48 | Name string `yaml:"name" validate:"required"`
49 | Kind string `yaml:"kind" validate:"required"`
50 | Source string `yaml:"source" validate:"required"`
51 | Description string `yaml:"description" validate:"required"`
52 | AuthRequired []string `yaml:"authRequired"`
53 | }
54 |
55 | // validate interface
56 | var _ tools.ToolConfig = Config{}
57 |
58 | func (cfg Config) ToolConfigKind() string {
59 | return kind
60 | }
61 |
62 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
63 | // verify source exists
64 | rawS, ok := srcs[cfg.Source]
65 | if !ok {
66 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
67 | }
68 |
69 | // verify the source is compatible
70 | s, ok := rawS.(*lookersrc.Source)
71 | if !ok {
72 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be `looker`", kind)
73 | }
74 |
75 | parameters := lookercommon.GetQueryParameters()
76 |
77 | mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters)
78 |
79 | // finish tool setup
80 | return Tool{
81 | Name: cfg.Name,
82 | Kind: kind,
83 | Parameters: parameters,
84 | AuthRequired: cfg.AuthRequired,
85 | UseClientOAuth: s.UseClientOAuth,
86 | Client: s.Client,
87 | ApiSettings: s.ApiSettings,
88 | manifest: tools.Manifest{
89 | Description: cfg.Description,
90 | Parameters: parameters.Manifest(),
91 | AuthRequired: cfg.AuthRequired,
92 | },
93 | mcpManifest: mcpManifest,
94 | }, nil
95 | }
96 |
97 | // validate interface
98 | var _ tools.Tool = Tool{}
99 |
100 | type Tool struct {
101 | Name string `yaml:"name"`
102 | Kind string `yaml:"kind"`
103 | UseClientOAuth bool
104 | Client *v4.LookerSDK
105 | ApiSettings *rtl.ApiSettings
106 | AuthRequired []string `yaml:"authRequired"`
107 | Parameters tools.Parameters `yaml:"parameters"`
108 | manifest tools.Manifest
109 | mcpManifest tools.McpManifest
110 | }
111 |
112 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
113 | logger, err := util.LoggerFromContext(ctx)
114 | if err != nil {
115 | return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
116 | }
117 | wq, err := lookercommon.ProcessQueryArgs(ctx, params)
118 | if err != nil {
119 | return nil, fmt.Errorf("error building query request: %w", err)
120 | }
121 | sdk, err := lookercommon.GetLookerSDK(t.UseClientOAuth, t.ApiSettings, t.Client, accessToken)
122 | if err != nil {
123 | return nil, fmt.Errorf("error getting sdk: %w", err)
124 | }
125 | resp, err := lookercommon.RunInlineQuery(ctx, sdk, wq, "sql", t.ApiSettings)
126 | if err != nil {
127 | return nil, fmt.Errorf("error making query request: %s", err)
128 | }
129 | logger.DebugContext(ctx, "resp = ", resp)
130 |
131 | return resp, nil
132 | }
133 |
134 | func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
135 | return tools.ParseParams(t.Parameters, data, claims)
136 | }
137 |
138 | func (t Tool) Manifest() tools.Manifest {
139 | return t.manifest
140 | }
141 |
142 | func (t Tool) McpManifest() tools.McpManifest {
143 | return t.mcpManifest
144 | }
145 |
146 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
147 | return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
148 | }
149 |
150 | func (t Tool) RequiresClientAuthorization() bool {
151 | return t.UseClientOAuth
152 | }
153 |
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestorequerycollection/firestorequerycollection_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package firestorequerycollection_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/testutils"
24 | "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestorequerycollection"
25 | )
26 |
27 | func TestParseFromYamlFirestoreQueryCollection(t *testing.T) {
28 | ctx, err := testutils.ContextWithNewLogger()
29 | if err != nil {
30 | t.Fatalf("unexpected error: %s", err)
31 | }
32 | tcs := []struct {
33 | desc string
34 | in string
35 | want server.ToolConfigs
36 | }{
37 | {
38 | desc: "basic example",
39 | in: `
40 | tools:
41 | query_users_tool:
42 | kind: firestore-query-collection
43 | source: my-firestore-instance
44 | description: Query users collection with filters and ordering
45 | `,
46 | want: server.ToolConfigs{
47 | "query_users_tool": firestorequerycollection.Config{
48 | Name: "query_users_tool",
49 | Kind: "firestore-query-collection",
50 | Source: "my-firestore-instance",
51 | Description: "Query users collection with filters and ordering",
52 | AuthRequired: []string{},
53 | },
54 | },
55 | },
56 | {
57 | desc: "with auth requirements",
58 | in: `
59 | tools:
60 | secure_query_tool:
61 | kind: firestore-query-collection
62 | source: prod-firestore
63 | description: Query collections with authentication
64 | authRequired:
65 | - google-auth-service
66 | - api-key-service
67 | `,
68 | want: server.ToolConfigs{
69 | "secure_query_tool": firestorequerycollection.Config{
70 | Name: "secure_query_tool",
71 | Kind: "firestore-query-collection",
72 | Source: "prod-firestore",
73 | Description: "Query collections with authentication",
74 | AuthRequired: []string{"google-auth-service", "api-key-service"},
75 | },
76 | },
77 | },
78 | }
79 | for _, tc := range tcs {
80 | t.Run(tc.desc, func(t *testing.T) {
81 | got := struct {
82 | Tools server.ToolConfigs `yaml:"tools"`
83 | }{}
84 | // Parse contents
85 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
86 | if err != nil {
87 | t.Fatalf("unable to unmarshal: %s", err)
88 | }
89 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
90 | t.Fatalf("incorrect parse: diff %v", diff)
91 | }
92 | })
93 | }
94 | }
95 |
96 | func TestParseFromYamlMultipleTools(t *testing.T) {
97 | ctx, err := testutils.ContextWithNewLogger()
98 | if err != nil {
99 | t.Fatalf("unexpected error: %s", err)
100 | }
101 | in := `
102 | tools:
103 | query_users:
104 | kind: firestore-query-collection
105 | source: users-firestore
106 | description: Query user documents with filtering
107 | authRequired:
108 | - user-auth
109 | query_products:
110 | kind: firestore-query-collection
111 | source: products-firestore
112 | description: Query product catalog
113 | query_orders:
114 | kind: firestore-query-collection
115 | source: orders-firestore
116 | description: Query customer orders with complex filters
117 | authRequired:
118 | - user-auth
119 | - admin-auth
120 | `
121 | want := server.ToolConfigs{
122 | "query_users": firestorequerycollection.Config{
123 | Name: "query_users",
124 | Kind: "firestore-query-collection",
125 | Source: "users-firestore",
126 | Description: "Query user documents with filtering",
127 | AuthRequired: []string{"user-auth"},
128 | },
129 | "query_products": firestorequerycollection.Config{
130 | Name: "query_products",
131 | Kind: "firestore-query-collection",
132 | Source: "products-firestore",
133 | Description: "Query product catalog",
134 | AuthRequired: []string{},
135 | },
136 | "query_orders": firestorequerycollection.Config{
137 | Name: "query_orders",
138 | Kind: "firestore-query-collection",
139 | Source: "orders-firestore",
140 | Description: "Query customer orders with complex filters",
141 | AuthRequired: []string{"user-auth", "admin-auth"},
142 | },
143 | }
144 |
145 | got := struct {
146 | Tools server.ToolConfigs `yaml:"tools"`
147 | }{}
148 | // Parse contents
149 | err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
150 | if err != nil {
151 | t.Fatalf("unable to unmarshal: %s", err)
152 | }
153 | if diff := cmp.Diff(want, got.Tools); diff != "" {
154 | t.Fatalf("incorrect parse: diff %v", diff)
155 | }
156 | }
157 |
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestorevalidaterules/firestorevalidaterules_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package firestorevalidaterules_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/testutils"
24 | "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestorevalidaterules"
25 | )
26 |
27 | func TestParseFromYamlFirestoreValidateRules(t *testing.T) {
28 | ctx, err := testutils.ContextWithNewLogger()
29 | if err != nil {
30 | t.Fatalf("unexpected error: %s", err)
31 | }
32 | tcs := []struct {
33 | desc string
34 | in string
35 | want server.ToolConfigs
36 | }{
37 | {
38 | desc: "basic example",
39 | in: `
40 | tools:
41 | validate_rules_tool:
42 | kind: firestore-validate-rules
43 | source: my-firestore-instance
44 | description: Validate Firestore security rules
45 | `,
46 | want: server.ToolConfigs{
47 | "validate_rules_tool": firestorevalidaterules.Config{
48 | Name: "validate_rules_tool",
49 | Kind: "firestore-validate-rules",
50 | Source: "my-firestore-instance",
51 | Description: "Validate Firestore security rules",
52 | AuthRequired: []string{},
53 | },
54 | },
55 | },
56 | {
57 | desc: "with auth requirements",
58 | in: `
59 | tools:
60 | secure_validate_rules:
61 | kind: firestore-validate-rules
62 | source: prod-firestore
63 | description: Validate rules with authentication
64 | authRequired:
65 | - google-auth-service
66 | - api-key-service
67 | `,
68 | want: server.ToolConfigs{
69 | "secure_validate_rules": firestorevalidaterules.Config{
70 | Name: "secure_validate_rules",
71 | Kind: "firestore-validate-rules",
72 | Source: "prod-firestore",
73 | Description: "Validate rules with authentication",
74 | AuthRequired: []string{"google-auth-service", "api-key-service"},
75 | },
76 | },
77 | },
78 | }
79 | for _, tc := range tcs {
80 | t.Run(tc.desc, func(t *testing.T) {
81 | got := struct {
82 | Tools server.ToolConfigs `yaml:"tools"`
83 | }{}
84 | // Parse contents
85 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
86 | if err != nil {
87 | t.Fatalf("unable to unmarshal: %s", err)
88 | }
89 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
90 | t.Fatalf("incorrect parse: diff %v", diff)
91 | }
92 | })
93 | }
94 | }
95 |
96 | func TestParseFromYamlMultipleTools(t *testing.T) {
97 | ctx, err := testutils.ContextWithNewLogger()
98 | if err != nil {
99 | t.Fatalf("unexpected error: %s", err)
100 | }
101 | in := `
102 | tools:
103 | validate_dev_rules:
104 | kind: firestore-validate-rules
105 | source: dev-firestore
106 | description: Validate development environment rules
107 | authRequired:
108 | - dev-auth
109 | validate_staging_rules:
110 | kind: firestore-validate-rules
111 | source: staging-firestore
112 | description: Validate staging environment rules
113 | validate_prod_rules:
114 | kind: firestore-validate-rules
115 | source: prod-firestore
116 | description: Validate production environment rules
117 | authRequired:
118 | - prod-auth
119 | - admin-auth
120 | `
121 | want := server.ToolConfigs{
122 | "validate_dev_rules": firestorevalidaterules.Config{
123 | Name: "validate_dev_rules",
124 | Kind: "firestore-validate-rules",
125 | Source: "dev-firestore",
126 | Description: "Validate development environment rules",
127 | AuthRequired: []string{"dev-auth"},
128 | },
129 | "validate_staging_rules": firestorevalidaterules.Config{
130 | Name: "validate_staging_rules",
131 | Kind: "firestore-validate-rules",
132 | Source: "staging-firestore",
133 | Description: "Validate staging environment rules",
134 | AuthRequired: []string{},
135 | },
136 | "validate_prod_rules": firestorevalidaterules.Config{
137 | Name: "validate_prod_rules",
138 | Kind: "firestore-validate-rules",
139 | Source: "prod-firestore",
140 | Description: "Validate production environment rules",
141 | AuthRequired: []string{"prod-auth", "admin-auth"},
142 | },
143 | }
144 |
145 | got := struct {
146 | Tools server.ToolConfigs `yaml:"tools"`
147 | }{}
148 | // Parse contents
149 | err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
150 | if err != nil {
151 | t.Fatalf("unable to unmarshal: %s", err)
152 | }
153 | if diff := cmp.Diff(want, got.Tools); diff != "" {
154 | t.Fatalf("incorrect parse: diff %v", diff)
155 | }
156 | }
157 |
```
--------------------------------------------------------------------------------
/internal/tools/postgres/postgreslistavailableextensions/postgreslistavailableextensions.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package postgreslistavailableextensions
16 |
17 | import (
18 | "context"
19 | "fmt"
20 |
21 | yaml "github.com/goccy/go-yaml"
22 | "github.com/googleapis/genai-toolbox/internal/sources"
23 | "github.com/googleapis/genai-toolbox/internal/sources/alloydbpg"
24 | "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlpg"
25 | "github.com/googleapis/genai-toolbox/internal/sources/postgres"
26 | "github.com/googleapis/genai-toolbox/internal/tools"
27 | "github.com/jackc/pgx/v5/pgxpool"
28 | )
29 |
30 | const kind string = "postgres-list-available-extensions"
31 |
32 | const listAvailableExtensionsQuery = `
33 | SELECT
34 | name,
35 | default_version,
36 | comment as description
37 | FROM
38 | pg_available_extensions
39 | ORDER BY name;
40 | `
41 |
42 | func init() {
43 | if !tools.Register(kind, newConfig) {
44 | panic(fmt.Sprintf("tool kind %q already registered", kind))
45 | }
46 | }
47 |
48 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
49 | actual := Config{Name: name}
50 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
51 | return nil, err
52 | }
53 | return actual, nil
54 | }
55 |
56 | type compatibleSource interface {
57 | PostgresPool() *pgxpool.Pool
58 | }
59 |
60 | // validate compatible sources are still compatible
61 | var _ compatibleSource = &alloydbpg.Source{}
62 | var _ compatibleSource = &cloudsqlpg.Source{}
63 | var _ compatibleSource = &postgres.Source{}
64 |
65 | var compatibleSources = [...]string{alloydbpg.SourceKind, cloudsqlpg.SourceKind, postgres.SourceKind}
66 |
67 | type Config struct {
68 | Name string `yaml:"name" validate:"required"`
69 | Kind string `yaml:"kind" validate:"required"`
70 | Source string `yaml:"source" validate:"required"`
71 | Description string `yaml:"description" validate:"required"`
72 | AuthRequired []string `yaml:"authRequired"`
73 | }
74 |
75 | // validate interface
76 | var _ tools.ToolConfig = Config{}
77 |
78 | func (cfg Config) ToolConfigKind() string {
79 | return kind
80 | }
81 |
82 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
83 | // verify source exists
84 | rawS, ok := srcs[cfg.Source]
85 | if !ok {
86 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
87 | }
88 |
89 | // verify the source is compatible
90 | s, ok := rawS.(compatibleSource)
91 | if !ok {
92 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
93 | }
94 |
95 | parameters := tools.Parameters{}
96 | mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters)
97 |
98 | // finish tool setup
99 | t := Tool{
100 | Name: cfg.Name,
101 | Kind: cfg.Kind,
102 | AuthRequired: cfg.AuthRequired,
103 | Pool: s.PostgresPool(),
104 | manifest: tools.Manifest{
105 | Description: cfg.Description,
106 | Parameters: parameters.Manifest(),
107 | AuthRequired: cfg.AuthRequired,
108 | },
109 | mcpManifest: mcpManifest,
110 | }
111 | return t, nil
112 | }
113 |
114 | // validate interface
115 | var _ tools.Tool = Tool{}
116 |
117 | type Tool struct {
118 | Name string `yaml:"name"`
119 | Kind string `yaml:"kind"`
120 | AuthRequired []string `yaml:"authRequired"`
121 | Pool *pgxpool.Pool
122 | manifest tools.Manifest
123 | mcpManifest tools.McpManifest
124 | }
125 |
126 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
127 | results, err := t.Pool.Query(ctx, listAvailableExtensionsQuery)
128 | if err != nil {
129 | return nil, fmt.Errorf("unable to execute query: %w", err)
130 | }
131 |
132 | fields := results.FieldDescriptions()
133 |
134 | var out []any
135 | for results.Next() {
136 | v, err := results.Values()
137 | if err != nil {
138 | return nil, fmt.Errorf("unable to parse row: %w", err)
139 | }
140 | vMap := make(map[string]any)
141 | for i, f := range fields {
142 | vMap[f.Name] = v[i]
143 | }
144 | out = append(out, vMap)
145 | }
146 |
147 | return out, nil
148 | }
149 |
150 | func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
151 | return tools.ParamValues{}, nil
152 | }
153 |
154 | func (t Tool) Manifest() tools.Manifest {
155 | return t.manifest
156 | }
157 |
158 | func (t Tool) McpManifest() tools.McpManifest {
159 | return t.mcpManifest
160 | }
161 |
162 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
163 | return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
164 | }
165 |
166 | func (t Tool) RequiresClientAuthorization() bool {
167 | return false
168 | }
169 |
```
--------------------------------------------------------------------------------
/internal/tools/dgraph/dgraph.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package dgraph
16 |
17 | import (
18 | "context"
19 | "encoding/json"
20 | "fmt"
21 |
22 | yaml "github.com/goccy/go-yaml"
23 | "github.com/googleapis/genai-toolbox/internal/sources"
24 | "github.com/googleapis/genai-toolbox/internal/sources/dgraph"
25 | "github.com/googleapis/genai-toolbox/internal/tools"
26 | )
27 |
28 | const kind string = "dgraph-dql"
29 |
30 | func init() {
31 | if !tools.Register(kind, newConfig) {
32 | panic(fmt.Sprintf("tool kind %q already registered", kind))
33 | }
34 | }
35 |
36 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
37 | actual := Config{Name: name}
38 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
39 | return nil, err
40 | }
41 | return actual, nil
42 | }
43 |
44 | type compatibleSource interface {
45 | DgraphClient() *dgraph.DgraphClient
46 | }
47 |
48 | // validate compatible sources are still compatible
49 | var _ compatibleSource = &dgraph.Source{}
50 |
51 | var compatibleSources = [...]string{dgraph.SourceKind}
52 |
53 | type Config struct {
54 | Name string `yaml:"name" validate:"required"`
55 | Kind string `yaml:"kind" validate:"required"`
56 | Source string `yaml:"source" validate:"required"`
57 | Description string `yaml:"description" validate:"required"`
58 | Statement string `yaml:"statement" validate:"required"`
59 | AuthRequired []string `yaml:"authRequired"`
60 | IsQuery bool `yaml:"isQuery"`
61 | Timeout string `yaml:"timeout"`
62 | Parameters tools.Parameters `yaml:"parameters"`
63 | }
64 |
65 | // validate interface
66 | var _ tools.ToolConfig = Config{}
67 |
68 | func (cfg Config) ToolConfigKind() string {
69 | return kind
70 | }
71 |
72 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
73 | // verify source exists
74 | rawS, ok := srcs[cfg.Source]
75 | if !ok {
76 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
77 | }
78 |
79 | // verify the source is compatible
80 | s, ok := rawS.(compatibleSource)
81 | if !ok {
82 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
83 | }
84 |
85 | mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, cfg.Parameters)
86 |
87 | // finish tool setup
88 | t := Tool{
89 | Name: cfg.Name,
90 | Kind: kind,
91 | Parameters: cfg.Parameters,
92 | Statement: cfg.Statement,
93 | AuthRequired: cfg.AuthRequired,
94 | DgraphClient: s.DgraphClient(),
95 | IsQuery: cfg.IsQuery,
96 | Timeout: cfg.Timeout,
97 | manifest: tools.Manifest{Description: cfg.Description, Parameters: cfg.Parameters.Manifest(), AuthRequired: cfg.AuthRequired},
98 | mcpManifest: mcpManifest,
99 | }
100 | return t, nil
101 | }
102 |
103 | // validate interface
104 | var _ tools.Tool = Tool{}
105 |
106 | type Tool struct {
107 | Name string `yaml:"name"`
108 | Kind string `yaml:"kind"`
109 | Parameters tools.Parameters `yaml:"parameters"`
110 | AuthRequired []string `yaml:"authRequired"`
111 | DgraphClient *dgraph.DgraphClient
112 | IsQuery bool
113 | Timeout string
114 | Statement string
115 | manifest tools.Manifest
116 | mcpManifest tools.McpManifest
117 | }
118 |
119 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
120 | paramsMap := params.AsMapWithDollarPrefix()
121 |
122 | resp, err := t.DgraphClient.ExecuteQuery(t.Statement, paramsMap, t.IsQuery, t.Timeout)
123 | if err != nil {
124 | return nil, err
125 | }
126 |
127 | if err := dgraph.CheckError(resp); err != nil {
128 | return nil, err
129 | }
130 |
131 | var result struct {
132 | Data map[string]interface{} `json:"data"`
133 | }
134 |
135 | if err := json.Unmarshal(resp, &result); err != nil {
136 | return nil, fmt.Errorf("error parsing JSON: %v", err)
137 | }
138 |
139 | return result.Data, nil
140 | }
141 |
142 | func (t Tool) ParseParams(data map[string]any, claimsMap map[string]map[string]any) (tools.ParamValues, error) {
143 | return tools.ParseParams(t.Parameters, data, claimsMap)
144 | }
145 |
146 | func (t Tool) Manifest() tools.Manifest {
147 | return t.manifest
148 | }
149 |
150 | func (t Tool) McpManifest() tools.McpManifest {
151 | return t.mcpManifest
152 | }
153 |
154 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
155 | return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
156 | }
157 |
158 | func (t Tool) RequiresClientAuthorization() bool {
159 | return false
160 | }
161 |
```
--------------------------------------------------------------------------------
/internal/tools/tools.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2024 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package tools
16 |
17 | import (
18 | "context"
19 | "errors"
20 | "fmt"
21 | "slices"
22 | "strings"
23 |
24 | yaml "github.com/goccy/go-yaml"
25 | "github.com/googleapis/genai-toolbox/internal/sources"
26 | )
27 |
28 | // ToolConfigFactory defines the signature for a function that creates and
29 | // decodes a specific tool's configuration. It takes the context, the tool's
30 | // name, and a YAML decoder to parse the config.
31 | type ToolConfigFactory func(ctx context.Context, name string, decoder *yaml.Decoder) (ToolConfig, error)
32 |
33 | var toolRegistry = make(map[string]ToolConfigFactory)
34 |
35 | // Register allows individual tool packages to register their configuration
36 | // factory function. This is typically called from an init() function in the
37 | // tool's package. It associates a 'kind' string with a function that can
38 | // produce the specific ToolConfig type. It returns true if the registration was
39 | // successful, and false if a tool with the same kind was already registered.
40 | func Register(kind string, factory ToolConfigFactory) bool {
41 | if _, exists := toolRegistry[kind]; exists {
42 | // Tool with this kind already exists, do not overwrite.
43 | return false
44 | }
45 | toolRegistry[kind] = factory
46 | return true
47 | }
48 |
49 | // DecodeConfig looks up the registered factory for the given kind and uses it
50 | // to decode the tool configuration.
51 | func DecodeConfig(ctx context.Context, kind string, name string, decoder *yaml.Decoder) (ToolConfig, error) {
52 | factory, found := toolRegistry[kind]
53 | if !found {
54 | return nil, fmt.Errorf("unknown tool kind: %q", kind)
55 | }
56 | toolConfig, err := factory(ctx, name, decoder)
57 | if err != nil {
58 | return nil, fmt.Errorf("unable to parse tool %q as kind %q: %w", name, kind, err)
59 | }
60 | return toolConfig, nil
61 | }
62 |
63 | type ToolConfig interface {
64 | ToolConfigKind() string
65 | Initialize(map[string]sources.Source) (Tool, error)
66 | }
67 |
68 | type AccessToken string
69 |
70 | func (token AccessToken) ParseBearerToken() (string, error) {
71 | headerParts := strings.Split(string(token), " ")
72 | if len(headerParts) != 2 || strings.ToLower(headerParts[0]) != "bearer" {
73 | return "", fmt.Errorf("authorization header must be in the format 'Bearer <token>': %w", ErrUnauthorized)
74 | }
75 | return headerParts[1], nil
76 | }
77 |
78 | type Tool interface {
79 | Invoke(context.Context, ParamValues, AccessToken) (any, error)
80 | ParseParams(map[string]any, map[string]map[string]any) (ParamValues, error)
81 | Manifest() Manifest
82 | McpManifest() McpManifest
83 | Authorized([]string) bool
84 | RequiresClientAuthorization() bool
85 | }
86 |
87 | // Manifest is the representation of tools sent to Client SDKs.
88 | type Manifest struct {
89 | Description string `json:"description"`
90 | Parameters []ParameterManifest `json:"parameters"`
91 | AuthRequired []string `json:"authRequired"`
92 | }
93 |
94 | // Definition for a tool the MCP client can call.
95 | type McpManifest struct {
96 | // The name of the tool.
97 | Name string `json:"name"`
98 | // A human-readable description of the tool.
99 | Description string `json:"description,omitempty"`
100 | // A JSON Schema object defining the expected parameters for the tool.
101 | InputSchema McpToolsSchema `json:"inputSchema,omitempty"`
102 | Metadata map[string]any `json:"_meta,omitempty"`
103 | }
104 |
105 | func GetMcpManifest(name, desc string, authInvoke []string, params Parameters) McpManifest {
106 | inputSchema, authParams := params.McpManifest()
107 | mcpManifest := McpManifest{
108 | Name: name,
109 | Description: desc,
110 | InputSchema: inputSchema,
111 | }
112 |
113 | // construct metadata, if applicable
114 | metadata := make(map[string]any)
115 | if len(authInvoke) > 0 {
116 | metadata["toolbox/authInvoke"] = authInvoke
117 | }
118 | if len(authParams) > 0 {
119 | metadata["toolbox/authParam"] = authParams
120 | }
121 | if len(metadata) > 0 {
122 | mcpManifest.Metadata = metadata
123 | }
124 | return mcpManifest
125 | }
126 |
127 | var ErrUnauthorized = errors.New("unauthorized")
128 |
129 | // Helper function that returns if a tool invocation request is authorized
130 | func IsAuthorized(authRequiredSources []string, verifiedAuthServices []string) bool {
131 | if len(authRequiredSources) == 0 {
132 | // no authorization requirement
133 | return true
134 | }
135 | for _, a := range authRequiredSources {
136 | if slices.Contains(verifiedAuthServices, a) {
137 | return true
138 | }
139 | }
140 | return false
141 | }
142 |
```
--------------------------------------------------------------------------------
/internal/tools/cloudsql/cloudsqlgetinstances/cloudsqlgetinstances.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package cloudsqlgetinstances
16 |
17 | import (
18 | "context"
19 | "fmt"
20 |
21 | yaml "github.com/goccy/go-yaml"
22 | "github.com/googleapis/genai-toolbox/internal/sources"
23 | "github.com/googleapis/genai-toolbox/internal/sources/cloudsqladmin"
24 | "github.com/googleapis/genai-toolbox/internal/tools"
25 | )
26 |
27 | const kind string = "cloud-sql-get-instance"
28 |
29 | func init() {
30 | if !tools.Register(kind, newConfig) {
31 | panic(fmt.Sprintf("tool kind %q already registered", kind))
32 | }
33 | }
34 |
35 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
36 | actual := Config{Name: name}
37 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
38 | return nil, err
39 | }
40 | return actual, nil
41 | }
42 |
43 | // Config defines the configuration for the get-instances tool.
44 | type Config struct {
45 | Name string `yaml:"name" validate:"required"`
46 | Kind string `yaml:"kind" validate:"required"`
47 | Description string `yaml:"description"`
48 | Source string `yaml:"source" validate:"required"`
49 | AuthRequired []string `yaml:"authRequired"`
50 | }
51 |
52 | // validate interface
53 | var _ tools.ToolConfig = Config{}
54 |
55 | // ToolConfigKind returns the kind of the tool.
56 | func (cfg Config) ToolConfigKind() string {
57 | return kind
58 | }
59 |
60 | // Initialize initializes the tool from the configuration.
61 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
62 | rawS, ok := srcs[cfg.Source]
63 | if !ok {
64 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
65 | }
66 |
67 | s, ok := rawS.(*cloudsqladmin.Source)
68 | if !ok {
69 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be `cloud-sql-admin`", kind)
70 | }
71 |
72 | allParameters := tools.Parameters{
73 | tools.NewStringParameter("projectId", "The project ID"),
74 | tools.NewStringParameter("instanceId", "The instance ID"),
75 | }
76 | paramManifest := allParameters.Manifest()
77 |
78 | description := cfg.Description
79 | if description == "" {
80 | description = "Gets a particular cloud sql instance."
81 | }
82 | mcpManifest := tools.GetMcpManifest(cfg.Name, description, cfg.AuthRequired, allParameters)
83 |
84 | return Tool{
85 | Name: cfg.Name,
86 | Kind: kind,
87 | AuthRequired: cfg.AuthRequired,
88 | Source: s,
89 | AllParams: allParameters,
90 | manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
91 | mcpManifest: mcpManifest,
92 | }, nil
93 | }
94 |
95 | // Tool represents the get-instances tool.
96 | type Tool struct {
97 | Name string `yaml:"name"`
98 | Kind string `yaml:"kind"`
99 | Description string `yaml:"description"`
100 | AuthRequired []string `yaml:"authRequired"`
101 |
102 | Source *cloudsqladmin.Source
103 | AllParams tools.Parameters `yaml:"allParams"`
104 | manifest tools.Manifest
105 | mcpManifest tools.McpManifest
106 | }
107 |
108 | // Invoke executes the tool's logic.
109 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
110 | paramsMap := params.AsMap()
111 |
112 | projectId, ok := paramsMap["projectId"].(string)
113 | if !ok {
114 | return nil, fmt.Errorf("missing 'projectId' parameter")
115 | }
116 | instanceId, ok := paramsMap["instanceId"].(string)
117 | if !ok {
118 | return nil, fmt.Errorf("missing 'instanceId' parameter")
119 | }
120 |
121 | service, err := t.Source.GetService(ctx, string(accessToken))
122 | if err != nil {
123 | return nil, err
124 | }
125 |
126 | resp, err := service.Instances.Get(projectId, instanceId).Do()
127 | if err != nil {
128 | return nil, fmt.Errorf("error getting instance: %w", err)
129 | }
130 |
131 | return resp, nil
132 | }
133 |
134 | // ParseParams parses the parameters for the tool.
135 | func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
136 | return tools.ParseParams(t.AllParams, data, claims)
137 | }
138 |
139 | // Manifest returns the tool's manifest.
140 | func (t Tool) Manifest() tools.Manifest {
141 | return t.manifest
142 | }
143 |
144 | // McpManifest returns the tool's MCP manifest.
145 | func (t Tool) McpManifest() tools.McpManifest {
146 | return t.mcpManifest
147 | }
148 |
149 | // Authorized checks if the tool is authorized.
150 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
151 | return true
152 | }
153 |
154 | func (t Tool) RequiresClientAuthorization() bool {
155 | return t.Source.UseClientAuthorization()
156 | }
157 |
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestorelistcollections/firestorelistcollections_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package firestorelistcollections_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/testutils"
24 | "github.com/googleapis/genai-toolbox/internal/tools/firestore/firestorelistcollections"
25 | )
26 |
27 | func TestParseFromYamlFirestoreListCollections(t *testing.T) {
28 | ctx, err := testutils.ContextWithNewLogger()
29 | if err != nil {
30 | t.Fatalf("unexpected error: %s", err)
31 | }
32 | tcs := []struct {
33 | desc string
34 | in string
35 | want server.ToolConfigs
36 | }{
37 | {
38 | desc: "basic example",
39 | in: `
40 | tools:
41 | list_collections_tool:
42 | kind: firestore-list-collections
43 | source: my-firestore-instance
44 | description: List collections in Firestore
45 | `,
46 | want: server.ToolConfigs{
47 | "list_collections_tool": firestorelistcollections.Config{
48 | Name: "list_collections_tool",
49 | Kind: "firestore-list-collections",
50 | Source: "my-firestore-instance",
51 | Description: "List collections in Firestore",
52 | AuthRequired: []string{},
53 | },
54 | },
55 | },
56 | {
57 | desc: "with auth requirements",
58 | in: `
59 | tools:
60 | secure_list_collections:
61 | kind: firestore-list-collections
62 | source: prod-firestore
63 | description: List collections with authentication
64 | authRequired:
65 | - google-auth-service
66 | - api-key-service
67 | `,
68 | want: server.ToolConfigs{
69 | "secure_list_collections": firestorelistcollections.Config{
70 | Name: "secure_list_collections",
71 | Kind: "firestore-list-collections",
72 | Source: "prod-firestore",
73 | Description: "List collections with authentication",
74 | AuthRequired: []string{"google-auth-service", "api-key-service"},
75 | },
76 | },
77 | },
78 | }
79 | for _, tc := range tcs {
80 | t.Run(tc.desc, func(t *testing.T) {
81 | got := struct {
82 | Tools server.ToolConfigs `yaml:"tools"`
83 | }{}
84 | // Parse contents
85 | err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got)
86 | if err != nil {
87 | t.Fatalf("unable to unmarshal: %s", err)
88 | }
89 | if diff := cmp.Diff(tc.want, got.Tools); diff != "" {
90 | t.Fatalf("incorrect parse: diff %v", diff)
91 | }
92 | })
93 | }
94 | }
95 |
96 | func TestParseFromYamlMultipleTools(t *testing.T) {
97 | ctx, err := testutils.ContextWithNewLogger()
98 | if err != nil {
99 | t.Fatalf("unexpected error: %s", err)
100 | }
101 | in := `
102 | tools:
103 | list_user_collections:
104 | kind: firestore-list-collections
105 | source: users-firestore
106 | description: List user-related collections
107 | authRequired:
108 | - user-auth
109 | list_product_collections:
110 | kind: firestore-list-collections
111 | source: products-firestore
112 | description: List product-related collections
113 | list_admin_collections:
114 | kind: firestore-list-collections
115 | source: admin-firestore
116 | description: List administrative collections
117 | authRequired:
118 | - user-auth
119 | - admin-auth
120 | `
121 | want := server.ToolConfigs{
122 | "list_user_collections": firestorelistcollections.Config{
123 | Name: "list_user_collections",
124 | Kind: "firestore-list-collections",
125 | Source: "users-firestore",
126 | Description: "List user-related collections",
127 | AuthRequired: []string{"user-auth"},
128 | },
129 | "list_product_collections": firestorelistcollections.Config{
130 | Name: "list_product_collections",
131 | Kind: "firestore-list-collections",
132 | Source: "products-firestore",
133 | Description: "List product-related collections",
134 | AuthRequired: []string{},
135 | },
136 | "list_admin_collections": firestorelistcollections.Config{
137 | Name: "list_admin_collections",
138 | Kind: "firestore-list-collections",
139 | Source: "admin-firestore",
140 | Description: "List administrative collections",
141 | AuthRequired: []string{"user-auth", "admin-auth"},
142 | },
143 | }
144 |
145 | got := struct {
146 | Tools server.ToolConfigs `yaml:"tools"`
147 | }{}
148 | // Parse contents
149 | err = yaml.UnmarshalContext(ctx, testutils.FormatYaml(in), &got)
150 | if err != nil {
151 | t.Fatalf("unable to unmarshal: %s", err)
152 | }
153 | if diff := cmp.Diff(want, got.Tools); diff != "" {
154 | t.Fatalf("incorrect parse: diff %v", diff)
155 | }
156 | }
157 |
```
--------------------------------------------------------------------------------
/internal/sources/cloudsqlmssql/cloud_sql_mssql_test.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package cloudsqlmssql_test
16 |
17 | import (
18 | "testing"
19 |
20 | yaml "github.com/goccy/go-yaml"
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/server"
23 | "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmssql"
24 | "github.com/googleapis/genai-toolbox/internal/testutils"
25 | )
26 |
27 | func TestParseFromYamlCloudSQLMssql(t *testing.T) {
28 | tcs := []struct {
29 | desc string
30 | in string
31 | want server.SourceConfigs
32 | }{
33 | {
34 | desc: "basic example",
35 | in: `
36 | sources:
37 | my-instance:
38 | kind: cloud-sql-mssql
39 | project: my-project
40 | region: my-region
41 | instance: my-instance
42 | database: my_db
43 | ipAddress: localhost
44 | user: my_user
45 | password: my_pass
46 | `,
47 | want: server.SourceConfigs{
48 | "my-instance": cloudsqlmssql.Config{
49 | Name: "my-instance",
50 | Kind: cloudsqlmssql.SourceKind,
51 | Project: "my-project",
52 | Region: "my-region",
53 | Instance: "my-instance",
54 | IPAddress: "localhost",
55 | IPType: "public",
56 | Database: "my_db",
57 | User: "my_user",
58 | Password: "my_pass",
59 | },
60 | },
61 | },
62 | {
63 | desc: "psc ipType",
64 | in: `
65 | sources:
66 | my-instance:
67 | kind: cloud-sql-mssql
68 | project: my-project
69 | region: my-region
70 | instance: my-instance
71 | database: my_db
72 | ipAddress: localhost
73 | user: my_user
74 | password: my_pass
75 | ipType: psc
76 | `,
77 | want: server.SourceConfigs{
78 | "my-instance": cloudsqlmssql.Config{
79 | Name: "my-instance",
80 | Kind: cloudsqlmssql.SourceKind,
81 | Project: "my-project",
82 | Region: "my-region",
83 | Instance: "my-instance",
84 | IPAddress: "localhost",
85 | IPType: "psc",
86 | Database: "my_db",
87 | User: "my_user",
88 | Password: "my_pass",
89 | },
90 | },
91 | },
92 | }
93 | for _, tc := range tcs {
94 | t.Run(tc.desc, func(t *testing.T) {
95 | got := struct {
96 | Sources server.SourceConfigs `yaml:"sources"`
97 | }{}
98 | // Parse contents
99 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
100 | if err != nil {
101 | t.Fatalf("unable to unmarshal: %s", err)
102 | }
103 | if !cmp.Equal(tc.want, got.Sources) {
104 | t.Fatalf("incorrect psarse: want %v, got %v", tc.want, got.Sources)
105 | }
106 | })
107 | }
108 |
109 | }
110 |
111 | func TestFailParseFromYaml(t *testing.T) {
112 | tcs := []struct {
113 | desc string
114 | in string
115 | err string
116 | }{
117 | {
118 | desc: "invalid ipType",
119 | in: `
120 | sources:
121 | my-instance:
122 | kind: cloud-sql-mssql
123 | project: my-project
124 | region: my-region
125 | instance: my-instance
126 | ipType: fail
127 | database: my_db
128 | ipAddress: localhost
129 | user: my_user
130 | password: my_pass
131 | `,
132 | err: "unable to parse source \"my-instance\" as \"cloud-sql-mssql\": ipType invalid: must be one of \"public\", \"private\", or \"psc\"",
133 | },
134 | {
135 | desc: "extra field",
136 | in: `
137 | sources:
138 | my-instance:
139 | kind: cloud-sql-mssql
140 | project: my-project
141 | region: my-region
142 | instance: my-instance
143 | database: my_db
144 | ipAddress: localhost
145 | user: my_user
146 | password: my_pass
147 | foo: bar
148 | `,
149 | err: "unable to parse source \"my-instance\" as \"cloud-sql-mssql\": [2:1] unknown field \"foo\"\n 1 | database: my_db\n> 2 | foo: bar\n ^\n 3 | instance: my-instance\n 4 | ipAddress: localhost\n 5 | kind: cloud-sql-mssql\n 6 | ",
150 | },
151 | {
152 | desc: "missing required field",
153 | in: `
154 | sources:
155 | my-instance:
156 | kind: cloud-sql-mssql
157 | region: my-region
158 | instance: my-instance
159 | database: my_db
160 | ipAddress: localhost
161 | user: my_user
162 | password: my_pass
163 | `,
164 | err: "unable to parse source \"my-instance\" as \"cloud-sql-mssql\": Key: 'Config.Project' Error:Field validation for 'Project' failed on the 'required' tag",
165 | },
166 | }
167 | for _, tc := range tcs {
168 | t.Run(tc.desc, func(t *testing.T) {
169 | got := struct {
170 | Sources server.SourceConfigs `yaml:"sources"`
171 | }{}
172 | // Parse contents
173 | err := yaml.Unmarshal(testutils.FormatYaml(tc.in), &got)
174 | if err == nil {
175 | t.Fatalf("expect parsing to fail")
176 | }
177 | errStr := err.Error()
178 | if errStr != tc.err {
179 | t.Fatalf("unexpected error: got %q, want %q", errStr, tc.err)
180 | }
181 | })
182 | }
183 | }
184 |
```