This is page 27 of 48. 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
│ │ │ ├── serverless-spark.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-create-project-file.md
│ │ │ ├── looker-delete-project-file.md
│ │ │ ├── looker-dev-mode.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-get-project-file.md
│ │ │ ├── looker-get-project-files.md
│ │ │ ├── looker-get-projects.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
│ │ │ └── looker-update-project-file.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
│ │ ├── serverless-spark
│ │ │ ├── _index.md
│ │ │ └── serverless-spark-list-batches.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
│ │ ├── serverless-spark.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
│ │ ├── serverlessspark
│ │ │ ├── serverlessspark_test.go
│ │ │ └── serverlessspark.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
│ │ │ ├── lookercreateprojectfile
│ │ │ │ ├── lookercreateprojectfile_test.go
│ │ │ │ └── lookercreateprojectfile.go
│ │ │ ├── lookerdeleteprojectfile
│ │ │ │ ├── lookerdeleteprojectfile_test.go
│ │ │ │ └── lookerdeleteprojectfile.go
│ │ │ ├── lookerdevmode
│ │ │ │ ├── lookerdevmode_test.go
│ │ │ │ └── lookerdevmode.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
│ │ │ ├── lookergetprojectfile
│ │ │ │ ├── lookergetprojectfile_test.go
│ │ │ │ └── lookergetprojectfile.go
│ │ │ ├── lookergetprojectfiles
│ │ │ │ ├── lookergetprojectfiles_test.go
│ │ │ │ └── lookergetprojectfiles.go
│ │ │ ├── lookergetprojects
│ │ │ │ ├── lookergetprojects_test.go
│ │ │ │ └── lookergetprojects.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
│ │ │ └── lookerupdateprojectfile
│ │ │ ├── lookerupdateprojectfile_test.go
│ │ │ └── lookerupdateprojectfile.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
│ │ ├── serverlessspark
│ │ │ └── serverlesssparklistbatches
│ │ │ ├── serverlesssparklistbatches_test.go
│ │ │ └── serverlesssparklistbatches.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
├── serverlessspark
│ └── serverless_spark_integration_test.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
--------------------------------------------------------------------------------
/tests/serverlessspark/serverless_spark_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 serverlessspark
16 |
17 | import (
18 | "bytes"
19 | "context"
20 | "encoding/json"
21 | "fmt"
22 | "io"
23 | "net/http"
24 | "os"
25 | "reflect"
26 | "regexp"
27 | "testing"
28 | "time"
29 |
30 | dataproc "cloud.google.com/go/dataproc/v2/apiv1"
31 | "cloud.google.com/go/dataproc/v2/apiv1/dataprocpb"
32 | "github.com/googleapis/genai-toolbox/internal/testutils"
33 | "github.com/googleapis/genai-toolbox/internal/tools/serverlessspark/serverlesssparklistbatches"
34 | "github.com/googleapis/genai-toolbox/tests"
35 | "google.golang.org/api/iterator"
36 | "google.golang.org/api/option"
37 | )
38 |
39 | var (
40 | serverlessSparkProject = os.Getenv("SERVERLESS_SPARK_PROJECT")
41 | serverlessSparkLocation = os.Getenv("SERVERLESS_SPARK_LOCATION")
42 | )
43 |
44 | func getServerlessSparkVars(t *testing.T) map[string]any {
45 | switch "" {
46 | case serverlessSparkProject:
47 | t.Fatal("'SERVERLESS_SPARK_PROJECT' not set")
48 | case serverlessSparkLocation:
49 | t.Fatal("'SERVERLESS_SPARK_LOCATION' not set")
50 | }
51 |
52 | return map[string]any{
53 | "kind": "serverless-spark",
54 | "project": serverlessSparkProject,
55 | "location": serverlessSparkLocation,
56 | }
57 | }
58 |
59 | func TestServerlessSparkToolEndpoints(t *testing.T) {
60 | sourceConfig := getServerlessSparkVars(t)
61 | ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
62 | defer cancel()
63 |
64 | toolsFile := map[string]any{
65 | "sources": map[string]any{
66 | "my-spark": sourceConfig,
67 | },
68 | "authServices": map[string]any{
69 | "my-google-auth": map[string]any{
70 | "kind": "google",
71 | "clientId": tests.ClientId,
72 | },
73 | },
74 | "tools": map[string]any{
75 | "list-batches": map[string]any{
76 | "kind": "serverless-spark-list-batches",
77 | "source": "my-spark",
78 | },
79 | "list-batches-with-auth": map[string]any{
80 | "kind": "serverless-spark-list-batches",
81 | "source": "my-spark",
82 | "authRequired": []string{"my-google-auth"},
83 | },
84 | },
85 | }
86 |
87 | cmd, cleanup, err := tests.StartCmd(ctx, toolsFile)
88 | if err != nil {
89 | t.Fatalf("command initialization returned an error: %s", err)
90 | }
91 | defer cleanup()
92 |
93 | waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
94 | defer cancel()
95 | out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
96 | if err != nil {
97 | t.Logf("toolbox command logs: \n%s", out)
98 | t.Fatalf("toolbox didn't start successfully: %s", err)
99 | }
100 |
101 | endpoint := fmt.Sprintf("%s-dataproc.googleapis.com:443", serverlessSparkLocation)
102 | client, err := dataproc.NewBatchControllerClient(ctx, option.WithEndpoint(endpoint))
103 | if err != nil {
104 | t.Fatalf("failed to create dataproc client: %v", err)
105 | }
106 | defer client.Close()
107 |
108 | runListBatchesTest(t, client, ctx)
109 | runListBatchesErrorTest(t)
110 | runListBatchesAuthTest(t)
111 | }
112 |
113 | // runListBatchesTest invokes the running list-batches tool and ensures it returns the correct
114 | // number of results. It can run successfully against any GCP project that has at least 2 succeeded
115 | // or failed Serverless Spark batches, of any age.
116 | func runListBatchesTest(t *testing.T, client *dataproc.BatchControllerClient, ctx context.Context) {
117 | batch2 := listBatchesRpc(t, client, ctx, "", 2, true)
118 | batch20 := listBatchesRpc(t, client, ctx, "", 20, false)
119 |
120 | tcs := []struct {
121 | name string
122 | filter string
123 | pageSize int
124 | numPages int
125 | want []serverlesssparklistbatches.Batch
126 | }{
127 | {name: "one page", pageSize: 2, numPages: 1, want: batch2},
128 | {name: "two pages", pageSize: 1, numPages: 2, want: batch2},
129 | {name: "20 batches", pageSize: 20, numPages: 1, want: batch20},
130 | {name: "omit page size", numPages: 1, want: batch20},
131 | {
132 | name: "filtered",
133 | filter: "state = SUCCEEDED",
134 | pageSize: 2,
135 | numPages: 1,
136 | want: listBatchesRpc(t, client, ctx, "state = SUCCEEDED", 2, true),
137 | },
138 | {
139 | name: "empty",
140 | filter: "state = SUCCEEDED AND state = FAILED",
141 | pageSize: 1,
142 | numPages: 1,
143 | want: nil,
144 | },
145 | }
146 |
147 | for _, tc := range tcs {
148 | t.Run(tc.name, func(t *testing.T) {
149 | var actual []serverlesssparklistbatches.Batch
150 | var pageToken string
151 | for i := 0; i < tc.numPages; i++ {
152 | request := map[string]any{
153 | "filter": tc.filter,
154 | "pageToken": pageToken,
155 | }
156 | if tc.pageSize > 0 {
157 | request["pageSize"] = tc.pageSize
158 | }
159 |
160 | resp, err := invokeListBatches("list-batches", request, nil)
161 | if err != nil {
162 | t.Fatalf("invokeListBatches failed: %v", err)
163 | }
164 | defer resp.Body.Close()
165 |
166 | if resp.StatusCode != http.StatusOK {
167 | bodyBytes, _ := io.ReadAll(resp.Body)
168 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
169 | }
170 |
171 | var body map[string]any
172 | if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
173 | t.Fatalf("error parsing response body: %v", err)
174 | }
175 |
176 | result, ok := body["result"].(string)
177 | if !ok {
178 | t.Fatalf("unable to find result in response body")
179 | }
180 |
181 | var listResponse serverlesssparklistbatches.ListBatchesResponse
182 | if err := json.Unmarshal([]byte(result), &listResponse); err != nil {
183 | t.Fatalf("error unmarshalling result: %s", err)
184 | }
185 | actual = append(actual, listResponse.Batches...)
186 | pageToken = listResponse.NextPageToken
187 | }
188 |
189 | if !reflect.DeepEqual(actual, tc.want) {
190 | t.Fatalf("unexpected batches: got %+v, want %+v", actual, tc.want)
191 | }
192 | })
193 | }
194 | }
195 |
196 | func listBatchesRpc(t *testing.T, client *dataproc.BatchControllerClient, ctx context.Context, filter string, n int, exact bool) []serverlesssparklistbatches.Batch {
197 | parent := fmt.Sprintf("projects/%s/locations/%s", serverlessSparkProject, serverlessSparkLocation)
198 | req := &dataprocpb.ListBatchesRequest{
199 | Parent: parent,
200 | PageSize: 2,
201 | OrderBy: "create_time desc",
202 | }
203 | if filter != "" {
204 | req.Filter = filter
205 | }
206 |
207 | it := client.ListBatches(ctx, req)
208 | pager := iterator.NewPager(it, n, "")
209 | var batchPbs []*dataprocpb.Batch
210 | _, err := pager.NextPage(&batchPbs)
211 | if err != nil {
212 | t.Fatalf("failed to list batches: %s", err)
213 | }
214 | if exact && len(batchPbs) != n {
215 | t.Fatalf("expected exactly %d batches, got %d", n, len(batchPbs))
216 | }
217 | if !exact && (len(batchPbs) == 0 || len(batchPbs) > n) {
218 | t.Fatalf("expected between 1 and %d batches, got %d", n, len(batchPbs))
219 | }
220 |
221 | return serverlesssparklistbatches.ToBatches(batchPbs)
222 | }
223 |
224 | func runListBatchesErrorTest(t *testing.T) {
225 | tcs := []struct {
226 | name string
227 | pageSize int
228 | wantCode int
229 | wantMsg string
230 | }{
231 | {
232 | name: "zero page size",
233 | pageSize: 0,
234 | wantCode: http.StatusBadRequest,
235 | wantMsg: "pageSize must be positive: 0",
236 | },
237 | {
238 | name: "negative page size",
239 | pageSize: -1,
240 | wantCode: http.StatusBadRequest,
241 | wantMsg: "pageSize must be positive: -1",
242 | },
243 | }
244 |
245 | for _, tc := range tcs {
246 | t.Run(tc.name, func(t *testing.T) {
247 | request := map[string]any{
248 | "pageSize": tc.pageSize,
249 | }
250 | resp, err := invokeListBatches("list-batches", request, nil)
251 | if err != nil {
252 | t.Fatalf("invokeListBatches failed: %v", err)
253 | }
254 | defer resp.Body.Close()
255 |
256 | if resp.StatusCode != tc.wantCode {
257 | bodyBytes, _ := io.ReadAll(resp.Body)
258 | t.Fatalf("response status code is not %d, got %d: %s", tc.wantCode, resp.StatusCode, string(bodyBytes))
259 | }
260 |
261 | bodyBytes, err := io.ReadAll(resp.Body)
262 | if err != nil {
263 | t.Fatalf("failed to read response body: %v", err)
264 | }
265 |
266 | if !bytes.Contains(bodyBytes, []byte(tc.wantMsg)) {
267 | t.Fatalf("response body does not contain %q: %s", tc.wantMsg, string(bodyBytes))
268 | }
269 | })
270 | }
271 | }
272 |
273 | func runListBatchesAuthTest(t *testing.T) {
274 | idToken, err := tests.GetGoogleIdToken(tests.ClientId)
275 | if err != nil {
276 | t.Fatalf("error getting Google ID token: %s", err)
277 | }
278 | tcs := []struct {
279 | name string
280 | toolName string
281 | headers map[string]string
282 | wantStatus int
283 | }{
284 | {
285 | name: "valid auth token",
286 | toolName: "list-batches-with-auth",
287 | headers: map[string]string{"my-google-auth_token": idToken},
288 | wantStatus: http.StatusOK,
289 | },
290 | {
291 | name: "invalid auth token",
292 | toolName: "list-batches-with-auth",
293 | headers: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
294 | wantStatus: http.StatusUnauthorized,
295 | },
296 | {
297 | name: "no auth token",
298 | toolName: "list-batches-with-auth",
299 | headers: nil,
300 | wantStatus: http.StatusUnauthorized,
301 | },
302 | }
303 |
304 | for _, tc := range tcs {
305 | t.Run(tc.name, func(t *testing.T) {
306 | request := map[string]any{
307 | "pageSize": 1,
308 | }
309 | resp, err := invokeListBatches(tc.toolName, request, tc.headers)
310 | if err != nil {
311 | t.Fatalf("invokeListBatches failed: %v", err)
312 | }
313 | defer resp.Body.Close()
314 |
315 | if resp.StatusCode != tc.wantStatus {
316 | bodyBytes, _ := io.ReadAll(resp.Body)
317 | t.Fatalf("response status code is not %d, got %d: %s", tc.wantStatus, resp.StatusCode, string(bodyBytes))
318 | }
319 | })
320 | }
321 | }
322 |
323 | func invokeListBatches(toolName string, request map[string]any, headers map[string]string) (*http.Response, error) {
324 | requestBytes, err := json.Marshal(request)
325 | if err != nil {
326 | return nil, fmt.Errorf("failed to marshal request: %w", err)
327 | }
328 |
329 | url := fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/invoke", toolName)
330 | req, err := http.NewRequest(http.MethodPost, url, bytes.NewBuffer(requestBytes))
331 | if err != nil {
332 | return nil, fmt.Errorf("unable to create request: %w", err)
333 | }
334 | req.Header.Add("Content-type", "application/json")
335 | for k, v := range headers {
336 | req.Header.Add(k, v)
337 | }
338 |
339 | return http.DefaultClient.Do(req)
340 | }
341 |
```
--------------------------------------------------------------------------------
/internal/prebuiltconfigs/tools/cloud-sql-mssql-observability.yaml:
--------------------------------------------------------------------------------
```yaml
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 | sources:
15 | cloud-monitoring-source:
16 | kind: cloud-monitoring
17 | tools:
18 | get_system_metrics:
19 | kind: cloud-monitoring-query-prometheus
20 | source: cloud-monitoring-source
21 | description: |
22 | Fetches system level cloudmonitoring data (timeseries metrics) for a SqlServer instance using a PromQL query. Take projectId and instanceId from the user for which the metrics timeseries data needs to be fetched.
23 | To use this tool, you must provide the Google Cloud `projectId` and a PromQL `query`.
24 |
25 | Generate PromQL `query` for SqlServer system metrics. Use the provided metrics and rules to construct queries, Get the labels like `instance_id` from user intent.
26 |
27 | Defaults:
28 | 1. Interval: Use a default interval of `5m` for `_over_time` aggregation functions unless a different window is specified by the user.
29 |
30 | PromQL Query Examples:
31 | 1. Basic Time Series: `avg_over_time({"__name__"="cloudsql.googleapis.com/database/cpu/utilization","monitored_resource"="cloudsql_database","project_id"="my-projectId","database_id"="my-projectId:my-instanceId"}[5m])`
32 | 2. Top K: `topk(30, avg_over_time({"__name__"="cloudsql.googleapis.com/database/cpu/utilization","monitored_resource"="cloudsql_database","project_id"="my-projectId","database_id"="my-projectId:my-instanceId"}[5m]))`
33 | 3. Mean: `avg(avg_over_time({"__name__"="cloudsql.googleapis.com/database/cpu/utilization","monitored_resource"="cloudsql_database","project_id"="my-projectId","database_id"="my-projectId:my-instanceId"}[5m]))`
34 | 4. Minimum: `min(min_over_time({"__name__"="cloudsql.googleapis.com/database/cpu/utilization","monitored_resource"="cloudsql_database","project_id"="my-projectId","database_id"="my-projectId:my-instanceId"}[5m]))`
35 | 5. Maximum: `max(max_over_time({"__name__"="cloudsql.googleapis.com/database/cpu/utilization","monitored_resource"="cloudsql_database","project_id"="my-projectId","database_id"="my-projectId:my-instanceId"}[5m]))`
36 | 6. Sum: `sum(avg_over_time({"__name__"="cloudsql.googleapis.com/database/cpu/utilization","monitored_resource"="cloudsql_database","project_id"="my-projectId","database_id"="my-projectId:my-instanceId"}[5m]))`
37 | 7. Count streams: `count(avg_over_time({"__name__"="cloudsql.googleapis.com/database/cpu/utilization","monitored_resource"="cloudsql_database","project_id"="my-projectId","database_id"="my-projectId:my-instanceId"}[5m]))`
38 | 8. Percentile with groupby on database_id: `quantile by ("database_id")(0.99,avg_over_time({"__name__"="cloudsql.googleapis.com/database/cpu/utilization","monitored_resource"="cloudsql_database","project_id"="my-projectId","database_id"="my-projectId:my-instanceId"}[5m]))`
39 |
40 | Available Metrics List: metricname. description. monitored resource. labels. database_id is actually the instance id and the format is `project_id:instance_id`.
41 | 1. `cloudsql.googleapis.com/database/cpu/utilization`: Current CPU utilization as a percentage of the reserved CPU. `cloudsql_database`. `database`, `project_id`, `database_id`.
42 | 2. `cloudsql.googleapis.com/database/memory/usage`: RAM usage in bytes, excluding buffer/cache. `cloudsql_database`. `database`, `project_id`, `database_id`.
43 | 3. `cloudsql.googleapis.com/database/memory/total_usage`: Total RAM usage in bytes, including buffer/cache. `cloudsql_database`. `database`, `project_id`, `database_id`.
44 | 4. `cloudsql.googleapis.com/database/disk/bytes_used`: Data utilization in bytes. `cloudsql_database`. `database`, `project_id`, `database_id`.
45 | 5. `cloudsql.googleapis.com/database/disk/quota`: Maximum data disk size in bytes. `cloudsql_database`. `database`, `project_id`, `database_id`.
46 | 6. `cloudsql.googleapis.com/database/disk/read_ops_count`: Delta count of data disk read IO operations. `cloudsql_database`. `database`, `project_id`, `database_id`.
47 | 7. `cloudsql.googleapis.com/database/disk/write_ops_count`: Delta count of data disk write IO operations. `cloudsql_database`. `database`, `project_id`, `database_id`.
48 | 8. `cloudsql.googleapis.com/database/network/received_bytes_count`: Delta count of bytes received through the network. `cloudsql_database`. `database`, `project_id`, `database_id`.
49 | 9. `cloudsql.googleapis.com/database/network/sent_bytes_count`: Delta count of bytes sent through the network. `cloudsql_database`. `destination`, `database`, `project_id`, `database_id`.
50 | 10. `cloudsql.googleapis.com/database/sqlserver/memory/buffer_cache_hit_ratio`: Current percentage of pages found in the buffer cache without reading from disk. `cloudsql_database`. `database`, `project_id`, `database_id`.
51 | 11. `cloudsql.googleapis.com/database/sqlserver/memory/memory_grants_pending`: Current number of processes waiting for a workspace memory grant. `cloudsql_database`. `database`, `project_id`, `database_id`.
52 | 12. `cloudsql.googleapis.com/database/sqlserver/memory/free_list_stall_count`: Total number of requests that waited for a free page. `cloudsql_database`. `database`, `project_id`, `database_id`.
53 | 13. `cloudsql.googleapis.com/database/swap/pages_swapped_in_count`: Total count of pages swapped in from disk since the system was booted. `cloudsql_database`. `database`, `project_id`, `database_id`.
54 | 14. `cloudsql.googleapis.com/database/swap/pages_swapped_out_count`: Total count of pages swapped out to disk since the system was booted. `cloudsql_database`. `database`, `project_id`, `database_id`.
55 | 15. `cloudsql.googleapis.com/database/sqlserver/memory/checkpoint_page_count`: Total number of pages flushed to disk by a checkpoint. `cloudsql_database`. `database`, `project_id`, `database_id`.
56 | 16. `cloudsql.googleapis.com/database/sqlserver/memory/lazy_write_count`: Total number of buffers written by the buffer manager's lazy writer. `cloudsql_database`. `database`, `project_id`, `database_id`.
57 | 17. `cloudsql.googleapis.com/database/sqlserver/memory/page_life_expectancy`: Current number of seconds a page will stay in the buffer pool. `cloudsql_database`. `database`, `project_id`, `database_id`.
58 | 18. `cloudsql.googleapis.com/database/sqlserver/memory/page_operation_count`: Total number of physical database page reads or writes. `cloudsql_database`. `operation`, `database`, `project_id`, `database_id`.
59 | 19. `cloudsql.googleapis.com/database/sqlserver/transactions/page_split_count`: Total number of page splits from overflowing index pages. `cloudsql_database`. `database`, `project_id`, `database_id`.
60 | 20. `cloudsql.googleapis.com/database/sqlserver/transactions/deadlock_count`: Total number of lock requests that resulted in a deadlock. `cloudsql_database`. `locked_resource`, `database`, `project_id`, `database_id`.
61 | 21. `cloudsql.googleapis.com/database/sqlserver/transactions/transaction_count`: Total number of transactions started. `cloudsql_database`. `database`, `project_id`, `database_id`.
62 | 22. `cloudsql.googleapis.com/database/sqlserver/transactions/batch_request_count`: Total number of Transact-SQL command batches received. `cloudsql_database`. `database`, `project_id`, `database_id`.
63 | 23. `cloudsql.googleapis.com/database/sqlserver/transactions/sql_compilation_count`: Total number of SQL compilations. `cloudsql_database`. `database`, `project_id`, `database_id`.
64 | 24. `cloudsql.googleapis.com/database/sqlserver/transactions/sql_recompilation_count`: Total number of SQL recompilations. `cloudsql_database`. `database`, `project_id`, `database_id`.
65 | 25. `cloudsql.googleapis.com/database/sqlserver/connections/processes_blocked`: Current number of blocked processes. `cloudsql_database`. `database`, `project_id`, `database_id`.
66 | 26. `cloudsql.googleapis.com/database/sqlserver/transactions/lock_wait_time`: Total time lock requests were waiting for locks. `cloudsql_database`. `locked_resource`, `database`, `project_id`, `database_id`.
67 | 27. `cloudsql.googleapis.com/database/sqlserver/transactions/lock_wait_count`: Total number of lock requests that required the caller to wait. `cloudsql_database`. `locked_resource`, `database`, `project_id`, `database_id`.
68 | 28. `cloudsql.googleapis.com/database/network/connections`: Number of connections to databases on the instance. `cloudsql_database`. `database`, `project_id`, `database_id`.
69 | 29. `cloudsql.googleapis.com/database/sqlserver/connections/login_attempt_count`: Total number of login attempts since the last server restart. `cloudsql_database`. `database`, `project_id`, `database_id`.
70 | 30. `cloudsql.googleapis.com/database/sqlserver/connections/logout_count`: Total number of logout operations since the last server restart. `cloudsql_database`. `database`, `project_id`, `database_id`.
71 | 31. `cloudsql.googleapis.com/database/sqlserver/connections/connection_reset_count`: Total number of logins started from the connection pool since the last server restart. `cloudsql_database`. `database`, `project_id`, `database_id`.
72 | 32. `cloudsql.googleapis.com/database/sqlserver/transactions/full_scan_count`: Total number of unrestricted full scans (base-table or full-index). `cloudsql_database`. `database`, `project_id`, `database_id`.
73 |
74 | toolsets:
75 | cloud_sql_mssql_cloud_monitoring_tools:
76 | - get_system_metrics
```
--------------------------------------------------------------------------------
/internal/tools/http/http.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 http
15 |
16 | import (
17 | "bytes"
18 | "context"
19 | "encoding/json"
20 | "fmt"
21 | "io"
22 | "net/http"
23 | "net/url"
24 | "slices"
25 | "strings"
26 |
27 | "maps"
28 | "text/template"
29 |
30 | yaml "github.com/goccy/go-yaml"
31 | "github.com/googleapis/genai-toolbox/internal/sources"
32 | httpsrc "github.com/googleapis/genai-toolbox/internal/sources/http"
33 | "github.com/googleapis/genai-toolbox/internal/tools"
34 | )
35 |
36 | const kind string = "http"
37 |
38 | func init() {
39 | if !tools.Register(kind, newConfig) {
40 | panic(fmt.Sprintf("tool kind %q already registered", kind))
41 | }
42 | }
43 |
44 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
45 | actual := Config{Name: name}
46 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
47 | return nil, err
48 | }
49 | return actual, nil
50 | }
51 |
52 | type Config struct {
53 | Name string `yaml:"name" validate:"required"`
54 | Kind string `yaml:"kind" validate:"required"`
55 | Source string `yaml:"source" validate:"required"`
56 | Description string `yaml:"description" validate:"required"`
57 | AuthRequired []string `yaml:"authRequired"`
58 | Path string `yaml:"path" validate:"required"`
59 | Method tools.HTTPMethod `yaml:"method" validate:"required"`
60 | Headers map[string]string `yaml:"headers"`
61 | RequestBody string `yaml:"requestBody"`
62 | PathParams tools.Parameters `yaml:"pathParams"`
63 | QueryParams tools.Parameters `yaml:"queryParams"`
64 | BodyParams tools.Parameters `yaml:"bodyParams"`
65 | HeaderParams tools.Parameters `yaml:"headerParams"`
66 | }
67 |
68 | // validate interface
69 | var _ tools.ToolConfig = Config{}
70 |
71 | func (cfg Config) ToolConfigKind() string {
72 | return kind
73 | }
74 |
75 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
76 | // verify source exists
77 | rawS, ok := srcs[cfg.Source]
78 | if !ok {
79 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
80 | }
81 |
82 | // verify the source is compatible
83 | s, ok := rawS.(*httpsrc.Source)
84 | if !ok {
85 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be `http`", kind)
86 | }
87 |
88 | // Combine Source and Tool headers.
89 | // In case of conflict, Tool header overrides Source header
90 | combinedHeaders := make(map[string]string)
91 | maps.Copy(combinedHeaders, s.DefaultHeaders)
92 | maps.Copy(combinedHeaders, cfg.Headers)
93 |
94 | // Create a slice for all parameters
95 | allParameters := slices.Concat(cfg.PathParams, cfg.BodyParams, cfg.HeaderParams, cfg.QueryParams)
96 |
97 | // Verify no duplicate parameter names
98 | err := tools.CheckDuplicateParameters(allParameters)
99 | if err != nil {
100 | return nil, err
101 | }
102 |
103 | // Create Toolbox manifest
104 | paramManifest := allParameters.Manifest()
105 |
106 | if paramManifest == nil {
107 | paramManifest = make([]tools.ParameterManifest, 0)
108 | }
109 |
110 | // Create MCP manifest
111 | mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters)
112 |
113 | // finish tool setup
114 | return Tool{
115 | Name: cfg.Name,
116 | Kind: kind,
117 | BaseURL: s.BaseURL,
118 | Path: cfg.Path,
119 | Method: cfg.Method,
120 | AuthRequired: cfg.AuthRequired,
121 | RequestBody: cfg.RequestBody,
122 | PathParams: cfg.PathParams,
123 | QueryParams: cfg.QueryParams,
124 | BodyParams: cfg.BodyParams,
125 | HeaderParams: cfg.HeaderParams,
126 | Headers: combinedHeaders,
127 | DefaultQueryParams: s.QueryParams,
128 | Client: s.Client,
129 | AllParams: allParameters,
130 | manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
131 | mcpManifest: mcpManifest,
132 | }, nil
133 | }
134 |
135 | // validate interface
136 | var _ tools.Tool = Tool{}
137 |
138 | type Tool struct {
139 | Name string `yaml:"name"`
140 | Kind string `yaml:"kind"`
141 | Description string `yaml:"description"`
142 | AuthRequired []string `yaml:"authRequired"`
143 |
144 | BaseURL string `yaml:"baseURL"`
145 | Path string `yaml:"path"`
146 | Method tools.HTTPMethod `yaml:"method"`
147 | Headers map[string]string `yaml:"headers"`
148 | DefaultQueryParams map[string]string `yaml:"defaultQueryParams"`
149 |
150 | RequestBody string `yaml:"requestBody"`
151 | PathParams tools.Parameters `yaml:"pathParams"`
152 | QueryParams tools.Parameters `yaml:"queryParams"`
153 | BodyParams tools.Parameters `yaml:"bodyParams"`
154 | HeaderParams tools.Parameters `yaml:"headerParams"`
155 | AllParams tools.Parameters `yaml:"allParams"`
156 |
157 | Client *http.Client
158 | manifest tools.Manifest
159 | mcpManifest tools.McpManifest
160 | }
161 |
162 | // Helper function to generate the HTTP request body upon Tool invocation.
163 | func getRequestBody(bodyParams tools.Parameters, requestBodyPayload string, paramsMap map[string]any) (string, error) {
164 | bodyParamValues, err := tools.GetParams(bodyParams, paramsMap)
165 | if err != nil {
166 | return "", err
167 | }
168 | bodyParamsMap := bodyParamValues.AsMap()
169 |
170 | requestBodyStr, err := tools.PopulateTemplateWithJSON("HTTPToolRequestBody", requestBodyPayload, bodyParamsMap)
171 | if err != nil {
172 | return "", err
173 | }
174 | return requestBodyStr, nil
175 | }
176 |
177 | // Helper function to generate the HTTP request URL upon Tool invocation.
178 | func getURL(baseURL, path string, pathParams, queryParams tools.Parameters, defaultQueryParams map[string]string, paramsMap map[string]any) (string, error) {
179 | // use Go template to replace path params
180 | pathParamValues, err := tools.GetParams(pathParams, paramsMap)
181 | if err != nil {
182 | return "", err
183 | }
184 | pathParamsMap := pathParamValues.AsMap()
185 |
186 | templ, err := template.New("url").Parse(path)
187 | if err != nil {
188 | return "", fmt.Errorf("error parsing URL: %s", err)
189 | }
190 | var templatedPath bytes.Buffer
191 | err = templ.Execute(&templatedPath, pathParamsMap)
192 | if err != nil {
193 | return "", fmt.Errorf("error replacing pathParams: %s", err)
194 | }
195 |
196 | // Create URL based on BaseURL and Path
197 | // Attach query parameters
198 | parsedURL, err := url.Parse(baseURL + templatedPath.String())
199 | if err != nil {
200 | return "", fmt.Errorf("error parsing URL: %s", err)
201 | }
202 |
203 | // Get existing query parameters from the URL
204 | queryParameters := parsedURL.Query()
205 | for key, value := range defaultQueryParams {
206 | queryParameters.Add(key, value)
207 | }
208 | parsedURL.RawQuery = queryParameters.Encode()
209 |
210 | // Set dynamic query parameters
211 | query := parsedURL.Query()
212 | for _, p := range queryParams {
213 | v, ok := paramsMap[p.GetName()]
214 | if !ok || v == nil {
215 | if !p.GetRequired(){
216 | // If the param is not required AND
217 | // Not provodid OR provided with a nil value
218 | // Omitted from the URL
219 | continue
220 | }
221 | v = ""
222 | }
223 | query.Add(p.GetName(), fmt.Sprintf("%v", v))
224 | }
225 | parsedURL.RawQuery = query.Encode()
226 | return parsedURL.String(), nil
227 | }
228 |
229 | // Helper function to generate the HTTP headers upon Tool invocation.
230 | func getHeaders(headerParams tools.Parameters, defaultHeaders map[string]string, paramsMap map[string]any) (map[string]string, error) {
231 | // Populate header params
232 | allHeaders := make(map[string]string)
233 | maps.Copy(allHeaders, defaultHeaders)
234 | for _, p := range headerParams {
235 | headerValue, ok := paramsMap[p.GetName()]
236 | if ok {
237 | if strValue, ok := headerValue.(string); ok {
238 | allHeaders[p.GetName()] = strValue
239 | } else {
240 | return nil, fmt.Errorf("header param %s got value of type %t, not string", p.GetName(), headerValue)
241 | }
242 | }
243 | }
244 | return allHeaders, nil
245 | }
246 |
247 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
248 | paramsMap := params.AsMap()
249 |
250 | // Calculate request body
251 | requestBody, err := getRequestBody(t.BodyParams, t.RequestBody, paramsMap)
252 | if err != nil {
253 | return nil, fmt.Errorf("error populating request body: %s", err)
254 | }
255 |
256 | // Calculate URL
257 | urlString, err := getURL(t.BaseURL, t.Path, t.PathParams, t.QueryParams, t.DefaultQueryParams, paramsMap)
258 | if err != nil {
259 | return nil, fmt.Errorf("error populating path parameters: %s", err)
260 | }
261 |
262 | req, _ := http.NewRequest(string(t.Method), urlString, strings.NewReader(requestBody))
263 |
264 | // Calculate request headers
265 | allHeaders, err := getHeaders(t.HeaderParams, t.Headers, paramsMap)
266 | if err != nil {
267 | return nil, fmt.Errorf("error populating request headers: %s", err)
268 | }
269 | // Set request headers
270 | for k, v := range allHeaders {
271 | req.Header.Set(k, v)
272 | }
273 |
274 | // Make request and fetch response
275 | resp, err := t.Client.Do(req)
276 | if err != nil {
277 | return nil, fmt.Errorf("error making HTTP request: %s", err)
278 | }
279 | defer resp.Body.Close()
280 |
281 | var body []byte
282 | body, err = io.ReadAll(resp.Body)
283 | if err != nil {
284 | return nil, err
285 | }
286 | if resp.StatusCode < 200 || resp.StatusCode > 299 {
287 | return nil, fmt.Errorf("unexpected status code: %d, response body: %s", resp.StatusCode, string(body))
288 | }
289 |
290 | var data any
291 | if err = json.Unmarshal(body, &data); err != nil {
292 | // if unable to unmarshal data, return result as string.
293 | return string(body), nil
294 | }
295 | return data, nil
296 | }
297 |
298 | func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
299 | return tools.ParseParams(t.AllParams, data, claims)
300 | }
301 |
302 | func (t Tool) Manifest() tools.Manifest {
303 | return t.manifest
304 | }
305 |
306 | func (t Tool) McpManifest() tools.McpManifest {
307 | return t.mcpManifest
308 | }
309 |
310 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
311 | return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
312 | }
313 |
314 | func (t Tool) RequiresClientAuthorization() bool {
315 | return false
316 | }
317 |
```
--------------------------------------------------------------------------------
/docs/en/how-to/connect-ide/looker_mcp.md:
--------------------------------------------------------------------------------
```markdown
1 | ---
2 | title: "Looker using MCP"
3 | type: docs
4 | weight: 2
5 | description: >
6 | Connect your IDE to Looker using Toolbox.
7 | ---
8 |
9 | [Model Context Protocol (MCP)](https://modelcontextprotocol.io/introduction) is
10 | an open protocol for connecting Large Language Models (LLMs) to data sources
11 | like Postgres. This guide covers how to use [MCP Toolbox for Databases][toolbox]
12 | to expose your developer assistant tools to a Looker instance:
13 |
14 | * [Gemini-CLI][gemini-cli]
15 | * [Cursor][cursor]
16 | * [Windsurf][windsurf] (Codium)
17 | * [Visual Studio Code][vscode] (Copilot)
18 | * [Cline][cline] (VS Code extension)
19 | * [Claude desktop][claudedesktop]
20 | * [Claude code][claudecode]
21 |
22 | [toolbox]: https://github.com/googleapis/genai-toolbox
23 | [gemini-cli]: #configure-your-mcp-client
24 | [cursor]: #configure-your-mcp-client
25 | [windsurf]: #configure-your-mcp-client
26 | [vscode]: #configure-your-mcp-client
27 | [cline]: #configure-your-mcp-client
28 | [claudedesktop]: #configure-your-mcp-client
29 | [claudecode]: #configure-your-mcp-client
30 |
31 | ## Set up Looker
32 |
33 | 1. Get a Looker Client ID and Client Secret. Follow the directions
34 | [here](https://cloud.google.com/looker/docs/api-auth#authentication_with_an_sdk).
35 |
36 | 1. Have the base URL of your Looker instance available. It is likely
37 | something like `https://looker.example.com`. In some cases the API is
38 | listening at a different port, and you will need to use
39 | `https://looker.example.com:19999` instead.
40 |
41 | ## Install MCP Toolbox
42 |
43 | 1. Download the latest version of Toolbox as a binary. Select the [correct
44 | binary](https://github.com/googleapis/genai-toolbox/releases) corresponding
45 | to your OS and CPU architecture. You are required to use Toolbox version
46 | v0.10.0+:
47 |
48 | <!-- {x-release-please-start-version} -->
49 | {{< tabpane persist=header >}}
50 | {{< tab header="linux/amd64" lang="bash" >}}
51 | curl -O https://storage.googleapis.com/genai-toolbox/v0.18.0/linux/amd64/toolbox
52 | {{< /tab >}}
53 |
54 | {{< tab header="darwin/arm64" lang="bash" >}}
55 | curl -O https://storage.googleapis.com/genai-toolbox/v0.18.0/darwin/arm64/toolbox
56 | {{< /tab >}}
57 |
58 | {{< tab header="darwin/amd64" lang="bash" >}}
59 | curl -O https://storage.googleapis.com/genai-toolbox/v0.18.0/darwin/amd64/toolbox
60 | {{< /tab >}}
61 |
62 | {{< tab header="windows/amd64" lang="bash" >}}
63 | curl -O https://storage.googleapis.com/genai-toolbox/v0.18.0/windows/amd64/toolbox.exe
64 | {{< /tab >}}
65 | {{< /tabpane >}}
66 | <!-- {x-release-please-end} -->
67 |
68 | 1. Make the binary executable:
69 |
70 | ```bash
71 | chmod +x toolbox
72 | ```
73 |
74 | 1. Verify the installation:
75 |
76 | ```bash
77 | ./toolbox --version
78 | ```
79 |
80 | ## Configure your MCP Client
81 |
82 | {{< tabpane text=true >}}
83 | {{% tab header="Gemini-CLI" lang="en" %}}
84 |
85 | 1. Install [Gemini-CLI](https://github.com/google-gemini/gemini-cli#install-globally-with-npm).
86 | 1. Create a directory `.gemini` in your home directory if it doesn't exist.
87 | 1. Create the file `.gemini/settings.json` if it doesn't exist.
88 | 1. Add the following configuration, or add the mcpServers stanza if you already
89 | have a `settings.json` with content. Replace the path to the toolbox
90 | executable and the environment variables with your values, and save:
91 |
92 | ```json
93 | {
94 | "mcpServers": {
95 | "looker-toolbox": {
96 | "command": "./PATH/TO/toolbox",
97 | "args": ["--stdio", "--prebuilt", "looker"],
98 | "env": {
99 | "LOOKER_BASE_URL": "https://looker.example.com",
100 | "LOOKER_CLIENT_ID": "",
101 | "LOOKER_CLIENT_SECRET": "",
102 | "LOOKER_VERIFY_SSL": "true"
103 | }
104 | }
105 | }
106 | }
107 | ```
108 |
109 | 1. Start Gemini-CLI with the `gemini` command and use the command `/mcp` to see
110 | the configured MCP tools.
111 | {{% /tab %}}
112 |
113 | {{% tab header="Claude code" lang="en" %}}
114 |
115 | 1. Install [Claude
116 | Code](https://docs.anthropic.com/en/docs/agents-and-tools/claude-code/overview).
117 | 1. Create a `.mcp.json` file in your project root if it doesn't exist.
118 | 1. Add the following configuration, replace the environment variables with your
119 | values, and save:
120 |
121 | ```json
122 | {
123 | "mcpServers": {
124 | "looker-toolbox": {
125 | "command": "./PATH/TO/toolbox",
126 | "args": ["--stdio", "--prebuilt", "looker"],
127 | "env": {
128 | "LOOKER_BASE_URL": "https://looker.example.com",
129 | "LOOKER_CLIENT_ID": "",
130 | "LOOKER_CLIENT_SECRET": "",
131 | "LOOKER_VERIFY_SSL": "true"
132 | }
133 | }
134 | }
135 | }
136 | ```
137 |
138 | 1. Restart Claude Code to apply the new configuration.
139 | {{% /tab %}}
140 |
141 | {{% tab header="Claude desktop" lang="en" %}}
142 |
143 | 1. Open [Claude desktop](https://claude.ai/download) and navigate to Settings.
144 | 1. Under the Developer tab, tap Edit Config to open the configuration file.
145 | 1. Add the following configuration, replace the environment variables with your
146 | values, and save:
147 |
148 | ```json
149 | {
150 | "mcpServers": {
151 | "looker-toolbox": {
152 | "command": "./PATH/TO/toolbox",
153 | "args": ["--stdio", "--prebuilt", "looker"],
154 | "env": {
155 | "LOOKER_BASE_URL": "https://looker.example.com",
156 | "LOOKER_CLIENT_ID": "",
157 | "LOOKER_CLIENT_SECRET": "",
158 | "LOOKER_VERIFY_SSL": "true"
159 | }
160 | }
161 | }
162 | }
163 | ```
164 |
165 | 1. Restart Claude desktop.
166 | 1. From the new chat screen, you should see a hammer (MCP) icon appear with the
167 | new MCP server available.
168 | {{% /tab %}}
169 |
170 | {{% tab header="Cline" lang="en" %}}
171 |
172 | 1. Open the [Cline](https://github.com/cline/cline) extension in VS Code and tap
173 | the **MCP Servers** icon.
174 | 1. Tap Configure MCP Servers to open the configuration file.
175 | 1. Add the following configuration, replace the environment variables with your
176 | values, and save:
177 |
178 | ```json
179 | {
180 | "mcpServers": {
181 | "looker-toolbox": {
182 | "command": "./PATH/TO/toolbox",
183 | "args": ["--stdio", "--prebuilt", "looker"],
184 | "env": {
185 | "LOOKER_BASE_URL": "https://looker.example.com",
186 | "LOOKER_CLIENT_ID": "",
187 | "LOOKER_CLIENT_SECRET": "",
188 | "LOOKER_VERIFY_SSL": "true"
189 | }
190 | }
191 | }
192 | }
193 | ```
194 |
195 | 1. You should see a green active status after the server is successfully
196 | connected.
197 | {{% /tab %}}
198 |
199 | {{% tab header="Cursor" lang="en" %}}
200 |
201 | 1. Create a `.cursor` directory in your project root if it doesn't exist.
202 | 1. Create a `.cursor/mcp.json` file if it doesn't exist and open it.
203 | 1. Add the following configuration, replace the environment variables with your
204 | values, and save:
205 |
206 | ```json
207 | {
208 | "mcpServers": {
209 | "looker-toolbox": {
210 | "command": "./PATH/TO/toolbox",
211 | "args": ["--stdio", "--prebuilt", "looker"],
212 | "env": {
213 | "LOOKER_BASE_URL": "https://looker.example.com",
214 | "LOOKER_CLIENT_ID": "",
215 | "LOOKER_CLIENT_SECRET": "",
216 | "LOOKER_VERIFY_SSL": "true"
217 | }
218 | }
219 | }
220 | }
221 | ```
222 |
223 | 1. Open [Cursor](https://www.cursor.com/) and navigate to **Settings > Cursor
224 | Settings > MCP**. You should see a green active status after the server is
225 | successfully connected.
226 | {{% /tab %}}
227 |
228 | {{% tab header="Visual Studio Code (Copilot)" lang="en" %}}
229 |
230 | 1. Open [VS Code](https://code.visualstudio.com/docs/copilot/overview) and
231 | create a `.vscode` directory in your project root if it doesn't exist.
232 | 1. Create a `.vscode/mcp.json` file if it doesn't exist and open it.
233 | 1. Add the following configuration, replace the environment variables with your
234 | values, and save:
235 |
236 | ```json
237 | {
238 | "servers": {
239 | "looker-toolbox": {
240 | "command": "./PATH/TO/toolbox",
241 | "args": ["--stdio", "--prebuilt", "looker"],
242 | "env": {
243 | "LOOKER_BASE_URL": "https://looker.example.com",
244 | "LOOKER_CLIENT_ID": "",
245 | "LOOKER_CLIENT_SECRET": "",
246 | "LOOKER_VERIFY_SSL": "true"
247 | }
248 | }
249 | }
250 | }
251 | ```
252 |
253 | {{% /tab %}}
254 |
255 | {{% tab header="Windsurf" lang="en" %}}
256 |
257 | 1. Open [Windsurf](https://docs.codeium.com/windsurf) and navigate to the
258 | Cascade assistant.
259 | 1. Tap on the hammer (MCP) icon, then Configure to open the configuration file.
260 | 1. Add the following configuration, replace the environment variables with your
261 | values, and save:
262 |
263 | ```json
264 | {
265 | "mcpServers": {
266 | "looker-toolbox": {
267 | "command": "./PATH/TO/toolbox",
268 | "args": ["--stdio", "--prebuilt", "looker"],
269 | "env": {
270 | "LOOKER_BASE_URL": "https://looker.example.com",
271 | "LOOKER_CLIENT_ID": "",
272 | "LOOKER_CLIENT_SECRET": "",
273 | "LOOKER_VERIFY_SSL": "true"
274 | }
275 | }
276 | }
277 | }
278 |
279 | ```
280 |
281 | {{% /tab %}}
282 | {{< /tabpane >}}
283 |
284 | ## Use Tools
285 |
286 | Your AI tool is now connected to Looker using MCP. Try asking your AI
287 | assistant to list models, explores, dimensions, and measures. Run a
288 | query, retrieve the SQL for a query, and run a saved Look.
289 |
290 | The following tools are available to the LLM:
291 |
292 | 1. **get_models**: list the LookML models in Looker
293 | 1. **get_explores**: list the explores in a given model
294 | 1. **get_dimensions**: list the dimensions in a given explore
295 | 1. **get_measures**: list the measures in a given explore
296 | 1. **get_filters**: list the filters in a given explore
297 | 1. **get_parameters**: list the parameters in a given explore
298 | 1. **query**: Run a query and return the data
299 | 1. **query_sql**: Return the SQL generated by Looker for a query
300 | 1. **query_url**: Return a link to the query in Looker for further exploration
301 | 1. **get_looks**: Return the saved Looks that match a title or description
302 | 1. **run_look**: Run a saved Look and return the data
303 | 1. **make_look**: Create a saved Look in Looker and return the URL
304 | 1. **get_dashboards**: Return the saved dashboards that match a title or description
305 | 1. **make_dashboard**: Create a saved dashboard in Looker and return the URL
306 | 1. **add_dashboard_element**: Add a tile to a dashboard
307 |
308 | {{< notice note >}}
309 | Prebuilt tools are pre-1.0, so expect some tool changes between versions. LLMs
310 | will adapt to the tools available, so this shouldn't affect most users.
311 | {{< /notice >}}
312 |
```
--------------------------------------------------------------------------------
/internal/sources/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 | "bytes"
19 | "context"
20 | "encoding/json"
21 | "fmt"
22 | "io"
23 | "net/http"
24 | "net/url"
25 | "strings"
26 |
27 | "github.com/goccy/go-yaml"
28 | "github.com/googleapis/genai-toolbox/internal/sources"
29 | "go.opentelemetry.io/otel/trace"
30 | )
31 |
32 | const SourceKind string = "dgraph"
33 |
34 | // validate interface
35 | var _ sources.SourceConfig = Config{}
36 |
37 | func init() {
38 | if !sources.Register(SourceKind, newConfig) {
39 | panic(fmt.Sprintf("source kind %q already registered", SourceKind))
40 | }
41 | }
42 |
43 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (sources.SourceConfig, error) {
44 | actual := Config{Name: name}
45 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
46 | return nil, err
47 | }
48 | return actual, nil
49 | }
50 |
51 | // HttpToken stores credentials for making HTTP request
52 | type HttpToken struct {
53 | UserId string
54 | Password string
55 | AccessJwt string
56 | RefreshToken string
57 | Namespace uint64
58 | }
59 |
60 | type DgraphClient struct {
61 | httpClient *http.Client
62 | *HttpToken
63 | baseUrl string
64 | apiKey string
65 | }
66 |
67 | type Config struct {
68 | Name string `yaml:"name" validate:"required"`
69 | Kind string `yaml:"kind" validate:"required"`
70 | DgraphUrl string `yaml:"dgraphUrl" validate:"required"`
71 | User string `yaml:"user"`
72 | Password string `yaml:"password"`
73 | Namespace uint64 `yaml:"namespace"`
74 | ApiKey string `yaml:"apiKey"`
75 | }
76 |
77 | func (r Config) SourceConfigKind() string {
78 | return SourceKind
79 | }
80 |
81 | func (r Config) Initialize(ctx context.Context, tracer trace.Tracer) (sources.Source, error) {
82 | hc, err := initDgraphHttpClient(ctx, tracer, r)
83 | if err != nil {
84 | return nil, err
85 | }
86 |
87 | if err := hc.healthCheck(); err != nil {
88 | return nil, err
89 | }
90 |
91 | s := &Source{
92 | Name: r.Name,
93 | Kind: SourceKind,
94 | Client: hc,
95 | }
96 | return s, nil
97 | }
98 |
99 | var _ sources.Source = &Source{}
100 |
101 | type Source struct {
102 | Name string `yaml:"name"`
103 | Kind string `yaml:"kind"`
104 | Client *DgraphClient `yaml:"client"`
105 | }
106 |
107 | func (s *Source) SourceKind() string {
108 | return SourceKind
109 | }
110 |
111 | func (s *Source) DgraphClient() *DgraphClient {
112 | return s.Client
113 | }
114 |
115 | func initDgraphHttpClient(ctx context.Context, tracer trace.Tracer, r Config) (*DgraphClient, error) {
116 | //nolint:all // Reassigned ctx
117 | ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, r.Name)
118 | defer span.End()
119 |
120 | if r.DgraphUrl == "" {
121 | return nil, fmt.Errorf("dgraph url should not be empty")
122 | }
123 |
124 | hc := &DgraphClient{
125 | httpClient: &http.Client{},
126 | baseUrl: r.DgraphUrl,
127 | HttpToken: &HttpToken{
128 | UserId: r.User,
129 | Namespace: r.Namespace,
130 | Password: r.Password,
131 | },
132 | apiKey: r.ApiKey,
133 | }
134 |
135 | if r.User != "" || r.Password != "" {
136 | if err := hc.loginWithCredentials(); err != nil {
137 | return nil, err
138 | }
139 | }
140 |
141 | return hc, nil
142 | }
143 |
144 | func (hc *DgraphClient) ExecuteQuery(query string, paramsMap map[string]interface{},
145 | isQuery bool, timeout string) ([]byte, error) {
146 | if isQuery {
147 | return hc.postDqlQuery(query, paramsMap, timeout)
148 | } else {
149 | return hc.mutate(query, paramsMap)
150 | }
151 | }
152 |
153 | // postDqlQuery sends a DQL query to the Dgraph server with query, parameters, and optional timeout.
154 | // Returns the response body ([]byte) and an error, if any.
155 | func (hc *DgraphClient) postDqlQuery(query string, paramsMap map[string]interface{}, timeout string) ([]byte, error) {
156 | urlParams := url.Values{}
157 | urlParams.Add("timeout", timeout)
158 | url, err := getUrl(hc.baseUrl, "/query", urlParams)
159 | if err != nil {
160 | return nil, err
161 | }
162 | p := struct {
163 | Query string `json:"query"`
164 | Variables map[string]interface{} `json:"variables"`
165 | }{
166 | Query: query,
167 | Variables: paramsMap,
168 | }
169 | body, err := json.Marshal(p)
170 | if err != nil {
171 | return nil, fmt.Errorf("error marshlling json: %v", err)
172 | }
173 |
174 | req, err := http.NewRequest(http.MethodPost, url, bytes.NewBuffer(body))
175 | if err != nil {
176 | return nil, fmt.Errorf("error building req for endpoint [%v] :%v", url, err)
177 | }
178 |
179 | req.Header.Add("Content-Type", "application/json")
180 |
181 | return hc.doReq(req)
182 | }
183 |
184 | // mutate sends an RDF mutation to the Dgraph server with "commitNow: true", embedding parameters.
185 | // Returns the server's response as a byte slice or an error if the mutation fails.
186 | func (hc *DgraphClient) mutate(mutation string, paramsMap map[string]interface{}) ([]byte, error) {
187 | mu := embedParamsIntoMutation(mutation, paramsMap)
188 | params := url.Values{}
189 | params.Add("commitNow", "true")
190 | url, err := getUrl(hc.baseUrl, "/mutate", params)
191 | if err != nil {
192 | return nil, err
193 | }
194 | req, err := http.NewRequest(http.MethodPost, url, bytes.NewBufferString(mu))
195 | if err != nil {
196 | return nil, fmt.Errorf("error building req for endpoint [%v] :%v", url, err)
197 | }
198 |
199 | req.Header.Add("Content-Type", "application/rdf")
200 |
201 | return hc.doReq(req)
202 | }
203 |
204 | func (hc *DgraphClient) doReq(req *http.Request) ([]byte, error) {
205 | if hc.HttpToken != nil {
206 | req.Header.Add("X-Dgraph-AccessToken", hc.AccessJwt)
207 | }
208 | if hc.apiKey != "" {
209 | req.Header.Set("Dg-Auth", hc.apiKey)
210 | }
211 |
212 | resp, err := hc.httpClient.Do(req)
213 |
214 | if err != nil && !strings.Contains(err.Error(), "Token is expired") {
215 | return nil, fmt.Errorf("error performing HTTP request: %w", err)
216 | } else if err != nil && strings.Contains(err.Error(), "Token is expired") {
217 | if errLogin := hc.loginWithToken(); errLogin != nil {
218 | return nil, errLogin
219 | }
220 | if hc.HttpToken != nil {
221 | req.Header.Add("X-Dgraph-AccessToken", hc.AccessJwt)
222 | }
223 | resp, err = hc.httpClient.Do(req)
224 | if err != nil {
225 | return nil, err
226 | }
227 | }
228 |
229 | defer resp.Body.Close()
230 |
231 | respBody, err := io.ReadAll(resp.Body)
232 | if err != nil {
233 | return nil, fmt.Errorf("error reading response body: url: [%v], err: [%v]", req.URL, err)
234 | }
235 | if resp.StatusCode != http.StatusOK {
236 | return nil, fmt.Errorf("got non 200 resp: %v", string(respBody))
237 | }
238 |
239 | return respBody, nil
240 | }
241 |
242 | func (hc *DgraphClient) loginWithCredentials() error {
243 | credentials := map[string]interface{}{
244 | "userid": hc.UserId,
245 | "password": hc.Password,
246 | "namespace": hc.Namespace,
247 | }
248 | return hc.doLogin(credentials)
249 | }
250 |
251 | func (hc *DgraphClient) loginWithToken() error {
252 | credentials := map[string]interface{}{
253 | "refreshJWT": hc.RefreshToken,
254 | "namespace": hc.Namespace,
255 | }
256 | return hc.doLogin(credentials)
257 | }
258 |
259 | func (hc *DgraphClient) doLogin(creds map[string]interface{}) error {
260 | url, err := getUrl(hc.baseUrl, "/login", nil)
261 | if err != nil {
262 | return err
263 | }
264 | payload, err := json.Marshal(creds)
265 | if err != nil {
266 | return fmt.Errorf("failed to marshal credentials: %v", err)
267 | }
268 | req, err := http.NewRequest(http.MethodPost, url, bytes.NewBuffer(payload))
269 | if err != nil {
270 | return fmt.Errorf("error building req for endpoint [%v] : %v", url, err)
271 | }
272 | req.Header.Add("Content-Type", "application/json")
273 | if hc.apiKey != "" {
274 | req.Header.Set("Dg-Auth", hc.apiKey)
275 | }
276 |
277 | resp, err := hc.doReq(req)
278 | if err != nil {
279 | if strings.Contains(err.Error(), "Token is expired") &&
280 | !strings.Contains(err.Error(), "unable to authenticate the refresh token") {
281 | return hc.loginWithToken()
282 | }
283 | return err
284 | }
285 |
286 | if err := CheckError(resp); err != nil {
287 | return err
288 | }
289 |
290 | var r struct {
291 | Data struct {
292 | AccessJWT string `json:"accessJWT"`
293 | RefreshJWT string `json:"refreshJWT"`
294 | } `json:"data"`
295 | }
296 |
297 | if err := json.Unmarshal(resp, &r); err != nil {
298 | return fmt.Errorf("failed to unmarshal response: %v", err)
299 | }
300 |
301 | if r.Data.AccessJWT == "" {
302 | return fmt.Errorf("no access JWT found in the response")
303 | }
304 | if r.Data.RefreshJWT == "" {
305 | return fmt.Errorf("no refresh JWT found in the response")
306 | }
307 |
308 | hc.AccessJwt = r.Data.AccessJWT
309 | hc.RefreshToken = r.Data.RefreshJWT
310 | return nil
311 | }
312 |
313 | func (hc *DgraphClient) healthCheck() error {
314 | url, err := getUrl(hc.baseUrl, "/health", nil)
315 | if err != nil {
316 | return err
317 | }
318 | req, err := http.NewRequest(http.MethodGet, url, nil)
319 | if err != nil {
320 | return fmt.Errorf("error creating request: %w", err)
321 | }
322 |
323 | resp, err := hc.httpClient.Do(req)
324 | if err != nil {
325 | return fmt.Errorf("error performing request: %w", err)
326 | }
327 |
328 | defer resp.Body.Close()
329 | data, err := io.ReadAll(resp.Body)
330 | if err != nil {
331 | return err
332 | }
333 | var result []struct {
334 | Instance string `json:"instance"`
335 | Address string `json:"address"`
336 | Status string `json:"status"`
337 | }
338 |
339 | // Unmarshal response into the struct
340 | if err := json.Unmarshal(data, &result); err != nil {
341 | return fmt.Errorf("failed to unmarshal json: %v", err)
342 | }
343 |
344 | if len(result) == 0 {
345 | return fmt.Errorf("health info should not empty for: %v", url)
346 | }
347 |
348 | var unhealthyErr error
349 | for _, info := range result {
350 | if info.Status != "healthy" {
351 | unhealthyErr = fmt.Errorf("dgraph instance [%v] is not in healthy state, address is %v",
352 | info.Instance, info.Address)
353 | } else {
354 | return nil
355 | }
356 | }
357 |
358 | return unhealthyErr
359 | }
360 |
361 | func getUrl(baseUrl, resource string, params url.Values) (string, error) {
362 | u, err := url.ParseRequestURI(baseUrl)
363 | if err != nil {
364 | return "", fmt.Errorf("failed to get url %v", err)
365 | }
366 | u.Path = resource
367 | u.RawQuery = params.Encode()
368 | return u.String(), nil
369 | }
370 |
371 | func CheckError(resp []byte) error {
372 | var errResp struct {
373 | Errors []struct {
374 | Message string `json:"message"`
375 | } `json:"errors"`
376 | }
377 |
378 | if err := json.Unmarshal(resp, &errResp); err != nil {
379 | return fmt.Errorf("failed to unmarshal json: %v", err)
380 | }
381 |
382 | if len(errResp.Errors) > 0 {
383 | return fmt.Errorf("error : %v", errResp.Errors)
384 | }
385 |
386 | return nil
387 | }
388 |
389 | func embedParamsIntoMutation(mutation string, paramsMap map[string]interface{}) string {
390 | for key, value := range paramsMap {
391 | mutation = strings.ReplaceAll(mutation, key, fmt.Sprintf(`"%v"`, value))
392 | }
393 | return mutation
394 | }
395 |
```
--------------------------------------------------------------------------------
/tests/option.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 tests
16 |
17 | /* Configurations for RunToolInvokeTest() */
18 |
19 | // InvokeTestConfig represents the various configuration options for RunToolInvokeTest()
20 | type InvokeTestConfig struct {
21 | myToolId3NameAliceWant string
22 | myToolById4Want string
23 | nullWant string
24 | myArrayToolWant string
25 | supportSelect1Want bool
26 | supportOptionalNullParam bool
27 | supportArrayParam bool
28 | supportClientAuth bool
29 | supportSelect1Auth bool
30 | }
31 |
32 | type InvokeTestOption func(*InvokeTestConfig)
33 |
34 | // WithMyToolId3NameAliceWant represents the response value for my-tool with id=3 and name=Alice.
35 | // e.g. tests.RunToolInvokeTest(t, select1Want, tests.WithMyToolId3NameAliceWant("custom"))
36 | func WithMyToolId3NameAliceWant(s string) InvokeTestOption {
37 | return func(c *InvokeTestConfig) {
38 | c.myToolId3NameAliceWant = s
39 | }
40 | }
41 |
42 | // WithMyArrayToolWant represents the response value for my-array-tool.
43 | // e.g. tests.RunToolInvokeTest(t, select1Want, tests.WithMyArrayToolWant("custom"))
44 | func WithMyArrayToolWant(s string) InvokeTestOption {
45 | return func(c *InvokeTestConfig) {
46 | c.myArrayToolWant = s
47 | }
48 | }
49 |
50 | // WithMyToolById4Want represents the response value for my-tool-by-id with id=4.
51 | // This response includes a null value column.
52 | // e.g. tests.RunToolInvokeTest(t, select1Want, tests.WithMyToolById4Want("custom"))
53 | func WithMyToolById4Want(s string) InvokeTestOption {
54 | return func(c *InvokeTestConfig) {
55 | c.myToolById4Want = s
56 | }
57 | }
58 |
59 | // WithNullWant represents a response value of null string.
60 | // e.g. tests.RunToolInvokeTest(t, select1Want, tests.WithNullWant("custom"))
61 | func WithNullWant(s string) InvokeTestOption {
62 | return func(c *InvokeTestConfig) {
63 | c.nullWant = s
64 | }
65 | }
66 |
67 | // DisableOptionalNullParamTest disables tests for optional null parameters.
68 | // e.g. tests.RunToolInvokeTest(t, select1Want, tests.DisableOptionalNullParamTest())
69 | func DisableOptionalNullParamTest() InvokeTestOption {
70 | return func(c *InvokeTestConfig) {
71 | c.supportOptionalNullParam = false
72 | }
73 | }
74 |
75 | // DisableArrayTest disables tests for sources that do not support array.
76 | // e.g. tests.RunToolInvokeTest(t, select1Want, tests.DisableArrayTest())
77 | func DisableArrayTest() InvokeTestOption {
78 | return func(c *InvokeTestConfig) {
79 | c.supportArrayParam = false
80 | }
81 | }
82 |
83 | // DisableSelect1Test disables tests for sources that do not support SELECT 1 query.
84 | // e.g. tests.RunToolInvokeTest(t, "", tests.DisableSelect1Test())
85 | func DisableSelect1Test() InvokeTestOption {
86 | return func(c *InvokeTestConfig) {
87 | c.supportSelect1Want = false
88 | }
89 | }
90 |
91 | // DisableSelect1AuthTest disables auth tests for sources that do not support SELECT 1 query.
92 | // e.g. tests.RunToolInvokeTest(t, "", tests.DisableSelect1AuthTest())
93 | func DisableSelect1AuthTest() InvokeTestOption {
94 | return func(c *InvokeTestConfig) {
95 | c.supportSelect1Auth = false
96 | }
97 | }
98 |
99 | // EnableClientAuthTest runs the client authorization tests.
100 | // Only enable it if your source supports the `useClientOAuth` configuration.
101 | // Currently, this should only be used with the BigQuery tests.
102 | func EnableClientAuthTest() InvokeTestOption {
103 | return func(c *InvokeTestConfig) {
104 | c.supportClientAuth = true
105 | }
106 | }
107 |
108 | /* Configurations for RunMCPToolCallMethod() */
109 |
110 | // MCPTestConfig represents the various configuration options for mcp tool call tests.
111 | type MCPTestConfig struct {
112 | myToolId3NameAliceWant string
113 | supportClientAuth bool
114 | supportSelect1Auth bool
115 | }
116 |
117 | type McpTestOption func(*MCPTestConfig)
118 |
119 | // WithMcpMyToolId3NameAliceWant represents the response value for my-tool with id=3 and name=Alice.
120 | // e.g. tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, tests.WithMcpMyToolId3NameAliceWant("custom"))
121 | func WithMcpMyToolId3NameAliceWant(s string) McpTestOption {
122 | return func(c *MCPTestConfig) {
123 | c.myToolId3NameAliceWant = s
124 | }
125 | }
126 |
127 | // EnableMcpClientAuthTest runs the client authorization tests.
128 | // Only enable it if your source supports the `useClientOAuth` configuration.
129 | // Currently, this should only be used with the BigQuery tests.
130 | func EnableMcpClientAuthTest() McpTestOption {
131 | return func(c *MCPTestConfig) {
132 | c.supportClientAuth = true
133 | }
134 | }
135 |
136 | // DisableMcpSelect1AuthTest disables the auth tool tests which use select 1.
137 | func DisableMcpSelect1AuthTest() McpTestOption {
138 | return func(c *MCPTestConfig) {
139 | c.supportSelect1Auth = false
140 | }
141 | }
142 |
143 | /* Configurations for RunExecuteSqlToolInvokeTest() */
144 |
145 | // ExecuteSqlTestConfig represents the various configuration options for RunExecuteSqlToolInvokeTest()
146 | type ExecuteSqlTestConfig struct {
147 | select1Statement string
148 | }
149 |
150 | type ExecuteSqlOption func(*ExecuteSqlTestConfig)
151 |
152 | // WithSelect1Statement represents the database's statement for `SELECT 1`.
153 | // e.g. tests.RunExecuteSqlToolInvokeTest(t, createTableStatement, select1Want, tests.WithSelect1Statement("custom"))
154 | func WithSelect1Statement(s string) ExecuteSqlOption {
155 | return func(c *ExecuteSqlTestConfig) {
156 | c.select1Statement = s
157 | }
158 | }
159 |
160 | /* Configurations for RunToolInvokeWithTemplateParameters() */
161 |
162 | // TemplateParameterTestConfig represents the various configuration options for template parameter tests.
163 | type TemplateParameterTestConfig struct {
164 | ddlWant string
165 | selectAllWant string
166 | selectId1Want string
167 | selectNameWant string
168 | selectEmptyWant string
169 | insert1Want string
170 |
171 | nameFieldArray string
172 | nameColFilter string
173 | createColArray string
174 |
175 | supportDdl bool
176 | supportInsert bool
177 | supportSelectFields bool
178 | }
179 |
180 | type TemplateParamOption func(*TemplateParameterTestConfig)
181 |
182 | // WithDdlWant represents the response value of ddl statements.
183 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithDdlWant("custom"))
184 | func WithDdlWant(s string) TemplateParamOption {
185 | return func(c *TemplateParameterTestConfig) {
186 | c.ddlWant = s
187 | }
188 | }
189 |
190 | // WithSelectAllWant represents the response value of select-templateParams-tool.
191 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithSelectAllWant("custom"))
192 | func WithSelectAllWant(s string) TemplateParamOption {
193 | return func(c *TemplateParameterTestConfig) {
194 | c.selectAllWant = s
195 | }
196 | }
197 |
198 | // WithTmplSelectId1Want represents the response value of select-templateParams-combined-tool with id=1.
199 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithTmplSelectId1Want("custom"))
200 | func WithTmplSelectId1Want(s string) TemplateParamOption {
201 | return func(c *TemplateParameterTestConfig) {
202 | c.selectId1Want = s
203 | }
204 | }
205 |
206 | // WithTmplSelectNameWant represents the response value of select-filter-templateParams-combined-tool with name.
207 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithTmplSelectNameWant("custom"))
208 | func WithTmplSelectNameWant(s string) TemplateParamOption {
209 | return func(c *TemplateParameterTestConfig) {
210 | c.selectNameWant = s
211 | }
212 | }
213 |
214 | // WithSelectEmptyWant represents the response value of select-templateParams-combined-tool with no results.
215 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithSelectEmptyWant("custom"))
216 | func WithSelectEmptyWant(s string) TemplateParamOption {
217 | return func(c *TemplateParameterTestConfig) {
218 | c.selectEmptyWant = s
219 | }
220 | }
221 |
222 | // WithInsert1Want represents the response value of insert-table-templateParams-tool.
223 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithInsert1Want("custom"))
224 | func WithInsert1Want(s string) TemplateParamOption {
225 | return func(c *TemplateParameterTestConfig) {
226 | c.insert1Want = s
227 | }
228 | }
229 |
230 | // WithNameFieldArray represents fields array parameter for select-fields-templateParams-tool.
231 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithNameFieldArray("custom"))
232 | func WithNameFieldArray(s string) TemplateParamOption {
233 | return func(c *TemplateParameterTestConfig) {
234 | c.nameFieldArray = s
235 | }
236 | }
237 |
238 | // WithNameColFilter represents the columnFilter parameter for select-filter-templateParams-combined-tool.
239 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithNameColFilter("custom"))
240 | func WithNameColFilter(s string) TemplateParamOption {
241 | return func(c *TemplateParameterTestConfig) {
242 | c.nameColFilter = s
243 | }
244 | }
245 |
246 | // WithCreateColArray represents the columns array parameter for create-table-templateParams-tool.
247 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.WithCreateColArray("custom"))
248 | func WithCreateColArray(s string) TemplateParamOption {
249 | return func(c *TemplateParameterTestConfig) {
250 | c.createColArray = s
251 | }
252 | }
253 |
254 | // DisableDdlTest disables tests for ddl statements for sources that do not support ddl.
255 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.DisableDdlTest())
256 | func DisableDdlTest() TemplateParamOption {
257 | return func(c *TemplateParameterTestConfig) {
258 | c.supportDdl = false
259 | }
260 | }
261 |
262 | // DisableInsertTest disables tests of insert statements for sources that do not support insert.
263 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.DisableInsertTest())
264 | func DisableInsertTest() TemplateParamOption {
265 | return func(c *TemplateParameterTestConfig) {
266 | c.supportInsert = false
267 | }
268 | }
269 |
270 | // DisableInsertTest disables tests of select-fields-templateParams-tool test.
271 | // e.g. tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam, tests.DisableSelectFilterTest())
272 | func DisableSelectFilterTest() TemplateParamOption {
273 | return func(c *TemplateParameterTestConfig) {
274 | c.supportSelectFields = false
275 | }
276 | }
277 |
```
--------------------------------------------------------------------------------
/tests/cassandra/cassandra_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 cassandra
16 |
17 | import (
18 | "context"
19 | "fmt"
20 | "log"
21 | "os"
22 | "regexp"
23 | "strings"
24 | "testing"
25 | "time"
26 |
27 | "github.com/gocql/gocql"
28 | "github.com/google/uuid"
29 | "github.com/googleapis/genai-toolbox/internal/testutils"
30 | "github.com/googleapis/genai-toolbox/tests"
31 | )
32 |
33 | var (
34 | CassandraSourceKind = "cassandra"
35 | CassandraToolKind = "cassandra-cql"
36 | Hosts = os.Getenv("CASSANDRA_HOST")
37 | Keyspace = "example_keyspace"
38 | Username = os.Getenv("CASSANDRA_USER")
39 | Password = os.Getenv("CASSANDRA_PASS")
40 | )
41 |
42 | func getCassandraVars(t *testing.T) map[string]any {
43 | switch "" {
44 | case Hosts:
45 | t.Fatal("'Hosts' not set")
46 | case Username:
47 | t.Fatal("'Username' not set")
48 | case Password:
49 | t.Fatal("'Password' not set")
50 | }
51 | return map[string]any{
52 | "kind": CassandraSourceKind,
53 | "hosts": strings.Split(Hosts, ","),
54 | "keyspace": Keyspace,
55 | "username": Username,
56 | "password": Password,
57 | }
58 | }
59 |
60 | func initCassandraSession() (*gocql.Session, error) {
61 | hostStrings := strings.Split(Hosts, ",")
62 |
63 | var hosts []string
64 | for _, h := range hostStrings {
65 | trimmedHost := strings.TrimSpace(h)
66 | if trimmedHost != "" {
67 | hosts = append(hosts, trimmedHost)
68 | }
69 | }
70 | if len(hosts) == 0 {
71 | return nil, fmt.Errorf("no valid hosts found in CASSANDRA_HOSTS env var")
72 | }
73 | // Configure cluster connection
74 | cluster := gocql.NewCluster(hosts...)
75 | cluster.Consistency = gocql.Quorum
76 | cluster.ProtoVersion = 4
77 | cluster.DisableInitialHostLookup = true
78 | cluster.ConnectTimeout = 10 * time.Second
79 | cluster.NumConns = 2
80 | cluster.Authenticator = gocql.PasswordAuthenticator{
81 | Username: Username,
82 | Password: Password,
83 | }
84 | cluster.RetryPolicy = &gocql.ExponentialBackoffRetryPolicy{
85 | NumRetries: 3,
86 | Min: 200 * time.Millisecond,
87 | Max: 2 * time.Second,
88 | }
89 |
90 | // Create session
91 | session, err := cluster.CreateSession()
92 | if err != nil {
93 | return nil, fmt.Errorf("Failed to create session: %v", err)
94 | }
95 |
96 | // Create keyspace
97 | err = session.Query(fmt.Sprintf(`
98 | CREATE KEYSPACE IF NOT EXISTS %s
99 | WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}
100 | `, Keyspace)).Exec()
101 | if err != nil {
102 | return nil, fmt.Errorf("Failed to create keyspace: %v", err)
103 | }
104 |
105 | return session, nil
106 | }
107 |
108 | func initTable(tableName string, session *gocql.Session) error {
109 |
110 | // Create table with additional columns
111 | err := session.Query(fmt.Sprintf(`
112 | CREATE TABLE IF NOT EXISTS %s.%s (
113 | id int PRIMARY KEY,
114 | name text,
115 | email text,
116 | age int,
117 | is_active boolean,
118 | created_at timestamp
119 | )
120 | `, Keyspace, tableName)).Exec()
121 | if err != nil {
122 | return fmt.Errorf("Failed to create table: %v", err)
123 | }
124 |
125 | // Use fixed timestamps for reproducibility
126 | fixedTime, _ := time.Parse(time.RFC3339, "2025-07-25T12:00:00Z")
127 | dayAgo := fixedTime.Add(-24 * time.Hour)
128 | twelveHoursAgo := fixedTime.Add(-12 * time.Hour)
129 |
130 | // Insert minimal diverse data with fixed time.Time for timestamps
131 | err = session.Query(fmt.Sprintf(`
132 | INSERT INTO %s.%s (id, name,email, age, is_active, created_at)
133 | VALUES (?, ?, ?, ?, ?, ?)`, Keyspace, tableName),
134 | 3, "Alice", tests.ServiceAccountEmail, 25, true, dayAgo,
135 | ).Exec()
136 | if err != nil {
137 | return fmt.Errorf("Failed to insert user: %v", err)
138 | }
139 | err = session.Query(fmt.Sprintf(`
140 | INSERT INTO %s.%s (id, name,email, age, is_active, created_at)
141 | VALUES (?, ?, ?, ?, ?, ?)`, Keyspace, tableName),
142 | 2, "Alex", "[email protected]", 30, false, twelveHoursAgo,
143 | ).Exec()
144 | if err != nil {
145 | return fmt.Errorf("Failed to insert user: %v", err)
146 | }
147 | err = session.Query(fmt.Sprintf(`
148 | INSERT INTO %s.%s (id, name,email, age, is_active, created_at)
149 | VALUES (?, ?, ?, ?, ?, ?)`, Keyspace, tableName),
150 | 1, "Sid", "[email protected]", 10, true, fixedTime,
151 | ).Exec()
152 | if err != nil {
153 | return fmt.Errorf("Failed to insert user: %v", err)
154 | }
155 | err = session.Query(fmt.Sprintf(`
156 | INSERT INTO %s.%s (id, name,email, age, is_active, created_at)
157 | VALUES (?, ?, ?, ?, ?, ?)`, Keyspace, tableName),
158 | 4, nil, "[email protected]", 40, false, fixedTime,
159 | ).Exec()
160 | if err != nil {
161 | return fmt.Errorf("Failed to insert user: %v", err)
162 | }
163 | return nil
164 | }
165 |
166 | func dropTable(session *gocql.Session, tableName string) {
167 | err := session.Query(fmt.Sprintf("drop table %s.%s", Keyspace, tableName)).Exec()
168 | if err != nil {
169 | log.Printf("Failed to drop table %s: %v", tableName, err)
170 | }
171 | }
172 |
173 | func TestCassandra(t *testing.T) {
174 | session, err := initCassandraSession()
175 | if err != nil {
176 | t.Fatal(err)
177 | }
178 | defer session.Close()
179 | sourceConfig := getCassandraVars(t)
180 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
181 | defer cancel()
182 |
183 | var args []string
184 | paramTableName := "param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
185 | tableNameAuth := "auth_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
186 | tableNameTemplateParam := "template_param_table_" + strings.ReplaceAll(uuid.New().String(), "-", "")
187 | err = initTable(paramTableName, session)
188 | if err != nil {
189 | t.Fatal(err)
190 | }
191 | defer dropTable(session, paramTableName)
192 |
193 | err = initTable(tableNameAuth, session)
194 | if err != nil {
195 | t.Fatal(err)
196 | }
197 | defer dropTable(session, tableNameAuth)
198 |
199 | err = initTable(tableNameTemplateParam, session)
200 | if err != nil {
201 | t.Fatal(err)
202 | }
203 | defer dropTable(session, tableNameTemplateParam)
204 |
205 | paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt := createParamToolInfo(paramTableName)
206 | _, _, authToolStmt := getCassandraAuthToolInfo(tableNameAuth)
207 | toolsFile := tests.GetToolsConfig(sourceConfig, CassandraToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt)
208 |
209 | tmplSelectCombined, tmplSelectFilterCombined := getCassandraTmplToolInfo()
210 | tmpSelectAll := "SELECT * FROM {{.tableName}} where id = 1"
211 |
212 | toolsFile = tests.AddTemplateParamConfig(t, toolsFile, CassandraToolKind, tmplSelectCombined, tmplSelectFilterCombined, tmpSelectAll)
213 |
214 | cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
215 | if err != nil {
216 | t.Fatalf("command initialization returned an error: %s", err)
217 | }
218 | defer cleanup()
219 |
220 | waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
221 | defer cancel()
222 | out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
223 | if err != nil {
224 | t.Logf("toolbox command logs: \n%s", out)
225 | t.Fatalf("toolbox didn't start successfully: %s", err)
226 | }
227 | selectIdNameWant, selectIdNullWant, selectArrayParamWant, mcpMyFailToolWant, mcpSelect1Want, mcpMyToolIdWant := getCassandraWants()
228 | selectAllWant, selectIdWant, selectNameWant := getCassandraTmplWants()
229 |
230 | tests.RunToolGetTest(t)
231 | tests.RunToolInvokeTest(t, "", tests.DisableSelect1Test(),
232 | tests.DisableOptionalNullParamTest(),
233 | tests.WithMyToolId3NameAliceWant(selectIdNameWant),
234 | tests.WithMyToolById4Want(selectIdNullWant),
235 | tests.WithMyArrayToolWant(selectArrayParamWant),
236 | tests.DisableSelect1AuthTest())
237 | tests.RunToolInvokeWithTemplateParameters(t, tableNameTemplateParam,
238 | tests.DisableSelectFilterTest(),
239 | tests.WithSelectAllWant(selectAllWant),
240 | tests.DisableDdlTest(), tests.DisableInsertTest(), tests.WithTmplSelectId1Want(selectIdWant), tests.WithTmplSelectNameWant(selectNameWant))
241 |
242 | tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want,
243 | tests.WithMcpMyToolId3NameAliceWant(mcpMyToolIdWant),
244 | tests.DisableMcpSelect1AuthTest())
245 |
246 | }
247 |
248 | func createParamToolInfo(tableName string) (string, string, string, string) {
249 | toolStatement := fmt.Sprintf("SELECT id, name FROM %s WHERE id = ? AND name = ? ALLOW FILTERING;", tableName)
250 | idParamStatement := fmt.Sprintf("SELECT id,name FROM %s WHERE id = ?;", tableName)
251 | nameParamStatement := fmt.Sprintf("SELECT id, name FROM %s WHERE name = ? ALLOW FILTERING;", tableName)
252 | arrayToolStatement := fmt.Sprintf("SELECT id, name FROM %s WHERE id IN ? AND name IN ? ALLOW FILTERING;", tableName)
253 | return toolStatement, idParamStatement, nameParamStatement, arrayToolStatement
254 |
255 | }
256 |
257 | func getCassandraAuthToolInfo(tableName string) (string, string, string) {
258 | createStatement := fmt.Sprintf("CREATE TABLE %s (id UUID PRIMARY KEY, name TEXT, email TEXT);", tableName)
259 | insertStatement := fmt.Sprintf("INSERT INTO %s (id, name, email) VALUES (uuid(), ?, ?), (uuid(), ?, ?);", tableName)
260 | toolStatement := fmt.Sprintf("SELECT name FROM %s WHERE email = ? ALLOW FILTERING;", tableName)
261 | return createStatement, insertStatement, toolStatement
262 | }
263 |
264 | func getCassandraTmplToolInfo() (string, string) {
265 | selectAllTemplateStmt := "SELECT age, id, name FROM {{.tableName}} where id = ?;"
266 | selectByIdTemplateStmt := "SELECT id, name FROM {{.tableName}} WHERE name = ? ALLOW FILTERING;"
267 | return selectAllTemplateStmt, selectByIdTemplateStmt
268 | }
269 |
270 | func getCassandraWants() (string, string, string, string, string, string) {
271 | selectIdNameWant := "[{\"id\":3,\"name\":\"Alice\"}]"
272 | selectIdNullWant := "[{\"id\":4,\"name\":\"\"}]"
273 | selectArrayParamWant := "[{\"id\":1,\"name\":\"Sid\"},{\"id\":3,\"name\":\"Alice\"}]"
274 | mcpMyFailToolWant := "{\"jsonrpc\":\"2.0\",\"id\":\"invoke-fail-tool\",\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"unable to parse rows: line 1:0 no viable alternative at input 'SELEC' ([SELEC]...)\"}],\"isError\":true}}"
275 | mcpMyToolIdWant := "{\"jsonrpc\":\"2.0\",\"id\":\"my-tool\",\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"[{\\\"id\\\":3,\\\"name\\\":\\\"Alice\\\"}]\"}]}}"
276 | return selectIdNameWant, selectIdNullWant, selectArrayParamWant, mcpMyFailToolWant, "nil", mcpMyToolIdWant
277 | }
278 |
279 | func getCassandraTmplWants() (string, string, string) {
280 | selectAllWant := "[{\"age\":10,\"created_at\":\"2025-07-25T12:00:00Z\",\"email\":\"[email protected]\",\"id\":1,\"is_active\":true,\"name\":\"Sid\"}]"
281 | selectIdWant := "[{\"age\":10,\"id\":1,\"name\":\"Sid\"}]"
282 | selectNameWant := "[{\"id\":2,\"name\":\"Alex\"}]"
283 | return selectAllWant, selectIdWant, selectNameWant
284 | }
285 |
```
--------------------------------------------------------------------------------
/tests/couchbase/couchbase_integration_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
16 |
17 | import (
18 | "context"
19 | "fmt"
20 | "os"
21 | "regexp"
22 | "strings"
23 | "testing"
24 | "time"
25 |
26 | "github.com/couchbase/gocb/v2"
27 | "github.com/google/uuid"
28 | "github.com/googleapis/genai-toolbox/internal/testutils"
29 | "github.com/googleapis/genai-toolbox/tests"
30 | )
31 |
32 | const (
33 | couchbaseSourceKind = "couchbase"
34 | couchbaseToolKind = "couchbase-sql"
35 | )
36 |
37 | var (
38 | couchbaseConnection = os.Getenv("COUCHBASE_CONNECTION")
39 | couchbaseBucket = os.Getenv("COUCHBASE_BUCKET")
40 | couchbaseScope = os.Getenv("COUCHBASE_SCOPE")
41 | couchbaseUser = os.Getenv("COUCHBASE_USER")
42 | couchbasePass = os.Getenv("COUCHBASE_PASS")
43 | )
44 |
45 | // getCouchbaseVars validates and returns Couchbase configuration variables
46 | func getCouchbaseVars(t *testing.T) map[string]any {
47 | switch "" {
48 | case couchbaseConnection:
49 | t.Fatal("'COUCHBASE_CONNECTION' not set")
50 | case couchbaseBucket:
51 | t.Fatal("'COUCHBASE_BUCKET' not set")
52 | case couchbaseScope:
53 | t.Fatal("'COUCHBASE_SCOPE' not set")
54 | case couchbaseUser:
55 | t.Fatal("'COUCHBASE_USER' not set")
56 | case couchbasePass:
57 | t.Fatal("'COUCHBASE_PASS' not set")
58 | }
59 |
60 | return map[string]any{
61 | "kind": couchbaseSourceKind,
62 | "connectionString": couchbaseConnection,
63 | "bucket": couchbaseBucket,
64 | "scope": couchbaseScope,
65 | "username": couchbaseUser,
66 | "password": couchbasePass,
67 | "queryScanConsistency": 2,
68 | }
69 | }
70 |
71 | // initCouchbaseCluster initializes a connection to the Couchbase cluster
72 | func initCouchbaseCluster(connectionString, username, password string) (*gocb.Cluster, error) {
73 | opts := gocb.ClusterOptions{
74 | Authenticator: gocb.PasswordAuthenticator{
75 | Username: username,
76 | Password: password,
77 | },
78 | }
79 |
80 | cluster, err := gocb.Connect(connectionString, opts)
81 | if err != nil {
82 | return nil, fmt.Errorf("gocb.Connect: %w", err)
83 | }
84 | return cluster, nil
85 | }
86 |
87 | func TestCouchbaseToolEndpoints(t *testing.T) {
88 | sourceConfig := getCouchbaseVars(t)
89 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
90 | defer cancel()
91 |
92 | var args []string
93 |
94 | cluster, err := initCouchbaseCluster(couchbaseConnection, couchbaseUser, couchbasePass)
95 | if err != nil {
96 | t.Fatalf("unable to create Couchbase connection: %s", err)
97 | }
98 | defer cluster.Close(nil)
99 |
100 | // Create collection names with UUID
101 | collectionNameParam := "param_" + strings.ReplaceAll(uuid.New().String(), "-", "")
102 | collectionNameAuth := "auth_" + strings.ReplaceAll(uuid.New().String(), "-", "")
103 | collectionNameTemplateParam := "template_param_" + strings.ReplaceAll(uuid.New().String(), "-", "")
104 |
105 | // Set up data for param tool
106 | paramToolStatement, idParamToolStmt, nameParamToolStmt, arrayToolStatement, paramTestParams := getCouchbaseParamToolInfo(collectionNameParam)
107 | teardownCollection1 := setupCouchbaseCollection(t, ctx, cluster, couchbaseBucket, couchbaseScope, collectionNameParam, paramTestParams)
108 | defer teardownCollection1(t)
109 |
110 | // Set up data for auth tool
111 | authToolStatement, authTestParams := getCouchbaseAuthToolInfo(collectionNameAuth)
112 | teardownCollection2 := setupCouchbaseCollection(t, ctx, cluster, couchbaseBucket, couchbaseScope, collectionNameAuth, authTestParams)
113 | defer teardownCollection2(t)
114 |
115 | // Setup up table for template param tool
116 | tmplSelectCombined, tmplSelectFilterCombined, tmplSelectAll, params3 := getCouchbaseTemplateParamToolInfo()
117 | teardownCollection3 := setupCouchbaseCollection(t, ctx, cluster, couchbaseBucket, couchbaseScope, collectionNameTemplateParam, params3)
118 | defer teardownCollection3(t)
119 |
120 | // Write config into a file and pass it to command
121 | toolsFile := tests.GetToolsConfig(sourceConfig, couchbaseToolKind, paramToolStatement, idParamToolStmt, nameParamToolStmt, arrayToolStatement, authToolStatement)
122 | toolsFile = tests.AddTemplateParamConfig(t, toolsFile, couchbaseToolKind, tmplSelectCombined, tmplSelectFilterCombined, tmplSelectAll)
123 |
124 | cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
125 | if err != nil {
126 | t.Fatalf("command initialization returned an error: %s", err)
127 | }
128 | defer cleanup()
129 |
130 | waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
131 | defer cancel()
132 | out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
133 | if err != nil {
134 | t.Logf("toolbox command logs: \n%s", out)
135 | t.Fatalf("toolbox didn't start successfully: %s", err)
136 | }
137 |
138 | // Get configs for tests
139 | select1Want := "[{\"$1\":1}]"
140 | mcpMyFailToolWant := `{"jsonrpc":"2.0","id":"invoke-fail-tool","result":{"content":[{"type":"text","text":"unable to execute query: parsing failure | {\"statement\":\"SELEC 1;\"`
141 | mcpSelect1Want := `{"jsonrpc":"2.0","id":"invoke my-auth-required-tool","result":{"content":[{"type":"text","text":"{\"$1\":1}"}]}}`
142 | tmplSelectId1Want := "[{\"age\":21,\"id\":1,\"name\":\"Alex\"}]"
143 | selectAllWant := "[{\"age\":21,\"id\":1,\"name\":\"Alex\"},{\"age\":100,\"id\":2,\"name\":\"Alice\"}]"
144 |
145 | // Run tests
146 | tests.RunToolGetTest(t)
147 | tests.RunToolInvokeTest(t, select1Want)
148 | tests.RunMCPToolCallMethod(t, mcpMyFailToolWant, mcpSelect1Want)
149 | tests.RunToolInvokeWithTemplateParameters(t, collectionNameTemplateParam,
150 | tests.WithTmplSelectId1Want(tmplSelectId1Want),
151 | tests.WithSelectAllWant(selectAllWant),
152 | tests.DisableDdlTest(),
153 | tests.DisableInsertTest(),
154 | )
155 | }
156 |
157 | // setupCouchbaseCollection creates a scope and collection and inserts test data
158 | func setupCouchbaseCollection(t *testing.T, ctx context.Context, cluster *gocb.Cluster,
159 | bucketName, scopeName, collectionName string, params []map[string]any) func(t *testing.T) {
160 |
161 | // Get bucket reference
162 | bucket := cluster.Bucket(bucketName)
163 |
164 | // Wait for bucket to be ready
165 | err := bucket.WaitUntilReady(5*time.Second, nil)
166 | if err != nil {
167 | t.Fatalf("failed to connect to bucket: %v", err)
168 | }
169 |
170 | // Create scope if it doesn't exist
171 | bucketMgr := bucket.CollectionsV2()
172 | err = bucketMgr.CreateScope(scopeName, nil)
173 | if err != nil && !strings.Contains(err.Error(), "already exists") {
174 | t.Logf("failed to create scope (might already exist): %v", err)
175 | }
176 |
177 | // Create a collection if it doesn't exist
178 | err = bucketMgr.CreateCollection(scopeName, collectionName, nil, nil)
179 | if err != nil && !strings.Contains(err.Error(), "already exists") {
180 | t.Fatalf("failed to create collection: %v", err)
181 | }
182 |
183 | // Get a reference to the collection
184 | collection := bucket.Scope(scopeName).Collection(collectionName)
185 |
186 | // Create primary index if it doesn't exist
187 | // Create primary index with retry logic
188 | maxRetries := 5
189 | retryDelay := 50 * time.Millisecond
190 | actualRetries := 0
191 | var lastErr error
192 | for attempt := 0; attempt < maxRetries; attempt++ {
193 | err = collection.QueryIndexes().CreatePrimaryIndex(
194 | &gocb.CreatePrimaryQueryIndexOptions{
195 | IgnoreIfExists: true,
196 | })
197 | if err == nil {
198 | lastErr = err // clear previous error
199 | break
200 | }
201 |
202 | lastErr = err
203 | t.Logf("Attempt %d: failed to create primary index: %v, retrying in %v", attempt+1, err, retryDelay)
204 | time.Sleep(retryDelay)
205 | // Exponential backoff
206 | retryDelay *= 2
207 | actualRetries += 1
208 | }
209 |
210 | if lastErr != nil {
211 | t.Fatalf("failed to create primary index collection after %d attempts: %v", actualRetries, lastErr)
212 | }
213 |
214 | // Insert test documents
215 | for i, param := range params {
216 | _, err = collection.Upsert(fmt.Sprintf("%d", i+1), param, &gocb.UpsertOptions{
217 | DurabilityLevel: gocb.DurabilityLevelMajority,
218 | })
219 | if err != nil {
220 | t.Fatalf("failed to insert test data: %v", err)
221 | }
222 | }
223 |
224 | // Return a cleanup function
225 | return func(t *testing.T) {
226 | // Drop the collection
227 | err := bucketMgr.DropCollection(scopeName, collectionName, nil)
228 | if err != nil {
229 | t.Logf("failed to drop collection: %v", err)
230 | }
231 | }
232 | }
233 |
234 | // getCouchbaseParamToolInfo returns statements and params for my-tool couchbase-sql kind
235 | func getCouchbaseParamToolInfo(collectionName string) (string, string, string, string, []map[string]any) {
236 | // N1QL uses positional or named parameters with $ prefix
237 | toolStatement := fmt.Sprintf("SELECT TONUMBER(meta().id) as id, "+
238 | "%s.* FROM %s WHERE meta().id = TOSTRING($id) OR name = $name order by meta().id",
239 | collectionName, collectionName)
240 | idToolStatement := fmt.Sprintf("SELECT TONUMBER(meta().id) as id, "+
241 | "%s.* FROM %s WHERE meta().id = TOSTRING($id) order by meta().id",
242 | collectionName, collectionName)
243 | nameToolStatement := fmt.Sprintf("SELECT TONUMBER(meta().id) as id, "+
244 | "%s.* FROM %s WHERE name = $name order by meta().id",
245 | collectionName, collectionName)
246 | arrayToolStatemnt := fmt.Sprintf("SELECT TONUMBER(meta().id) as id, "+
247 | "%s.* FROM %s WHERE TONUMBER(meta().id) IN $idArray AND name IN $nameArray order by meta().id", collectionName, collectionName)
248 | params := []map[string]any{
249 | {"name": "Alice"},
250 | {"name": "Jane"},
251 | {"name": "Sid"},
252 | {"name": nil},
253 | }
254 | return toolStatement, idToolStatement, nameToolStatement, arrayToolStatemnt, params
255 | }
256 |
257 | // getCouchbaseAuthToolInfo returns statements and param of my-auth-tool for couchbase-sql kind
258 | func getCouchbaseAuthToolInfo(collectionName string) (string, []map[string]any) {
259 | toolStatement := fmt.Sprintf("SELECT name FROM %s WHERE email = $email", collectionName)
260 |
261 | params := []map[string]any{
262 | {"name": "Alice", "email": tests.ServiceAccountEmail},
263 | {"name": "Jane", "email": "[email protected]"},
264 | }
265 | return toolStatement, params
266 | }
267 |
268 | func getCouchbaseTemplateParamToolInfo() (string, string, string, []map[string]any) {
269 | tmplSelectCombined := "SELECT {{.tableName}}.* FROM {{.tableName}} WHERE id = $id"
270 | tmplSelectFilterCombined := "SELECT {{.tableName}}.* FROM {{.tableName}} WHERE {{.columnFilter}} = $name"
271 | tmplSelectAll := "SELECT {{.tableName}}.* FROM {{.tableName}}"
272 |
273 | params := []map[string]any{
274 | {"name": "Alex", "id": 1, "age": 21},
275 | {"name": "Alice", "id": 2, "age": 100},
276 | }
277 | return tmplSelectCombined, tmplSelectFilterCombined, tmplSelectAll, params
278 | }
279 |
```
--------------------------------------------------------------------------------
/internal/tools/firestore/firestoreupdatedocument/firestoreupdatedocument.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 firestoreupdatedocument
16 |
17 | import (
18 | "context"
19 | "fmt"
20 | "strings"
21 |
22 | firestoreapi "cloud.google.com/go/firestore"
23 | yaml "github.com/goccy/go-yaml"
24 | "github.com/googleapis/genai-toolbox/internal/sources"
25 | firestoreds "github.com/googleapis/genai-toolbox/internal/sources/firestore"
26 | "github.com/googleapis/genai-toolbox/internal/tools"
27 | "github.com/googleapis/genai-toolbox/internal/tools/firestore/util"
28 | )
29 |
30 | const kind string = "firestore-update-document"
31 | const documentPathKey string = "documentPath"
32 | const documentDataKey string = "documentData"
33 | const updateMaskKey string = "updateMask"
34 | const returnDocumentDataKey string = "returnData"
35 |
36 | func init() {
37 | if !tools.Register(kind, newConfig) {
38 | panic(fmt.Sprintf("tool kind %q already registered", kind))
39 | }
40 | }
41 |
42 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
43 | actual := Config{Name: name}
44 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
45 | return nil, err
46 | }
47 | return actual, nil
48 | }
49 |
50 | type compatibleSource interface {
51 | FirestoreClient() *firestoreapi.Client
52 | }
53 |
54 | // validate compatible sources are still compatible
55 | var _ compatibleSource = &firestoreds.Source{}
56 |
57 | var compatibleSources = [...]string{firestoreds.SourceKind}
58 |
59 | type Config struct {
60 | Name string `yaml:"name" validate:"required"`
61 | Kind string `yaml:"kind" validate:"required"`
62 | Source string `yaml:"source" validate:"required"`
63 | Description string `yaml:"description" validate:"required"`
64 | AuthRequired []string `yaml:"authRequired"`
65 | }
66 |
67 | // validate interface
68 | var _ tools.ToolConfig = Config{}
69 |
70 | func (cfg Config) ToolConfigKind() string {
71 | return kind
72 | }
73 |
74 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
75 | // verify source exists
76 | rawS, ok := srcs[cfg.Source]
77 | if !ok {
78 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
79 | }
80 |
81 | // verify the source is compatible
82 | s, ok := rawS.(compatibleSource)
83 | if !ok {
84 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
85 | }
86 |
87 | // Create parameters
88 | documentPathParameter := tools.NewStringParameter(
89 | documentPathKey,
90 | "The relative path of the document which needs to be updated (e.g., 'users/userId' or 'users/userId/posts/postId'). Note: This is a relative path, NOT an absolute path like 'projects/{project_id}/databases/{database_id}/documents/...'",
91 | )
92 |
93 | documentDataParameter := tools.NewMapParameter(
94 | documentDataKey,
95 | `The document data in Firestore's native JSON format. Each field must be wrapped with a type indicator:
96 | - Strings: {"stringValue": "text"}
97 | - Integers: {"integerValue": "123"} or {"integerValue": 123}
98 | - Doubles: {"doubleValue": 123.45}
99 | - Booleans: {"booleanValue": true}
100 | - Timestamps: {"timestampValue": "2025-01-07T10:00:00Z"}
101 | - GeoPoints: {"geoPointValue": {"latitude": 34.05, "longitude": -118.24}}
102 | - Arrays: {"arrayValue": {"values": [{"stringValue": "item1"}, {"integerValue": "2"}]}}
103 | - Maps: {"mapValue": {"fields": {"key1": {"stringValue": "value1"}, "key2": {"booleanValue": true}}}}
104 | - Null: {"nullValue": null}
105 | - Bytes: {"bytesValue": "base64EncodedString"}
106 | - References: {"referenceValue": "collection/document"}`,
107 | "", // Empty string for generic map that accepts any value type
108 | )
109 |
110 | updateMaskParameter := tools.NewArrayParameterWithRequired(
111 | updateMaskKey,
112 | "The selective fields to update. If not provided, all fields in documentData will be updated. When provided, only the specified fields will be updated. Fields referenced in the mask but not present in documentData will be deleted from the document",
113 | false, // not required
114 | tools.NewStringParameter("field", "Field path to update or delete. Use dot notation to access nested fields within maps (e.g., 'address.city' to update the city field within an address map, or 'user.profile.name' for deeply nested fields). To delete a field, include it in the mask but omit it from documentData. Note: You cannot update individual array elements; you must update the entire array field"),
115 | )
116 |
117 | returnDataParameter := tools.NewBooleanParameterWithDefault(
118 | returnDocumentDataKey,
119 | false,
120 | "If set to true the output will have the data of the updated document. This flag if set to false will help avoid overloading the context of the agent.",
121 | )
122 |
123 | parameters := tools.Parameters{
124 | documentPathParameter,
125 | documentDataParameter,
126 | updateMaskParameter,
127 | returnDataParameter,
128 | }
129 |
130 | mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters)
131 |
132 | // finish tool setup
133 | t := Tool{
134 | Name: cfg.Name,
135 | Kind: kind,
136 | Parameters: parameters,
137 | AuthRequired: cfg.AuthRequired,
138 | Client: s.FirestoreClient(),
139 | manifest: tools.Manifest{Description: cfg.Description, Parameters: parameters.Manifest(), AuthRequired: cfg.AuthRequired},
140 | mcpManifest: mcpManifest,
141 | }
142 | return t, nil
143 | }
144 |
145 | // validate interface
146 | var _ tools.Tool = Tool{}
147 |
148 | type Tool struct {
149 | Name string `yaml:"name"`
150 | Kind string `yaml:"kind"`
151 | AuthRequired []string `yaml:"authRequired"`
152 | Parameters tools.Parameters `yaml:"parameters"`
153 |
154 | Client *firestoreapi.Client
155 | manifest tools.Manifest
156 | mcpManifest tools.McpManifest
157 | }
158 |
159 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
160 | mapParams := params.AsMap()
161 |
162 | // Get document path
163 | documentPath, ok := mapParams[documentPathKey].(string)
164 | if !ok || documentPath == "" {
165 | return nil, fmt.Errorf("invalid or missing '%s' parameter", documentPathKey)
166 | }
167 |
168 | // Validate document path
169 | if err := util.ValidateDocumentPath(documentPath); err != nil {
170 | return nil, fmt.Errorf("invalid document path: %w", err)
171 | }
172 |
173 | // Get document data
174 | documentDataRaw, ok := mapParams[documentDataKey]
175 | if !ok {
176 | return nil, fmt.Errorf("invalid or missing '%s' parameter", documentDataKey)
177 | }
178 |
179 | // Get update mask if provided
180 | var updatePaths []string
181 | if updateMaskRaw, ok := mapParams[updateMaskKey]; ok && updateMaskRaw != nil {
182 | if updateMaskArray, ok := updateMaskRaw.([]any); ok {
183 | // Use ConvertAnySliceToTyped to convert the slice
184 | typedSlice, err := tools.ConvertAnySliceToTyped(updateMaskArray, "string")
185 | if err != nil {
186 | return nil, fmt.Errorf("failed to convert update mask: %w", err)
187 | }
188 | updatePaths, ok = typedSlice.([]string)
189 | if !ok {
190 | return nil, fmt.Errorf("unexpected type conversion error for update mask")
191 | }
192 | }
193 | }
194 |
195 | // Get return document data flag
196 | returnData := false
197 | if val, ok := mapParams[returnDocumentDataKey].(bool); ok {
198 | returnData = val
199 | }
200 |
201 | // Get the document reference
202 | docRef := t.Client.Doc(documentPath)
203 |
204 | // Prepare update data
205 | var writeResult *firestoreapi.WriteResult
206 | var writeErr error
207 |
208 | if len(updatePaths) > 0 {
209 | // Use selective field update with update mask
210 | updates := make([]firestoreapi.Update, 0, len(updatePaths))
211 |
212 | // Convert document data without delete markers
213 | dataMap, err := util.JSONToFirestoreValue(documentDataRaw, t.Client)
214 | if err != nil {
215 | return nil, fmt.Errorf("failed to convert document data: %w", err)
216 | }
217 |
218 | // Ensure it's a map
219 | dataMapTyped, ok := dataMap.(map[string]interface{})
220 | if !ok {
221 | return nil, fmt.Errorf("document data must be a map")
222 | }
223 |
224 | for _, path := range updatePaths {
225 | // Get the value for this path from the document data
226 | value, exists := getFieldValue(dataMapTyped, path)
227 | if !exists {
228 | // Field not in document data but in mask - delete it
229 | value = firestoreapi.Delete
230 | }
231 |
232 | updates = append(updates, firestoreapi.Update{
233 | Path: path,
234 | Value: value,
235 | })
236 | }
237 |
238 | writeResult, writeErr = docRef.Update(ctx, updates)
239 | } else {
240 | // Update all fields in the document data (merge)
241 | documentData, err := util.JSONToFirestoreValue(documentDataRaw, t.Client)
242 | if err != nil {
243 | return nil, fmt.Errorf("failed to convert document data: %w", err)
244 | }
245 | writeResult, writeErr = docRef.Set(ctx, documentData, firestoreapi.MergeAll)
246 | }
247 |
248 | if writeErr != nil {
249 | return nil, fmt.Errorf("failed to update document: %w", writeErr)
250 | }
251 |
252 | // Build the response
253 | response := map[string]any{
254 | "documentPath": docRef.Path,
255 | "updateTime": writeResult.UpdateTime.Format("2006-01-02T15:04:05.999999999Z"),
256 | }
257 |
258 | // Add document data if requested
259 | if returnData {
260 | // Fetch the updated document to return the current state
261 | snapshot, err := docRef.Get(ctx)
262 | if err != nil {
263 | return nil, fmt.Errorf("failed to retrieve updated document: %w", err)
264 | }
265 |
266 | // Convert the document data to simple JSON format
267 | simplifiedData := util.FirestoreValueToJSON(snapshot.Data())
268 | response["documentData"] = simplifiedData
269 | }
270 |
271 | return response, nil
272 | }
273 |
274 | // getFieldValue retrieves a value from a nested map using a dot-separated path
275 | func getFieldValue(data map[string]interface{}, path string) (interface{}, bool) {
276 | // Split the path by dots for nested field access
277 | parts := strings.Split(path, ".")
278 |
279 | current := data
280 | for i, part := range parts {
281 | if i == len(parts)-1 {
282 | // Last part - return the value
283 | if value, exists := current[part]; exists {
284 | return value, true
285 | }
286 | return nil, false
287 | }
288 |
289 | // Navigate deeper into the structure
290 | if next, ok := current[part].(map[string]interface{}); ok {
291 | current = next
292 | } else {
293 | return nil, false
294 | }
295 | }
296 |
297 | return nil, false
298 | }
299 |
300 | func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
301 | return tools.ParseParams(t.Parameters, data, claims)
302 | }
303 |
304 | func (t Tool) Manifest() tools.Manifest {
305 | return t.manifest
306 | }
307 |
308 | func (t Tool) McpManifest() tools.McpManifest {
309 | return t.mcpManifest
310 | }
311 |
312 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
313 | return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
314 | }
315 |
316 | func (t Tool) RequiresClientAuthorization() bool {
317 | return false
318 | }
319 |
```
--------------------------------------------------------------------------------
/internal/log/log_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 log
16 |
17 | import (
18 | "bytes"
19 | "context"
20 | "encoding/json"
21 | "log/slog"
22 | "strings"
23 | "testing"
24 |
25 | "github.com/google/go-cmp/cmp"
26 | )
27 |
28 | func TestSeverityToLevel(t *testing.T) {
29 | tcs := []struct {
30 | name string
31 | in string
32 | want slog.Level
33 | }{
34 | {
35 | name: "test debug",
36 | in: "Debug",
37 | want: slog.LevelDebug,
38 | },
39 | {
40 | name: "test info",
41 | in: "Info",
42 | want: slog.LevelInfo,
43 | },
44 | {
45 | name: "test warn",
46 | in: "Warn",
47 | want: slog.LevelWarn,
48 | },
49 | {
50 | name: "test error",
51 | in: "Error",
52 | want: slog.LevelError,
53 | },
54 | }
55 | for _, tc := range tcs {
56 | t.Run(tc.name, func(t *testing.T) {
57 | got, err := SeverityToLevel(tc.in)
58 | if err != nil {
59 | t.Fatalf("unexpected error: %s", err)
60 | }
61 | if got != tc.want {
62 | t.Fatalf("incorrect level to severity: got %v, want %v", got, tc.want)
63 | }
64 |
65 | })
66 | }
67 | }
68 |
69 | func TestSeverityToLevelError(t *testing.T) {
70 | _, err := SeverityToLevel("fail")
71 | if err == nil {
72 | t.Fatalf("expected error on incorrect level")
73 | }
74 | }
75 |
76 | func TestLevelToSeverity(t *testing.T) {
77 | tcs := []struct {
78 | name string
79 | in string
80 | want string
81 | }{
82 | {
83 | name: "test debug",
84 | in: slog.LevelDebug.String(),
85 | want: "DEBUG",
86 | },
87 | {
88 | name: "test info",
89 | in: slog.LevelInfo.String(),
90 | want: "INFO",
91 | },
92 | {
93 | name: "test warn",
94 | in: slog.LevelWarn.String(),
95 | want: "WARN",
96 | },
97 | {
98 | name: "test error",
99 | in: slog.LevelError.String(),
100 | want: "ERROR",
101 | },
102 | }
103 | for _, tc := range tcs {
104 | t.Run(tc.name, func(t *testing.T) {
105 | got, err := levelToSeverity(tc.in)
106 | if err != nil {
107 | t.Fatalf("unexpected error: %s", err)
108 | }
109 | if got != tc.want {
110 | t.Fatalf("incorrect level to severity: got %v, want %v", got, tc.want)
111 | }
112 |
113 | })
114 | }
115 | }
116 |
117 | func TestLevelToSeverityError(t *testing.T) {
118 | _, err := levelToSeverity("fail")
119 | if err == nil {
120 | t.Fatalf("expected error on incorrect slog level")
121 | }
122 | }
123 |
124 | func runLogger(logger Logger, logMsg string) {
125 | ctx := context.Background()
126 | switch logMsg {
127 | case "info":
128 | logger.InfoContext(ctx, "log info")
129 | case "debug":
130 | logger.DebugContext(ctx, "log debug")
131 | case "warn":
132 | logger.WarnContext(ctx, "log warn")
133 | case "error":
134 | logger.ErrorContext(ctx, "log error")
135 | }
136 | }
137 |
138 | func TestStdLogger(t *testing.T) {
139 | tcs := []struct {
140 | name string
141 | logLevel string
142 | logMsg string
143 | wantOut string
144 | wantErr string
145 | }{
146 | {
147 | name: "debug logger logging debug",
148 | logLevel: "debug",
149 | logMsg: "debug",
150 | wantOut: "DEBUG \"log debug\" \n",
151 | wantErr: "",
152 | },
153 | {
154 | name: "info logger logging debug",
155 | logLevel: "info",
156 | logMsg: "debug",
157 | wantOut: "",
158 | wantErr: "",
159 | },
160 | {
161 | name: "warn logger logging debug",
162 | logLevel: "warn",
163 | logMsg: "debug",
164 | wantOut: "",
165 | wantErr: "",
166 | },
167 | {
168 | name: "error logger logging debug",
169 | logLevel: "error",
170 | logMsg: "debug",
171 | wantOut: "",
172 | wantErr: "",
173 | },
174 | {
175 | name: "debug logger logging info",
176 | logLevel: "debug",
177 | logMsg: "info",
178 | wantOut: "INFO \"log info\" \n",
179 | wantErr: "",
180 | },
181 | {
182 | name: "info logger logging info",
183 | logLevel: "info",
184 | logMsg: "info",
185 | wantOut: "INFO \"log info\" \n",
186 | wantErr: "",
187 | },
188 | {
189 | name: "warn logger logging info",
190 | logLevel: "warn",
191 | logMsg: "info",
192 | wantOut: "",
193 | wantErr: "",
194 | },
195 | {
196 | name: "error logger logging info",
197 | logLevel: "error",
198 | logMsg: "info",
199 | wantOut: "",
200 | wantErr: "",
201 | },
202 | {
203 | name: "debug logger logging warn",
204 | logLevel: "debug",
205 | logMsg: "warn",
206 | wantOut: "",
207 | wantErr: "WARN \"log warn\" \n",
208 | },
209 | {
210 | name: "info logger logging warn",
211 | logLevel: "info",
212 | logMsg: "warn",
213 | wantOut: "",
214 | wantErr: "WARN \"log warn\" \n",
215 | },
216 | {
217 | name: "warn logger logging warn",
218 | logLevel: "warn",
219 | logMsg: "warn",
220 | wantOut: "",
221 | wantErr: "WARN \"log warn\" \n",
222 | },
223 | {
224 | name: "error logger logging warn",
225 | logLevel: "error",
226 | logMsg: "warn",
227 | wantOut: "",
228 | wantErr: "",
229 | },
230 | {
231 | name: "debug logger logging error",
232 | logLevel: "debug",
233 | logMsg: "error",
234 | wantOut: "",
235 | wantErr: "ERROR \"log error\" \n",
236 | },
237 | {
238 | name: "info logger logging error",
239 | logLevel: "info",
240 | logMsg: "error",
241 | wantOut: "",
242 | wantErr: "ERROR \"log error\" \n",
243 | },
244 | {
245 | name: "warn logger logging error",
246 | logLevel: "warn",
247 | logMsg: "error",
248 | wantOut: "",
249 | wantErr: "ERROR \"log error\" \n",
250 | },
251 | {
252 | name: "error logger logging error",
253 | logLevel: "error",
254 | logMsg: "error",
255 | wantOut: "",
256 | wantErr: "ERROR \"log error\" \n",
257 | },
258 | }
259 | for _, tc := range tcs {
260 | t.Run(tc.name, func(t *testing.T) {
261 | outW := new(bytes.Buffer)
262 | errW := new(bytes.Buffer)
263 |
264 | logger, err := NewStdLogger(outW, errW, tc.logLevel)
265 | if err != nil {
266 | t.Fatalf("unexpected error: %s", err)
267 | }
268 | runLogger(logger, tc.logMsg)
269 |
270 | outWString := outW.String()
271 | spaceIndexOut := strings.Index(outWString, " ")
272 | gotOut := outWString[spaceIndexOut+1:]
273 |
274 | errWString := errW.String()
275 | spaceIndexErr := strings.Index(errWString, " ")
276 | gotErr := errWString[spaceIndexErr+1:]
277 |
278 | if diff := cmp.Diff(gotOut, tc.wantOut); diff != "" {
279 | t.Fatalf("incorrect log: diff %v", diff)
280 | }
281 | if diff := cmp.Diff(gotErr, tc.wantErr); diff != "" {
282 | t.Fatalf("incorrect log: diff %v", diff)
283 | }
284 | })
285 | }
286 | }
287 |
288 | func TestStructuredLoggerDebugLog(t *testing.T) {
289 | tcs := []struct {
290 | name string
291 | logLevel string
292 | logMsg string
293 | wantOut map[string]string
294 | wantErr map[string]string
295 | }{
296 | {
297 | name: "debug logger logging debug",
298 | logLevel: "debug",
299 | logMsg: "debug",
300 | wantOut: map[string]string{
301 | "severity": "DEBUG",
302 | "message": "log debug",
303 | },
304 | wantErr: map[string]string{},
305 | },
306 | {
307 | name: "info logger logging debug",
308 | logLevel: "info",
309 | logMsg: "debug",
310 | wantOut: map[string]string{},
311 | wantErr: map[string]string{},
312 | },
313 | {
314 | name: "warn logger logging debug",
315 | logLevel: "warn",
316 | logMsg: "debug",
317 | wantOut: map[string]string{},
318 | wantErr: map[string]string{},
319 | },
320 | {
321 | name: "error logger logging debug",
322 | logLevel: "error",
323 | logMsg: "debug",
324 | wantOut: map[string]string{},
325 | wantErr: map[string]string{},
326 | },
327 | {
328 | name: "debug logger logging info",
329 | logLevel: "debug",
330 | logMsg: "info",
331 | wantOut: map[string]string{
332 | "severity": "INFO",
333 | "message": "log info",
334 | },
335 | wantErr: map[string]string{},
336 | },
337 | {
338 | name: "info logger logging info",
339 | logLevel: "info",
340 | logMsg: "info",
341 | wantOut: map[string]string{
342 | "severity": "INFO",
343 | "message": "log info",
344 | },
345 | wantErr: map[string]string{},
346 | },
347 | {
348 | name: "warn logger logging info",
349 | logLevel: "warn",
350 | logMsg: "info",
351 | wantOut: map[string]string{},
352 | wantErr: map[string]string{},
353 | },
354 | {
355 | name: "error logger logging info",
356 | logLevel: "error",
357 | logMsg: "info",
358 | wantOut: map[string]string{},
359 | wantErr: map[string]string{},
360 | },
361 | {
362 | name: "debug logger logging warn",
363 | logLevel: "debug",
364 | logMsg: "warn",
365 | wantOut: map[string]string{},
366 | wantErr: map[string]string{
367 | "severity": "WARN",
368 | "message": "log warn",
369 | },
370 | },
371 | {
372 | name: "info logger logging warn",
373 | logLevel: "info",
374 | logMsg: "warn",
375 | wantOut: map[string]string{},
376 | wantErr: map[string]string{
377 | "severity": "WARN",
378 | "message": "log warn",
379 | },
380 | },
381 | {
382 | name: "warn logger logging warn",
383 | logLevel: "warn",
384 | logMsg: "warn",
385 | wantOut: map[string]string{},
386 | wantErr: map[string]string{
387 | "severity": "WARN",
388 | "message": "log warn",
389 | },
390 | },
391 | {
392 | name: "error logger logging warn",
393 | logLevel: "error",
394 | logMsg: "warn",
395 | wantOut: map[string]string{},
396 | wantErr: map[string]string{},
397 | },
398 | {
399 | name: "debug logger logging error",
400 | logLevel: "debug",
401 | logMsg: "error",
402 | wantOut: map[string]string{},
403 | wantErr: map[string]string{
404 | "severity": "ERROR",
405 | "message": "log error",
406 | },
407 | },
408 | {
409 | name: "info logger logging error",
410 | logLevel: "info",
411 | logMsg: "error",
412 | wantOut: map[string]string{},
413 | wantErr: map[string]string{
414 | "severity": "ERROR",
415 | "message": "log error",
416 | },
417 | },
418 | {
419 | name: "warn logger logging error",
420 | logLevel: "warn",
421 | logMsg: "error",
422 | wantOut: map[string]string{},
423 | wantErr: map[string]string{
424 | "severity": "ERROR",
425 | "message": "log error",
426 | },
427 | },
428 | {
429 | name: "error logger logging error",
430 | logLevel: "error",
431 | logMsg: "error",
432 | wantOut: map[string]string{},
433 | wantErr: map[string]string{
434 | "severity": "ERROR",
435 | "message": "log error",
436 | },
437 | },
438 | }
439 | for _, tc := range tcs {
440 | t.Run(tc.name, func(t *testing.T) {
441 | outW := new(bytes.Buffer)
442 | errW := new(bytes.Buffer)
443 |
444 | logger, err := NewStructuredLogger(outW, errW, tc.logLevel)
445 | if err != nil {
446 | t.Fatalf("unexpected error: %s", err)
447 | }
448 | runLogger(logger, tc.logMsg)
449 |
450 | if len(tc.wantOut) != 0 {
451 | got := make(map[string]interface{})
452 |
453 | if err := json.Unmarshal(outW.Bytes(), &got); err != nil {
454 | t.Fatalf("failed to parse writer")
455 | }
456 |
457 | if got["severity"] != tc.wantOut["severity"] {
458 | t.Fatalf("incorrect severity: got %v, want %v", got["severity"], tc.wantOut["severity"])
459 | }
460 |
461 | } else {
462 | if outW.String() != "" {
463 | t.Fatalf("incorrect log. got %v, want %v", outW.String(), "")
464 | }
465 | }
466 |
467 | if len(tc.wantErr) != 0 {
468 | got := make(map[string]interface{})
469 |
470 | if err := json.Unmarshal(errW.Bytes(), &got); err != nil {
471 | t.Fatalf("failed to parse writer")
472 | }
473 |
474 | if got["severity"] != tc.wantErr["severity"] {
475 | t.Fatalf("incorrect severity: got %v, want %v", got["severity"], tc.wantErr["severity"])
476 | }
477 |
478 | } else {
479 | if errW.String() != "" {
480 | t.Fatalf("incorrect log. got %v, want %v", errW.String(), "")
481 | }
482 | }
483 | })
484 | }
485 | }
486 |
```