This is page 29 of 45. Use http://codebase.md/googleapis/genai-toolbox?lines=true&page={x} to view the full context.
# Directory Structure
```
├── .ci
│ ├── continuous.release.cloudbuild.yaml
│ ├── generate_release_table.sh
│ ├── integration.cloudbuild.yaml
│ ├── quickstart_test
│ │ ├── go.integration.cloudbuild.yaml
│ │ ├── js.integration.cloudbuild.yaml
│ │ ├── py.integration.cloudbuild.yaml
│ │ ├── run_go_tests.sh
│ │ ├── run_js_tests.sh
│ │ ├── run_py_tests.sh
│ │ └── setup_hotels_sample.sql
│ ├── test_with_coverage.sh
│ └── versioned.release.cloudbuild.yaml
├── .github
│ ├── auto-label.yaml
│ ├── blunderbuss.yml
│ ├── CODEOWNERS
│ ├── header-checker-lint.yml
│ ├── ISSUE_TEMPLATE
│ │ ├── bug_report.yml
│ │ ├── config.yml
│ │ ├── feature_request.yml
│ │ └── question.yml
│ ├── label-sync.yml
│ ├── labels.yaml
│ ├── PULL_REQUEST_TEMPLATE.md
│ ├── release-please.yml
│ ├── renovate.json5
│ ├── sync-repo-settings.yaml
│ └── workflows
│ ├── cloud_build_failure_reporter.yml
│ ├── deploy_dev_docs.yaml
│ ├── deploy_previous_version_docs.yaml
│ ├── deploy_versioned_docs.yaml
│ ├── docs_deploy.yaml
│ ├── docs_preview_clean.yaml
│ ├── docs_preview_deploy.yaml
│ ├── lint.yaml
│ ├── schedule_reporter.yml
│ ├── sync-labels.yaml
│ └── tests.yaml
├── .gitignore
├── .gitmodules
├── .golangci.yaml
├── .hugo
│ ├── archetypes
│ │ └── default.md
│ ├── assets
│ │ ├── icons
│ │ │ └── logo.svg
│ │ └── scss
│ │ ├── _styles_project.scss
│ │ └── _variables_project.scss
│ ├── go.mod
│ ├── go.sum
│ ├── hugo.toml
│ ├── layouts
│ │ ├── _default
│ │ │ └── home.releases.releases
│ │ ├── index.llms-full.txt
│ │ ├── index.llms.txt
│ │ ├── partials
│ │ │ ├── hooks
│ │ │ │ └── head-end.html
│ │ │ ├── navbar-version-selector.html
│ │ │ ├── page-meta-links.html
│ │ │ └── td
│ │ │ └── render-heading.html
│ │ ├── robot.txt
│ │ └── shortcodes
│ │ ├── include.html
│ │ ├── ipynb.html
│ │ └── regionInclude.html
│ ├── package-lock.json
│ ├── package.json
│ └── static
│ ├── favicons
│ │ ├── android-chrome-192x192.png
│ │ ├── android-chrome-512x512.png
│ │ ├── apple-touch-icon.png
│ │ ├── favicon-16x16.png
│ │ ├── favicon-32x32.png
│ │ └── favicon.ico
│ └── js
│ └── w3.js
├── CHANGELOG.md
├── cmd
│ ├── options_test.go
│ ├── options.go
│ ├── root_test.go
│ ├── root.go
│ └── version.txt
├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── DEVELOPER.md
├── Dockerfile
├── docs
│ └── en
│ ├── _index.md
│ ├── about
│ │ ├── _index.md
│ │ └── faq.md
│ ├── concepts
│ │ ├── _index.md
│ │ └── telemetry
│ │ ├── index.md
│ │ ├── telemetry_flow.png
│ │ └── telemetry_traces.png
│ ├── getting-started
│ │ ├── _index.md
│ │ ├── colab_quickstart.ipynb
│ │ ├── configure.md
│ │ ├── introduction
│ │ │ ├── _index.md
│ │ │ └── architecture.png
│ │ ├── local_quickstart_go.md
│ │ ├── local_quickstart_js.md
│ │ ├── local_quickstart.md
│ │ ├── mcp_quickstart
│ │ │ ├── _index.md
│ │ │ ├── inspector_tools.png
│ │ │ └── inspector.png
│ │ └── quickstart
│ │ ├── go
│ │ │ ├── genAI
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ ├── genkit
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ ├── langchain
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ ├── openAI
│ │ │ │ ├── go.mod
│ │ │ │ ├── go.sum
│ │ │ │ └── quickstart.go
│ │ │ └── quickstart_test.go
│ │ ├── golden.txt
│ │ ├── js
│ │ │ ├── genAI
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ ├── genkit
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ ├── langchain
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ ├── llamaindex
│ │ │ │ ├── package-lock.json
│ │ │ │ ├── package.json
│ │ │ │ └── quickstart.js
│ │ │ └── quickstart.test.js
│ │ ├── python
│ │ │ ├── __init__.py
│ │ │ ├── adk
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ ├── core
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ ├── langchain
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ ├── llamaindex
│ │ │ │ ├── quickstart.py
│ │ │ │ └── requirements.txt
│ │ │ └── quickstart_test.py
│ │ └── shared
│ │ ├── cloud_setup.md
│ │ ├── configure_toolbox.md
│ │ └── database_setup.md
│ ├── how-to
│ │ ├── _index.md
│ │ ├── connect_via_geminicli.md
│ │ ├── connect_via_mcp.md
│ │ ├── connect-ide
│ │ │ ├── _index.md
│ │ │ ├── alloydb_pg_admin_mcp.md
│ │ │ ├── alloydb_pg_mcp.md
│ │ │ ├── bigquery_mcp.md
│ │ │ ├── cloud_sql_mssql_admin_mcp.md
│ │ │ ├── cloud_sql_mssql_mcp.md
│ │ │ ├── cloud_sql_mysql_admin_mcp.md
│ │ │ ├── cloud_sql_mysql_mcp.md
│ │ │ ├── cloud_sql_pg_admin_mcp.md
│ │ │ ├── cloud_sql_pg_mcp.md
│ │ │ ├── firestore_mcp.md
│ │ │ ├── looker_mcp.md
│ │ │ ├── mssql_mcp.md
│ │ │ ├── mysql_mcp.md
│ │ │ ├── neo4j_mcp.md
│ │ │ ├── postgres_mcp.md
│ │ │ ├── spanner_mcp.md
│ │ │ └── sqlite_mcp.md
│ │ ├── deploy_docker.md
│ │ ├── deploy_gke.md
│ │ ├── deploy_toolbox.md
│ │ ├── export_telemetry.md
│ │ └── toolbox-ui
│ │ ├── edit-headers.gif
│ │ ├── edit-headers.png
│ │ ├── index.md
│ │ ├── optional-param-checked.png
│ │ ├── optional-param-unchecked.png
│ │ ├── run-tool.gif
│ │ ├── tools.png
│ │ └── toolsets.png
│ ├── reference
│ │ ├── _index.md
│ │ ├── cli.md
│ │ └── prebuilt-tools.md
│ ├── resources
│ │ ├── _index.md
│ │ ├── authServices
│ │ │ ├── _index.md
│ │ │ └── google.md
│ │ ├── sources
│ │ │ ├── _index.md
│ │ │ ├── alloydb-admin.md
│ │ │ ├── alloydb-pg.md
│ │ │ ├── bigquery.md
│ │ │ ├── bigtable.md
│ │ │ ├── cassandra.md
│ │ │ ├── clickhouse.md
│ │ │ ├── cloud-monitoring.md
│ │ │ ├── cloud-sql-admin.md
│ │ │ ├── cloud-sql-mssql.md
│ │ │ ├── cloud-sql-mysql.md
│ │ │ ├── cloud-sql-pg.md
│ │ │ ├── couchbase.md
│ │ │ ├── dataplex.md
│ │ │ ├── dgraph.md
│ │ │ ├── firebird.md
│ │ │ ├── firestore.md
│ │ │ ├── http.md
│ │ │ ├── looker.md
│ │ │ ├── mongodb.md
│ │ │ ├── mssql.md
│ │ │ ├── mysql.md
│ │ │ ├── neo4j.md
│ │ │ ├── oceanbase.md
│ │ │ ├── oracle.md
│ │ │ ├── postgres.md
│ │ │ ├── redis.md
│ │ │ ├── spanner.md
│ │ │ ├── sqlite.md
│ │ │ ├── tidb.md
│ │ │ ├── trino.md
│ │ │ ├── valkey.md
│ │ │ └── yugabytedb.md
│ │ └── tools
│ │ ├── _index.md
│ │ ├── alloydb
│ │ │ ├── _index.md
│ │ │ ├── alloydb-create-cluster.md
│ │ │ ├── alloydb-create-instance.md
│ │ │ ├── alloydb-create-user.md
│ │ │ ├── alloydb-get-cluster.md
│ │ │ ├── alloydb-get-instance.md
│ │ │ ├── alloydb-get-user.md
│ │ │ ├── alloydb-list-clusters.md
│ │ │ ├── alloydb-list-instances.md
│ │ │ ├── alloydb-list-users.md
│ │ │ └── alloydb-wait-for-operation.md
│ │ ├── alloydbainl
│ │ │ ├── _index.md
│ │ │ └── alloydb-ai-nl.md
│ │ ├── bigquery
│ │ │ ├── _index.md
│ │ │ ├── bigquery-analyze-contribution.md
│ │ │ ├── bigquery-conversational-analytics.md
│ │ │ ├── bigquery-execute-sql.md
│ │ │ ├── bigquery-forecast.md
│ │ │ ├── bigquery-get-dataset-info.md
│ │ │ ├── bigquery-get-table-info.md
│ │ │ ├── bigquery-list-dataset-ids.md
│ │ │ ├── bigquery-list-table-ids.md
│ │ │ ├── bigquery-search-catalog.md
│ │ │ └── bigquery-sql.md
│ │ ├── bigtable
│ │ │ ├── _index.md
│ │ │ └── bigtable-sql.md
│ │ ├── cassandra
│ │ │ ├── _index.md
│ │ │ └── cassandra-cql.md
│ │ ├── clickhouse
│ │ │ ├── _index.md
│ │ │ ├── clickhouse-execute-sql.md
│ │ │ ├── clickhouse-list-databases.md
│ │ │ ├── clickhouse-list-tables.md
│ │ │ └── clickhouse-sql.md
│ │ ├── cloudmonitoring
│ │ │ ├── _index.md
│ │ │ └── cloud-monitoring-query-prometheus.md
│ │ ├── cloudsql
│ │ │ ├── _index.md
│ │ │ ├── cloudsqlcreatedatabase.md
│ │ │ ├── cloudsqlcreateusers.md
│ │ │ ├── cloudsqlgetinstances.md
│ │ │ ├── cloudsqllistdatabases.md
│ │ │ ├── cloudsqllistinstances.md
│ │ │ ├── cloudsqlmssqlcreateinstance.md
│ │ │ ├── cloudsqlmysqlcreateinstance.md
│ │ │ ├── cloudsqlpgcreateinstances.md
│ │ │ └── cloudsqlwaitforoperation.md
│ │ ├── couchbase
│ │ │ ├── _index.md
│ │ │ └── couchbase-sql.md
│ │ ├── dataform
│ │ │ ├── _index.md
│ │ │ └── dataform-compile-local.md
│ │ ├── dataplex
│ │ │ ├── _index.md
│ │ │ ├── dataplex-lookup-entry.md
│ │ │ ├── dataplex-search-aspect-types.md
│ │ │ └── dataplex-search-entries.md
│ │ ├── dgraph
│ │ │ ├── _index.md
│ │ │ └── dgraph-dql.md
│ │ ├── firebird
│ │ │ ├── _index.md
│ │ │ ├── firebird-execute-sql.md
│ │ │ └── firebird-sql.md
│ │ ├── firestore
│ │ │ ├── _index.md
│ │ │ ├── firestore-add-documents.md
│ │ │ ├── firestore-delete-documents.md
│ │ │ ├── firestore-get-documents.md
│ │ │ ├── firestore-get-rules.md
│ │ │ ├── firestore-list-collections.md
│ │ │ ├── firestore-query-collection.md
│ │ │ ├── firestore-query.md
│ │ │ ├── firestore-update-document.md
│ │ │ └── firestore-validate-rules.md
│ │ ├── http
│ │ │ ├── _index.md
│ │ │ └── http.md
│ │ ├── looker
│ │ │ ├── _index.md
│ │ │ ├── looker-add-dashboard-element.md
│ │ │ ├── looker-conversational-analytics.md
│ │ │ ├── looker-get-dashboards.md
│ │ │ ├── looker-get-dimensions.md
│ │ │ ├── looker-get-explores.md
│ │ │ ├── looker-get-filters.md
│ │ │ ├── looker-get-looks.md
│ │ │ ├── looker-get-measures.md
│ │ │ ├── looker-get-models.md
│ │ │ ├── looker-get-parameters.md
│ │ │ ├── looker-health-analyze.md
│ │ │ ├── looker-health-pulse.md
│ │ │ ├── looker-health-vacuum.md
│ │ │ ├── looker-make-dashboard.md
│ │ │ ├── looker-make-look.md
│ │ │ ├── looker-query-sql.md
│ │ │ ├── looker-query-url.md
│ │ │ ├── looker-query.md
│ │ │ └── looker-run-look.md
│ │ ├── mongodb
│ │ │ ├── _index.md
│ │ │ ├── mongodb-aggregate.md
│ │ │ ├── mongodb-delete-many.md
│ │ │ ├── mongodb-delete-one.md
│ │ │ ├── mongodb-find-one.md
│ │ │ ├── mongodb-find.md
│ │ │ ├── mongodb-insert-many.md
│ │ │ ├── mongodb-insert-one.md
│ │ │ ├── mongodb-update-many.md
│ │ │ └── mongodb-update-one.md
│ │ ├── mssql
│ │ │ ├── _index.md
│ │ │ ├── mssql-execute-sql.md
│ │ │ ├── mssql-list-tables.md
│ │ │ └── mssql-sql.md
│ │ ├── mysql
│ │ │ ├── _index.md
│ │ │ ├── mysql-execute-sql.md
│ │ │ ├── mysql-list-active-queries.md
│ │ │ ├── mysql-list-table-fragmentation.md
│ │ │ ├── mysql-list-tables-missing-unique-indexes.md
│ │ │ ├── mysql-list-tables.md
│ │ │ └── mysql-sql.md
│ │ ├── neo4j
│ │ │ ├── _index.md
│ │ │ ├── neo4j-cypher.md
│ │ │ ├── neo4j-execute-cypher.md
│ │ │ └── neo4j-schema.md
│ │ ├── oceanbase
│ │ │ ├── _index.md
│ │ │ ├── oceanbase-execute-sql.md
│ │ │ └── oceanbase-sql.md
│ │ ├── oracle
│ │ │ ├── _index.md
│ │ │ ├── oracle-execute-sql.md
│ │ │ └── oracle-sql.md
│ │ ├── postgres
│ │ │ ├── _index.md
│ │ │ ├── postgres-execute-sql.md
│ │ │ ├── postgres-list-active-queries.md
│ │ │ ├── postgres-list-available-extensions.md
│ │ │ ├── postgres-list-installed-extensions.md
│ │ │ ├── postgres-list-tables.md
│ │ │ └── postgres-sql.md
│ │ ├── redis
│ │ │ ├── _index.md
│ │ │ └── redis.md
│ │ ├── spanner
│ │ │ ├── _index.md
│ │ │ ├── spanner-execute-sql.md
│ │ │ ├── spanner-list-tables.md
│ │ │ └── spanner-sql.md
│ │ ├── sqlite
│ │ │ ├── _index.md
│ │ │ ├── sqlite-execute-sql.md
│ │ │ └── sqlite-sql.md
│ │ ├── tidb
│ │ │ ├── _index.md
│ │ │ ├── tidb-execute-sql.md
│ │ │ └── tidb-sql.md
│ │ ├── trino
│ │ │ ├── _index.md
│ │ │ ├── trino-execute-sql.md
│ │ │ └── trino-sql.md
│ │ ├── utility
│ │ │ ├── _index.md
│ │ │ └── wait.md
│ │ ├── valkey
│ │ │ ├── _index.md
│ │ │ └── valkey.md
│ │ └── yuagbytedb
│ │ ├── _index.md
│ │ └── yugabytedb-sql.md
│ ├── samples
│ │ ├── _index.md
│ │ ├── alloydb
│ │ │ ├── _index.md
│ │ │ ├── ai-nl
│ │ │ │ ├── alloydb_ai_nl.ipynb
│ │ │ │ └── index.md
│ │ │ └── mcp_quickstart.md
│ │ ├── bigquery
│ │ │ ├── _index.md
│ │ │ ├── colab_quickstart_bigquery.ipynb
│ │ │ ├── local_quickstart.md
│ │ │ └── mcp_quickstart
│ │ │ ├── _index.md
│ │ │ ├── inspector_tools.png
│ │ │ └── inspector.png
│ │ └── looker
│ │ ├── _index.md
│ │ ├── looker_gemini_oauth
│ │ │ ├── _index.md
│ │ │ ├── authenticated.png
│ │ │ ├── authorize.png
│ │ │ └── registration.png
│ │ ├── looker_gemini.md
│ │ └── looker_mcp_inspector
│ │ ├── _index.md
│ │ ├── inspector_tools.png
│ │ └── inspector.png
│ └── sdks
│ ├── _index.md
│ ├── go-sdk.md
│ ├── js-sdk.md
│ └── python-sdk.md
├── gemini-extension.json
├── go.mod
├── go.sum
├── internal
│ ├── auth
│ │ ├── auth.go
│ │ └── google
│ │ └── google.go
│ ├── log
│ │ ├── handler.go
│ │ ├── log_test.go
│ │ ├── log.go
│ │ └── logger.go
│ ├── prebuiltconfigs
│ │ ├── prebuiltconfigs_test.go
│ │ ├── prebuiltconfigs.go
│ │ └── tools
│ │ ├── alloydb-postgres-admin.yaml
│ │ ├── alloydb-postgres-observability.yaml
│ │ ├── alloydb-postgres.yaml
│ │ ├── bigquery.yaml
│ │ ├── clickhouse.yaml
│ │ ├── cloud-sql-mssql-admin.yaml
│ │ ├── cloud-sql-mssql-observability.yaml
│ │ ├── cloud-sql-mssql.yaml
│ │ ├── cloud-sql-mysql-admin.yaml
│ │ ├── cloud-sql-mysql-observability.yaml
│ │ ├── cloud-sql-mysql.yaml
│ │ ├── cloud-sql-postgres-admin.yaml
│ │ ├── cloud-sql-postgres-observability.yaml
│ │ ├── cloud-sql-postgres.yaml
│ │ ├── dataplex.yaml
│ │ ├── firestore.yaml
│ │ ├── looker-conversational-analytics.yaml
│ │ ├── looker.yaml
│ │ ├── mssql.yaml
│ │ ├── mysql.yaml
│ │ ├── neo4j.yaml
│ │ ├── oceanbase.yaml
│ │ ├── postgres.yaml
│ │ ├── spanner-postgres.yaml
│ │ ├── spanner.yaml
│ │ └── sqlite.yaml
│ ├── server
│ │ ├── api_test.go
│ │ ├── api.go
│ │ ├── common_test.go
│ │ ├── config.go
│ │ ├── mcp
│ │ │ ├── jsonrpc
│ │ │ │ ├── jsonrpc_test.go
│ │ │ │ └── jsonrpc.go
│ │ │ ├── mcp.go
│ │ │ ├── util
│ │ │ │ └── lifecycle.go
│ │ │ ├── v20241105
│ │ │ │ ├── method.go
│ │ │ │ └── types.go
│ │ │ ├── v20250326
│ │ │ │ ├── method.go
│ │ │ │ └── types.go
│ │ │ └── v20250618
│ │ │ ├── method.go
│ │ │ └── types.go
│ │ ├── mcp_test.go
│ │ ├── mcp.go
│ │ ├── server_test.go
│ │ ├── server.go
│ │ ├── static
│ │ │ ├── assets
│ │ │ │ └── mcptoolboxlogo.png
│ │ │ ├── css
│ │ │ │ └── style.css
│ │ │ ├── index.html
│ │ │ ├── js
│ │ │ │ ├── auth.js
│ │ │ │ ├── loadTools.js
│ │ │ │ ├── mainContent.js
│ │ │ │ ├── navbar.js
│ │ │ │ ├── runTool.js
│ │ │ │ ├── toolDisplay.js
│ │ │ │ ├── tools.js
│ │ │ │ └── toolsets.js
│ │ │ ├── tools.html
│ │ │ └── toolsets.html
│ │ ├── web_test.go
│ │ └── web.go
│ ├── sources
│ │ ├── alloydbadmin
│ │ │ ├── alloydbadmin_test.go
│ │ │ └── alloydbadmin.go
│ │ ├── alloydbpg
│ │ │ ├── alloydb_pg_test.go
│ │ │ └── alloydb_pg.go
│ │ ├── bigquery
│ │ │ ├── bigquery_test.go
│ │ │ └── bigquery.go
│ │ ├── bigtable
│ │ │ ├── bigtable_test.go
│ │ │ └── bigtable.go
│ │ ├── cassandra
│ │ │ ├── cassandra_test.go
│ │ │ └── cassandra.go
│ │ ├── clickhouse
│ │ │ ├── clickhouse_test.go
│ │ │ └── clickhouse.go
│ │ ├── cloudmonitoring
│ │ │ ├── cloud_monitoring_test.go
│ │ │ └── cloud_monitoring.go
│ │ ├── cloudsqladmin
│ │ │ ├── cloud_sql_admin_test.go
│ │ │ └── cloud_sql_admin.go
│ │ ├── cloudsqlmssql
│ │ │ ├── cloud_sql_mssql_test.go
│ │ │ └── cloud_sql_mssql.go
│ │ ├── cloudsqlmysql
│ │ │ ├── cloud_sql_mysql_test.go
│ │ │ └── cloud_sql_mysql.go
│ │ ├── cloudsqlpg
│ │ │ ├── cloud_sql_pg_test.go
│ │ │ └── cloud_sql_pg.go
│ │ ├── couchbase
│ │ │ ├── couchbase_test.go
│ │ │ └── couchbase.go
│ │ ├── dataplex
│ │ │ ├── dataplex_test.go
│ │ │ └── dataplex.go
│ │ ├── dgraph
│ │ │ ├── dgraph_test.go
│ │ │ └── dgraph.go
│ │ ├── dialect.go
│ │ ├── firebird
│ │ │ ├── firebird_test.go
│ │ │ └── firebird.go
│ │ ├── firestore
│ │ │ ├── firestore_test.go
│ │ │ └── firestore.go
│ │ ├── http
│ │ │ ├── http_test.go
│ │ │ └── http.go
│ │ ├── ip_type.go
│ │ ├── looker
│ │ │ ├── looker_test.go
│ │ │ └── looker.go
│ │ ├── mongodb
│ │ │ ├── mongodb_test.go
│ │ │ └── mongodb.go
│ │ ├── mssql
│ │ │ ├── mssql_test.go
│ │ │ └── mssql.go
│ │ ├── mysql
│ │ │ ├── mysql_test.go
│ │ │ └── mysql.go
│ │ ├── neo4j
│ │ │ ├── neo4j_test.go
│ │ │ └── neo4j.go
│ │ ├── oceanbase
│ │ │ ├── oceanbase_test.go
│ │ │ └── oceanbase.go
│ │ ├── oracle
│ │ │ └── oracle.go
│ │ ├── postgres
│ │ │ ├── postgres_test.go
│ │ │ └── postgres.go
│ │ ├── redis
│ │ │ ├── redis_test.go
│ │ │ └── redis.go
│ │ ├── sources.go
│ │ ├── spanner
│ │ │ ├── spanner_test.go
│ │ │ └── spanner.go
│ │ ├── sqlite
│ │ │ ├── sqlite_test.go
│ │ │ └── sqlite.go
│ │ ├── tidb
│ │ │ ├── tidb_test.go
│ │ │ └── tidb.go
│ │ ├── trino
│ │ │ ├── trino_test.go
│ │ │ └── trino.go
│ │ ├── util.go
│ │ ├── valkey
│ │ │ ├── valkey_test.go
│ │ │ └── valkey.go
│ │ └── yugabytedb
│ │ ├── yugabytedb_test.go
│ │ └── yugabytedb.go
│ ├── telemetry
│ │ ├── instrumentation.go
│ │ └── telemetry.go
│ ├── testutils
│ │ └── testutils.go
│ ├── tools
│ │ ├── alloydb
│ │ │ ├── alloydbcreatecluster
│ │ │ │ ├── alloydbcreatecluster_test.go
│ │ │ │ └── alloydbcreatecluster.go
│ │ │ ├── alloydbcreateinstance
│ │ │ │ ├── alloydbcreateinstance_test.go
│ │ │ │ └── alloydbcreateinstance.go
│ │ │ ├── alloydbcreateuser
│ │ │ │ ├── alloydbcreateuser_test.go
│ │ │ │ └── alloydbcreateuser.go
│ │ │ ├── alloydbgetcluster
│ │ │ │ ├── alloydbgetcluster_test.go
│ │ │ │ └── alloydbgetcluster.go
│ │ │ ├── alloydbgetinstance
│ │ │ │ ├── alloydbgetinstance_test.go
│ │ │ │ └── alloydbgetinstance.go
│ │ │ ├── alloydbgetuser
│ │ │ │ ├── alloydbgetuser_test.go
│ │ │ │ └── alloydbgetuser.go
│ │ │ ├── alloydblistclusters
│ │ │ │ ├── alloydblistclusters_test.go
│ │ │ │ └── alloydblistclusters.go
│ │ │ ├── alloydblistinstances
│ │ │ │ ├── alloydblistinstances_test.go
│ │ │ │ └── alloydblistinstances.go
│ │ │ ├── alloydblistusers
│ │ │ │ ├── alloydblistusers_test.go
│ │ │ │ └── alloydblistusers.go
│ │ │ └── alloydbwaitforoperation
│ │ │ ├── alloydbwaitforoperation_test.go
│ │ │ └── alloydbwaitforoperation.go
│ │ ├── alloydbainl
│ │ │ ├── alloydbainl_test.go
│ │ │ └── alloydbainl.go
│ │ ├── bigquery
│ │ │ ├── bigqueryanalyzecontribution
│ │ │ │ ├── bigqueryanalyzecontribution_test.go
│ │ │ │ └── bigqueryanalyzecontribution.go
│ │ │ ├── bigquerycommon
│ │ │ │ ├── table_name_parser_test.go
│ │ │ │ ├── table_name_parser.go
│ │ │ │ └── util.go
│ │ │ ├── bigqueryconversationalanalytics
│ │ │ │ ├── bigqueryconversationalanalytics_test.go
│ │ │ │ └── bigqueryconversationalanalytics.go
│ │ │ ├── bigqueryexecutesql
│ │ │ │ ├── bigqueryexecutesql_test.go
│ │ │ │ └── bigqueryexecutesql.go
│ │ │ ├── bigqueryforecast
│ │ │ │ ├── bigqueryforecast_test.go
│ │ │ │ └── bigqueryforecast.go
│ │ │ ├── bigquerygetdatasetinfo
│ │ │ │ ├── bigquerygetdatasetinfo_test.go
│ │ │ │ └── bigquerygetdatasetinfo.go
│ │ │ ├── bigquerygettableinfo
│ │ │ │ ├── bigquerygettableinfo_test.go
│ │ │ │ └── bigquerygettableinfo.go
│ │ │ ├── bigquerylistdatasetids
│ │ │ │ ├── bigquerylistdatasetids_test.go
│ │ │ │ └── bigquerylistdatasetids.go
│ │ │ ├── bigquerylisttableids
│ │ │ │ ├── bigquerylisttableids_test.go
│ │ │ │ └── bigquerylisttableids.go
│ │ │ ├── bigquerysearchcatalog
│ │ │ │ ├── bigquerysearchcatalog_test.go
│ │ │ │ └── bigquerysearchcatalog.go
│ │ │ └── bigquerysql
│ │ │ ├── bigquerysql_test.go
│ │ │ └── bigquerysql.go
│ │ ├── bigtable
│ │ │ ├── bigtable_test.go
│ │ │ └── bigtable.go
│ │ ├── cassandra
│ │ │ └── cassandracql
│ │ │ ├── cassandracql_test.go
│ │ │ └── cassandracql.go
│ │ ├── clickhouse
│ │ │ ├── clickhouseexecutesql
│ │ │ │ ├── clickhouseexecutesql_test.go
│ │ │ │ └── clickhouseexecutesql.go
│ │ │ ├── clickhouselistdatabases
│ │ │ │ ├── clickhouselistdatabases_test.go
│ │ │ │ └── clickhouselistdatabases.go
│ │ │ ├── clickhouselisttables
│ │ │ │ ├── clickhouselisttables_test.go
│ │ │ │ └── clickhouselisttables.go
│ │ │ └── clickhousesql
│ │ │ ├── clickhousesql_test.go
│ │ │ └── clickhousesql.go
│ │ ├── cloudmonitoring
│ │ │ ├── cloudmonitoring_test.go
│ │ │ └── cloudmonitoring.go
│ │ ├── cloudsql
│ │ │ ├── cloudsqlcreatedatabase
│ │ │ │ ├── cloudsqlcreatedatabase_test.go
│ │ │ │ └── cloudsqlcreatedatabase.go
│ │ │ ├── cloudsqlcreateusers
│ │ │ │ ├── cloudsqlcreateusers_test.go
│ │ │ │ └── cloudsqlcreateusers.go
│ │ │ ├── cloudsqlgetinstances
│ │ │ │ ├── cloudsqlgetinstances_test.go
│ │ │ │ └── cloudsqlgetinstances.go
│ │ │ ├── cloudsqllistdatabases
│ │ │ │ ├── cloudsqllistdatabases_test.go
│ │ │ │ └── cloudsqllistdatabases.go
│ │ │ ├── cloudsqllistinstances
│ │ │ │ ├── cloudsqllistinstances_test.go
│ │ │ │ └── cloudsqllistinstances.go
│ │ │ └── cloudsqlwaitforoperation
│ │ │ ├── cloudsqlwaitforoperation_test.go
│ │ │ └── cloudsqlwaitforoperation.go
│ │ ├── cloudsqlmssql
│ │ │ └── cloudsqlmssqlcreateinstance
│ │ │ ├── cloudsqlmssqlcreateinstance_test.go
│ │ │ └── cloudsqlmssqlcreateinstance.go
│ │ ├── cloudsqlmysql
│ │ │ └── cloudsqlmysqlcreateinstance
│ │ │ ├── cloudsqlmysqlcreateinstance_test.go
│ │ │ └── cloudsqlmysqlcreateinstance.go
│ │ ├── cloudsqlpg
│ │ │ └── cloudsqlpgcreateinstances
│ │ │ ├── cloudsqlpgcreateinstances_test.go
│ │ │ └── cloudsqlpgcreateinstances.go
│ │ ├── common_test.go
│ │ ├── common.go
│ │ ├── couchbase
│ │ │ ├── couchbase_test.go
│ │ │ └── couchbase.go
│ │ ├── dataform
│ │ │ └── dataformcompilelocal
│ │ │ ├── dataformcompilelocal_test.go
│ │ │ └── dataformcompilelocal.go
│ │ ├── dataplex
│ │ │ ├── dataplexlookupentry
│ │ │ │ ├── dataplexlookupentry_test.go
│ │ │ │ └── dataplexlookupentry.go
│ │ │ ├── dataplexsearchaspecttypes
│ │ │ │ ├── dataplexsearchaspecttypes_test.go
│ │ │ │ └── dataplexsearchaspecttypes.go
│ │ │ └── dataplexsearchentries
│ │ │ ├── dataplexsearchentries_test.go
│ │ │ └── dataplexsearchentries.go
│ │ ├── dgraph
│ │ │ ├── dgraph_test.go
│ │ │ └── dgraph.go
│ │ ├── firebird
│ │ │ ├── firebirdexecutesql
│ │ │ │ ├── firebirdexecutesql_test.go
│ │ │ │ └── firebirdexecutesql.go
│ │ │ └── firebirdsql
│ │ │ ├── firebirdsql_test.go
│ │ │ └── firebirdsql.go
│ │ ├── firestore
│ │ │ ├── firestoreadddocuments
│ │ │ │ ├── firestoreadddocuments_test.go
│ │ │ │ └── firestoreadddocuments.go
│ │ │ ├── firestoredeletedocuments
│ │ │ │ ├── firestoredeletedocuments_test.go
│ │ │ │ └── firestoredeletedocuments.go
│ │ │ ├── firestoregetdocuments
│ │ │ │ ├── firestoregetdocuments_test.go
│ │ │ │ └── firestoregetdocuments.go
│ │ │ ├── firestoregetrules
│ │ │ │ ├── firestoregetrules_test.go
│ │ │ │ └── firestoregetrules.go
│ │ │ ├── firestorelistcollections
│ │ │ │ ├── firestorelistcollections_test.go
│ │ │ │ └── firestorelistcollections.go
│ │ │ ├── firestorequery
│ │ │ │ ├── firestorequery_test.go
│ │ │ │ └── firestorequery.go
│ │ │ ├── firestorequerycollection
│ │ │ │ ├── firestorequerycollection_test.go
│ │ │ │ └── firestorequerycollection.go
│ │ │ ├── firestoreupdatedocument
│ │ │ │ ├── firestoreupdatedocument_test.go
│ │ │ │ └── firestoreupdatedocument.go
│ │ │ ├── firestorevalidaterules
│ │ │ │ ├── firestorevalidaterules_test.go
│ │ │ │ └── firestorevalidaterules.go
│ │ │ └── util
│ │ │ ├── converter_test.go
│ │ │ ├── converter.go
│ │ │ ├── validator_test.go
│ │ │ └── validator.go
│ │ ├── http
│ │ │ ├── http_test.go
│ │ │ └── http.go
│ │ ├── http_method.go
│ │ ├── looker
│ │ │ ├── lookeradddashboardelement
│ │ │ │ ├── lookeradddashboardelement_test.go
│ │ │ │ └── lookeradddashboardelement.go
│ │ │ ├── lookercommon
│ │ │ │ ├── lookercommon_test.go
│ │ │ │ └── lookercommon.go
│ │ │ ├── lookerconversationalanalytics
│ │ │ │ ├── lookerconversationalanalytics_test.go
│ │ │ │ └── lookerconversationalanalytics.go
│ │ │ ├── lookergetdashboards
│ │ │ │ ├── lookergetdashboards_test.go
│ │ │ │ └── lookergetdashboards.go
│ │ │ ├── lookergetdimensions
│ │ │ │ ├── lookergetdimensions_test.go
│ │ │ │ └── lookergetdimensions.go
│ │ │ ├── lookergetexplores
│ │ │ │ ├── lookergetexplores_test.go
│ │ │ │ └── lookergetexplores.go
│ │ │ ├── lookergetfilters
│ │ │ │ ├── lookergetfilters_test.go
│ │ │ │ └── lookergetfilters.go
│ │ │ ├── lookergetlooks
│ │ │ │ ├── lookergetlooks_test.go
│ │ │ │ └── lookergetlooks.go
│ │ │ ├── lookergetmeasures
│ │ │ │ ├── lookergetmeasures_test.go
│ │ │ │ └── lookergetmeasures.go
│ │ │ ├── lookergetmodels
│ │ │ │ ├── lookergetmodels_test.go
│ │ │ │ └── lookergetmodels.go
│ │ │ ├── lookergetparameters
│ │ │ │ ├── lookergetparameters_test.go
│ │ │ │ └── lookergetparameters.go
│ │ │ ├── lookerhealthanalyze
│ │ │ │ ├── lookerhealthanalyze_test.go
│ │ │ │ └── lookerhealthanalyze.go
│ │ │ ├── lookerhealthpulse
│ │ │ │ ├── lookerhealthpulse_test.go
│ │ │ │ └── lookerhealthpulse.go
│ │ │ ├── lookerhealthvacuum
│ │ │ │ ├── lookerhealthvacuum_test.go
│ │ │ │ └── lookerhealthvacuum.go
│ │ │ ├── lookermakedashboard
│ │ │ │ ├── lookermakedashboard_test.go
│ │ │ │ └── lookermakedashboard.go
│ │ │ ├── lookermakelook
│ │ │ │ ├── lookermakelook_test.go
│ │ │ │ └── lookermakelook.go
│ │ │ ├── lookerquery
│ │ │ │ ├── lookerquery_test.go
│ │ │ │ └── lookerquery.go
│ │ │ ├── lookerquerysql
│ │ │ │ ├── lookerquerysql_test.go
│ │ │ │ └── lookerquerysql.go
│ │ │ ├── lookerqueryurl
│ │ │ │ ├── lookerqueryurl_test.go
│ │ │ │ └── lookerqueryurl.go
│ │ │ └── lookerrunlook
│ │ │ ├── lookerrunlook_test.go
│ │ │ └── lookerrunlook.go
│ │ ├── mongodb
│ │ │ ├── mongodbaggregate
│ │ │ │ ├── mongodbaggregate_test.go
│ │ │ │ └── mongodbaggregate.go
│ │ │ ├── mongodbdeletemany
│ │ │ │ ├── mongodbdeletemany_test.go
│ │ │ │ └── mongodbdeletemany.go
│ │ │ ├── mongodbdeleteone
│ │ │ │ ├── mongodbdeleteone_test.go
│ │ │ │ └── mongodbdeleteone.go
│ │ │ ├── mongodbfind
│ │ │ │ ├── mongodbfind_test.go
│ │ │ │ └── mongodbfind.go
│ │ │ ├── mongodbfindone
│ │ │ │ ├── mongodbfindone_test.go
│ │ │ │ └── mongodbfindone.go
│ │ │ ├── mongodbinsertmany
│ │ │ │ ├── mongodbinsertmany_test.go
│ │ │ │ └── mongodbinsertmany.go
│ │ │ ├── mongodbinsertone
│ │ │ │ ├── mongodbinsertone_test.go
│ │ │ │ └── mongodbinsertone.go
│ │ │ ├── mongodbupdatemany
│ │ │ │ ├── mongodbupdatemany_test.go
│ │ │ │ └── mongodbupdatemany.go
│ │ │ └── mongodbupdateone
│ │ │ ├── mongodbupdateone_test.go
│ │ │ └── mongodbupdateone.go
│ │ ├── mssql
│ │ │ ├── mssqlexecutesql
│ │ │ │ ├── mssqlexecutesql_test.go
│ │ │ │ └── mssqlexecutesql.go
│ │ │ ├── mssqllisttables
│ │ │ │ ├── mssqllisttables_test.go
│ │ │ │ └── mssqllisttables.go
│ │ │ └── mssqlsql
│ │ │ ├── mssqlsql_test.go
│ │ │ └── mssqlsql.go
│ │ ├── mysql
│ │ │ ├── mysqlcommon
│ │ │ │ └── mysqlcommon.go
│ │ │ ├── mysqlexecutesql
│ │ │ │ ├── mysqlexecutesql_test.go
│ │ │ │ └── mysqlexecutesql.go
│ │ │ ├── mysqllistactivequeries
│ │ │ │ ├── mysqllistactivequeries_test.go
│ │ │ │ └── mysqllistactivequeries.go
│ │ │ ├── mysqllisttablefragmentation
│ │ │ │ ├── mysqllisttablefragmentation_test.go
│ │ │ │ └── mysqllisttablefragmentation.go
│ │ │ ├── mysqllisttables
│ │ │ │ ├── mysqllisttables_test.go
│ │ │ │ └── mysqllisttables.go
│ │ │ ├── mysqllisttablesmissinguniqueindexes
│ │ │ │ ├── mysqllisttablesmissinguniqueindexes_test.go
│ │ │ │ └── mysqllisttablesmissinguniqueindexes.go
│ │ │ └── mysqlsql
│ │ │ ├── mysqlsql_test.go
│ │ │ └── mysqlsql.go
│ │ ├── neo4j
│ │ │ ├── neo4jcypher
│ │ │ │ ├── neo4jcypher_test.go
│ │ │ │ └── neo4jcypher.go
│ │ │ ├── neo4jexecutecypher
│ │ │ │ ├── classifier
│ │ │ │ │ ├── classifier_test.go
│ │ │ │ │ └── classifier.go
│ │ │ │ ├── neo4jexecutecypher_test.go
│ │ │ │ └── neo4jexecutecypher.go
│ │ │ └── neo4jschema
│ │ │ ├── cache
│ │ │ │ ├── cache_test.go
│ │ │ │ └── cache.go
│ │ │ ├── helpers
│ │ │ │ ├── helpers_test.go
│ │ │ │ └── helpers.go
│ │ │ ├── neo4jschema_test.go
│ │ │ ├── neo4jschema.go
│ │ │ └── types
│ │ │ └── types.go
│ │ ├── oceanbase
│ │ │ ├── oceanbaseexecutesql
│ │ │ │ ├── oceanbaseexecutesql_test.go
│ │ │ │ └── oceanbaseexecutesql.go
│ │ │ └── oceanbasesql
│ │ │ ├── oceanbasesql_test.go
│ │ │ └── oceanbasesql.go
│ │ ├── oracle
│ │ │ ├── oracleexecutesql
│ │ │ │ └── oracleexecutesql.go
│ │ │ └── oraclesql
│ │ │ └── oraclesql.go
│ │ ├── parameters_test.go
│ │ ├── parameters.go
│ │ ├── postgres
│ │ │ ├── postgresexecutesql
│ │ │ │ ├── postgresexecutesql_test.go
│ │ │ │ └── postgresexecutesql.go
│ │ │ ├── postgreslistactivequeries
│ │ │ │ ├── postgreslistactivequeries_test.go
│ │ │ │ └── postgreslistactivequeries.go
│ │ │ ├── postgreslistavailableextensions
│ │ │ │ ├── postgreslistavailableextensions_test.go
│ │ │ │ └── postgreslistavailableextensions.go
│ │ │ ├── postgreslistinstalledextensions
│ │ │ │ ├── postgreslistinstalledextensions_test.go
│ │ │ │ └── postgreslistinstalledextensions.go
│ │ │ ├── postgreslisttables
│ │ │ │ ├── postgreslisttables_test.go
│ │ │ │ └── postgreslisttables.go
│ │ │ └── postgressql
│ │ │ ├── postgressql_test.go
│ │ │ └── postgressql.go
│ │ ├── redis
│ │ │ ├── redis_test.go
│ │ │ └── redis.go
│ │ ├── spanner
│ │ │ ├── spannerexecutesql
│ │ │ │ ├── spannerexecutesql_test.go
│ │ │ │ └── spannerexecutesql.go
│ │ │ ├── spannerlisttables
│ │ │ │ ├── spannerlisttables_test.go
│ │ │ │ └── spannerlisttables.go
│ │ │ └── spannersql
│ │ │ ├── spanner_test.go
│ │ │ └── spannersql.go
│ │ ├── sqlite
│ │ │ ├── sqliteexecutesql
│ │ │ │ ├── sqliteexecutesql_test.go
│ │ │ │ └── sqliteexecutesql.go
│ │ │ └── sqlitesql
│ │ │ ├── sqlitesql_test.go
│ │ │ └── sqlitesql.go
│ │ ├── tidb
│ │ │ ├── tidbexecutesql
│ │ │ │ ├── tidbexecutesql_test.go
│ │ │ │ └── tidbexecutesql.go
│ │ │ └── tidbsql
│ │ │ ├── tidbsql_test.go
│ │ │ └── tidbsql.go
│ │ ├── tools_test.go
│ │ ├── tools.go
│ │ ├── toolsets.go
│ │ ├── trino
│ │ │ ├── trinoexecutesql
│ │ │ │ ├── trinoexecutesql_test.go
│ │ │ │ └── trinoexecutesql.go
│ │ │ └── trinosql
│ │ │ ├── trinosql_test.go
│ │ │ └── trinosql.go
│ │ ├── utility
│ │ │ └── wait
│ │ │ ├── wait_test.go
│ │ │ └── wait.go
│ │ ├── valkey
│ │ │ ├── valkey_test.go
│ │ │ └── valkey.go
│ │ └── yugabytedbsql
│ │ ├── yugabytedbsql_test.go
│ │ └── yugabytedbsql.go
│ └── util
│ └── util.go
├── LICENSE
├── logo.png
├── main.go
├── MCP-TOOLBOX-EXTENSION.md
├── README.md
└── tests
├── alloydb
│ ├── alloydb_integration_test.go
│ └── alloydb_wait_for_operation_test.go
├── alloydbainl
│ └── alloydb_ai_nl_integration_test.go
├── alloydbpg
│ └── alloydb_pg_integration_test.go
├── auth.go
├── bigquery
│ └── bigquery_integration_test.go
├── bigtable
│ └── bigtable_integration_test.go
├── cassandra
│ └── cassandra_integration_test.go
├── clickhouse
│ └── clickhouse_integration_test.go
├── cloudmonitoring
│ └── cloud_monitoring_integration_test.go
├── cloudsql
│ ├── cloud_sql_create_database_test.go
│ ├── cloud_sql_create_users_test.go
│ ├── cloud_sql_get_instances_test.go
│ ├── cloud_sql_list_databases_test.go
│ ├── cloudsql_list_instances_test.go
│ └── cloudsql_wait_for_operation_test.go
├── cloudsqlmssql
│ ├── cloud_sql_mssql_create_instance_integration_test.go
│ └── cloud_sql_mssql_integration_test.go
├── cloudsqlmysql
│ ├── cloud_sql_mysql_create_instance_integration_test.go
│ └── cloud_sql_mysql_integration_test.go
├── cloudsqlpg
│ ├── cloud_sql_pg_create_instances_test.go
│ └── cloud_sql_pg_integration_test.go
├── common.go
├── couchbase
│ └── couchbase_integration_test.go
├── dataform
│ └── dataform_integration_test.go
├── dataplex
│ └── dataplex_integration_test.go
├── dgraph
│ └── dgraph_integration_test.go
├── firebird
│ └── firebird_integration_test.go
├── firestore
│ └── firestore_integration_test.go
├── http
│ └── http_integration_test.go
├── looker
│ └── looker_integration_test.go
├── mongodb
│ └── mongodb_integration_test.go
├── mssql
│ └── mssql_integration_test.go
├── mysql
│ └── mysql_integration_test.go
├── neo4j
│ └── neo4j_integration_test.go
├── oceanbase
│ └── oceanbase_integration_test.go
├── option.go
├── oracle
│ └── oracle_integration_test.go
├── postgres
│ └── postgres_integration_test.go
├── redis
│ └── redis_test.go
├── server.go
├── source.go
├── spanner
│ └── spanner_integration_test.go
├── sqlite
│ └── sqlite_integration_test.go
├── tidb
│ └── tidb_integration_test.go
├── tool.go
├── trino
│ └── trino_integration_test.go
├── utility
│ └── wait_integration_test.go
├── valkey
│ └── valkey_test.go
└── yugabytedb
└── yugabytedb_integration_test.go
```
# Files
--------------------------------------------------------------------------------
/tests/alloydbainl/alloydb_ai_nl_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 alloydbainl
16 |
17 | import (
18 | "bytes"
19 | "context"
20 | "encoding/json"
21 | "io"
22 | "net/http"
23 | "os"
24 | "reflect"
25 | "regexp"
26 | "strings"
27 | "testing"
28 | "time"
29 |
30 | "github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
31 | "github.com/googleapis/genai-toolbox/internal/testutils"
32 | "github.com/googleapis/genai-toolbox/tests"
33 | )
34 |
35 | var (
36 | AlloyDBAINLSourceKind = "alloydb-postgres"
37 | AlloyDBAINLToolKind = "alloydb-ai-nl"
38 | AlloyDBAINLProject = os.Getenv("ALLOYDB_AI_NL_PROJECT")
39 | AlloyDBAINLRegion = os.Getenv("ALLOYDB_AI_NL_REGION")
40 | AlloyDBAINLCluster = os.Getenv("ALLOYDB_AI_NL_CLUSTER")
41 | AlloyDBAINLInstance = os.Getenv("ALLOYDB_AI_NL_INSTANCE")
42 | AlloyDBAINLDatabase = os.Getenv("ALLOYDB_AI_NL_DATABASE")
43 | AlloyDBAINLUser = os.Getenv("ALLOYDB_AI_NL_USER")
44 | AlloyDBAINLPass = os.Getenv("ALLOYDB_AI_NL_PASS")
45 | )
46 |
47 | func getAlloyDBAINLVars(t *testing.T) map[string]any {
48 | switch "" {
49 | case AlloyDBAINLProject:
50 | t.Fatal("'ALLOYDB_AI_NL_PROJECT' not set")
51 | case AlloyDBAINLRegion:
52 | t.Fatal("'ALLOYDB_AI_NL_REGION' not set")
53 | case AlloyDBAINLCluster:
54 | t.Fatal("'ALLOYDB_AI_NL_CLUSTER' not set")
55 | case AlloyDBAINLInstance:
56 | t.Fatal("'ALLOYDB_AI_NL_INSTANCE' not set")
57 | case AlloyDBAINLDatabase:
58 | t.Fatal("'ALLOYDB_AI_NL_DATABASE' not set")
59 | case AlloyDBAINLUser:
60 | t.Fatal("'ALLOYDB_AI_NL_USER' not set")
61 | case AlloyDBAINLPass:
62 | t.Fatal("'ALLOYDB_AI_NL_PASS' not set")
63 | }
64 | return map[string]any{
65 | "kind": AlloyDBAINLSourceKind,
66 | "project": AlloyDBAINLProject,
67 | "cluster": AlloyDBAINLCluster,
68 | "instance": AlloyDBAINLInstance,
69 | "region": AlloyDBAINLRegion,
70 | "database": AlloyDBAINLDatabase,
71 | "user": AlloyDBAINLUser,
72 | "password": AlloyDBAINLPass,
73 | }
74 | }
75 |
76 | func TestAlloyDBAINLToolEndpoints(t *testing.T) {
77 | sourceConfig := getAlloyDBAINLVars(t)
78 | ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
79 | defer cancel()
80 |
81 | var args []string
82 |
83 | // Write config into a file and pass it to command
84 | toolsFile := getAINLToolsConfig(sourceConfig)
85 |
86 | cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
87 | if err != nil {
88 | t.Fatalf("command initialization returned an error: %s", err)
89 | }
90 | defer cleanup()
91 |
92 | waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
93 | defer cancel()
94 | out, err := testutils.WaitForString(waitCtx, regexp.MustCompile(`Server ready to serve`), cmd.Out)
95 | if err != nil {
96 | t.Logf("toolbox command logs: \n%s", out)
97 | t.Fatalf("toolbox didn't start successfully: %s", err)
98 | }
99 |
100 | runAINLToolGetTest(t)
101 | runAINLToolInvokeTest(t)
102 | runAINLMCPToolCallMethod(t)
103 | }
104 |
105 | func runAINLToolGetTest(t *testing.T) {
106 | // Test tool get endpoint
107 | tcs := []struct {
108 | name string
109 | api string
110 | want map[string]any
111 | }{
112 | {
113 | name: "get my-simple-tool",
114 | api: "http://127.0.0.1:5000/api/tool/my-simple-tool/",
115 | want: map[string]any{
116 | "my-simple-tool": map[string]any{
117 | "description": "Simple tool to test end to end functionality.",
118 | "parameters": []any{
119 | map[string]any{
120 | "name": "question",
121 | "type": "string",
122 | "required": true,
123 | "description": "The natural language question to ask.",
124 | "authSources": []any{},
125 | },
126 | },
127 | "authRequired": []any{},
128 | },
129 | },
130 | },
131 | }
132 | for _, tc := range tcs {
133 | t.Run(tc.name, func(t *testing.T) {
134 | resp, err := http.Get(tc.api)
135 | if err != nil {
136 | t.Fatalf("error when sending a request: %s", err)
137 | }
138 | defer resp.Body.Close()
139 | if resp.StatusCode != 200 {
140 | t.Fatalf("response status code is not 200")
141 | }
142 |
143 | var body map[string]interface{}
144 | err = json.NewDecoder(resp.Body).Decode(&body)
145 | if err != nil {
146 | t.Fatalf("error parsing response body")
147 | }
148 |
149 | got, ok := body["tools"]
150 | if !ok {
151 | t.Fatalf("unable to find tools in response body")
152 | }
153 | if !reflect.DeepEqual(got, tc.want) {
154 | t.Fatalf("got %q, want %q", got, tc.want)
155 | }
156 | })
157 | }
158 | }
159 |
160 | func runAINLToolInvokeTest(t *testing.T) {
161 | // Get ID token
162 | idToken, err := tests.GetGoogleIdToken(tests.ClientId)
163 | if err != nil {
164 | t.Fatalf("error getting Google ID token: %s", err)
165 | }
166 |
167 | // Test tool invoke endpoint
168 | invokeTcs := []struct {
169 | name string
170 | api string
171 | requestHeader map[string]string
172 | requestBody io.Reader
173 | want string
174 | isErr bool
175 | }{
176 | {
177 | name: "invoke my-simple-tool",
178 | api: "http://127.0.0.1:5000/api/tool/my-simple-tool/invoke",
179 | requestHeader: map[string]string{},
180 | requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)),
181 | want: "[{\"execute_nl_query\":{\"?column?\":1}}]",
182 | isErr: false,
183 | },
184 | {
185 | name: "Invoke my-tool without parameters",
186 | api: "http://127.0.0.1:5000/api/tool/my-tool/invoke",
187 | requestHeader: map[string]string{},
188 | requestBody: bytes.NewBuffer([]byte(`{}`)),
189 | isErr: true,
190 | },
191 | {
192 | name: "Invoke my-auth-tool with auth token",
193 | api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke",
194 | requestHeader: map[string]string{"my-google-auth_token": idToken},
195 | requestBody: bytes.NewBuffer([]byte(`{"question": "can you show me the name of this user?"}`)),
196 | want: "[{\"execute_nl_query\":{\"name\":\"Alice\"}}]",
197 | isErr: false,
198 | },
199 | {
200 | name: "Invoke my-auth-tool with invalid auth token",
201 | api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke",
202 | requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
203 | requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)),
204 | isErr: true,
205 | },
206 | {
207 | name: "Invoke my-auth-tool without auth token",
208 | api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke",
209 | requestHeader: map[string]string{},
210 | requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)),
211 | isErr: true,
212 | },
213 | {
214 | name: "Invoke my-auth-required-tool with auth token",
215 | api: "http://127.0.0.1:5000/api/tool/my-auth-required-tool/invoke",
216 | requestHeader: map[string]string{"my-google-auth_token": idToken},
217 | requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)),
218 | isErr: false,
219 | want: "[{\"execute_nl_query\":{\"?column?\":1}}]",
220 | },
221 | {
222 | name: "Invoke my-auth-required-tool with invalid auth token",
223 | api: "http://127.0.0.1:5000/api/tool/my-auth-required-tool/invoke",
224 | requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
225 | requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)),
226 | isErr: true,
227 | },
228 | {
229 | name: "Invoke my-auth-required-tool without auth token",
230 | api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke",
231 | requestHeader: map[string]string{},
232 | requestBody: bytes.NewBuffer([]byte(`{"question": "return 1"}`)),
233 | isErr: true,
234 | },
235 | }
236 | for _, tc := range invokeTcs {
237 | t.Run(tc.name, func(t *testing.T) {
238 | // Send Tool invocation request
239 | req, err := http.NewRequest(http.MethodPost, tc.api, tc.requestBody)
240 | if err != nil {
241 | t.Fatalf("unable to create request: %s", err)
242 | }
243 | req.Header.Add("Content-type", "application/json")
244 | for k, v := range tc.requestHeader {
245 | req.Header.Add(k, v)
246 | }
247 | resp, err := http.DefaultClient.Do(req)
248 | if err != nil {
249 | t.Fatalf("unable to send request: %s", err)
250 | }
251 | defer resp.Body.Close()
252 |
253 | if resp.StatusCode != http.StatusOK {
254 | if tc.isErr == true {
255 | return
256 | }
257 | bodyBytes, _ := io.ReadAll(resp.Body)
258 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
259 | }
260 |
261 | // Check response body
262 | var body map[string]interface{}
263 | err = json.NewDecoder(resp.Body).Decode(&body)
264 | if err != nil {
265 | t.Fatalf("error parsing response body")
266 | }
267 | got, ok := body["result"].(string)
268 | if !ok {
269 | t.Fatalf("unable to find result in response body")
270 | }
271 |
272 | if got != tc.want {
273 | t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
274 | }
275 | })
276 | }
277 |
278 | }
279 |
280 | func getAINLToolsConfig(sourceConfig map[string]any) map[string]any {
281 | // Write config into a file and pass it to command
282 | toolsFile := map[string]any{
283 | "sources": map[string]any{
284 | "my-instance": sourceConfig,
285 | },
286 | "authServices": map[string]any{
287 | "my-google-auth": map[string]any{
288 | "kind": "google",
289 | "clientId": tests.ClientId,
290 | },
291 | },
292 | "tools": map[string]any{
293 | "my-simple-tool": map[string]any{
294 | "kind": AlloyDBAINLToolKind,
295 | "source": "my-instance",
296 | "description": "Simple tool to test end to end functionality.",
297 | "nlConfig": "my_nl_config",
298 | },
299 | "my-auth-tool": map[string]any{
300 | "kind": AlloyDBAINLToolKind,
301 | "source": "my-instance",
302 | "description": "Tool to test authenticated parameters.",
303 | "nlConfig": "my_nl_config",
304 | "nlConfigParameters": []map[string]any{
305 | {
306 | "name": "email",
307 | "type": "string",
308 | "description": "user email",
309 | "authServices": []map[string]string{
310 | {
311 | "name": "my-google-auth",
312 | "field": "email",
313 | },
314 | },
315 | },
316 | },
317 | },
318 | "my-auth-required-tool": map[string]any{
319 | "kind": AlloyDBAINLToolKind,
320 | "source": "my-instance",
321 | "description": "Tool to test auth required invocation.",
322 | "nlConfig": "my_nl_config",
323 | "authRequired": []string{
324 | "my-google-auth",
325 | },
326 | },
327 | },
328 | }
329 |
330 | return toolsFile
331 | }
332 |
333 | func runAINLMCPToolCallMethod(t *testing.T) {
334 | sessionId := tests.RunInitialize(t, "2024-11-05")
335 | header := map[string]string{}
336 | if sessionId != "" {
337 | header["Mcp-Session-Id"] = sessionId
338 | }
339 |
340 | // Test tool invoke endpoint
341 | invokeTcs := []struct {
342 | name string
343 | api string
344 | requestBody jsonrpc.JSONRPCRequest
345 | requestHeader map[string]string
346 | want string
347 | }{
348 | {
349 | name: "MCP Invoke my-simple-tool",
350 | api: "http://127.0.0.1:5000/mcp",
351 | requestHeader: map[string]string{},
352 | requestBody: jsonrpc.JSONRPCRequest{
353 | Jsonrpc: "2.0",
354 | Id: "my-simple-tool",
355 | Request: jsonrpc.Request{
356 | Method: "tools/call",
357 | },
358 | Params: map[string]any{
359 | "name": "my-simple-tool",
360 | "arguments": map[string]any{
361 | "question": "return 1",
362 | },
363 | },
364 | },
365 | want: `{"jsonrpc":"2.0","id":"my-simple-tool","result":{"content":[{"type":"text","text":"{\"execute_nl_query\":{\"?column?\":1}}"}]}}`,
366 | },
367 | {
368 | name: "MCP Invoke invalid tool",
369 | api: "http://127.0.0.1:5000/mcp",
370 | requestHeader: map[string]string{},
371 | requestBody: jsonrpc.JSONRPCRequest{
372 | Jsonrpc: "2.0",
373 | Id: "invalid-tool",
374 | Request: jsonrpc.Request{
375 | Method: "tools/call",
376 | },
377 | Params: map[string]any{
378 | "name": "foo",
379 | "arguments": map[string]any{},
380 | },
381 | },
382 | want: `{"jsonrpc":"2.0","id":"invalid-tool","error":{"code":-32602,"message":"invalid tool name: tool with name \"foo\" does not exist"}}`,
383 | },
384 | {
385 | name: "MCP Invoke my-auth-tool without parameters",
386 | api: "http://127.0.0.1:5000/mcp",
387 | requestHeader: map[string]string{},
388 | requestBody: jsonrpc.JSONRPCRequest{
389 | Jsonrpc: "2.0",
390 | Id: "invoke-without-parameter",
391 | Request: jsonrpc.Request{
392 | Method: "tools/call",
393 | },
394 | Params: map[string]any{
395 | "name": "my-auth-tool",
396 | "arguments": map[string]any{},
397 | },
398 | },
399 | want: `{"jsonrpc":"2.0","id":"invoke-without-parameter","error":{"code":-32602,"message":"provided parameters were invalid: parameter question is required"}}`,
400 | },
401 | }
402 | for _, tc := range invokeTcs {
403 | t.Run(tc.name, func(t *testing.T) {
404 | reqMarshal, err := json.Marshal(tc.requestBody)
405 | if err != nil {
406 | t.Fatalf("unexpected error during marshaling of request body")
407 | }
408 | // Send Tool invocation request
409 | req, err := http.NewRequest(http.MethodPost, tc.api, bytes.NewBuffer(reqMarshal))
410 | if err != nil {
411 | t.Fatalf("unable to create request: %s", err)
412 | }
413 | req.Header.Add("Content-type", "application/json")
414 | for k, v := range header {
415 | req.Header.Add(k, v)
416 | }
417 | resp, err := http.DefaultClient.Do(req)
418 | if err != nil {
419 | t.Fatalf("unable to send request: %s", err)
420 | }
421 | respBody, err := io.ReadAll(resp.Body)
422 | if err != nil {
423 | t.Fatalf("unable to read request body: %s", err)
424 | }
425 | defer resp.Body.Close()
426 | got := string(bytes.TrimSpace(respBody))
427 |
428 | // Remove `\` and `"` for string comparison
429 | got = strings.ReplaceAll(got, "\\", "")
430 | want := strings.ReplaceAll(tc.want, "\\", "")
431 | got = strings.ReplaceAll(got, "\"", "")
432 | want = strings.ReplaceAll(want, "\"", "")
433 |
434 | if !strings.Contains(got, want) {
435 | t.Fatalf("Expected substring not found:\ngot: %q\nwant: %q (to be contained within got)", got, want)
436 | }
437 | })
438 | }
439 | }
440 |
```
--------------------------------------------------------------------------------
/internal/tools/bigquery/bigqueryexecutesql/bigqueryexecutesql.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 bigqueryexecutesql
16 |
17 | import (
18 | "context"
19 | "encoding/json"
20 | "fmt"
21 | "strings"
22 |
23 | bigqueryapi "cloud.google.com/go/bigquery"
24 | yaml "github.com/goccy/go-yaml"
25 | "github.com/googleapis/genai-toolbox/internal/sources"
26 | bigqueryds "github.com/googleapis/genai-toolbox/internal/sources/bigquery"
27 | "github.com/googleapis/genai-toolbox/internal/tools"
28 | bqutil "github.com/googleapis/genai-toolbox/internal/tools/bigquery/bigquerycommon"
29 | "github.com/googleapis/genai-toolbox/internal/util"
30 | bigqueryrestapi "google.golang.org/api/bigquery/v2"
31 | "google.golang.org/api/iterator"
32 | )
33 |
34 | const kind string = "bigquery-execute-sql"
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 | BigQueryClient() *bigqueryapi.Client
52 | BigQuerySession() bigqueryds.BigQuerySessionProvider
53 | BigQueryWriteMode() string
54 | BigQueryRestService() *bigqueryrestapi.Service
55 | BigQueryClientCreator() bigqueryds.BigqueryClientCreator
56 | UseClientAuthorization() bool
57 | IsDatasetAllowed(projectID, datasetID string) bool
58 | BigQueryAllowedDatasets() []string
59 | }
60 |
61 | // validate compatible sources are still compatible
62 | var _ compatibleSource = &bigqueryds.Source{}
63 |
64 | var compatibleSources = [...]string{bigqueryds.SourceKind}
65 |
66 | type Config struct {
67 | Name string `yaml:"name" validate:"required"`
68 | Kind string `yaml:"kind" validate:"required"`
69 | Source string `yaml:"source" validate:"required"`
70 | Description string `yaml:"description" validate:"required"`
71 | AuthRequired []string `yaml:"authRequired"`
72 | }
73 |
74 | // validate interface
75 | var _ tools.ToolConfig = Config{}
76 |
77 | func (cfg Config) ToolConfigKind() string {
78 | return kind
79 | }
80 |
81 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
82 | // verify source exists
83 | rawS, ok := srcs[cfg.Source]
84 | if !ok {
85 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
86 | }
87 |
88 | // verify the source is compatible
89 | s, ok := rawS.(compatibleSource)
90 | if !ok {
91 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
92 | }
93 |
94 | var sqlDescriptionBuilder strings.Builder
95 | switch s.BigQueryWriteMode() {
96 | case bigqueryds.WriteModeBlocked:
97 | sqlDescriptionBuilder.WriteString("The SQL to execute. In 'blocked' mode, only SELECT statements are allowed; other statement types will fail.")
98 | case bigqueryds.WriteModeProtected:
99 | sqlDescriptionBuilder.WriteString("The SQL to execute. Only SELECT statements and writes to the session's temporary dataset are allowed (e.g., `CREATE TEMP TABLE ...`).")
100 | default: // WriteModeAllowed
101 | sqlDescriptionBuilder.WriteString("The SQL to execute.")
102 | }
103 |
104 | allowedDatasets := s.BigQueryAllowedDatasets()
105 | if len(allowedDatasets) > 0 {
106 | if len(allowedDatasets) == 1 {
107 | datasetFQN := allowedDatasets[0]
108 | parts := strings.Split(datasetFQN, ".")
109 | if len(parts) < 2 {
110 | return nil, fmt.Errorf("expected allowedDataset to have at least 2 parts (project.dataset): %s", datasetFQN)
111 | }
112 | datasetID := parts[1]
113 | sqlDescriptionBuilder.WriteString(fmt.Sprintf(" The query must only access the `%s` dataset. "+
114 | "To query a table within this dataset (e.g., `my_table`), "+
115 | "qualify it with the dataset id (e.g., `%s.my_table`).", datasetFQN, datasetID))
116 | } else {
117 | datasetIDs := []string{}
118 | for _, ds := range allowedDatasets {
119 | datasetIDs = append(datasetIDs, fmt.Sprintf("`%s`", ds))
120 | }
121 | sqlDescriptionBuilder.WriteString(fmt.Sprintf(" The query must only access datasets from the following list: %s.", strings.Join(datasetIDs, ", ")))
122 | }
123 | }
124 |
125 | sqlParameter := tools.NewStringParameter("sql", sqlDescriptionBuilder.String())
126 | dryRunParameter := tools.NewBooleanParameterWithDefault(
127 | "dry_run",
128 | false,
129 | "If set to true, the query will be validated and information about the execution will be returned "+
130 | "without running the query. Defaults to false.",
131 | )
132 | parameters := tools.Parameters{sqlParameter, dryRunParameter}
133 | mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters)
134 |
135 | // finish tool setup
136 | t := Tool{
137 | Name: cfg.Name,
138 | Kind: kind,
139 | Parameters: parameters,
140 | AuthRequired: cfg.AuthRequired,
141 | UseClientOAuth: s.UseClientAuthorization(),
142 | ClientCreator: s.BigQueryClientCreator(),
143 | Client: s.BigQueryClient(),
144 | RestService: s.BigQueryRestService(),
145 | WriteMode: s.BigQueryWriteMode(),
146 | SessionProvider: s.BigQuerySession(),
147 | IsDatasetAllowed: s.IsDatasetAllowed,
148 | AllowedDatasets: allowedDatasets,
149 | manifest: tools.Manifest{Description: cfg.Description, Parameters: parameters.Manifest(), AuthRequired: cfg.AuthRequired},
150 | mcpManifest: mcpManifest,
151 | }
152 | return t, nil
153 | }
154 |
155 | // validate interface
156 | var _ tools.Tool = Tool{}
157 |
158 | type Tool struct {
159 | Name string `yaml:"name"`
160 | Kind string `yaml:"kind"`
161 | AuthRequired []string `yaml:"authRequired"`
162 | UseClientOAuth bool `yaml:"useClientOAuth"`
163 | Parameters tools.Parameters `yaml:"parameters"`
164 |
165 | Client *bigqueryapi.Client
166 | RestService *bigqueryrestapi.Service
167 | WriteMode string
168 | SessionProvider bigqueryds.BigQuerySessionProvider
169 | ClientCreator bigqueryds.BigqueryClientCreator
170 | IsDatasetAllowed func(projectID, datasetID string) bool
171 | AllowedDatasets []string
172 | manifest tools.Manifest
173 | mcpManifest tools.McpManifest
174 | }
175 |
176 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
177 | paramsMap := params.AsMap()
178 | sql, ok := paramsMap["sql"].(string)
179 | if !ok {
180 | return nil, fmt.Errorf("unable to cast sql parameter %s", paramsMap["sql"])
181 | }
182 | dryRun, ok := paramsMap["dry_run"].(bool)
183 | if !ok {
184 | return nil, fmt.Errorf("unable to cast dry_run parameter %s", paramsMap["dry_run"])
185 | }
186 |
187 | bqClient := t.Client
188 | restService := t.RestService
189 |
190 | var err error
191 | // Initialize new client if using user OAuth token
192 | if t.UseClientOAuth {
193 | tokenStr, err := accessToken.ParseBearerToken()
194 | if err != nil {
195 | return nil, fmt.Errorf("error parsing access token: %w", err)
196 | }
197 | bqClient, restService, err = t.ClientCreator(tokenStr, true)
198 | if err != nil {
199 | return nil, fmt.Errorf("error creating client from OAuth access token: %w", err)
200 | }
201 | }
202 |
203 | var connProps []*bigqueryapi.ConnectionProperty
204 | var session *bigqueryds.Session
205 | if t.WriteMode == bigqueryds.WriteModeProtected {
206 | session, err = t.SessionProvider(ctx)
207 | if err != nil {
208 | return nil, fmt.Errorf("failed to get BigQuery session for protected mode: %w", err)
209 | }
210 | connProps = []*bigqueryapi.ConnectionProperty{
211 | {Key: "session_id", Value: session.ID},
212 | }
213 | }
214 |
215 | dryRunJob, err := bqutil.DryRunQuery(ctx, restService, bqClient.Project(), bqClient.Location, sql, nil, connProps)
216 | if err != nil {
217 | return nil, fmt.Errorf("query validation failed: %w", err)
218 | }
219 |
220 | statementType := dryRunJob.Statistics.Query.StatementType
221 |
222 | switch t.WriteMode {
223 | case bigqueryds.WriteModeBlocked:
224 | if statementType != "SELECT" {
225 | return nil, fmt.Errorf("write mode is 'blocked', only SELECT statements are allowed")
226 | }
227 | case bigqueryds.WriteModeProtected:
228 | if dryRunJob.Configuration != nil && dryRunJob.Configuration.Query != nil {
229 | if dest := dryRunJob.Configuration.Query.DestinationTable; dest != nil && dest.DatasetId != session.DatasetID {
230 | return nil, fmt.Errorf("protected write mode only supports SELECT statements, or write operations in the anonymous "+
231 | "dataset of a BigQuery session, but destination was %q", dest.DatasetId)
232 | }
233 | }
234 | }
235 |
236 | if len(t.AllowedDatasets) > 0 {
237 | switch statementType {
238 | case "CREATE_SCHEMA", "DROP_SCHEMA", "ALTER_SCHEMA":
239 | return nil, fmt.Errorf("dataset-level operations like '%s' are not allowed when dataset restrictions are in place", statementType)
240 | case "CREATE_FUNCTION", "CREATE_TABLE_FUNCTION", "CREATE_PROCEDURE":
241 | return nil, fmt.Errorf("creating stored routines ('%s') is not allowed when dataset restrictions are in place, as their contents cannot be safely analyzed", statementType)
242 | case "CALL":
243 | return nil, fmt.Errorf("calling stored procedures ('%s') is not allowed when dataset restrictions are in place, as their contents cannot be safely analyzed", statementType)
244 | }
245 |
246 | // Use a map to avoid duplicate table names.
247 | tableIDSet := make(map[string]struct{})
248 |
249 | // Get all tables from the dry run result. This is the most reliable method.
250 | queryStats := dryRunJob.Statistics.Query
251 | if queryStats != nil {
252 | for _, tableRef := range queryStats.ReferencedTables {
253 | tableIDSet[fmt.Sprintf("%s.%s.%s", tableRef.ProjectId, tableRef.DatasetId, tableRef.TableId)] = struct{}{}
254 | }
255 | if tableRef := queryStats.DdlTargetTable; tableRef != nil {
256 | tableIDSet[fmt.Sprintf("%s.%s.%s", tableRef.ProjectId, tableRef.DatasetId, tableRef.TableId)] = struct{}{}
257 | }
258 | if tableRef := queryStats.DdlDestinationTable; tableRef != nil {
259 | tableIDSet[fmt.Sprintf("%s.%s.%s", tableRef.ProjectId, tableRef.DatasetId, tableRef.TableId)] = struct{}{}
260 | }
261 | }
262 |
263 | var tableNames []string
264 | if len(tableIDSet) > 0 {
265 | for tableID := range tableIDSet {
266 | tableNames = append(tableNames, tableID)
267 | }
268 | } else if statementType != "SELECT" {
269 | // If dry run yields no tables, fall back to the parser for non-SELECT statements
270 | // to catch unsafe operations like EXECUTE IMMEDIATE.
271 | parsedTables, parseErr := bqutil.TableParser(sql, t.Client.Project())
272 | if parseErr != nil {
273 | // If parsing fails (e.g., EXECUTE IMMEDIATE), we cannot guarantee safety, so we must fail.
274 | return nil, fmt.Errorf("could not parse tables from query to validate against allowed datasets: %w", parseErr)
275 | }
276 | tableNames = parsedTables
277 | }
278 |
279 | for _, tableID := range tableNames {
280 | parts := strings.Split(tableID, ".")
281 | if len(parts) == 3 {
282 | projectID, datasetID := parts[0], parts[1]
283 | if !t.IsDatasetAllowed(projectID, datasetID) {
284 | return nil, fmt.Errorf("query accesses dataset '%s.%s', which is not in the allowed list", projectID, datasetID)
285 | }
286 | }
287 | }
288 | }
289 |
290 | if dryRun {
291 | if dryRunJob != nil {
292 | jobJSON, err := json.MarshalIndent(dryRunJob, "", " ")
293 | if err != nil {
294 | return nil, fmt.Errorf("failed to marshal dry run job to JSON: %w", err)
295 | }
296 | return string(jobJSON), nil
297 | }
298 | // This case should not be reached, but as a fallback, we return a message.
299 | return "Dry run was requested, but no job information was returned.", nil
300 | }
301 |
302 | query := bqClient.Query(sql)
303 | query.Location = bqClient.Location
304 |
305 | query.ConnectionProperties = connProps
306 |
307 | // Log the query executed for debugging.
308 | logger, err := util.LoggerFromContext(ctx)
309 | if err != nil {
310 | return nil, fmt.Errorf("error getting logger: %s", err)
311 | }
312 | logger.DebugContext(ctx, "executing `%s` tool query: %s", kind, sql)
313 |
314 | // This block handles SELECT statements, which return a row set.
315 | // We iterate through the results, convert each row into a map of
316 | // column names to values, and return the collection of rows.
317 | var out []any
318 | job, err := query.Run(ctx)
319 | if err != nil {
320 | return nil, fmt.Errorf("unable to execute query: %w", err)
321 | }
322 | it, err := job.Read(ctx)
323 | if err != nil {
324 | return nil, fmt.Errorf("unable to read query results: %w", err)
325 | }
326 | for {
327 | var row map[string]bigqueryapi.Value
328 | err = it.Next(&row)
329 | if err == iterator.Done {
330 | break
331 | }
332 | if err != nil {
333 | return nil, fmt.Errorf("unable to iterate through query results: %w", err)
334 | }
335 | vMap := make(map[string]any)
336 | for key, value := range row {
337 | vMap[key] = value
338 | }
339 | out = append(out, vMap)
340 | }
341 | // If the query returned any rows, return them directly.
342 | if len(out) > 0 {
343 | return out, nil
344 | }
345 |
346 | // This handles the standard case for a SELECT query that successfully
347 | // executes but returns zero rows.
348 | if statementType == "SELECT" {
349 | return "The query returned 0 rows.", nil
350 | }
351 | // This is the fallback for a successful query that doesn't return content.
352 | // In most cases, this will be for DML/DDL statements like INSERT, UPDATE, CREATE, etc.
353 | // However, it is also possible that this was a query that was expected to return rows
354 | // but returned none, a case that we cannot distinguish here.
355 | return "Query executed successfully and returned no content.", nil
356 | }
357 |
358 | func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
359 | return tools.ParseParams(t.Parameters, data, claims)
360 | }
361 |
362 | func (t Tool) Manifest() tools.Manifest {
363 | return t.manifest
364 | }
365 |
366 | func (t Tool) McpManifest() tools.McpManifest {
367 | return t.mcpManifest
368 | }
369 |
370 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
371 | return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
372 | }
373 |
374 | func (t Tool) RequiresClientAuthorization() bool {
375 | return t.UseClientOAuth
376 | }
377 |
```
--------------------------------------------------------------------------------
/internal/tools/looker/lookerhealthvacuum/lookerhealthvacuum.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 lookerhealthvacuum
15 |
16 | import (
17 | "context"
18 | "encoding/json"
19 | "fmt"
20 | "regexp"
21 | "strings"
22 |
23 | yaml "github.com/goccy/go-yaml"
24 | "github.com/googleapis/genai-toolbox/internal/sources"
25 | lookersrc "github.com/googleapis/genai-toolbox/internal/sources/looker"
26 | "github.com/googleapis/genai-toolbox/internal/tools"
27 | "github.com/googleapis/genai-toolbox/internal/tools/looker/lookercommon"
28 | "github.com/googleapis/genai-toolbox/internal/util"
29 | "github.com/looker-open-source/sdk-codegen/go/rtl"
30 | v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4"
31 | )
32 |
33 | // =================================================================================================================
34 | // START MCP SERVER CORE LOGIC
35 | // =================================================================================================================
36 | const kind string = "looker-health-vacuum"
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 | Parameters map[string]any `yaml:"parameters"`
59 | }
60 |
61 | var _ tools.ToolConfig = Config{}
62 |
63 | func (cfg Config) ToolConfigKind() string {
64 | return kind
65 | }
66 |
67 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
68 | rawS, ok := srcs[cfg.Source]
69 | if !ok {
70 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
71 | }
72 |
73 | s, ok := rawS.(*lookersrc.Source)
74 | if !ok {
75 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be `looker`", kind)
76 | }
77 |
78 | actionParameter := tools.NewStringParameterWithRequired("action", "The vacuum action to run. Can be 'models', or 'explores'.", true)
79 | projectParameter := tools.NewStringParameterWithDefault("project", "", "The Looker project to vacuum (optional).")
80 | modelParameter := tools.NewStringParameterWithDefault("model", "", "The Looker model to vacuum (optional).")
81 | exploreParameter := tools.NewStringParameterWithDefault("explore", "", "The Looker explore to vacuum (optional).")
82 | timeframeParameter := tools.NewIntParameterWithDefault("timeframe", 90, "The timeframe in days to analyze.")
83 | minQueriesParameter := tools.NewIntParameterWithDefault("min_queries", 1, "The minimum number of queries for a model or explore to be considered used.")
84 |
85 | parameters := tools.Parameters{
86 | actionParameter,
87 | projectParameter,
88 | modelParameter,
89 | exploreParameter,
90 | timeframeParameter,
91 | minQueriesParameter,
92 | }
93 |
94 | mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters)
95 |
96 | return Tool{
97 | Name: cfg.Name,
98 | Kind: kind,
99 | Parameters: parameters,
100 | AuthRequired: cfg.AuthRequired,
101 | UseClientOAuth: s.UseClientOAuth,
102 | Client: s.Client,
103 | ApiSettings: s.ApiSettings,
104 | manifest: tools.Manifest{
105 | Description: cfg.Description,
106 | Parameters: parameters.Manifest(),
107 | AuthRequired: cfg.AuthRequired,
108 | },
109 | mcpManifest: mcpManifest,
110 | }, nil
111 | }
112 |
113 | var _ tools.Tool = Tool{}
114 |
115 | type Tool struct {
116 | Name string `yaml:"name"`
117 | Kind string `yaml:"kind"`
118 | UseClientOAuth bool
119 | Client *v4.LookerSDK
120 | ApiSettings *rtl.ApiSettings
121 | AuthRequired []string `yaml:"authRequired"`
122 | Parameters tools.Parameters
123 | manifest tools.Manifest
124 | mcpManifest tools.McpManifest
125 | }
126 |
127 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
128 | sdk, err := lookercommon.GetLookerSDK(t.UseClientOAuth, t.ApiSettings, t.Client, accessToken)
129 | if err != nil {
130 | return nil, fmt.Errorf("error getting sdk: %w", err)
131 | }
132 |
133 | paramsMap := params.AsMap()
134 | timeframe, _ := paramsMap["timeframe"].(int)
135 | if timeframe == 0 {
136 | timeframe = 90
137 | }
138 | minQueries, _ := paramsMap["min_queries"].(int)
139 | if minQueries == 0 {
140 | minQueries = 1
141 | }
142 |
143 | vacuumTool := &vacuumTool{
144 | SdkClient: sdk,
145 | timeframe: timeframe,
146 | minQueries: minQueries,
147 | }
148 |
149 | action, ok := paramsMap["action"].(string)
150 | if !ok {
151 | return nil, fmt.Errorf("action parameter not found")
152 | }
153 |
154 | switch action {
155 | case "models":
156 | project, _ := paramsMap["project"].(string)
157 | model, _ := paramsMap["model"].(string)
158 | return vacuumTool.models(ctx, project, model)
159 | case "explores":
160 | model, _ := paramsMap["model"].(string)
161 | explore, _ := paramsMap["explore"].(string)
162 | return vacuumTool.explores(ctx, model, explore)
163 | default:
164 | return nil, fmt.Errorf("unknown action: %s", action)
165 | }
166 | }
167 |
168 | func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
169 | return tools.ParseParams(t.Parameters, data, claims)
170 | }
171 |
172 | func (t Tool) Manifest() tools.Manifest {
173 | return t.manifest
174 | }
175 |
176 | func (t Tool) McpManifest() tools.McpManifest {
177 | return t.mcpManifest
178 | }
179 |
180 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
181 | return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
182 | }
183 |
184 | func (t Tool) RequiresClientAuthorization() bool {
185 | return t.UseClientOAuth
186 | }
187 |
188 | // =================================================================================================================
189 | // END MCP SERVER CORE LOGIC
190 | // =================================================================================================================
191 |
192 | // =================================================================================================================
193 | // START LOOKER HEALTH VACUUM CORE LOGIC
194 | // =================================================================================================================
195 | type vacuumTool struct {
196 | SdkClient *v4.LookerSDK
197 | timeframe int
198 | minQueries int
199 | }
200 |
201 | func (t *vacuumTool) models(ctx context.Context, project, model string) ([]map[string]interface{}, error) {
202 | logger, err := util.LoggerFromContext(ctx)
203 | if err != nil {
204 | return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
205 | }
206 | logger.InfoContext(ctx, "Vacuuming models...")
207 |
208 | usedModels, err := t.getUsedModels(ctx)
209 | if err != nil {
210 | return nil, err
211 | }
212 |
213 | lookmlModels, err := t.SdkClient.AllLookmlModels(v4.RequestAllLookmlModels{}, nil)
214 | if err != nil {
215 | return nil, fmt.Errorf("error fetching LookML models: %w", err)
216 | }
217 |
218 | var results []map[string]interface{}
219 | for _, m := range lookmlModels {
220 | if (project == "" || (m.ProjectName != nil && *m.ProjectName == project)) &&
221 | (model == "" || (m.Name != nil && *m.Name == model)) {
222 |
223 | queryCount := 0
224 | if qc, ok := usedModels[*m.Name]; ok {
225 | queryCount = qc
226 | }
227 |
228 | unusedExplores, err := t.getUnusedExplores(ctx, *m.Name)
229 | if err != nil {
230 | return nil, err
231 | }
232 |
233 | results = append(results, map[string]interface{}{
234 | "Model": *m.Name,
235 | "Unused Explores": unusedExplores,
236 | "Model Query Count": queryCount,
237 | })
238 | }
239 | }
240 | return results, nil
241 | }
242 |
243 | func (t *vacuumTool) explores(ctx context.Context, model, explore string) ([]map[string]interface{}, error) {
244 | logger, err := util.LoggerFromContext(ctx)
245 | if err != nil {
246 | return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
247 | }
248 | logger.InfoContext(ctx, "Vacuuming explores...")
249 |
250 | lookmlModels, err := t.SdkClient.AllLookmlModels(v4.RequestAllLookmlModels{}, nil)
251 | if err != nil {
252 | return nil, fmt.Errorf("error fetching LookML models: %w", err)
253 | }
254 |
255 | var results []map[string]interface{}
256 | for _, m := range lookmlModels {
257 | if model != "" && (m.Name == nil || *m.Name != model) {
258 | continue
259 | }
260 | if m.Explores == nil {
261 | continue
262 | }
263 |
264 | for _, e := range *m.Explores {
265 | if explore != "" && (e.Name == nil || *e.Name != explore) {
266 | continue
267 | }
268 | if e.Name == nil {
269 | continue
270 | }
271 |
272 | exploreDetail, err := t.SdkClient.LookmlModelExplore(v4.RequestLookmlModelExplore{
273 | LookmlModelName: *m.Name,
274 | ExploreName: *e.Name,
275 | }, nil)
276 | if err != nil {
277 | logger.ErrorContext(ctx, fmt.Sprintf("Error fetching detail for explore %s.%s: %v", *m.Name, *e.Name, err))
278 | continue
279 | }
280 |
281 | usedFields, err := t.getUsedExploreFields(ctx, *m.Name, *e.Name)
282 | if err != nil {
283 | logger.ErrorContext(ctx, fmt.Sprintf("Error fetching used fields for explore %s.%s: %v", *m.Name, *e.Name, err))
284 | continue
285 | }
286 |
287 | var allFields []string
288 | if exploreDetail.Fields != nil {
289 | for _, d := range *exploreDetail.Fields.Dimensions {
290 | if !*d.Hidden {
291 | allFields = append(allFields, *d.Name)
292 | }
293 | }
294 | for _, ms := range *exploreDetail.Fields.Measures {
295 | if !*ms.Hidden {
296 | allFields = append(allFields, *ms.Name)
297 | }
298 | }
299 | }
300 |
301 | var unusedFields []string
302 | for _, field := range allFields {
303 | if _, ok := usedFields[field]; !ok {
304 | unusedFields = append(unusedFields, field)
305 | }
306 | }
307 |
308 | joinStats := make(map[string]int)
309 | if exploreDetail.Joins != nil {
310 | for field, queryCount := range usedFields {
311 | join := strings.Split(field, ".")[0]
312 | joinStats[join] += queryCount
313 | }
314 | for _, join := range *exploreDetail.Joins {
315 | if _, ok := joinStats[*join.Name]; !ok {
316 | joinStats[*join.Name] = 0
317 | }
318 | }
319 | }
320 |
321 | var unusedJoins []string
322 | for join, count := range joinStats {
323 | if count == 0 {
324 | unusedJoins = append(unusedJoins, join)
325 | }
326 | }
327 |
328 | results = append(results, map[string]interface{}{
329 | "Model": *m.Name,
330 | "Explore": *e.Name,
331 | "Unused Joins": unusedJoins,
332 | "Unused Fields": unusedFields,
333 | })
334 | }
335 | }
336 | return results, nil
337 | }
338 |
339 | func (t *vacuumTool) getUsedModels(ctx context.Context) (map[string]int, error) {
340 | limit := "5000"
341 | query := &v4.WriteQuery{
342 | Model: "system__activity",
343 | View: "history",
344 | Fields: &[]string{"history.query_run_count", "query.model"},
345 | Filters: &map[string]any{
346 | "history.created_date": fmt.Sprintf("%d days", t.timeframe),
347 | "query.model": "-system__activity, -i__looker",
348 | "history.query_run_count": fmt.Sprintf(">%d", t.minQueries-1),
349 | "user.dev_branch_name": "NULL",
350 | },
351 | Limit: &limit,
352 | }
353 | raw, err := lookercommon.RunInlineQuery(ctx, t.SdkClient, query, "json", nil)
354 | if err != nil {
355 | return nil, err
356 | }
357 |
358 | var data []map[string]interface{}
359 | _ = json.Unmarshal([]byte(raw), &data)
360 |
361 | results := make(map[string]int)
362 | for _, row := range data {
363 | model, _ := row["query.model"].(string)
364 | count, _ := row["history.query_run_count"].(float64)
365 | results[model] = int(count)
366 | }
367 | return results, nil
368 | }
369 |
370 | func (t *vacuumTool) getUnusedExplores(ctx context.Context, modelName string) ([]string, error) {
371 | lookmlModel, err := t.SdkClient.LookmlModel(modelName, "", nil)
372 | if err != nil {
373 | return nil, fmt.Errorf("error fetching LookML model %s: %w", modelName, err)
374 | }
375 |
376 | var unusedExplores []string
377 | if lookmlModel.Explores != nil {
378 | for _, e := range *lookmlModel.Explores {
379 | limit := "1"
380 | queryCountQueryBody := &v4.WriteQuery{
381 | Model: "system__activity",
382 | View: "history",
383 | Fields: &[]string{"history.query_run_count"},
384 | Filters: &map[string]any{
385 | "query.model": modelName,
386 | "query.view": *e.Name,
387 | "history.created_date": fmt.Sprintf("%d days", t.timeframe),
388 | "history.query_run_count": fmt.Sprintf(">%d", t.minQueries-1),
389 | "user.dev_branch_name": "NULL",
390 | },
391 | Limit: &limit,
392 | }
393 |
394 | rawQueryCount, err := lookercommon.RunInlineQuery(ctx, t.SdkClient, queryCountQueryBody, "json", nil)
395 | if err != nil {
396 | // Log the error but continue
397 | continue
398 | }
399 |
400 | var data []map[string]interface{}
401 | _ = json.Unmarshal([]byte(rawQueryCount), &data)
402 | if len(data) == 0 {
403 | unusedExplores = append(unusedExplores, *e.Name)
404 | }
405 | }
406 | }
407 | return unusedExplores, nil
408 | }
409 |
410 | func (t *vacuumTool) getUsedExploreFields(ctx context.Context, model, explore string) (map[string]int, error) {
411 | limit := "5000"
412 | query := &v4.WriteQuery{
413 | Model: "system__activity",
414 | View: "history",
415 | Fields: &[]string{"query.formatted_fields", "query.filters", "history.query_run_count"},
416 | Filters: &map[string]any{
417 | "history.created_date": fmt.Sprintf("%d days", t.timeframe),
418 | "query.model": strings.ReplaceAll(model, "_", "^_"),
419 | "query.view": strings.ReplaceAll(explore, "_", "^_"),
420 | "query.formatted_fields": "-NULL",
421 | "history.workspace_id": "production",
422 | },
423 | Limit: &limit,
424 | }
425 | raw, err := lookercommon.RunInlineQuery(ctx, t.SdkClient, query, "json", nil)
426 | if err != nil {
427 | return nil, err
428 | }
429 |
430 | var data []map[string]interface{}
431 | _ = json.Unmarshal([]byte(raw), &data)
432 |
433 | results := make(map[string]int)
434 | fieldRegex := regexp.MustCompile(`(\w+\.\w+)`)
435 |
436 | for _, row := range data {
437 | count, _ := row["history.query_run_count"].(float64)
438 | formattedFields, _ := row["query.formatted_fields"].(string)
439 | filters, _ := row["query.filters"].(string)
440 |
441 | usedFields := make(map[string]bool)
442 |
443 | for _, field := range fieldRegex.FindAllString(formattedFields, -1) {
444 | results[field] += int(count)
445 | usedFields[field] = true
446 | }
447 |
448 | for _, field := range fieldRegex.FindAllString(filters, -1) {
449 | if _, ok := usedFields[field]; !ok {
450 | results[field] += int(count)
451 | }
452 | }
453 | }
454 | return results, nil
455 | }
456 |
457 | // =================================================================================================================
458 | // END LOOKER HEALTH VACUUM CORE LOGIC
459 | // =================================================================================================================
460 |
```
--------------------------------------------------------------------------------
/internal/server/static/css/style.css:
--------------------------------------------------------------------------------
```css
1 | :root {
2 | --toolbox-blue: #4285f4;
3 | --text-primary-gray: #444444;
4 | --text-secondary-gray: #6e6e6e;
5 | --button-primary: var(--toolbox-blue);
6 | --button-secondary: #616161;
7 | --section-border: #e0e0e0;
8 | }
9 |
10 | body {
11 | display: flex;
12 | height: 100vh;
13 | margin: 0;
14 | font-family: 'Trebuchet MS';
15 | background-color: #f8f9fa;
16 | box-sizing: border-box;
17 | }
18 |
19 | *, *:before, *:after {
20 | box-sizing: inherit;
21 | }
22 |
23 | #navbar-container {
24 | flex: 0 0 250px;
25 | height: 100%;
26 | position: relative;
27 | z-index: 10;
28 | }
29 |
30 | #main-content-container {
31 | flex: 1;
32 | display: flex;
33 | flex-direction: column;
34 | min-width: 0;
35 | overflow-x: hidden;
36 | }
37 |
38 | .left-nav {
39 | background-color: #fff;
40 | box-shadow: 4px 0px 12px rgba(0, 0, 0, 0.15);
41 | display: flex;
42 | flex-direction: column;
43 | padding: 15px;
44 | align-items: center;
45 | width: 100%;
46 | height: 100%;
47 | z-index: 3;
48 |
49 | ul {
50 | font-family: 'Verdana';
51 | list-style: none;
52 | padding: 0;
53 | margin: 0;
54 | width: 100%;
55 |
56 | li {
57 | margin-bottom: 5px;
58 |
59 | a {
60 | display: flex;
61 | align-items: center;
62 | padding: 12px;
63 | text-decoration: none;
64 | color: #333;
65 | border-radius: 0;
66 |
67 | &:hover {
68 | background-color: #e9e9e9;
69 | border-radius: 35px;
70 | }
71 |
72 | &.active {
73 | background-color: #d0d0d0;
74 | font-weight: bold;
75 | border-radius: 35px;
76 | }
77 | }
78 | }
79 | }
80 | }
81 |
82 | .second-nav {
83 | flex: 0 0 250px;
84 | background-color: #fff;
85 | box-shadow: 4px 0px 12px rgba(0, 0, 0, 0.15);
86 | z-index: 2;
87 | display: flex;
88 | flex-direction: column;
89 | padding: 15px;
90 | align-items: center;
91 | position: relative;
92 | }
93 |
94 | .nav-logo {
95 | width: 90%;
96 | margin-bottom: 40px;
97 | flex-shrink: 0;
98 |
99 | img {
100 | max-width: 100%;
101 | height: auto;
102 | display: block;
103 | cursor: pointer;
104 | }
105 | }
106 |
107 | .main-content-area {
108 | flex: 1;
109 | display: flex;
110 | flex-direction: column;
111 | min-width: 0;
112 | overflow-x: hidden;
113 | }
114 |
115 | .top-bar {
116 | background-color: #fff;
117 | padding: 30px 30px;
118 | display: flex;
119 | justify-content: flex-end;
120 | align-items: center;
121 | border-bottom: 1px solid #eee;
122 | }
123 |
124 | .content {
125 | padding: 20px;
126 | flex-grow: 1;
127 | overflow-y: auto;
128 | display: flex;
129 | flex-direction: column;
130 | }
131 |
132 | .btn {
133 | display: flex;
134 | align-items: center;
135 | justify-content: center;
136 | padding: 10px 20px;
137 | color: white;
138 | border: none;
139 | border-radius: 30px;
140 | font: inherit;
141 | font-size: 1em;
142 | font-weight: bolder;
143 | cursor: pointer;
144 |
145 | &:hover {
146 | opacity: 0.8;
147 | }
148 | }
149 |
150 | .btn--run {
151 | background-color: var(--button-primary);
152 | }
153 |
154 | .btn--editHeaders {
155 | background-color: var(--button-secondary)
156 | }
157 |
158 | .btn--saveHeaders {
159 | background-color: var(--button-primary)
160 | }
161 |
162 | .btn--closeHeaders {
163 | background-color: var(--button-secondary)
164 | }
165 |
166 | .btn--setup-gis {
167 | background-color: white;
168 | color: var(--text-primary-gray);
169 | border: 2px solid var(--text-primary-gray);
170 | }
171 |
172 | .btn--externalDocs {
173 | background-color: var(--button-secondary);
174 | text-decoration: none;
175 | display: inline-flex;
176 | }
177 |
178 | .tool-button {
179 | display: flex;
180 | align-items: center;
181 | padding: 12px;
182 | text-decoration: none;
183 | color: #333;
184 | background-color: transparent;
185 | border: none;
186 | border-radius: 0;
187 | width: 100%;
188 | text-align: left;
189 | cursor: pointer;
190 | font-family: inherit;
191 | font-size: inherit;
192 |
193 | transition: background-color 0.1s ease-in-out, border-radius 0.1s ease-in-out;
194 |
195 | &:hover {
196 | background-color: #e9e9e9;
197 | border-radius: 35px;
198 | }
199 |
200 | &:focus {
201 | outline: none;
202 | box-shadow: 0 0 0 2px rgba(208, 208, 208, 0.5);
203 | }
204 |
205 | &.active {
206 | background-color: #d0d0d0;
207 | font-weight: bold;
208 | border-radius: 35px;
209 |
210 | &:hover {
211 | background-color: #d0d0d0;
212 | }
213 | }
214 | }
215 |
216 | #secondary-panel-content {
217 | flex: 1;
218 | overflow-y: auto;
219 | width: 100%;
220 | min-height: 0;
221 |
222 | ul {
223 | list-style: none;
224 | padding: 0;
225 | margin: 0;
226 | width: 100%;
227 | }
228 | }
229 |
230 | .tool-details-grid {
231 | display: grid;
232 | grid-template-columns: 1fr 2fr;
233 | gap: 20px;
234 | margin: 0 0 20px 0;
235 | align-items: start;
236 | flex-shrink: 0;
237 | }
238 |
239 | .tool-info {
240 | display: flex;
241 | flex-direction: column;
242 | gap: 15px;
243 | }
244 |
245 | .tool-execution-area {
246 | display: flex;
247 | flex-direction: column;
248 | gap: 12px;
249 | }
250 |
251 | .tool-params {
252 | background-color: #ffffff;
253 | padding: 15px;
254 | border-radius: 4px;
255 | border: 1px solid #ddd;
256 |
257 | h5 {
258 | margin-bottom: 0;
259 | }
260 | }
261 |
262 | .tool-box {
263 | background-color: #ffffff;
264 | padding: 15px;
265 | border-radius: 4px;
266 | border: 1px solid #eee;
267 |
268 | h5 {
269 | color: var(--toolbox-blue);
270 | margin-top: 0;
271 | font-weight: bold;
272 | }
273 | }
274 |
275 | .params-header {
276 | display: flex;
277 | justify-content: flex-end;
278 | margin-bottom: 8px;
279 | padding-right: 6px;
280 | font-weight: bold;
281 | font-size: 0.9em;
282 | color: var(--text-secondary-gray);
283 | }
284 |
285 | .params-disclaimer {
286 | font-style: italic;
287 | color: var(--text-secondary-gray);
288 | font-size: 0.8em;
289 | margin-bottom: 10px;
290 | width: 100%;
291 | word-wrap: break-word;
292 | }
293 |
294 | .param-item {
295 | margin-bottom: 12px;
296 |
297 | label {
298 | display: block;
299 | margin-bottom: 4px;
300 | font-family: inherit;
301 | }
302 |
303 | &.disabled-param {
304 | > label {
305 | color: #888;
306 | text-decoration: line-through;
307 | }
308 |
309 | .param-input-element {
310 | background-color: #f5f5f5;
311 | opacity: 0.6;
312 | }
313 | }
314 |
315 | input[type="text"],
316 | input[type="number"],
317 | select,
318 | textarea {
319 | width: calc(100% - 12px);
320 | padding: 6px;
321 | border: 1px solid #ccc;
322 | border-radius: 4px;
323 | font-family: inherit;
324 | }
325 |
326 | input[type="checkbox"].param-input-element {
327 | width: auto;
328 | padding: 0;
329 | border: initial;
330 | border-radius: initial;
331 | vertical-align: middle;
332 | margin-right: 4px;
333 | accent-color: var(--toolbox-blue);
334 | flex-grow: 0;
335 | }
336 | }
337 |
338 | .input-checkbox-wrapper {
339 | display: flex;
340 | align-items: center;
341 | gap: 10px;
342 | }
343 |
344 | .param-input-element-container {
345 | flex-grow: 1;
346 | }
347 |
348 | .param-input-element {
349 | box-sizing: border-box;
350 | }
351 |
352 | .include-param-container {
353 | display: flex;
354 | align-items: center;
355 | white-space: nowrap;
356 |
357 | input[type="checkbox"] {
358 | width: auto;
359 | padding: 0;
360 | border: initial;
361 | border-radius: initial;
362 | vertical-align: middle;
363 | margin-right: 0;
364 | accent-color: var(--toolbox-blue);
365 | }
366 | }
367 |
368 | .include-param-container input[type="checkbox"] {
369 | width: auto;
370 | padding: 0;
371 | border: initial;
372 | border-radius: initial;
373 | vertical-align: middle;
374 | margin: 0;
375 | accent-color: var(--toolbox-blue);
376 | }
377 |
378 | .checkbox-bool-label {
379 | margin-left: 5px;
380 | font-style: italic;
381 | color: var(--text-primary-gray);
382 | }
383 |
384 | .checkbox-bool-label.disabled {
385 | color: #aaa;
386 | cursor: not-allowed;
387 | }
388 |
389 | .param-label-extras {
390 | font-style: italic;
391 | font-weight: lighter;
392 | color: var(--text-secondary-gray);
393 | }
394 |
395 | .auth-param-input {
396 | background-color: var(--section-border);
397 | cursor: not-allowed;
398 | }
399 |
400 | .run-button-container {
401 | display: flex;
402 | justify-content: flex-end;
403 | gap: 20px;
404 | }
405 |
406 | .header-modal {
407 | display: none;
408 | position: fixed;
409 | z-index: 1000;
410 | left: 0;
411 | top: 0;
412 | width: 100%;
413 | height: 100%;
414 | overflow: auto;
415 | background-color: rgba(0,0,0,0.4);
416 |
417 | li {
418 | margin-bottom: 10px;
419 | }
420 |
421 | .header-modal-content {
422 | background-color: #fefefe;
423 | margin: 10% auto;
424 | padding: 20px;
425 | border: 1px solid #888;
426 | width: 80%;
427 | max-width: 50%;
428 | border-radius: 8px;
429 | display: flex;
430 | flex-direction: column;
431 | gap: 15px;
432 | align-items: center;
433 |
434 | h5 {
435 | margin-top: 0;
436 | font-size: 1.2em;
437 | }
438 |
439 | .headers-textarea {
440 | width: calc(100% - 16px);
441 | padding: 8px;
442 | font-family: monospace;
443 | border: 1px solid #ccc;
444 | border-radius: 4px;
445 | min-height: 150px;
446 | }
447 |
448 | .header-modal-actions {
449 | display: flex;
450 | justify-content: center;
451 | gap: 30px;
452 | width: 100%;
453 | }
454 |
455 | .auth-token-details {
456 | width: 100%;
457 | max-width: calc(100% - 16px);
458 | margin-left: 8px;
459 | margin-right: 8px;
460 |
461 | summary {
462 | cursor: pointer;
463 | text-align: left;
464 | padding: 5px 0;
465 | }
466 |
467 | .auth-token-content {
468 | padding: 10px;
469 | border: 1px solid #eee;
470 | margin-top: 5px;
471 | background-color: #f9f9f9;
472 | text-align: left;
473 | max-width: 100%;
474 | overflow-wrap: break-word;
475 |
476 | .auth-tab-group {
477 | display: flex;
478 | border-bottom: 1px solid #ccc;
479 | margin-bottom: 10px;
480 | }
481 |
482 | .auth-tab-picker {
483 | padding: 8px 12px;
484 | cursor: pointer;
485 | border: 1px solid transparent;
486 | border-bottom: 1px solid transparent;
487 | margin-bottom: -1px;
488 | background-color: #f0f0f0;
489 |
490 | &.active {
491 | background-color: #fff;
492 | border-color: #ccc;
493 | border-bottom-color: #fff;
494 | font-weight: bold;
495 | }
496 | }
497 |
498 | .auth-tab-content {
499 | display: none;
500 | overflow-wrap: break-word;
501 | word-wrap: break-word;
502 | max-width: 100%;
503 |
504 | &.active {
505 | display: block;
506 | }
507 |
508 | pre {
509 | white-space: pre-wrap;
510 | word-wrap: break-word;
511 | overflow-x: auto;
512 | background-color: #f5f5f5;
513 | padding: 10px;
514 | border: 1px solid #ccc;
515 | border-radius: 4px;
516 | max-width: 100%;
517 |
518 | code {
519 | display: block;
520 | word-wrap: break-word;
521 | color: inherit;
522 | }
523 | }
524 | }
525 | }
526 | }
527 | }
528 | }
529 |
530 | .auth-method-header {
531 | display: flex;
532 | justify-content: space-between;
533 | align-items: center;
534 | padding: 8px 12px;
535 | }
536 |
537 | .auth-method-label {
538 | font-weight: 500;
539 | color: var(--text-primary-gray);
540 | word-break: break-word;
541 | }
542 |
543 | .auth-helper-section {
544 | border: 1px solid var(--section-border);
545 | background-color: transparent;
546 | padding: 16px;
547 | border-radius: 8px;
548 | margin-top: 20px;
549 | width: 80%;
550 | }
551 |
552 | .auth-method-list {
553 | display: flex;
554 | flex-direction: column;
555 | gap: 12px;
556 | }
557 |
558 | .auth-method-details {
559 | padding: 16px;
560 | border: 1px solid var(--section-border);
561 | border-radius: 4px;
562 | margin-bottom: 16px;
563 | background-color: #fff;
564 | overflow-x: auto;
565 | }
566 |
567 | .auth-controls {
568 | display: flex;
569 | flex-direction: column;
570 | gap: 16px;
571 | align-items: flex-start;
572 | }
573 |
574 | .auth-input-row {
575 | display: flex;
576 | flex-direction: column;
577 | gap: 6px;
578 |
579 | & label {
580 | font-size: 14px;
581 | color: var(--text-primary-gray);
582 | margin-bottom: 4px;
583 | }
584 | }
585 |
586 | .auth-input {
587 | padding: 8px 8px;
588 | border: 1px solid #bdc1c6;
589 | border-radius: 4px;
590 | font-size: 14px;
591 | width: 100%;
592 | box-sizing: border-box;
593 |
594 | &:focus {
595 | outline: none;
596 | border-color: var(--toolbox-blue);
597 | box-shadow: 0 0 0 1px #1a73e8;
598 | }
599 | }
600 |
601 | .auth-method-actions {
602 | display: flex;
603 | align-items: center;
604 | gap: 12px;
605 | }
606 |
607 | .auth-instructions {
608 | font-size: 0.8em;
609 | margin-top: 5px;
610 | color: var(--text-secondary-gray);
611 | }
612 |
613 | .tool-response {
614 | margin: 20px 0 0 0;
615 |
616 | textarea {
617 | width: 100%;
618 | min-height: 150px;
619 | padding: 12px;
620 | border: 1px solid #ddd;
621 | border-radius: 4px;
622 | font-family: monospace;
623 | }
624 | }
625 |
626 | .search-container {
627 | display: flex;
628 | width: 100%;
629 | margin-bottom: 15px;
630 |
631 | #toolset-search-input {
632 | flex-grow: 1;
633 | padding: 10px 12px;
634 | border: 1px solid #ccc;
635 | border-radius: 20px 0 0 20px;
636 | border-right: none;
637 | font-family: inherit;
638 | font-size: 0.9em;
639 | color: var(--text-primary-gray);
640 |
641 | &:focus {
642 | outline: none;
643 | border-color: var(--toolbox-blue);
644 | box-shadow: 0 0 0 2px rgba(66, 133, 244, 0.3);
645 | }
646 |
647 | &::placeholder {
648 | color: var(--text-secondary-gray);
649 | }
650 | }
651 |
652 | #toolset-search-button {
653 | padding: 10px 15px;
654 | border: 1px solid var(--button-primary);
655 | background-color: var(--button-primary);
656 | color: white;
657 | border-radius: 0 20px 20px 0;
658 | cursor: pointer;
659 | font-family: inherit;
660 | font-size: 0.9em;
661 | font-weight: bold;
662 | transition: opacity 0.2s ease-in-out;
663 | flex-shrink: 0;
664 | line-height: 1;
665 |
666 | &:hover {
667 | opacity: 0.8;
668 | }
669 |
670 | &:focus {
671 | outline: none;
672 | box-shadow: 0 0 0 2px rgba(66, 133, 244, 0.3);
673 | }
674 | }
675 | }
676 |
677 | .toggle-details-tab {
678 | background-color: transparent;
679 | color: var(--toolbox-blue);
680 | border: none;
681 | padding: 8px 12px;
682 | border-radius: 4px;
683 | cursor: pointer;
684 | font-size: 1em;
685 | font-weight: bold;
686 |
687 | &:hover {
688 | opacity: 0.8;
689 | }
690 | }
691 |
692 | .resource-instructions {
693 | font-family: inherit;
694 | color: var(--text-primary-gray);
695 | padding: 24px;
696 | background-color: #ffffff;
697 | border-radius: 4px;
698 | box-shadow: 0 1px 3px rgba(0,0,0,0.12), 0 1px 2px rgba(0,0,0,0.24);
699 | margin: 16px 0;
700 |
701 | li {
702 | margin-bottom: 10px;
703 | color: var(--text-secondary-gray);
704 | }
705 | }
706 |
707 | .resource-title {
708 | color: var(--toolbox-blue);
709 | font-size: 24px;
710 | font-weight: bold;
711 | margin-top: 0;
712 | margin-bottom: 16px;
713 | }
714 |
715 | .resource-subtitle {
716 | color: var(--text-primary-gray);
717 | font-size: 20px;
718 | font-weight: normal;
719 | margin-top: 24px;
720 | margin-bottom: 12px;
721 | }
722 |
723 | .resource-intro,
724 | .resource-description {
725 | font-size: 16px;
726 | line-height: 1.6;
727 | color: var(--text-secondary-gray);
728 | margin-bottom: 16px;
729 |
730 | code {
731 | background-color: #e9e9e9;
732 | color: #c0392b;
733 | padding: 2px 4px;
734 | border-radius: 4px;
735 | font-family: monospace;
736 | font-size: 90%;
737 | vertical-align: baseline;
738 | }
739 | }
740 |
```
--------------------------------------------------------------------------------
/internal/tools/neo4j/neo4jschema/helpers/helpers_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 helpers
16 |
17 | import (
18 | "testing"
19 | "time"
20 |
21 | "github.com/google/go-cmp/cmp"
22 | "github.com/googleapis/genai-toolbox/internal/tools/neo4j/neo4jschema/types"
23 | "github.com/neo4j/neo4j-go-driver/v5/neo4j"
24 | )
25 |
26 | func TestHelperFunctions(t *testing.T) {
27 | t.Run("ConvertToStringSlice", func(t *testing.T) {
28 | tests := []struct {
29 | name string
30 | input []any
31 | want []string
32 | }{
33 | {
34 | name: "empty slice",
35 | input: []any{},
36 | want: []string{},
37 | },
38 | {
39 | name: "string values",
40 | input: []any{"a", "b", "c"},
41 | want: []string{"a", "b", "c"},
42 | },
43 | {
44 | name: "mixed types",
45 | input: []any{"string", 123, true, 45.67},
46 | want: []string{"string", "123", "true", "45.67"},
47 | },
48 | }
49 |
50 | for _, tt := range tests {
51 | t.Run(tt.name, func(t *testing.T) {
52 | got := ConvertToStringSlice(tt.input)
53 | if diff := cmp.Diff(tt.want, got); diff != "" {
54 | t.Errorf("ConvertToStringSlice() mismatch (-want +got):\n%s", diff)
55 | }
56 | })
57 | }
58 | })
59 |
60 | t.Run("GetStringValue", func(t *testing.T) {
61 | tests := []struct {
62 | name string
63 | input any
64 | want string
65 | }{
66 | {
67 | name: "nil value",
68 | input: nil,
69 | want: "",
70 | },
71 | {
72 | name: "string value",
73 | input: "test",
74 | want: "test",
75 | },
76 | {
77 | name: "int value",
78 | input: 42,
79 | want: "42",
80 | },
81 | {
82 | name: "bool value",
83 | input: true,
84 | want: "true",
85 | },
86 | }
87 |
88 | for _, tt := range tests {
89 | t.Run(tt.name, func(t *testing.T) {
90 | got := GetStringValue(tt.input)
91 | if got != tt.want {
92 | t.Errorf("GetStringValue() got %q, want %q", got, tt.want)
93 | }
94 | })
95 | }
96 | })
97 | }
98 |
99 | func TestMapToAPOCSchema(t *testing.T) {
100 | tests := []struct {
101 | name string
102 | input map[string]any
103 | want *types.APOCSchemaResult
104 | wantErr bool
105 | }{
106 | {
107 | name: "simple node schema",
108 | input: map[string]any{
109 | "Person": map[string]any{
110 | "type": "node",
111 | "count": int64(150),
112 | "properties": map[string]any{
113 | "name": map[string]any{
114 | "type": "STRING",
115 | "unique": false,
116 | "indexed": true,
117 | "existence": false,
118 | },
119 | },
120 | },
121 | },
122 | want: &types.APOCSchemaResult{
123 | Value: map[string]types.APOCEntity{
124 | "Person": {
125 | Type: "node",
126 | Count: 150,
127 | Properties: map[string]types.APOCProperty{
128 | "name": {
129 | Type: "STRING",
130 | Unique: false,
131 | Indexed: true,
132 | Existence: false,
133 | },
134 | },
135 | },
136 | },
137 | },
138 | wantErr: false,
139 | },
140 | {
141 | name: "empty input",
142 | input: map[string]any{},
143 | want: &types.APOCSchemaResult{Value: map[string]types.APOCEntity{}},
144 | wantErr: false,
145 | },
146 | }
147 |
148 | for _, tt := range tests {
149 | t.Run(tt.name, func(t *testing.T) {
150 | got, err := MapToAPOCSchema(tt.input)
151 | if (err != nil) != tt.wantErr {
152 | t.Errorf("MapToAPOCSchema() error = %v, wantErr %v", err, tt.wantErr)
153 | return
154 | }
155 | if diff := cmp.Diff(tt.want, got); diff != "" {
156 | t.Errorf("MapToAPOCSchema() mismatch (-want +got):\n%s", diff)
157 | }
158 | })
159 | }
160 | }
161 |
162 | func TestProcessAPOCSchema(t *testing.T) {
163 | tests := []struct {
164 | name string
165 | input *types.APOCSchemaResult
166 | wantNodes []types.NodeLabel
167 | wantRels []types.Relationship
168 | wantStats *types.Statistics
169 | statsAreEmpty bool
170 | }{
171 | {
172 | name: "empty schema",
173 | input: &types.APOCSchemaResult{
174 | Value: map[string]types.APOCEntity{},
175 | },
176 | wantNodes: nil,
177 | wantRels: nil,
178 | statsAreEmpty: true,
179 | },
180 | {
181 | name: "simple node only",
182 | input: &types.APOCSchemaResult{
183 | Value: map[string]types.APOCEntity{
184 | "Person": {
185 | Type: "node",
186 | Count: 100,
187 | Properties: map[string]types.APOCProperty{
188 | "name": {Type: "STRING", Indexed: true},
189 | "age": {Type: "INTEGER"},
190 | },
191 | },
192 | },
193 | },
194 | wantNodes: []types.NodeLabel{
195 | {
196 | Name: "Person",
197 | Count: 100,
198 | Properties: []types.PropertyInfo{
199 | {Name: "age", Types: []string{"INTEGER"}},
200 | {Name: "name", Types: []string{"STRING"}, Indexed: true},
201 | },
202 | },
203 | },
204 | wantRels: nil,
205 | wantStats: &types.Statistics{
206 | NodesByLabel: map[string]int64{"Person": 100},
207 | PropertiesByLabel: map[string]int64{"Person": 2},
208 | TotalNodes: 100,
209 | TotalProperties: 200,
210 | },
211 | },
212 | {
213 | name: "nodes and relationships",
214 | input: &types.APOCSchemaResult{
215 | Value: map[string]types.APOCEntity{
216 | "Person": {
217 | Type: "node",
218 | Count: 100,
219 | Properties: map[string]types.APOCProperty{
220 | "name": {Type: "STRING", Unique: true, Indexed: true, Existence: true},
221 | },
222 | Relationships: map[string]types.APOCRelationshipInfo{
223 | "KNOWS": {
224 | Direction: "out",
225 | Count: 50,
226 | Labels: []string{"Person"},
227 | Properties: map[string]types.APOCProperty{
228 | "since": {Type: "INTEGER"},
229 | },
230 | },
231 | },
232 | },
233 | "Post": {
234 | Type: "node",
235 | Count: 200,
236 | Properties: map[string]types.APOCProperty{"content": {Type: "STRING"}},
237 | },
238 | "FOLLOWS": {Type: "relationship", Count: 80},
239 | },
240 | },
241 | wantNodes: []types.NodeLabel{
242 | {
243 | Name: "Post",
244 | Count: 200,
245 | Properties: []types.PropertyInfo{
246 | {Name: "content", Types: []string{"STRING"}},
247 | },
248 | },
249 | {
250 | Name: "Person",
251 | Count: 100,
252 | Properties: []types.PropertyInfo{
253 | {Name: "name", Types: []string{"STRING"}, Unique: true, Indexed: true, Mandatory: true},
254 | },
255 | },
256 | },
257 | wantRels: []types.Relationship{
258 | {
259 | Type: "KNOWS",
260 | StartNode: "Person",
261 | EndNode: "Person",
262 | Count: 50,
263 | Properties: []types.PropertyInfo{
264 | {Name: "since", Types: []string{"INTEGER"}},
265 | },
266 | },
267 | },
268 | wantStats: &types.Statistics{
269 | NodesByLabel: map[string]int64{"Person": 100, "Post": 200},
270 | RelationshipsByType: map[string]int64{"KNOWS": 50},
271 | PropertiesByLabel: map[string]int64{"Person": 1, "Post": 1},
272 | PropertiesByRelType: map[string]int64{"KNOWS": 1},
273 | TotalNodes: 300,
274 | TotalRelationships: 50,
275 | TotalProperties: 350, // (100*1 + 200*1) for nodes + (50*1) for rels
276 | },
277 | },
278 | }
279 |
280 | for _, tt := range tests {
281 | t.Run(tt.name, func(t *testing.T) {
282 | gotNodes, gotRels, gotStats := ProcessAPOCSchema(tt.input)
283 |
284 | if diff := cmp.Diff(tt.wantNodes, gotNodes); diff != "" {
285 | t.Errorf("ProcessAPOCSchema() node labels mismatch (-want +got):\n%s", diff)
286 | }
287 | if diff := cmp.Diff(tt.wantRels, gotRels); diff != "" {
288 | t.Errorf("ProcessAPOCSchema() relationships mismatch (-want +got):\n%s", diff)
289 | }
290 | if tt.statsAreEmpty {
291 | tt.wantStats = &types.Statistics{}
292 | }
293 |
294 | if diff := cmp.Diff(tt.wantStats, gotStats); diff != "" {
295 | t.Errorf("ProcessAPOCSchema() statistics mismatch (-want +got):\n%s", diff)
296 | }
297 | })
298 | }
299 | }
300 |
301 | func TestProcessNonAPOCSchema(t *testing.T) {
302 | t.Run("full schema processing", func(t *testing.T) {
303 | nodeCounts := map[string]int64{"Person": 10, "City": 5}
304 | nodePropsMap := map[string]map[string]map[string]bool{
305 | "Person": {"name": {"STRING": true}, "age": {"INTEGER": true}},
306 | "City": {"name": {"STRING": true, "TEXT": true}},
307 | }
308 | relCounts := map[string]int64{"LIVES_IN": 8}
309 | relPropsMap := map[string]map[string]map[string]bool{
310 | "LIVES_IN": {"since": {"DATE": true}},
311 | }
312 | relConnectivity := map[string]types.RelConnectivityInfo{
313 | "LIVES_IN": {StartNode: "Person", EndNode: "City", Count: 8},
314 | }
315 |
316 | wantNodes := []types.NodeLabel{
317 | {
318 | Name: "Person",
319 | Count: 10,
320 | Properties: []types.PropertyInfo{
321 | {Name: "age", Types: []string{"INTEGER"}},
322 | {Name: "name", Types: []string{"STRING"}},
323 | },
324 | },
325 | {
326 | Name: "City",
327 | Count: 5,
328 | Properties: []types.PropertyInfo{
329 | {Name: "name", Types: []string{"STRING", "TEXT"}},
330 | },
331 | },
332 | }
333 | wantRels := []types.Relationship{
334 | {
335 | Type: "LIVES_IN",
336 | Count: 8,
337 | StartNode: "Person",
338 | EndNode: "City",
339 | Properties: []types.PropertyInfo{
340 | {Name: "since", Types: []string{"DATE"}},
341 | },
342 | },
343 | }
344 | wantStats := &types.Statistics{
345 | TotalNodes: 15,
346 | TotalRelationships: 8,
347 | TotalProperties: 33, // (10*2 + 5*1) for nodes + (8*1) for rels
348 | NodesByLabel: map[string]int64{"Person": 10, "City": 5},
349 | RelationshipsByType: map[string]int64{"LIVES_IN": 8},
350 | PropertiesByLabel: map[string]int64{"Person": 2, "City": 1},
351 | PropertiesByRelType: map[string]int64{"LIVES_IN": 1},
352 | }
353 |
354 | gotNodes, gotRels, gotStats := ProcessNonAPOCSchema(nodeCounts, nodePropsMap, relCounts, relPropsMap, relConnectivity)
355 |
356 | if diff := cmp.Diff(wantNodes, gotNodes); diff != "" {
357 | t.Errorf("ProcessNonAPOCSchema() nodes mismatch (-want +got):\n%s", diff)
358 | }
359 | if diff := cmp.Diff(wantRels, gotRels); diff != "" {
360 | t.Errorf("ProcessNonAPOCSchema() relationships mismatch (-want +got):\n%s", diff)
361 | }
362 | if diff := cmp.Diff(wantStats, gotStats); diff != "" {
363 | t.Errorf("ProcessNonAPOCSchema() stats mismatch (-want +got):\n%s", diff)
364 | }
365 | })
366 |
367 | t.Run("empty schema", func(t *testing.T) {
368 | gotNodes, gotRels, gotStats := ProcessNonAPOCSchema(
369 | map[string]int64{},
370 | map[string]map[string]map[string]bool{},
371 | map[string]int64{},
372 | map[string]map[string]map[string]bool{},
373 | map[string]types.RelConnectivityInfo{},
374 | )
375 |
376 | if len(gotNodes) != 0 {
377 | t.Errorf("expected 0 nodes, got %d", len(gotNodes))
378 | }
379 | if len(gotRels) != 0 {
380 | t.Errorf("expected 0 relationships, got %d", len(gotRels))
381 | }
382 | if diff := cmp.Diff(&types.Statistics{}, gotStats); diff != "" {
383 | t.Errorf("ProcessNonAPOCSchema() stats mismatch (-want +got):\n%s", diff)
384 | }
385 | })
386 | }
387 |
388 | func TestConvertValue(t *testing.T) {
389 | tests := []struct {
390 | name string
391 | input any
392 | want any
393 | }{
394 | {
395 | name: "nil value",
396 | input: nil,
397 | want: nil,
398 | },
399 | {
400 | name: "neo4j.InvalidValue",
401 | input: neo4j.InvalidValue{},
402 | want: nil,
403 | },
404 | {
405 | name: "primitive bool",
406 | input: true,
407 | want: true,
408 | },
409 | {
410 | name: "primitive int",
411 | input: int64(42),
412 | want: int64(42),
413 | },
414 | {
415 | name: "primitive float",
416 | input: 3.14,
417 | want: 3.14,
418 | },
419 | {
420 | name: "primitive string",
421 | input: "hello",
422 | want: "hello",
423 | },
424 | {
425 | name: "neo4j.Date",
426 | input: neo4j.Date(time.Date(2024, 6, 1, 0, 0, 0, 0, time.UTC)),
427 | want: "2024-06-01",
428 | },
429 | {
430 | name: "neo4j.LocalTime",
431 | input: neo4j.LocalTime(time.Date(0, 0, 0, 12, 34, 56, 0, time.Local)),
432 | want: "12:34:56",
433 | },
434 | {
435 | name: "neo4j.Time",
436 | input: neo4j.Time(time.Date(0, 0, 0, 1, 2, 3, 0, time.UTC)),
437 | want: "01:02:03Z",
438 | },
439 | {
440 | name: "neo4j.LocalDateTime",
441 | input: neo4j.LocalDateTime(time.Date(2024, 6, 1, 10, 20, 30, 0, time.Local)),
442 | want: "2024-06-01T10:20:30",
443 | },
444 | {
445 | name: "neo4j.Duration",
446 | input: neo4j.Duration{Months: 1, Days: 2, Seconds: 3, Nanos: 4},
447 | want: "P1M2DT3.000000004S",
448 | },
449 | {
450 | name: "neo4j.Point2D",
451 | input: neo4j.Point2D{X: 1.1, Y: 2.2, SpatialRefId: 1234},
452 | want: map[string]any{"x": 1.1, "y": 2.2, "srid": uint32(1234)},
453 | },
454 | {
455 | name: "neo4j.Point3D",
456 | input: neo4j.Point3D{X: 1.1, Y: 2.2, Z: 3.3, SpatialRefId: 5467},
457 | want: map[string]any{"x": 1.1, "y": 2.2, "z": 3.3, "srid": uint32(5467)},
458 | },
459 | {
460 | name: "neo4j.Node (handled by Entity case, losing labels)",
461 | input: neo4j.Node{
462 | ElementId: "element-1",
463 | Labels: []string{"Person"},
464 | Props: map[string]any{"name": "Alice"},
465 | },
466 | want: map[string]any{
467 | "elementId": "element-1",
468 | "labels": []string{"Person"},
469 | "properties": map[string]any{"name": "Alice"},
470 | },
471 | },
472 | {
473 | name: "neo4j.Relationship (handled by Entity case, losing type/endpoints)",
474 | input: neo4j.Relationship{
475 | ElementId: "element-2",
476 | StartElementId: "start-1",
477 | EndElementId: "end-1",
478 | Type: "KNOWS",
479 | Props: map[string]any{"since": 2024},
480 | },
481 | want: map[string]any{
482 | "elementId": "element-2",
483 | "properties": map[string]any{"since": 2024},
484 | "startElementId": "start-1",
485 | "endElementId": "end-1",
486 | "type": "KNOWS",
487 | },
488 | },
489 | {
490 | name: "neo4j.Path (elements handled by Entity case)",
491 | input: func() neo4j.Path {
492 | node1 := neo4j.Node{ElementId: "n10", Labels: []string{"A"}, Props: map[string]any{"p1": "v1"}}
493 | node2 := neo4j.Node{ElementId: "n11", Labels: []string{"B"}, Props: map[string]any{"p2": "v2"}}
494 | rel1 := neo4j.Relationship{ElementId: "r12", StartElementId: "n10", EndElementId: "n11", Type: "REL", Props: map[string]any{"p3": "v3"}}
495 | return neo4j.Path{
496 | Nodes: []neo4j.Node{node1, node2},
497 | Relationships: []neo4j.Relationship{rel1},
498 | }
499 | }(),
500 | want: map[string]any{
501 | "nodes": []any{
502 | map[string]any{
503 | "elementId": "n10",
504 | "properties": map[string]any{"p1": "v1"},
505 | "labels": []string{"A"},
506 | },
507 | map[string]any{
508 | "elementId": "n11",
509 | "properties": map[string]any{"p2": "v2"},
510 | "labels": []string{"B"},
511 | },
512 | },
513 | "relationships": []any{
514 | map[string]any{
515 | "elementId": "r12",
516 | "properties": map[string]any{"p3": "v3"},
517 | "startElementId": "n10",
518 | "endElementId": "n11",
519 | "type": "REL",
520 | },
521 | },
522 | },
523 | },
524 | {
525 | name: "slice of primitives",
526 | input: []any{"a", 1, true},
527 | want: []any{"a", 1, true},
528 | },
529 | {
530 | name: "slice of mixed types",
531 | input: []any{"a", neo4j.Date(time.Date(2024, 6, 1, 0, 0, 0, 0, time.UTC))},
532 | want: []any{"a", "2024-06-01"},
533 | },
534 | {
535 | name: "map of primitives",
536 | input: map[string]any{"foo": 1, "bar": "baz"},
537 | want: map[string]any{"foo": 1, "bar": "baz"},
538 | },
539 | {
540 | name: "map with nested neo4j type",
541 | input: map[string]any{"date": neo4j.Date(time.Date(2024, 6, 1, 0, 0, 0, 0, time.UTC))},
542 | want: map[string]any{"date": "2024-06-01"},
543 | },
544 | {
545 | name: "unhandled type",
546 | input: struct{ X int }{X: 5},
547 | want: "{5}",
548 | },
549 | }
550 |
551 | for _, tt := range tests {
552 | t.Run(tt.name, func(t *testing.T) {
553 | got := ConvertValue(tt.input)
554 | if !cmp.Equal(got, tt.want) {
555 | t.Errorf("ConvertValue() mismatch (-want +got):\n%s", cmp.Diff(tt.want, got))
556 | }
557 | })
558 | }
559 | }
560 |
```
--------------------------------------------------------------------------------
/internal/tools/mysql/mysqllisttables/mysqllisttables.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 mysqllisttables
16 |
17 | import (
18 | "context"
19 | "database/sql"
20 | "fmt"
21 |
22 | yaml "github.com/goccy/go-yaml"
23 | "github.com/googleapis/genai-toolbox/internal/sources"
24 | "github.com/googleapis/genai-toolbox/internal/sources/cloudsqlmysql"
25 | "github.com/googleapis/genai-toolbox/internal/sources/mysql"
26 | "github.com/googleapis/genai-toolbox/internal/tools"
27 | "github.com/googleapis/genai-toolbox/internal/tools/mysql/mysqlcommon"
28 | )
29 |
30 | const kind string = "mysql-list-tables"
31 |
32 | const listTablesStatement = `
33 | SELECT
34 | T.TABLE_SCHEMA AS schema_name,
35 | T.TABLE_NAME AS object_name,
36 | CASE
37 | WHEN @output_format = 'simple' THEN
38 | JSON_OBJECT('name', T.TABLE_NAME)
39 | ELSE
40 | CONVERT(
41 | JSON_OBJECT(
42 | 'schema_name', T.TABLE_SCHEMA,
43 | 'object_name', T.TABLE_NAME,
44 | 'object_type', 'TABLE',
45 | 'owner', (
46 | SELECT
47 | IFNULL(U.GRANTEE, 'N/A')
48 | FROM
49 | INFORMATION_SCHEMA.SCHEMA_PRIVILEGES U
50 | WHERE
51 | U.TABLE_SCHEMA = T.TABLE_SCHEMA
52 | LIMIT 1
53 | ),
54 | 'comment', IFNULL(T.TABLE_COMMENT, ''),
55 | 'columns', (
56 | SELECT
57 | IFNULL(
58 | JSON_ARRAYAGG(
59 | JSON_OBJECT(
60 | 'column_name', C.COLUMN_NAME,
61 | 'data_type', C.COLUMN_TYPE,
62 | 'ordinal_position', C.ORDINAL_POSITION,
63 | 'is_not_nullable', IF(C.IS_NULLABLE = 'NO', TRUE, FALSE),
64 | 'column_default', C.COLUMN_DEFAULT,
65 | 'column_comment', IFNULL(C.COLUMN_COMMENT, '')
66 | )
67 | ),
68 | JSON_ARRAY()
69 | )
70 | FROM
71 | INFORMATION_SCHEMA.COLUMNS C
72 | WHERE
73 | C.TABLE_SCHEMA = T.TABLE_SCHEMA AND C.TABLE_NAME = T.TABLE_NAME
74 | ORDER BY C.ORDINAL_POSITION
75 | ),
76 | 'constraints', (
77 | SELECT
78 | IFNULL(
79 | JSON_ARRAYAGG(
80 | JSON_OBJECT(
81 | 'constraint_name', TC.CONSTRAINT_NAME,
82 | 'constraint_type',
83 | CASE TC.CONSTRAINT_TYPE
84 | WHEN 'PRIMARY KEY' THEN 'PRIMARY KEY'
85 | WHEN 'FOREIGN KEY' THEN 'FOREIGN KEY'
86 | WHEN 'UNIQUE' THEN 'UNIQUE'
87 | ELSE TC.CONSTRAINT_TYPE
88 | END,
89 | 'constraint_definition', '',
90 | 'constraint_columns', (
91 | SELECT
92 | IFNULL(JSON_ARRAYAGG(KCU.COLUMN_NAME), JSON_ARRAY())
93 | FROM
94 | INFORMATION_SCHEMA.KEY_COLUMN_USAGE KCU
95 | WHERE
96 | KCU.CONSTRAINT_SCHEMA = TC.CONSTRAINT_SCHEMA
97 | AND KCU.CONSTRAINT_NAME = TC.CONSTRAINT_NAME
98 | AND KCU.TABLE_NAME = TC.TABLE_NAME
99 | ORDER BY KCU.ORDINAL_POSITION
100 | ),
101 | 'foreign_key_referenced_table', IF(TC.CONSTRAINT_TYPE = 'FOREIGN KEY', RC.REFERENCED_TABLE_NAME, NULL),
102 | 'foreign_key_referenced_columns', IF(TC.CONSTRAINT_TYPE = 'FOREIGN KEY',
103 | (SELECT IFNULL(JSON_ARRAYAGG(FKCU.REFERENCED_COLUMN_NAME), JSON_ARRAY())
104 | FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE FKCU
105 | WHERE FKCU.CONSTRAINT_SCHEMA = TC.CONSTRAINT_SCHEMA
106 | AND FKCU.CONSTRAINT_NAME = TC.CONSTRAINT_NAME
107 | AND FKCU.TABLE_NAME = TC.TABLE_NAME
108 | AND FKCU.REFERENCED_TABLE_NAME IS NOT NULL
109 | ORDER BY FKCU.ORDINAL_POSITION),
110 | NULL
111 | )
112 | )
113 | ),
114 | JSON_ARRAY()
115 | )
116 | FROM
117 | INFORMATION_SCHEMA.TABLE_CONSTRAINTS TC
118 | LEFT JOIN
119 | INFORMATION_SCHEMA.REFERENTIAL_CONSTRAINTS RC
120 | ON TC.CONSTRAINT_SCHEMA = RC.CONSTRAINT_SCHEMA
121 | AND TC.CONSTRAINT_NAME = RC.CONSTRAINT_NAME
122 | AND TC.TABLE_NAME = RC.TABLE_NAME
123 | WHERE
124 | TC.TABLE_SCHEMA = T.TABLE_SCHEMA AND TC.TABLE_NAME = T.TABLE_NAME
125 | ),
126 | 'indexes', (
127 | SELECT
128 | IFNULL(
129 | JSON_ARRAYAGG(
130 | JSON_OBJECT(
131 | 'index_name', IndexData.INDEX_NAME,
132 | 'is_unique', IF(IndexData.NON_UNIQUE = 0, TRUE, FALSE),
133 | 'is_primary', IF(IndexData.INDEX_NAME = 'PRIMARY', TRUE, FALSE),
134 | 'index_columns', IFNULL(IndexData.INDEX_COLUMNS_ARRAY, JSON_ARRAY())
135 | )
136 | ),
137 | JSON_ARRAY()
138 | )
139 | FROM (
140 | SELECT
141 | S.TABLE_SCHEMA,
142 | S.TABLE_NAME,
143 | S.INDEX_NAME,
144 | MIN(S.NON_UNIQUE) AS NON_UNIQUE,
145 | JSON_ARRAYAGG(S.COLUMN_NAME) AS INDEX_COLUMNS_ARRAY
146 | FROM
147 | INFORMATION_SCHEMA.STATISTICS S
148 | GROUP BY
149 | S.TABLE_SCHEMA, S.TABLE_NAME, S.INDEX_NAME
150 | ) AS IndexData
151 | WHERE IndexData.TABLE_SCHEMA = T.TABLE_SCHEMA AND IndexData.TABLE_NAME = T.TABLE_NAME
152 | ORDER BY IndexData.INDEX_NAME
153 | ),
154 | 'triggers', (
155 | SELECT
156 | IFNULL(
157 | JSON_ARRAYAGG(
158 | JSON_OBJECT(
159 | 'trigger_name', TR.TRIGGER_NAME,
160 | 'trigger_definition', TR.ACTION_STATEMENT
161 | )
162 | ),
163 | JSON_ARRAY()
164 | )
165 | FROM
166 | INFORMATION_SCHEMA.TRIGGERS TR
167 | WHERE
168 | TR.EVENT_OBJECT_SCHEMA = T.TABLE_SCHEMA AND TR.EVENT_OBJECT_TABLE = T.TABLE_NAME
169 | ORDER BY TR.TRIGGER_NAME
170 | )
171 | )
172 | USING utf8mb4)
173 | END AS object_details
174 | FROM
175 | INFORMATION_SCHEMA.TABLES T
176 | CROSS JOIN (SELECT @table_names := ?, @output_format := ?) AS variables
177 | WHERE
178 | T.TABLE_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')
179 | AND (NULLIF(TRIM(@table_names), '') IS NULL OR FIND_IN_SET(T.TABLE_NAME, @table_names))
180 | AND T.TABLE_TYPE = 'BASE TABLE'
181 | ORDER BY
182 | T.TABLE_SCHEMA, T.TABLE_NAME;
183 | `
184 |
185 | func init() {
186 | if !tools.Register(kind, newConfig) {
187 | panic(fmt.Sprintf("tool kind %q already registered", kind))
188 | }
189 | }
190 |
191 | func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
192 | actual := Config{Name: name}
193 | if err := decoder.DecodeContext(ctx, &actual); err != nil {
194 | return nil, err
195 | }
196 | return actual, nil
197 | }
198 |
199 | type compatibleSource interface {
200 | MySQLPool() *sql.DB
201 | }
202 |
203 | // validate compatible sources are still compatible
204 | var _ compatibleSource = &cloudsqlmysql.Source{}
205 | var _ compatibleSource = &mysql.Source{}
206 |
207 | var compatibleSources = [...]string{cloudsqlmysql.SourceKind, mysql.SourceKind}
208 |
209 | type Config struct {
210 | Name string `yaml:"name" validate:"required"`
211 | Kind string `yaml:"kind" validate:"required"`
212 | Source string `yaml:"source" validate:"required"`
213 | Description string `yaml:"description" validate:"required"`
214 | AuthRequired []string `yaml:"authRequired"`
215 | }
216 |
217 | // validate interface
218 | var _ tools.ToolConfig = Config{}
219 |
220 | func (cfg Config) ToolConfigKind() string {
221 | return kind
222 | }
223 |
224 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
225 | // verify source exists
226 | rawS, ok := srcs[cfg.Source]
227 | if !ok {
228 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
229 | }
230 |
231 | // verify the source is compatible
232 | s, ok := rawS.(compatibleSource)
233 | if !ok {
234 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources)
235 | }
236 |
237 | allParameters := tools.Parameters{
238 | tools.NewStringParameterWithDefault("table_names", "", "Optional: A comma-separated list of table names. If empty, details for all tables will be listed."),
239 | tools.NewStringParameterWithDefault("output_format", "detailed", "Optional: Use 'simple' for names only or 'detailed' for full info."),
240 | }
241 | paramManifest := allParameters.Manifest()
242 | mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, allParameters)
243 |
244 | // finish tool setup
245 | t := Tool{
246 | Name: cfg.Name,
247 | Kind: kind,
248 | AllParams: allParameters,
249 | AuthRequired: cfg.AuthRequired,
250 | Pool: s.MySQLPool(),
251 | manifest: tools.Manifest{Description: cfg.Description, Parameters: paramManifest, AuthRequired: cfg.AuthRequired},
252 | mcpManifest: mcpManifest,
253 | }
254 | return t, nil
255 | }
256 |
257 | // validate interface
258 | var _ tools.Tool = Tool{}
259 |
260 | type Tool struct {
261 | Name string `yaml:"name"`
262 | Kind string `yaml:"kind"`
263 | AuthRequired []string `yaml:"authRequired"`
264 | AllParams tools.Parameters `yaml:"allParams"`
265 |
266 | Pool *sql.DB
267 | manifest tools.Manifest
268 | mcpManifest tools.McpManifest
269 | }
270 |
271 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
272 | paramsMap := params.AsMap()
273 |
274 | tableNames, ok := paramsMap["table_names"].(string)
275 | if !ok {
276 | return nil, fmt.Errorf("invalid '%s' parameter; expected a string", tableNames)
277 | }
278 | outputFormat, _ := paramsMap["output_format"].(string)
279 | if outputFormat != "simple" && outputFormat != "detailed" {
280 | return nil, fmt.Errorf("invalid value for output_format: must be 'simple' or 'detailed', but got %q", outputFormat)
281 | }
282 |
283 | results, err := t.Pool.QueryContext(ctx, listTablesStatement, tableNames, outputFormat)
284 | if err != nil {
285 | return nil, fmt.Errorf("unable to execute query: %w", err)
286 | }
287 |
288 | cols, err := results.Columns()
289 | if err != nil {
290 | return nil, fmt.Errorf("unable to retrieve rows column name: %w", err)
291 | }
292 |
293 | // create an array of values for each column, which can be re-used to scan each row
294 | rawValues := make([]any, len(cols))
295 | values := make([]any, len(cols))
296 | for i := range rawValues {
297 | values[i] = &rawValues[i]
298 | }
299 | defer results.Close()
300 |
301 | colTypes, err := results.ColumnTypes()
302 | if err != nil {
303 | return nil, fmt.Errorf("unable to get column types: %w", err)
304 | }
305 |
306 | var out []any
307 | for results.Next() {
308 | err := results.Scan(values...)
309 | if err != nil {
310 | return nil, fmt.Errorf("unable to parse row: %w", err)
311 | }
312 | vMap := make(map[string]any)
313 | for i, name := range cols {
314 | val := rawValues[i]
315 | if val == nil {
316 | vMap[name] = nil
317 | continue
318 | }
319 |
320 | vMap[name], err = mysqlcommon.ConvertToType(colTypes[i], val)
321 | if err != nil {
322 | return nil, fmt.Errorf("errors encountered when converting values: %w", err)
323 | }
324 | }
325 | out = append(out, vMap)
326 | }
327 |
328 | if err := results.Err(); err != nil {
329 | return nil, fmt.Errorf("errors encountered during row iteration: %w", err)
330 | }
331 |
332 | return out, nil
333 | }
334 |
335 | func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
336 | return tools.ParseParams(t.AllParams, data, claims)
337 | }
338 |
339 | func (t Tool) Manifest() tools.Manifest {
340 | return t.manifest
341 | }
342 |
343 | func (t Tool) McpManifest() tools.McpManifest {
344 | return t.mcpManifest
345 | }
346 |
347 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
348 | return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
349 | }
350 |
351 | func (t Tool) RequiresClientAuthorization() bool {
352 | return false
353 | }
354 |
```
--------------------------------------------------------------------------------
/internal/tools/looker/lookerhealthpulse/lookerhealthpulse.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 lookerhealthpulse
15 |
16 | import (
17 | "context"
18 | "encoding/json"
19 | "fmt"
20 | "strings"
21 |
22 | yaml "github.com/goccy/go-yaml"
23 | "github.com/googleapis/genai-toolbox/internal/sources"
24 | lookersrc "github.com/googleapis/genai-toolbox/internal/sources/looker"
25 | "github.com/googleapis/genai-toolbox/internal/tools"
26 | "github.com/googleapis/genai-toolbox/internal/tools/looker/lookercommon"
27 | "github.com/googleapis/genai-toolbox/internal/util"
28 |
29 | "github.com/looker-open-source/sdk-codegen/go/rtl"
30 | v4 "github.com/looker-open-source/sdk-codegen/go/sdk/v4"
31 | )
32 |
33 | // =================================================================================================================
34 | // START MCP SERVER CORE LOGIC
35 | // =================================================================================================================
36 | const kind string = "looker-health-pulse"
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 | Parameters map[string]any `yaml:"parameters"`
59 | }
60 |
61 | // validate interface
62 | var _ tools.ToolConfig = Config{}
63 |
64 | func (cfg Config) ToolConfigKind() string {
65 | return kind
66 | }
67 |
68 | func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) {
69 | // verify source exists
70 | rawS, ok := srcs[cfg.Source]
71 | if !ok {
72 | return nil, fmt.Errorf("no source named %q configured", cfg.Source)
73 | }
74 |
75 | // verify the source is compatible
76 | s, ok := rawS.(*lookersrc.Source)
77 | if !ok {
78 | return nil, fmt.Errorf("invalid source for %q tool: source kind must be `looker`", kind)
79 | }
80 |
81 | actionParameter := tools.NewStringParameterWithRequired("action", "The health check to run. Can be either: `check_db_connections`, `check_dashboard_performance`,`check_dashboard_errors`,`check_explore_performance`,`check_schedule_failures`, or `check_legacy_features`", true)
82 |
83 | parameters := tools.Parameters{
84 | actionParameter,
85 | }
86 |
87 | mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters)
88 |
89 | // finish tool setup
90 | return Tool{
91 | Name: cfg.Name,
92 | Kind: kind,
93 | Parameters: parameters,
94 | AuthRequired: cfg.AuthRequired,
95 | UseClientOAuth: s.UseClientOAuth,
96 | Client: s.Client,
97 | ApiSettings: s.ApiSettings,
98 | manifest: tools.Manifest{
99 | Description: cfg.Description,
100 | Parameters: parameters.Manifest(),
101 | AuthRequired: cfg.AuthRequired,
102 | },
103 | mcpManifest: mcpManifest,
104 | }, nil
105 | }
106 |
107 | // validate interface
108 | var _ tools.Tool = Tool{}
109 |
110 | type Tool struct {
111 | Name string `yaml:"name"`
112 | Kind string `yaml:"kind"`
113 | UseClientOAuth bool
114 | Client *v4.LookerSDK
115 | ApiSettings *rtl.ApiSettings
116 | AuthRequired []string `yaml:"authRequired"`
117 | Parameters tools.Parameters `yaml:"parameters"`
118 | manifest tools.Manifest
119 | mcpManifest tools.McpManifest
120 | }
121 |
122 | func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken tools.AccessToken) (any, error) {
123 | logger, err := util.LoggerFromContext(ctx)
124 | if err != nil {
125 | return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
126 | }
127 |
128 | sdk, err := lookercommon.GetLookerSDK(t.UseClientOAuth, t.ApiSettings, t.Client, accessToken)
129 | if err != nil {
130 | return nil, fmt.Errorf("error getting sdk: %w", err)
131 | }
132 |
133 | pulseTool := &pulseTool{
134 | ApiSettings: t.ApiSettings,
135 | SdkClient: sdk,
136 | }
137 |
138 | paramsMap := params.AsMap()
139 | action, ok := paramsMap["action"].(string)
140 | if !ok {
141 | return nil, fmt.Errorf("action parameter not found")
142 | }
143 |
144 | pulseParams := PulseParams{
145 | Action: action,
146 | }
147 |
148 | result, err := pulseTool.RunPulse(ctx, pulseParams)
149 | if err != nil {
150 | return nil, fmt.Errorf("error running pulse: %w", err)
151 | }
152 |
153 | logger.DebugContext(ctx, "result = ", result)
154 |
155 | return result, nil
156 | }
157 |
158 | func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (tools.ParamValues, error) {
159 | return tools.ParseParams(t.Parameters, data, claims)
160 | }
161 |
162 | func (t Tool) Manifest() tools.Manifest {
163 | return t.manifest
164 | }
165 |
166 | func (t Tool) McpManifest() tools.McpManifest {
167 | return t.mcpManifest
168 | }
169 |
170 | func (t Tool) Authorized(verifiedAuthServices []string) bool {
171 | return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices)
172 | }
173 |
174 | func (t Tool) RequiresClientAuthorization() bool {
175 | return t.UseClientOAuth
176 | }
177 |
178 | // =================================================================================================================
179 | // END MCP SERVER CORE LOGIC
180 | // =================================================================================================================
181 |
182 | // =================================================================================================================
183 | // START LOOKER HEALTH PULSE CORE LOGIC
184 | // =================================================================================================================
185 | type PulseParams struct {
186 | Action string
187 | // Optionally add more parameters if needed
188 | }
189 |
190 | // pulseTool holds Looker API settings and client
191 | type pulseTool struct {
192 | ApiSettings *rtl.ApiSettings
193 | SdkClient *v4.LookerSDK
194 | }
195 |
196 | func (t *pulseTool) RunPulse(ctx context.Context, params PulseParams) (interface{}, error) {
197 | switch params.Action {
198 | case "check_db_connections":
199 | return t.checkDBConnections(ctx)
200 | case "check_dashboard_performance":
201 | return t.checkDashboardPerformance(ctx)
202 | case "check_dashboard_errors":
203 | return t.checkDashboardErrors(ctx)
204 | case "check_explore_performance":
205 | return t.checkExplorePerformance(ctx)
206 | case "check_schedule_failures":
207 | return t.checkScheduleFailures(ctx)
208 | case "check_legacy_features":
209 | return t.checkLegacyFeatures(ctx)
210 | default:
211 | return nil, fmt.Errorf("unknown action: %s", params.Action)
212 | }
213 | }
214 |
215 | // Check DB connections and run tests
216 | func (t *pulseTool) checkDBConnections(ctx context.Context) (interface{}, error) {
217 | logger, err := util.LoggerFromContext(ctx)
218 | if err != nil {
219 | return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
220 | }
221 | logger.InfoContext(ctx, "Test 1/6: Checking connections")
222 |
223 | reservedNames := map[string]struct{}{
224 | "looker__internal__analytics__replica": {},
225 | "looker__internal__analytics": {},
226 | "looker": {},
227 | "looker__ilooker": {},
228 | }
229 |
230 | connections, err := t.SdkClient.AllConnections("", t.ApiSettings)
231 | if err != nil {
232 | return nil, fmt.Errorf("error fetching connections: %w", err)
233 | }
234 |
235 | var filteredConnections []v4.DBConnection
236 | for _, c := range connections {
237 | if _, reserved := reservedNames[*c.Name]; !reserved {
238 | filteredConnections = append(filteredConnections, c)
239 | }
240 | }
241 | if len(filteredConnections) == 0 {
242 | return nil, fmt.Errorf("no connections found")
243 | }
244 |
245 | var results []map[string]interface{}
246 | for _, conn := range filteredConnections {
247 | var errors []string
248 | // Test connection (simulate test_connection endpoint)
249 | resp, err := t.SdkClient.TestConnection(*conn.Name, nil, t.ApiSettings)
250 | if err != nil {
251 | errors = append(errors, "API JSONDecode Error")
252 | } else {
253 | for _, r := range resp {
254 | if *r.Status == "error" {
255 | errors = append(errors, *r.Message)
256 | }
257 | }
258 | }
259 |
260 | // Run inline query for connection activity
261 | limit := "1"
262 | query := &v4.WriteQuery{
263 | Model: "system__activity",
264 | View: "history",
265 | Fields: &[]string{"history.query_run_count"},
266 | Filters: &map[string]any{
267 | "history.connection_name": *conn.Name,
268 | "history.created_date": "90 days",
269 | "user.dev_branch_name": "NULL",
270 | },
271 | Limit: &limit,
272 | }
273 | raw, err := lookercommon.RunInlineQuery(ctx, t.SdkClient, query, "json", t.ApiSettings)
274 | if err != nil {
275 | return nil, err
276 | }
277 | var queryRunCount interface{}
278 | var data []map[string]interface{}
279 | _ = json.Unmarshal([]byte(raw), &data)
280 | if len(data) > 0 {
281 | queryRunCount = data[0]["history.query_run_count"]
282 | }
283 |
284 | results = append(results, map[string]interface{}{
285 | "Connection": *conn.Name,
286 | "Status": "OK",
287 | "Errors": errors,
288 | "Query Count": queryRunCount,
289 | })
290 | }
291 | return results, nil
292 | }
293 |
294 | func (t *pulseTool) checkDashboardPerformance(ctx context.Context) (interface{}, error) {
295 | logger, err := util.LoggerFromContext(ctx)
296 | if err != nil {
297 | return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
298 | }
299 | logger.InfoContext(ctx, "Test 2/6: Checking for dashboards with queries slower than 30 seconds in the last 7 days")
300 |
301 | limit := "20"
302 | query := &v4.WriteQuery{
303 | Model: "system__activity",
304 | View: "history",
305 | Fields: &[]string{"dashboard.title", "query.count"},
306 | Filters: &map[string]any{
307 | "history.created_date": "7 days",
308 | "history.real_dash_id": "-NULL",
309 | "history.runtime": ">30",
310 | "history.status": "complete",
311 | },
312 | Sorts: &[]string{"query.count desc"},
313 | Limit: &limit,
314 | }
315 | raw, err := lookercommon.RunInlineQuery(ctx, t.SdkClient, query, "json", t.ApiSettings)
316 | if err != nil {
317 | return nil, err
318 | }
319 | var dashboards []map[string]interface{}
320 | if err := json.Unmarshal([]byte(raw), &dashboards); err != nil {
321 | return nil, err
322 | }
323 | return dashboards, nil
324 | }
325 |
326 | func (t *pulseTool) checkDashboardErrors(ctx context.Context) (interface{}, error) {
327 | logger, err := util.LoggerFromContext(ctx)
328 | if err != nil {
329 | return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
330 | }
331 | logger.InfoContext(ctx, "Test 3/6: Checking for dashboards with erroring queries in the last 7 days")
332 |
333 | limit := "20"
334 | query := &v4.WriteQuery{
335 | Model: "system__activity",
336 | View: "history",
337 | Fields: &[]string{"dashboard.title", "history.query_run_count"},
338 | Filters: &map[string]any{
339 | "dashboard.title": "-NULL",
340 | "history.created_date": "7 days",
341 | "history.dashboard_session": "-NULL",
342 | "history.status": "error",
343 | },
344 | Sorts: &[]string{"history.query_run_count desc"},
345 | Limit: &limit,
346 | }
347 | raw, err := lookercommon.RunInlineQuery(ctx, t.SdkClient, query, "json", t.ApiSettings)
348 | if err != nil {
349 | return nil, err
350 | }
351 | var dashboards []map[string]interface{}
352 | if err := json.Unmarshal([]byte(raw), &dashboards); err != nil {
353 | return nil, err
354 | }
355 | return dashboards, nil
356 | }
357 |
358 | func (t *pulseTool) checkExplorePerformance(ctx context.Context) (interface{}, error) {
359 | logger, err := util.LoggerFromContext(ctx)
360 | if err != nil {
361 | return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
362 | }
363 | logger.InfoContext(ctx, "Test 4/6: Checking for the slowest explores in the past 7 days")
364 |
365 | limit := "20"
366 | query := &v4.WriteQuery{
367 | Model: "system__activity",
368 | View: "history",
369 | Fields: &[]string{"query.model", "query.view", "history.average_runtime"},
370 | Filters: &map[string]any{
371 | "history.created_date": "7 days",
372 | "query.model": "-NULL, -system^_^_activity",
373 | },
374 | Sorts: &[]string{"history.average_runtime desc"},
375 | Limit: &limit,
376 | }
377 | raw, err := lookercommon.RunInlineQuery(ctx, t.SdkClient, query, "json", t.ApiSettings)
378 | if err != nil {
379 | return nil, err
380 | }
381 | var explores []map[string]interface{}
382 | if err := json.Unmarshal([]byte(raw), &explores); err != nil {
383 | return nil, err
384 | }
385 |
386 | // Average query runtime
387 | query.Fields = &[]string{"history.average_runtime"}
388 | rawAvg, err := lookercommon.RunInlineQuery(ctx, t.SdkClient, query, "json", t.ApiSettings)
389 | if err != nil {
390 | return nil, err
391 | }
392 | var avgData []map[string]interface{}
393 | if err := json.Unmarshal([]byte(rawAvg), &avgData); err == nil {
394 | if len(avgData) > 0 {
395 | if avgRuntime, ok := avgData[0]["history.average_runtime"].(float64); ok {
396 | logger.InfoContext(ctx, fmt.Sprintf("For context, the average query runtime is %.4fs", avgRuntime))
397 | }
398 | }
399 | }
400 | return explores, nil
401 | }
402 |
403 | func (t *pulseTool) checkScheduleFailures(ctx context.Context) (interface{}, error) {
404 | logger, err := util.LoggerFromContext(ctx)
405 | if err != nil {
406 | return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
407 | }
408 | logger.InfoContext(ctx, "Test 5/6: Checking for failing schedules")
409 |
410 | limit := "500"
411 | query := &v4.WriteQuery{
412 | Model: "system__activity",
413 | View: "scheduled_plan",
414 | Fields: &[]string{"scheduled_job.name", "scheduled_job.count"},
415 | Filters: &map[string]any{
416 | "scheduled_job.created_date": "7 days",
417 | "scheduled_job.status": "failure",
418 | },
419 | Sorts: &[]string{"scheduled_job.count desc"},
420 | Limit: &limit,
421 | }
422 | raw, err := lookercommon.RunInlineQuery(ctx, t.SdkClient, query, "json", t.ApiSettings)
423 | if err != nil {
424 | return nil, err
425 | }
426 | var schedules []map[string]interface{}
427 | if err := json.Unmarshal([]byte(raw), &schedules); err != nil {
428 | return nil, err
429 | }
430 | return schedules, nil
431 | }
432 |
433 | func (t *pulseTool) checkLegacyFeatures(ctx context.Context) (interface{}, error) {
434 | logger, err := util.LoggerFromContext(ctx)
435 | if err != nil {
436 | return nil, fmt.Errorf("unable to get logger from ctx: %s", err)
437 | }
438 | logger.InfoContext(ctx, "Test 6/6: Checking for enabled legacy features")
439 |
440 | features, err := t.SdkClient.AllLegacyFeatures(t.ApiSettings)
441 | if err != nil {
442 | if strings.Contains(err.Error(), "Unsupported in Looker (Google Cloud core)") {
443 | return []map[string]string{{"Feature": "Unsupported in Looker (Google Cloud core)"}}, nil
444 | }
445 | logger.ErrorContext(ctx, err.Error())
446 | return []map[string]string{{"Feature": "Unable to pull legacy features due to SDK error"}}, nil
447 | }
448 | var legacyFeatures []map[string]string
449 | for _, f := range features {
450 | if *f.Enabled {
451 | legacyFeatures = append(legacyFeatures, map[string]string{"Feature": *f.Name})
452 | }
453 | }
454 | return legacyFeatures, nil
455 | }
456 |
457 | // =================================================================================================================
458 | // END LOOKER HEALTH PULSE CORE LOGIC
459 | // =================================================================================================================
460 |
```