This is page 76 of 76. 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_prompts_with_coverage.sh
│ ├── test_with_coverage.sh
│ └── versioned.release.cloudbuild.yaml
├── .gemini
│ └── config.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
│ ├── trusted-contribution.yml
│ └── workflows
│ ├── cloud_build_failure_reporter.yml
│ ├── deploy_dev_docs.yaml
│ ├── deploy_previous_version_docs.yaml
│ ├── deploy_versioned_docs.yaml
│ ├── docs_preview_clean.yaml
│ ├── docs_preview_deploy.yaml
│ ├── link_checker_workflow.yaml
│ ├── lint.yaml
│ ├── publish-mcp.yml
│ ├── 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
├── .lycheeignore
├── 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
│ ├── ALLOYDBADMIN_README.md
│ ├── ALLOYDBPG_README.md
│ ├── BIGQUERY_README.md
│ ├── CLOUDSQLMSSQL_README.md
│ ├── CLOUDSQLMSSQLADMIN_README.md
│ ├── CLOUDSQLMYSQL_README.md
│ ├── CLOUDSQLMYSQLADMIN_README.md
│ ├── CLOUDSQLPG_README.md
│ ├── CLOUDSQLPGADMIN_README.md
│ ├── DATAPLEX_README.md
│ ├── en
│ │ ├── _index.md
│ │ ├── about
│ │ │ ├── _index.md
│ │ │ └── faq.md
│ │ ├── blogs
│ │ │ └── _index.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
│ │ │ ├── prompts_quickstart_gemini_cli.md
│ │ │ └── quickstart
│ │ │ ├── go
│ │ │ │ ├── adkgo
│ │ │ │ │ ├── go.mod
│ │ │ │ │ ├── go.sum
│ │ │ │ │ └── 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
│ │ │ │ ├── adk
│ │ │ │ │ ├── package-lock.json
│ │ │ │ │ ├── package.json
│ │ │ │ │ └── quickstart.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_adk_agent.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
│ │ │ ├── embeddingModels
│ │ │ │ ├── _index.md
│ │ │ │ └── gemini.md
│ │ │ ├── prompts
│ │ │ │ ├── _index.md
│ │ │ │ └── custom
│ │ │ │ └── _index.md
│ │ │ ├── sources
│ │ │ │ ├── _index.md
│ │ │ │ ├── alloydb-admin.md
│ │ │ │ ├── alloydb-pg.md
│ │ │ │ ├── bigquery.md
│ │ │ │ ├── bigtable.md
│ │ │ │ ├── cassandra.md
│ │ │ │ ├── clickhouse.md
│ │ │ │ ├── cloud-gda.md
│ │ │ │ ├── cloud-healthcare.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
│ │ │ │ ├── elasticsearch.md
│ │ │ │ ├── firebird.md
│ │ │ │ ├── firestore.md
│ │ │ │ ├── http.md
│ │ │ │ ├── looker.md
│ │ │ │ ├── mariadb.md
│ │ │ │ ├── mindsdb.md
│ │ │ │ ├── mongodb.md
│ │ │ │ ├── mssql.md
│ │ │ │ ├── mysql.md
│ │ │ │ ├── neo4j.md
│ │ │ │ ├── oceanbase.md
│ │ │ │ ├── oracle.md
│ │ │ │ ├── postgres.md
│ │ │ │ ├── redis.md
│ │ │ │ ├── serverless-spark.md
│ │ │ │ ├── singlestore.md
│ │ │ │ ├── snowflake.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
│ │ │ ├── cloudgda
│ │ │ │ ├── _index.md
│ │ │ │ └── cloud-gda-query.md
│ │ │ ├── cloudhealthcare
│ │ │ │ ├── _index.md
│ │ │ │ ├── cloud-healthcare-fhir-fetch-page.md
│ │ │ │ ├── cloud-healthcare-fhir-patient-everything.md
│ │ │ │ ├── cloud-healthcare-fhir-patient-search.md
│ │ │ │ ├── cloud-healthcare-get-dataset.md
│ │ │ │ ├── cloud-healthcare-get-dicom-store-metrics.md
│ │ │ │ ├── cloud-healthcare-get-dicom-store.md
│ │ │ │ ├── cloud-healthcare-get-fhir-resource.md
│ │ │ │ ├── cloud-healthcare-get-fhir-store-metrics.md
│ │ │ │ ├── cloud-healthcare-get-fhir-store.md
│ │ │ │ ├── cloud-healthcare-list-dicom-stores.md
│ │ │ │ ├── cloud-healthcare-list-fhir-stores.md
│ │ │ │ ├── cloud-healthcare-retrieve-rendered-dicom-instance.md
│ │ │ │ ├── cloud-healthcare-search-dicom-instances.md
│ │ │ │ ├── cloud-healthcare-search-dicom-series.md
│ │ │ │ └── cloud-healthcare-search-dicom-studies.md
│ │ │ ├── cloudmonitoring
│ │ │ │ ├── _index.md
│ │ │ │ └── cloud-monitoring-query-prometheus.md
│ │ │ ├── cloudsql
│ │ │ │ ├── _index.md
│ │ │ │ ├── cloudsqlcloneinstance.md
│ │ │ │ ├── cloudsqlcreatedatabase.md
│ │ │ │ ├── cloudsqlcreateusers.md
│ │ │ │ ├── cloudsqlgetinstances.md
│ │ │ │ ├── cloudsqllistdatabases.md
│ │ │ │ ├── cloudsqllistinstances.md
│ │ │ │ ├── cloudsqlmssqlcreateinstance.md
│ │ │ │ ├── cloudsqlmysqlcreateinstance.md
│ │ │ │ ├── cloudsqlpgcreateinstances.md
│ │ │ │ ├── cloudsqlpgupgradeprecheck.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
│ │ │ ├── elasticsearch
│ │ │ │ ├── _index.md
│ │ │ │ └── elasticsearch-esql.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-add-dashboard-filter.md
│ │ │ │ ├── looker-conversational-analytics.md
│ │ │ │ ├── looker-create-project-file.md
│ │ │ │ ├── looker-delete-project-file.md
│ │ │ │ ├── looker-dev-mode.md
│ │ │ │ ├── looker-generate-embed-url.md
│ │ │ │ ├── looker-get-connection-databases.md
│ │ │ │ ├── looker-get-connection-schemas.md
│ │ │ │ ├── looker-get-connection-table-columns.md
│ │ │ │ ├── looker-get-connection-tables.md
│ │ │ │ ├── looker-get-connections.md
│ │ │ │ ├── looker-get-dashboards.md
│ │ │ │ ├── looker-get-dimensions.md
│ │ │ │ ├── looker-get-explores.md
│ │ │ │ ├── looker-get-filters.md
│ │ │ │ ├── looker-get-looks.md
│ │ │ │ ├── looker-get-measures.md
│ │ │ │ ├── looker-get-models.md
│ │ │ │ ├── looker-get-parameters.md
│ │ │ │ ├── looker-get-project-file.md
│ │ │ │ ├── looker-get-project-files.md
│ │ │ │ ├── looker-get-projects.md
│ │ │ │ ├── looker-health-analyze.md
│ │ │ │ ├── looker-health-pulse.md
│ │ │ │ ├── looker-health-vacuum.md
│ │ │ │ ├── looker-make-dashboard.md
│ │ │ │ ├── looker-make-look.md
│ │ │ │ ├── looker-query-sql.md
│ │ │ │ ├── looker-query-url.md
│ │ │ │ ├── looker-query.md
│ │ │ │ ├── looker-run-dashboard.md
│ │ │ │ ├── looker-run-look.md
│ │ │ │ └── looker-update-project-file.md
│ │ │ ├── mindsdb
│ │ │ │ ├── _index.md
│ │ │ │ ├── mindsdb-execute-sql.md
│ │ │ │ └── mindsdb-sql.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-get-query-plan.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-database-overview.md
│ │ │ │ ├── postgres-execute-sql.md
│ │ │ │ ├── postgres-get-column-cardinality.md
│ │ │ │ ├── postgres-list-active-queries.md
│ │ │ │ ├── postgres-list-available-extensions.md
│ │ │ │ ├── postgres-list-database-stats.md
│ │ │ │ ├── postgres-list-indexes.md
│ │ │ │ ├── postgres-list-installed-extensions.md
│ │ │ │ ├── postgres-list-locks.md
│ │ │ │ ├── postgres-list-pg-settings.md
│ │ │ │ ├── postgres-list-publication-tables.md
│ │ │ │ ├── postgres-list-query-stats.md
│ │ │ │ ├── postgres-list-roles.md
│ │ │ │ ├── postgres-list-schemas.md
│ │ │ │ ├── postgres-list-sequences.md
│ │ │ │ ├── postgres-list-stored-procedure.md
│ │ │ │ ├── postgres-list-table-stats.md
│ │ │ │ ├── postgres-list-tables.md
│ │ │ │ ├── postgres-list-tablespaces.md
│ │ │ │ ├── postgres-list-triggers.md
│ │ │ │ ├── postgres-list-views.md
│ │ │ │ ├── postgres-long-running-transactions.md
│ │ │ │ ├── postgres-replication-stats.md
│ │ │ │ └── postgres-sql.md
│ │ │ ├── redis
│ │ │ │ ├── _index.md
│ │ │ │ └── redis.md
│ │ │ ├── serverless-spark
│ │ │ │ ├── _index.md
│ │ │ │ ├── serverless-spark-cancel-batch.md
│ │ │ │ ├── serverless-spark-create-pyspark-batch.md
│ │ │ │ ├── serverless-spark-create-spark-batch.md
│ │ │ │ ├── serverless-spark-get-batch.md
│ │ │ │ └── serverless-spark-list-batches.md
│ │ │ ├── singlestore
│ │ │ │ ├── _index.md
│ │ │ │ ├── singlestore-execute-sql.md
│ │ │ │ └── singlestore-sql.md
│ │ │ ├── snowflake
│ │ │ │ ├── _index.md
│ │ │ │ ├── snowflake-execute-sql.md
│ │ │ │ └── snowflake-sql.md
│ │ │ ├── spanner
│ │ │ │ ├── _index.md
│ │ │ │ ├── spanner-execute-sql.md
│ │ │ │ ├── spanner-list-graphs.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
│ │ │ └── snowflake
│ │ │ ├── _index.md
│ │ │ ├── runme.py
│ │ │ ├── snowflake-config.yaml
│ │ │ ├── snowflake-env.sh
│ │ │ └── test-snowflake.sh
│ │ └── sdks
│ │ ├── _index.md
│ │ ├── go-sdk.md
│ │ ├── js-sdk.md
│ │ └── python-sdk.md
│ ├── LOOKER_README.md
│ ├── SPANNER_README.md
│ └── TOOLBOX_README.md
├── gemini-extension.json
├── go.mod
├── go.sum
├── internal
│ ├── auth
│ │ ├── auth.go
│ │ └── google
│ │ └── google.go
│ ├── embeddingmodels
│ │ ├── embeddingmodels.go
│ │ └── gemini
│ │ ├── gemini_test.go
│ │ └── gemini.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-healthcare.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
│ │ ├── elasticsearch.yaml
│ │ ├── firestore.yaml
│ │ ├── looker-conversational-analytics.yaml
│ │ ├── looker.yaml
│ │ ├── mindsdb.yaml
│ │ ├── mssql.yaml
│ │ ├── mysql.yaml
│ │ ├── neo4j.yaml
│ │ ├── oceanbase.yaml
│ │ ├── postgres.yaml
│ │ ├── serverless-spark.yaml
│ │ ├── singlestore.yaml
│ │ ├── snowflake.yaml
│ │ ├── spanner-postgres.yaml
│ │ ├── spanner.yaml
│ │ └── sqlite.yaml
│ ├── prompts
│ │ ├── arguments_test.go
│ │ ├── arguments.go
│ │ ├── custom
│ │ │ ├── custom_test.go
│ │ │ └── custom.go
│ │ ├── messages_test.go
│ │ ├── messages.go
│ │ ├── prompts_test.go
│ │ ├── prompts.go
│ │ ├── promptsets_test.go
│ │ └── promptsets.go
│ ├── 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
│ │ ├── resources
│ │ │ ├── resources_test.go
│ │ │ └── resources.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
│ │ │ └── cache.go
│ │ ├── bigtable
│ │ │ ├── bigtable_test.go
│ │ │ └── bigtable.go
│ │ ├── cassandra
│ │ │ ├── cassandra_test.go
│ │ │ └── cassandra.go
│ │ ├── clickhouse
│ │ │ ├── clickhouse_test.go
│ │ │ └── clickhouse.go
│ │ ├── cloudgda
│ │ │ ├── cloud_gda_test.go
│ │ │ └── cloud_gda.go
│ │ ├── cloudhealthcare
│ │ │ ├── cloud_healthcare_test.go
│ │ │ └── cloud_healthcare.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
│ │ ├── elasticsearch
│ │ │ ├── elasticsearch_test.go
│ │ │ └── elasticsearch.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
│ │ ├── mindsdb
│ │ │ ├── mindsdb_test.go
│ │ │ └── mindsdb.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_test.go
│ │ │ └── oracle.go
│ │ ├── postgres
│ │ │ ├── postgres_test.go
│ │ │ └── postgres.go
│ │ ├── redis
│ │ │ ├── redis_test.go
│ │ │ └── redis.go
│ │ ├── serverlessspark
│ │ │ ├── serverlessspark_test.go
│ │ │ ├── serverlessspark.go
│ │ │ ├── url_test.go
│ │ │ └── url.go
│ │ ├── singlestore
│ │ │ ├── singlestore_test.go
│ │ │ └── singlestore.go
│ │ ├── snowflake
│ │ │ ├── snowflake_test.go
│ │ │ └── snowflake.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
│ │ ├── cloudgda
│ │ │ ├── cloudgda_test.go
│ │ │ ├── cloudgda.go
│ │ │ └── types.go
│ │ ├── cloudhealthcare
│ │ │ ├── cloudhealthcarefhirfetchpage
│ │ │ │ ├── cloudhealthcarefhirfetchpage_test.go
│ │ │ │ └── cloudhealthcarefhirfetchpage.go
│ │ │ ├── cloudhealthcarefhirpatienteverything
│ │ │ │ ├── cloudhealthcarefhirpatienteverything_test.go
│ │ │ │ └── cloudhealthcarefhirpatienteverything.go
│ │ │ ├── cloudhealthcarefhirpatientsearch
│ │ │ │ ├── cloudhealthcarefhirpatientsearch_test.go
│ │ │ │ └── cloudhealthcarefhirpatientsearch.go
│ │ │ ├── cloudhealthcaregetdataset
│ │ │ │ ├── cloudhealthcaregetdataset_test.go
│ │ │ │ └── cloudhealthcaregetdataset.go
│ │ │ ├── cloudhealthcaregetdicomstore
│ │ │ │ ├── cloudhealthcaregetdicomstore_test.go
│ │ │ │ └── cloudhealthcaregetdicomstore.go
│ │ │ ├── cloudhealthcaregetdicomstoremetrics
│ │ │ │ ├── cloudhealthcaregetdicomstoremetrics_test.go
│ │ │ │ └── cloudhealthcaregetdicomstoremetrics.go
│ │ │ ├── cloudhealthcaregetfhirresource
│ │ │ │ ├── cloudhealthcaregetfhirresource_test.go
│ │ │ │ └── cloudhealthcaregetfhirresource.go
│ │ │ ├── cloudhealthcaregetfhirstore
│ │ │ │ ├── cloudhealthcaregetfhirstore_test.go
│ │ │ │ └── cloudhealthcaregetfhirstore.go
│ │ │ ├── cloudhealthcaregetfhirstoremetrics
│ │ │ │ ├── cloudhealthcaregetfhirstoremetrics_test.go
│ │ │ │ └── cloudhealthcaregetfhirstoremetrics.go
│ │ │ ├── cloudhealthcarelistdicomstores
│ │ │ │ ├── cloudhealthcarelistdicomstores_test.go
│ │ │ │ └── cloudhealthcarelistdicomstores.go
│ │ │ ├── cloudhealthcarelistfhirstores
│ │ │ │ ├── cloudhealthcarelistfhirstores_test.go
│ │ │ │ └── cloudhealthcarelistfhirstores.go
│ │ │ ├── cloudhealthcareretrieverendereddicominstance
│ │ │ │ ├── cloudhealthcareretrieverendereddicominstance_test.go
│ │ │ │ └── cloudhealthcareretrieverendereddicominstance.go
│ │ │ ├── cloudhealthcaresearchdicominstances
│ │ │ │ ├── cloudhealthcaresearchdicominstances_test.go
│ │ │ │ └── cloudhealthcaresearchdicominstances.go
│ │ │ ├── cloudhealthcaresearchdicomseries
│ │ │ │ ├── cloudhealthcaresearchdicomseries_test.go
│ │ │ │ └── cloudhealthcaresearchdicomseries.go
│ │ │ ├── cloudhealthcaresearchdicomstudies
│ │ │ │ ├── cloudhealthcaresearchdicomstudies_test.go
│ │ │ │ └── cloudhealthcaresearchdicomstudies.go
│ │ │ └── common
│ │ │ └── util.go
│ │ ├── cloudmonitoring
│ │ │ ├── cloudmonitoring_test.go
│ │ │ └── cloudmonitoring.go
│ │ ├── cloudsql
│ │ │ ├── cloudsqlcloneinstance
│ │ │ │ ├── cloudsqlcloneinstance_test.go
│ │ │ │ └── cloudsqlcloneinstance.go
│ │ │ ├── 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
│ │ │ └── cloudsqlpgupgradeprecheck
│ │ │ ├── cloudsqlpgupgradeprecheck_test.go
│ │ │ └── cloudsqlpgupgradeprecheck.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
│ │ ├── elasticsearch
│ │ │ └── elasticsearchesql
│ │ │ ├── elasticsearchesql_test.go
│ │ │ └── elasticsearchesql.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
│ │ │ ├── lookeradddashboardfilter
│ │ │ │ ├── lookeradddashboardfilter_test.go
│ │ │ │ └── lookeradddashboardfilter.go
│ │ │ ├── lookercommon
│ │ │ │ ├── lookercommon_test.go
│ │ │ │ └── lookercommon.go
│ │ │ ├── lookerconversationalanalytics
│ │ │ │ ├── lookerconversationalanalytics_test.go
│ │ │ │ └── lookerconversationalanalytics.go
│ │ │ ├── lookercreateprojectfile
│ │ │ │ ├── lookercreateprojectfile_test.go
│ │ │ │ └── lookercreateprojectfile.go
│ │ │ ├── lookerdeleteprojectfile
│ │ │ │ ├── lookerdeleteprojectfile_test.go
│ │ │ │ └── lookerdeleteprojectfile.go
│ │ │ ├── lookerdevmode
│ │ │ │ ├── lookerdevmode_test.go
│ │ │ │ └── lookerdevmode.go
│ │ │ ├── lookergenerateembedurl
│ │ │ │ ├── lookergenerateembedurl_test.go
│ │ │ │ └── lookergenerateembedurl.go
│ │ │ ├── lookergetconnectiondatabases
│ │ │ │ ├── lookergetconnectiondatabases_test.go
│ │ │ │ └── lookergetconnectiondatabases.go
│ │ │ ├── lookergetconnections
│ │ │ │ ├── lookergetconnections_test.go
│ │ │ │ └── lookergetconnections.go
│ │ │ ├── lookergetconnectionschemas
│ │ │ │ ├── lookergetconnectionschemas_test.go
│ │ │ │ └── lookergetconnectionschemas.go
│ │ │ ├── lookergetconnectiontablecolumns
│ │ │ │ ├── lookergetconnectiontablecolumns_test.go
│ │ │ │ └── lookergetconnectiontablecolumns.go
│ │ │ ├── lookergetconnectiontables
│ │ │ │ ├── lookergetconnectiontables_test.go
│ │ │ │ └── lookergetconnectiontables.go
│ │ │ ├── lookergetdashboards
│ │ │ │ ├── lookergetdashboards_test.go
│ │ │ │ └── lookergetdashboards.go
│ │ │ ├── lookergetdimensions
│ │ │ │ ├── lookergetdimensions_test.go
│ │ │ │ └── lookergetdimensions.go
│ │ │ ├── lookergetexplores
│ │ │ │ ├── lookergetexplores_test.go
│ │ │ │ └── lookergetexplores.go
│ │ │ ├── lookergetfilters
│ │ │ │ ├── lookergetfilters_test.go
│ │ │ │ └── lookergetfilters.go
│ │ │ ├── lookergetlooks
│ │ │ │ ├── lookergetlooks_test.go
│ │ │ │ └── lookergetlooks.go
│ │ │ ├── lookergetmeasures
│ │ │ │ ├── lookergetmeasures_test.go
│ │ │ │ └── lookergetmeasures.go
│ │ │ ├── lookergetmodels
│ │ │ │ ├── lookergetmodels_test.go
│ │ │ │ └── lookergetmodels.go
│ │ │ ├── lookergetparameters
│ │ │ │ ├── lookergetparameters_test.go
│ │ │ │ └── lookergetparameters.go
│ │ │ ├── lookergetprojectfile
│ │ │ │ ├── lookergetprojectfile_test.go
│ │ │ │ └── lookergetprojectfile.go
│ │ │ ├── lookergetprojectfiles
│ │ │ │ ├── lookergetprojectfiles_test.go
│ │ │ │ └── lookergetprojectfiles.go
│ │ │ ├── lookergetprojects
│ │ │ │ ├── lookergetprojects_test.go
│ │ │ │ └── lookergetprojects.go
│ │ │ ├── lookerhealthanalyze
│ │ │ │ ├── lookerhealthanalyze_test.go
│ │ │ │ └── lookerhealthanalyze.go
│ │ │ ├── lookerhealthpulse
│ │ │ │ ├── lookerhealthpulse_test.go
│ │ │ │ └── lookerhealthpulse.go
│ │ │ ├── lookerhealthvacuum
│ │ │ │ ├── lookerhealthvacuum_test.go
│ │ │ │ └── lookerhealthvacuum.go
│ │ │ ├── lookermakedashboard
│ │ │ │ ├── lookermakedashboard_test.go
│ │ │ │ └── lookermakedashboard.go
│ │ │ ├── lookermakelook
│ │ │ │ ├── lookermakelook_test.go
│ │ │ │ └── lookermakelook.go
│ │ │ ├── lookerquery
│ │ │ │ ├── lookerquery_test.go
│ │ │ │ └── lookerquery.go
│ │ │ ├── lookerquerysql
│ │ │ │ ├── lookerquerysql_test.go
│ │ │ │ └── lookerquerysql.go
│ │ │ ├── lookerqueryurl
│ │ │ │ ├── lookerqueryurl_test.go
│ │ │ │ └── lookerqueryurl.go
│ │ │ ├── lookerrundashboard
│ │ │ │ ├── lookerrundashboard_test.go
│ │ │ │ └── lookerrundashboard.go
│ │ │ ├── lookerrunlook
│ │ │ │ ├── lookerrunlook_test.go
│ │ │ │ └── lookerrunlook.go
│ │ │ └── lookerupdateprojectfile
│ │ │ ├── lookerupdateprojectfile_test.go
│ │ │ └── lookerupdateprojectfile.go
│ │ ├── mindsdb
│ │ │ ├── mindsdbexecutesql
│ │ │ │ ├── mindsdbexecutesql_test.go
│ │ │ │ └── mindsdbexecutesql.go
│ │ │ └── mindsdbsql
│ │ │ ├── mindsdbsql_test.go
│ │ │ └── mindsdbsql.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
│ │ │ ├── mysqlgetqueryplan
│ │ │ │ ├── mysqlgetqueryplan_test.go
│ │ │ │ └── mysqlgetqueryplan.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_test.go
│ │ │ │ └── oracleexecutesql.go
│ │ │ └── oraclesql
│ │ │ ├── oraclesql_test.go
│ │ │ └── oraclesql.go
│ │ ├── postgres
│ │ │ ├── postgresdatabaseoverview
│ │ │ │ ├── postgresdatabaseoverview_test.go
│ │ │ │ └── postgresdatabaseoverview.go
│ │ │ ├── postgresexecutesql
│ │ │ │ ├── postgresexecutesql_test.go
│ │ │ │ └── postgresexecutesql.go
│ │ │ ├── postgresgetcolumncardinality
│ │ │ │ ├── postgresgetcolumncardinality_test.go
│ │ │ │ └── postgresgetcolumncardinality.go
│ │ │ ├── postgreslistactivequeries
│ │ │ │ ├── postgreslistactivequeries_test.go
│ │ │ │ └── postgreslistactivequeries.go
│ │ │ ├── postgreslistavailableextensions
│ │ │ │ ├── postgreslistavailableextensions_test.go
│ │ │ │ └── postgreslistavailableextensions.go
│ │ │ ├── postgreslistdatabasestats
│ │ │ │ ├── postgreslistdatabasestats_test.go
│ │ │ │ └── postgreslistdatabasestats.go
│ │ │ ├── postgreslistindexes
│ │ │ │ ├── postgreslistindexes_test.go
│ │ │ │ └── postgreslistindexes.go
│ │ │ ├── postgreslistinstalledextensions
│ │ │ │ ├── postgreslistinstalledextensions_test.go
│ │ │ │ └── postgreslistinstalledextensions.go
│ │ │ ├── postgreslistlocks
│ │ │ │ ├── postgreslistlocks_test.go
│ │ │ │ └── postgreslistlocks.go
│ │ │ ├── postgreslistpgsettings
│ │ │ │ ├── postgreslistpgsettings_test.go
│ │ │ │ └── postgreslistpgsettings.go
│ │ │ ├── postgreslistpublicationtables
│ │ │ │ ├── postgreslistpublicationtables_test.go
│ │ │ │ └── postgreslistpublicationtables.go
│ │ │ ├── postgreslistquerystats
│ │ │ │ ├── postgreslistquerystats_test.go
│ │ │ │ └── postgreslistquerystats.go
│ │ │ ├── postgreslistroles
│ │ │ │ ├── postgreslistroles_test.go
│ │ │ │ └── postgreslistroles.go
│ │ │ ├── postgreslistschemas
│ │ │ │ ├── postgreslistschemas_test.go
│ │ │ │ └── postgreslistschemas.go
│ │ │ ├── postgreslistsequences
│ │ │ │ ├── postgreslistsequences_test.go
│ │ │ │ └── postgreslistsequences.go
│ │ │ ├── postgresliststoredprocedure
│ │ │ │ ├── postgresliststoredprocedure_test.go
│ │ │ │ └── postgresliststoredprocedure.go
│ │ │ ├── postgreslisttables
│ │ │ │ ├── postgreslisttables_test.go
│ │ │ │ └── postgreslisttables.go
│ │ │ ├── postgreslisttablespaces
│ │ │ │ ├── postgreslisttablespaces_test.go
│ │ │ │ └── postgreslisttablespaces.go
│ │ │ ├── postgreslisttablestats
│ │ │ │ ├── postgreslisttablestats_test.go
│ │ │ │ └── postgreslisttablestats.go
│ │ │ ├── postgreslisttriggers
│ │ │ │ ├── postgreslisttriggers_test.go
│ │ │ │ └── postgreslisttriggers.go
│ │ │ ├── postgreslistviews
│ │ │ │ ├── postgreslistviews_test.go
│ │ │ │ └── postgreslistviews.go
│ │ │ ├── postgreslongrunningtransactions
│ │ │ │ ├── postgreslongrunningtransactions_test.go
│ │ │ │ └── postgreslongrunningtransactions.go
│ │ │ ├── postgresreplicationstats
│ │ │ │ ├── postgresreplicationstats_test.go
│ │ │ │ └── postgresreplicationstats.go
│ │ │ └── postgressql
│ │ │ ├── postgressql_test.go
│ │ │ └── postgressql.go
│ │ ├── redis
│ │ │ ├── redis_test.go
│ │ │ └── redis.go
│ │ ├── serverlessspark
│ │ │ ├── createbatch
│ │ │ │ ├── config.go
│ │ │ │ └── tool.go
│ │ │ ├── serverlesssparkcancelbatch
│ │ │ │ ├── serverlesssparkcancelbatch_test.go
│ │ │ │ └── serverlesssparkcancelbatch.go
│ │ │ ├── serverlesssparkcreatepysparkbatch
│ │ │ │ ├── serverlesssparkcreatepysparkbatch_test.go
│ │ │ │ └── serverlesssparkcreatepysparkbatch.go
│ │ │ ├── serverlesssparkcreatesparkbatch
│ │ │ │ ├── serverlesssparkcreatesparkbatch_test.go
│ │ │ │ └── serverlesssparkcreatesparkbatch.go
│ │ │ ├── serverlesssparkgetbatch
│ │ │ │ ├── serverlesssparkgetbatch_test.go
│ │ │ │ └── serverlesssparkgetbatch.go
│ │ │ ├── serverlesssparklistbatches
│ │ │ │ ├── serverlesssparklistbatches_test.go
│ │ │ │ └── serverlesssparklistbatches.go
│ │ │ └── testutils
│ │ │ └── testutils.go
│ │ ├── singlestore
│ │ │ ├── singlestoreexecutesql
│ │ │ │ ├── singlestoreexecutesql_test.go
│ │ │ │ └── singlestoreexecutesql.go
│ │ │ └── singlestoresql
│ │ │ ├── singlestoresql_test.go
│ │ │ └── singlestoresql.go
│ │ ├── snowflake
│ │ │ ├── snowflakeexecutesql
│ │ │ │ ├── snowflakeexecutesql_test.go
│ │ │ │ └── snowflakeexecutesql.go
│ │ │ └── snowflakesql
│ │ │ ├── snowflakesql_test.go
│ │ │ └── snowflakesql.go
│ │ ├── spanner
│ │ │ ├── spannerexecutesql
│ │ │ │ ├── spannerexecutesql_test.go
│ │ │ │ └── spannerexecutesql.go
│ │ │ ├── spannerlistgraphs
│ │ │ │ ├── spannerlistgraphs_test.go
│ │ │ │ └── spannerlistgraphs.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
│ ├── orderedmap
│ │ ├── orderedmap_test.go
│ │ └── orderedmap.go
│ ├── parameters
│ │ ├── common_test.go
│ │ ├── common.go
│ │ ├── parameters_test.go
│ │ └── parameters.go
│ └── util.go
├── LICENSE
├── logo.png
├── main.go
├── MCP-TOOLBOX-EXTENSION.md
├── README.md
├── server.json
└── 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
├── cloudgda
│ └── cloud_gda_integration_test.go
├── cloudhealthcare
│ └── cloud_healthcare_integration_test.go
├── cloudmonitoring
│ └── cloud_monitoring_integration_test.go
├── cloudsql
│ ├── cloud_sql_clone_instance_test.go
│ ├── 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
│ └── cloud_sql_pg_upgrade_precheck_test.go
├── common.go
├── couchbase
│ └── couchbase_integration_test.go
├── dataform
│ └── dataform_integration_test.go
├── dataplex
│ └── dataplex_integration_test.go
├── dgraph
│ └── dgraph_integration_test.go
├── elasticsearch
│ └── elasticsearch_integration_test.go
├── firebird
│ └── firebird_integration_test.go
├── firestore
│ └── firestore_integration_test.go
├── http
│ └── http_integration_test.go
├── looker
│ └── looker_integration_test.go
├── mariadb
│ └── mariadb_integration_test.go
├── mindsdb
│ └── mindsdb_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
├── prompts
│ └── custom
│ └── prompts_integration_test.go
├── redis
│ └── redis_test.go
├── server.go
├── serverlessspark
│ └── serverless_spark_integration_test.go
├── singlestore
│ └── singlestore_integration_test.go
├── snowflake
│ └── snowflake_integration_test.go
├── source.go
├── spanner
│ └── spanner_integration_test.go
├── sqlite
│ └── sqlite_integration_test.go
├── tidb
│ └── tidb_integration_test.go
├── tool.go
├── trino
│ └── trino_integration_test.go
├── utility
│ └── wait_integration_test.go
├── valkey
│ └── valkey_test.go
└── yugabytedb
└── yugabytedb_integration_test.go
```
# Files
--------------------------------------------------------------------------------
/tests/tool.go:
--------------------------------------------------------------------------------
```go
1 | // Copyright 2025 Google LLC
2 | //
3 | // Licensed under the Apache License, Version 2.0 (the "License");
4 | // you may not use this file except in compliance with the License.
5 | // You may obtain a copy of the License at
6 | //
7 | // http://www.apache.org/licenses/LICENSE-2.0
8 | //
9 | // Unless required by applicable law or agreed to in writing, software
10 | // distributed under the License is distributed on an "AS IS" BASIS,
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | // See the License for the specific language governing permissions and
13 | // limitations under the License.
14 |
15 | package tests
16 |
17 | import (
18 | "bytes"
19 | "context"
20 | "database/sql"
21 | "encoding/json"
22 | "fmt"
23 | "io"
24 | "net/http"
25 | "reflect"
26 | "sort"
27 | "strings"
28 | "sync"
29 | "testing"
30 | "time"
31 |
32 | "github.com/google/go-cmp/cmp"
33 | "github.com/google/go-cmp/cmp/cmpopts"
34 | "github.com/google/uuid"
35 | "github.com/googleapis/genai-toolbox/internal/server/mcp/jsonrpc"
36 | "github.com/googleapis/genai-toolbox/internal/sources"
37 | "github.com/jackc/pgx/v5/pgxpool"
38 | )
39 |
40 | // RunToolGet runs the tool get endpoint
41 | func RunToolGetTest(t *testing.T) {
42 | // Test tool get endpoint
43 | tcs := []struct {
44 | name string
45 | api string
46 | want map[string]any
47 | }{
48 | {
49 | name: "get my-simple-tool",
50 | api: "http://127.0.0.1:5000/api/tool/my-simple-tool/",
51 | want: map[string]any{
52 | "my-simple-tool": map[string]any{
53 | "description": "Simple tool to test end to end functionality.",
54 | "parameters": []any{},
55 | "authRequired": []any{},
56 | },
57 | },
58 | },
59 | }
60 | for _, tc := range tcs {
61 | t.Run(tc.name, func(t *testing.T) {
62 | resp, err := http.Get(tc.api)
63 | if err != nil {
64 | t.Fatalf("error when sending a request: %s", err)
65 | }
66 | defer resp.Body.Close()
67 | if resp.StatusCode != 200 {
68 | t.Fatalf("response status code is not 200")
69 | }
70 |
71 | var body map[string]interface{}
72 | err = json.NewDecoder(resp.Body).Decode(&body)
73 | if err != nil {
74 | t.Fatalf("error parsing response body")
75 | }
76 |
77 | got, ok := body["tools"]
78 | if !ok {
79 | t.Fatalf("unable to find tools in response body")
80 | }
81 | if !reflect.DeepEqual(got, tc.want) {
82 | t.Fatalf("got %q, want %q", got, tc.want)
83 | }
84 | })
85 | }
86 | }
87 |
88 | func RunToolGetTestByName(t *testing.T, name string, want map[string]any) {
89 | // Test tool get endpoint
90 | tcs := []struct {
91 | name string
92 | api string
93 | want map[string]any
94 | }{
95 | {
96 | name: fmt.Sprintf("get %s", name),
97 | api: fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/", name),
98 | want: want,
99 | },
100 | }
101 | for _, tc := range tcs {
102 | t.Run(tc.name, func(t *testing.T) {
103 | resp, err := http.Get(tc.api)
104 | if err != nil {
105 | t.Fatalf("error when sending a request: %s", err)
106 | }
107 | defer resp.Body.Close()
108 | if resp.StatusCode != 200 {
109 | t.Fatalf("response status code is not 200")
110 | }
111 |
112 | var body map[string]interface{}
113 | err = json.NewDecoder(resp.Body).Decode(&body)
114 | if err != nil {
115 | t.Fatalf("error parsing response body")
116 | }
117 |
118 | got, ok := body["tools"]
119 | if !ok {
120 | t.Fatalf("unable to find tools in response body")
121 | }
122 | if !reflect.DeepEqual(got, tc.want) {
123 | t.Fatalf("got %q, want %q", got, tc.want)
124 | }
125 | })
126 | }
127 | }
128 |
129 | // RunToolInvokeSimpleTest runs the tool invoke endpoint with no parameters
130 | func RunToolInvokeSimpleTest(t *testing.T, name string, simpleWant string) {
131 | // Test tool invoke endpoint
132 | invokeTcs := []struct {
133 | name string
134 | api string
135 | requestHeader map[string]string
136 | requestBody io.Reader
137 | want string
138 | isErr bool
139 | }{
140 | {
141 | name: fmt.Sprintf("invoke %s", name),
142 | api: fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/invoke", name),
143 | requestHeader: map[string]string{},
144 | requestBody: bytes.NewBuffer([]byte(`{}`)),
145 | want: simpleWant,
146 | isErr: false,
147 | },
148 | }
149 | for _, tc := range invokeTcs {
150 | t.Run(tc.name, func(t *testing.T) {
151 | // Send Tool invocation request
152 | resp, respBody := RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
153 | if resp.StatusCode != http.StatusOK {
154 | if tc.isErr {
155 | return
156 | }
157 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBody))
158 | }
159 |
160 | // Check response body
161 | var body map[string]interface{}
162 | err := json.Unmarshal(respBody, &body)
163 | if err != nil {
164 | t.Fatalf("error parsing response body")
165 | }
166 |
167 | got, ok := body["result"].(string)
168 | if !ok {
169 | t.Fatalf("unable to find result in response body")
170 | }
171 |
172 | if !strings.Contains(got, tc.want) {
173 | t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
174 | }
175 | })
176 | }
177 | }
178 |
179 | func RunToolInvokeParametersTest(t *testing.T, name string, params []byte, simpleWant string) {
180 | // Test tool invoke endpoint
181 | invokeTcs := []struct {
182 | name string
183 | api string
184 | requestHeader map[string]string
185 | requestBody io.Reader
186 | want string
187 | isErr bool
188 | }{
189 | {
190 | name: fmt.Sprintf("invoke %s", name),
191 | api: fmt.Sprintf("http://127.0.0.1:5000/api/tool/%s/invoke", name),
192 | requestHeader: map[string]string{},
193 | requestBody: bytes.NewBuffer(params),
194 | want: simpleWant,
195 | isErr: false,
196 | },
197 | }
198 | for _, tc := range invokeTcs {
199 | t.Run(tc.name, func(t *testing.T) {
200 | // Send Tool invocation request
201 | resp, respBody := RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
202 | if resp.StatusCode != http.StatusOK {
203 | if tc.isErr {
204 | return
205 | }
206 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBody))
207 | }
208 |
209 | // Check response body
210 | var body map[string]interface{}
211 | err := json.Unmarshal(respBody, &body)
212 | if err != nil {
213 | t.Fatalf("error parsing response body")
214 | }
215 |
216 | got, ok := body["result"].(string)
217 | if !ok {
218 | t.Fatalf("unable to find result in response body")
219 | }
220 |
221 | if !strings.Contains(got, tc.want) {
222 | t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
223 | }
224 | })
225 | }
226 | }
227 |
228 | // RunToolInvoke runs the tool invoke endpoint
229 | func RunToolInvokeTest(t *testing.T, select1Want string, options ...InvokeTestOption) {
230 | // Resolve options
231 | // Default values for InvokeTestConfig
232 | configs := &InvokeTestConfig{
233 | myToolId3NameAliceWant: "[{\"id\":1,\"name\":\"Alice\"},{\"id\":3,\"name\":\"Sid\"}]",
234 | myToolById4Want: "[{\"id\":4,\"name\":null}]",
235 | myArrayToolWant: "[{\"id\":1,\"name\":\"Alice\"},{\"id\":3,\"name\":\"Sid\"}]",
236 | nullWant: "null",
237 | supportOptionalNullParam: true,
238 | supportArrayParam: true,
239 | supportClientAuth: false,
240 | supportSelect1Want: true,
241 | supportSelect1Auth: true,
242 | }
243 |
244 | // Apply provided options
245 | for _, option := range options {
246 | option(configs)
247 | }
248 |
249 | // Get ID token
250 | idToken, err := GetGoogleIdToken(ClientId)
251 | if err != nil {
252 | t.Fatalf("error getting Google ID token: %s", err)
253 | }
254 |
255 | // Get access token
256 | accessToken, err := sources.GetIAMAccessToken(t.Context())
257 | if err != nil {
258 | t.Fatalf("error getting access token from ADC: %s", err)
259 | }
260 | accessToken = "Bearer " + accessToken
261 |
262 | // Test tool invoke endpoint
263 | invokeTcs := []struct {
264 | name string
265 | api string
266 | enabled bool
267 | requestHeader map[string]string
268 | requestBody io.Reader
269 | wantStatusCode int
270 | wantBody string
271 | }{
272 | {
273 | name: "invoke my-simple-tool",
274 | api: "http://127.0.0.1:5000/api/tool/my-simple-tool/invoke",
275 | enabled: configs.supportSelect1Want,
276 | requestHeader: map[string]string{},
277 | requestBody: bytes.NewBuffer([]byte(`{}`)),
278 | wantBody: select1Want,
279 | wantStatusCode: http.StatusOK,
280 | },
281 | {
282 | name: "invoke my-tool",
283 | api: "http://127.0.0.1:5000/api/tool/my-tool/invoke",
284 | enabled: true,
285 | requestHeader: map[string]string{},
286 | requestBody: bytes.NewBuffer([]byte(`{"id": 3, "name": "Alice"}`)),
287 | wantBody: configs.myToolId3NameAliceWant,
288 | wantStatusCode: http.StatusOK,
289 | },
290 | {
291 | name: "invoke my-tool-by-id with nil response",
292 | api: "http://127.0.0.1:5000/api/tool/my-tool-by-id/invoke",
293 | enabled: true,
294 | requestHeader: map[string]string{},
295 | requestBody: bytes.NewBuffer([]byte(`{"id": 4}`)),
296 | wantBody: configs.myToolById4Want,
297 | wantStatusCode: http.StatusOK,
298 | },
299 | {
300 | name: "invoke my-tool-by-name with nil response",
301 | api: "http://127.0.0.1:5000/api/tool/my-tool-by-name/invoke",
302 | enabled: configs.supportOptionalNullParam,
303 | requestHeader: map[string]string{},
304 | requestBody: bytes.NewBuffer([]byte(`{}`)),
305 | wantBody: configs.nullWant,
306 | wantStatusCode: http.StatusOK,
307 | },
308 | {
309 | name: "Invoke my-tool without parameters",
310 | api: "http://127.0.0.1:5000/api/tool/my-tool/invoke",
311 | enabled: true,
312 | requestHeader: map[string]string{},
313 | requestBody: bytes.NewBuffer([]byte(`{}`)),
314 | wantBody: "",
315 | wantStatusCode: http.StatusBadRequest,
316 | },
317 | {
318 | name: "Invoke my-tool with insufficient parameters",
319 | api: "http://127.0.0.1:5000/api/tool/my-tool/invoke",
320 | enabled: true,
321 | requestHeader: map[string]string{},
322 | requestBody: bytes.NewBuffer([]byte(`{"id": 1}`)),
323 | wantBody: "",
324 | wantStatusCode: http.StatusBadRequest,
325 | },
326 | {
327 | name: "invoke my-array-tool",
328 | api: "http://127.0.0.1:5000/api/tool/my-array-tool/invoke",
329 | enabled: configs.supportArrayParam,
330 | requestHeader: map[string]string{},
331 | requestBody: bytes.NewBuffer([]byte(`{"idArray": [1,2,3], "nameArray": ["Alice", "Sid", "RandomName"], "cmdArray": ["HGETALL", "row3"]}`)),
332 | wantBody: configs.myArrayToolWant,
333 | wantStatusCode: http.StatusOK,
334 | },
335 | {
336 | name: "Invoke my-auth-tool with auth token",
337 | api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke",
338 | enabled: configs.supportSelect1Auth,
339 | requestHeader: map[string]string{"my-google-auth_token": idToken},
340 | requestBody: bytes.NewBuffer([]byte(`{}`)),
341 | wantBody: configs.myAuthToolWant,
342 | wantStatusCode: http.StatusOK,
343 | },
344 | {
345 | name: "Invoke my-auth-tool with invalid auth token",
346 | api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke",
347 | enabled: configs.supportSelect1Auth,
348 | requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
349 | requestBody: bytes.NewBuffer([]byte(`{}`)),
350 | wantBody: "",
351 | wantStatusCode: http.StatusUnauthorized,
352 | },
353 | {
354 | name: "Invoke my-auth-tool without auth token",
355 | api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke",
356 | enabled: true,
357 | requestHeader: map[string]string{},
358 | requestBody: bytes.NewBuffer([]byte(`{}`)),
359 | wantBody: "",
360 | wantStatusCode: http.StatusUnauthorized,
361 | },
362 | {
363 | name: "Invoke my-auth-required-tool with auth token",
364 | api: "http://127.0.0.1:5000/api/tool/my-auth-required-tool/invoke",
365 | enabled: configs.supportSelect1Auth,
366 | requestHeader: map[string]string{"my-google-auth_token": idToken},
367 | requestBody: bytes.NewBuffer([]byte(`{}`)),
368 | wantBody: select1Want,
369 | wantStatusCode: http.StatusOK,
370 | },
371 | {
372 | name: "Invoke my-auth-required-tool with invalid auth token",
373 | api: "http://127.0.0.1:5000/api/tool/my-auth-required-tool/invoke",
374 | enabled: true,
375 | requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
376 | requestBody: bytes.NewBuffer([]byte(`{}`)),
377 | wantBody: "",
378 | wantStatusCode: http.StatusUnauthorized,
379 | },
380 | {
381 | name: "Invoke my-auth-required-tool without auth token",
382 | api: "http://127.0.0.1:5000/api/tool/my-auth-tool/invoke",
383 | enabled: true,
384 | requestHeader: map[string]string{},
385 | requestBody: bytes.NewBuffer([]byte(`{}`)),
386 | wantBody: "",
387 | wantStatusCode: http.StatusUnauthorized,
388 | },
389 | {
390 | name: "Invoke my-client-auth-tool with auth token",
391 | api: "http://127.0.0.1:5000/api/tool/my-client-auth-tool/invoke",
392 | enabled: configs.supportClientAuth,
393 | requestHeader: map[string]string{"Authorization": accessToken},
394 | requestBody: bytes.NewBuffer([]byte(`{}`)),
395 | wantBody: select1Want,
396 | wantStatusCode: http.StatusOK,
397 | },
398 | {
399 | name: "Invoke my-client-auth-tool without auth token",
400 | api: "http://127.0.0.1:5000/api/tool/my-client-auth-tool/invoke",
401 | enabled: configs.supportClientAuth,
402 | requestHeader: map[string]string{},
403 | requestBody: bytes.NewBuffer([]byte(`{}`)),
404 | wantStatusCode: http.StatusUnauthorized,
405 | },
406 | {
407 |
408 | name: "Invoke my-client-auth-tool with invalid auth token",
409 | api: "http://127.0.0.1:5000/api/tool/my-client-auth-tool/invoke",
410 | enabled: configs.supportClientAuth,
411 | requestHeader: map[string]string{"Authorization": "Bearer invalid-token"},
412 | requestBody: bytes.NewBuffer([]byte(`{}`)),
413 | wantStatusCode: http.StatusUnauthorized,
414 | },
415 | }
416 | for _, tc := range invokeTcs {
417 | t.Run(tc.name, func(t *testing.T) {
418 | if !tc.enabled {
419 | return
420 | }
421 | // Send Tool invocation request
422 | resp, respBody := RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
423 |
424 | // Check status code
425 | if resp.StatusCode != tc.wantStatusCode {
426 | t.Errorf("StatusCode mismatch: got %d, want %d. Response body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
427 | }
428 |
429 | // skip response body check
430 | if tc.wantBody == "" {
431 | return
432 | }
433 |
434 | // Check response body
435 | var body map[string]interface{}
436 | err = json.Unmarshal(respBody, &body)
437 | if err != nil {
438 | t.Fatalf("error parsing response body: %s", err)
439 | }
440 |
441 | got, ok := body["result"].(string)
442 | if !ok {
443 | t.Fatalf("unable to find result in response body")
444 | }
445 |
446 | if got != tc.wantBody {
447 | t.Fatalf("unexpected value: got %q, want %q", got, tc.wantBody)
448 | }
449 | })
450 | }
451 | }
452 |
453 | // RunToolInvokeWithTemplateParameters runs tool invoke test cases with template parameters.
454 | func RunToolInvokeWithTemplateParameters(t *testing.T, tableName string, options ...TemplateParamOption) {
455 | // Resolve options
456 | // Default values for TemplateParameterTestConfig
457 | configs := &TemplateParameterTestConfig{
458 | ddlWant: "null",
459 | selectAllWant: "[{\"age\":21,\"id\":1,\"name\":\"Alex\"},{\"age\":100,\"id\":2,\"name\":\"Alice\"}]",
460 | selectId1Want: "[{\"age\":21,\"id\":1,\"name\":\"Alex\"}]",
461 | selectNameWant: "[{\"age\":21,\"id\":1,\"name\":\"Alex\"}]",
462 | selectEmptyWant: "null",
463 | insert1Want: "null",
464 |
465 | nameFieldArray: `["name"]`,
466 | nameColFilter: "name",
467 | createColArray: `["id INT","name VARCHAR(20)","age INT"]`,
468 |
469 | supportDdl: true,
470 | supportInsert: true,
471 | }
472 |
473 | // Apply provided options
474 | for _, option := range options {
475 | option(configs)
476 | }
477 |
478 | selectOnlyNamesWant := "[{\"name\":\"Alex\"},{\"name\":\"Alice\"}]"
479 |
480 | // Test tool invoke endpoint
481 | invokeTcs := []struct {
482 | name string
483 | enabled bool
484 | ddl bool
485 | insert bool
486 | api string
487 | requestHeader map[string]string
488 | requestBody io.Reader
489 | want string
490 | isErr bool
491 | }{
492 | {
493 | name: "invoke create-table-templateParams-tool",
494 | ddl: true,
495 | api: "http://127.0.0.1:5000/api/tool/create-table-templateParams-tool/invoke",
496 | requestHeader: map[string]string{},
497 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s", "columns":%s}`, tableName, configs.createColArray))),
498 | want: configs.ddlWant,
499 | isErr: false,
500 | },
501 | {
502 | name: "invoke insert-table-templateParams-tool",
503 | insert: true,
504 | api: "http://127.0.0.1:5000/api/tool/insert-table-templateParams-tool/invoke",
505 | requestHeader: map[string]string{},
506 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s", "columns":["id","name","age"], "values":"1, 'Alex', 21"}`, tableName))),
507 | want: configs.insert1Want,
508 | isErr: false,
509 | },
510 | {
511 | name: "invoke insert-table-templateParams-tool",
512 | insert: true,
513 | api: "http://127.0.0.1:5000/api/tool/insert-table-templateParams-tool/invoke",
514 | requestHeader: map[string]string{},
515 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s", "columns":["id","name","age"], "values":"2, 'Alice', 100"}`, tableName))),
516 | want: configs.insert1Want,
517 | isErr: false,
518 | },
519 | {
520 | name: "invoke select-templateParams-tool",
521 | api: "http://127.0.0.1:5000/api/tool/select-templateParams-tool/invoke",
522 | requestHeader: map[string]string{},
523 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s"}`, tableName))),
524 | want: configs.selectAllWant,
525 | isErr: false,
526 | },
527 | {
528 | name: "invoke select-templateParams-combined-tool",
529 | api: "http://127.0.0.1:5000/api/tool/select-templateParams-combined-tool/invoke",
530 | requestHeader: map[string]string{},
531 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"id": 1, "tableName": "%s"}`, tableName))),
532 | want: configs.selectId1Want,
533 | isErr: false,
534 | },
535 | {
536 | name: "invoke select-templateParams-combined-tool with no results",
537 | api: "http://127.0.0.1:5000/api/tool/select-templateParams-combined-tool/invoke",
538 | requestHeader: map[string]string{},
539 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"id": 999, "tableName": "%s"}`, tableName))),
540 | want: configs.selectEmptyWant,
541 | isErr: false,
542 | },
543 | {
544 | name: "invoke select-fields-templateParams-tool",
545 | enabled: configs.supportSelectFields,
546 | api: "http://127.0.0.1:5000/api/tool/select-fields-templateParams-tool/invoke",
547 | requestHeader: map[string]string{},
548 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s", "fields":%s}`, tableName, configs.nameFieldArray))),
549 | want: selectOnlyNamesWant,
550 | isErr: false,
551 | },
552 | {
553 | name: "invoke select-filter-templateParams-combined-tool",
554 | api: "http://127.0.0.1:5000/api/tool/select-filter-templateParams-combined-tool/invoke",
555 | requestHeader: map[string]string{},
556 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"name": "Alex", "tableName": "%s", "columnFilter": "%s"}`, tableName, configs.nameColFilter))),
557 | want: configs.selectNameWant,
558 | isErr: false,
559 | },
560 | {
561 | name: "invoke drop-table-templateParams-tool",
562 | ddl: true,
563 | api: "http://127.0.0.1:5000/api/tool/drop-table-templateParams-tool/invoke",
564 | requestHeader: map[string]string{},
565 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"tableName": "%s"}`, tableName))),
566 | want: configs.ddlWant,
567 | isErr: false,
568 | },
569 | }
570 | for _, tc := range invokeTcs {
571 | t.Run(tc.name, func(t *testing.T) {
572 | if !tc.enabled {
573 | return
574 | }
575 | // if test case is DDL and source support ddl test cases
576 | ddlAllow := !tc.ddl || (tc.ddl && configs.supportDdl)
577 | // if test case is insert statement and source support insert test cases
578 | insertAllow := !tc.insert || (tc.insert && configs.supportInsert)
579 | if ddlAllow && insertAllow {
580 | // Send Tool invocation request
581 | resp, respBody := RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
582 | if resp.StatusCode != http.StatusOK {
583 | if tc.isErr {
584 | return
585 | }
586 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBody))
587 | }
588 |
589 | // Check response body
590 | var body map[string]interface{}
591 | err := json.Unmarshal(respBody, &body)
592 | if err != nil {
593 | t.Fatalf("error parsing response body")
594 | }
595 |
596 | got, ok := body["result"].(string)
597 | if !ok {
598 | t.Fatalf("unable to find result in response body")
599 | }
600 |
601 | if got != tc.want {
602 | t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
603 | }
604 | }
605 | })
606 | }
607 | }
608 |
609 | func RunExecuteSqlToolInvokeTest(t *testing.T, createTableStatement, select1Want string, options ...ExecuteSqlOption) {
610 | // Resolve options
611 | // Default values for ExecuteSqlTestConfig
612 | configs := &ExecuteSqlTestConfig{
613 | select1Statement: `"SELECT 1"`,
614 | createWant: "null",
615 | dropWant: "null",
616 | selectEmptyWant: "null",
617 | }
618 |
619 | // Apply provided options
620 | for _, option := range options {
621 | option(configs)
622 | }
623 |
624 | // Get ID token
625 | idToken, err := GetGoogleIdToken(ClientId)
626 | if err != nil {
627 | t.Fatalf("error getting Google ID token: %s", err)
628 | }
629 |
630 | // Test tool invoke endpoint
631 | invokeTcs := []struct {
632 | name string
633 | api string
634 | requestHeader map[string]string
635 | requestBody io.Reader
636 | want string
637 | isErr bool
638 | }{
639 | {
640 | name: "invoke my-exec-sql-tool",
641 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
642 | requestHeader: map[string]string{},
643 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sql": %s}`, configs.select1Statement))),
644 | want: select1Want,
645 | isErr: false,
646 | },
647 | {
648 | name: "invoke my-exec-sql-tool create table",
649 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
650 | requestHeader: map[string]string{},
651 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sql": %s}`, createTableStatement))),
652 | want: configs.createWant,
653 | isErr: false,
654 | },
655 | {
656 | name: "invoke my-exec-sql-tool select table",
657 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
658 | requestHeader: map[string]string{},
659 | requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT * FROM t"}`)),
660 | want: configs.selectEmptyWant,
661 | isErr: false,
662 | },
663 | {
664 | name: "invoke my-exec-sql-tool drop table",
665 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
666 | requestHeader: map[string]string{},
667 | requestBody: bytes.NewBuffer([]byte(`{"sql":"DROP TABLE t"}`)),
668 | want: configs.dropWant,
669 | isErr: false,
670 | },
671 | {
672 | name: "invoke my-exec-sql-tool without body",
673 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
674 | requestHeader: map[string]string{},
675 | requestBody: bytes.NewBuffer([]byte(`{}`)),
676 | isErr: true,
677 | },
678 | {
679 | name: "Invoke my-auth-exec-sql-tool with auth token",
680 | api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke",
681 | requestHeader: map[string]string{"my-google-auth_token": idToken},
682 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sql": %s}`, configs.select1Statement))),
683 | isErr: false,
684 | want: select1Want,
685 | },
686 | {
687 | name: "Invoke my-auth-exec-sql-tool with invalid auth token",
688 | api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke",
689 | requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
690 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sql": %s}`, configs.select1Statement))),
691 | isErr: true,
692 | },
693 | {
694 | name: "Invoke my-auth-exec-sql-tool without auth token",
695 | api: "http://127.0.0.1:5000/api/tool/my-auth-exec-sql-tool/invoke",
696 | requestHeader: map[string]string{},
697 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sql": %s}`, configs.select1Statement))),
698 | isErr: true,
699 | },
700 | {
701 | name: "invoke my-exec-sql-tool with invalid SELECT SQL",
702 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
703 | requestHeader: map[string]string{},
704 | requestBody: bytes.NewBuffer([]byte(`{"sql":"SELECT * FROM non_existent_table"}`)),
705 | isErr: true,
706 | },
707 | {
708 | name: "invoke my-exec-sql-tool with invalid ALTER SQL",
709 | api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
710 | requestHeader: map[string]string{},
711 | requestBody: bytes.NewBuffer([]byte(`{"sql":"ALTER TALE t ALTER COLUMN id DROP NOT NULL"}`)),
712 | isErr: true,
713 | },
714 | }
715 | for _, tc := range invokeTcs {
716 | t.Run(tc.name, func(t *testing.T) {
717 | // Send Tool invocation request
718 | resp, respBody := RunRequest(t, http.MethodPost, tc.api, tc.requestBody, tc.requestHeader)
719 | if resp.StatusCode != http.StatusOK {
720 | if tc.isErr {
721 | return
722 | }
723 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBody))
724 | }
725 |
726 | // Check response body
727 | var body map[string]interface{}
728 | err = json.Unmarshal(respBody, &body)
729 | if err != nil {
730 | t.Fatalf("error parsing response body")
731 | }
732 |
733 | got, ok := body["result"].(string)
734 | if !ok {
735 | t.Fatalf("unable to find result in response body")
736 | }
737 |
738 | if got != tc.want {
739 | t.Fatalf("unexpected value: got %q, want %q", got, tc.want)
740 | }
741 | })
742 | }
743 | }
744 |
745 | // RunInitialize runs the initialize lifecycle for mcp to set up client-server connection
746 | func RunInitialize(t *testing.T, protocolVersion string) string {
747 | url := "http://127.0.0.1:5000/mcp"
748 |
749 | initializeRequestBody := map[string]any{
750 | "jsonrpc": "2.0",
751 | "id": "mcp-initialize",
752 | "method": "initialize",
753 | "params": map[string]any{
754 | "protocolVersion": protocolVersion,
755 | },
756 | }
757 | reqMarshal, err := json.Marshal(initializeRequestBody)
758 | if err != nil {
759 | t.Fatalf("unexpected error during marshaling of body")
760 | }
761 |
762 | resp, _ := RunRequest(t, http.MethodPost, url, bytes.NewBuffer(reqMarshal), nil)
763 | if resp.StatusCode != 200 {
764 | t.Fatalf("response status code is not 200")
765 | }
766 |
767 | if contentType := resp.Header.Get("Content-type"); contentType != "application/json" {
768 | t.Fatalf("unexpected content-type header: want %s, got %s", "application/json", contentType)
769 | }
770 |
771 | sessionId := resp.Header.Get("Mcp-Session-Id")
772 |
773 | header := map[string]string{}
774 | if sessionId != "" {
775 | header["Mcp-Session-Id"] = sessionId
776 | }
777 |
778 | initializeNotificationBody := map[string]any{
779 | "jsonrpc": "2.0",
780 | "method": "notifications/initialized",
781 | }
782 | notiMarshal, err := json.Marshal(initializeNotificationBody)
783 | if err != nil {
784 | t.Fatalf("unexpected error during marshaling of notifications body")
785 | }
786 |
787 | _, _ = RunRequest(t, http.MethodPost, url, bytes.NewBuffer(notiMarshal), header)
788 | return sessionId
789 | }
790 |
791 | // RunMCPToolCallMethod runs the tool/call for mcp endpoint
792 | func RunMCPToolCallMethod(t *testing.T, myFailToolWant, select1Want string, options ...McpTestOption) {
793 | // Resolve options
794 | // Default values for MCPTestConfig
795 | configs := &MCPTestConfig{
796 | myToolId3NameAliceWant: `{"jsonrpc":"2.0","id":"my-tool","result":{"content":[{"type":"text","text":"{\"id\":1,\"name\":\"Alice\"}"},{"type":"text","text":"{\"id\":3,\"name\":\"Sid\"}"}]}}`,
797 | mcpSelect1Want: select1Want,
798 | supportClientAuth: false,
799 | supportSelect1Auth: true,
800 | }
801 |
802 | // Apply provided options
803 | for _, option := range options {
804 | option(configs)
805 | }
806 |
807 | sessionId := RunInitialize(t, "2024-11-05")
808 |
809 | // Get access token
810 | accessToken, err := sources.GetIAMAccessToken(t.Context())
811 | if err != nil {
812 | t.Fatalf("error getting access token from ADC: %s", err)
813 | }
814 | accessToken = "Bearer " + accessToken
815 |
816 | idToken, err := GetGoogleIdToken(ClientId)
817 | if err != nil {
818 | t.Fatalf("error getting Google ID token: %s", err)
819 | }
820 |
821 | // Test tool invoke endpoint
822 | invokeTcs := []struct {
823 | name string
824 | api string
825 | enabled bool // switch to turn on/off the test case
826 | requestBody jsonrpc.JSONRPCRequest
827 | requestHeader map[string]string
828 | wantStatusCode int
829 | wantBody string
830 | }{
831 | {
832 | name: "MCP Invoke my-tool",
833 | api: "http://127.0.0.1:5000/mcp",
834 | enabled: true,
835 | requestHeader: map[string]string{},
836 | requestBody: jsonrpc.JSONRPCRequest{
837 | Jsonrpc: "2.0",
838 | Id: "my-tool",
839 | Request: jsonrpc.Request{
840 | Method: "tools/call",
841 | },
842 | Params: map[string]any{
843 | "name": "my-tool",
844 | "arguments": map[string]any{
845 | "id": int(3),
846 | "name": "Alice",
847 | },
848 | },
849 | },
850 | wantStatusCode: http.StatusOK,
851 | wantBody: configs.myToolId3NameAliceWant,
852 | },
853 | {
854 | name: "MCP Invoke invalid tool",
855 | api: "http://127.0.0.1:5000/mcp",
856 | enabled: true,
857 | requestHeader: map[string]string{},
858 | requestBody: jsonrpc.JSONRPCRequest{
859 | Jsonrpc: "2.0",
860 | Id: "invalid-tool",
861 | Request: jsonrpc.Request{
862 | Method: "tools/call",
863 | },
864 | Params: map[string]any{
865 | "name": "foo",
866 | "arguments": map[string]any{},
867 | },
868 | },
869 | wantStatusCode: http.StatusOK,
870 | wantBody: `{"jsonrpc":"2.0","id":"invalid-tool","error":{"code":-32602,"message":"invalid tool name: tool with name \"foo\" does not exist"}}`,
871 | },
872 | {
873 | name: "MCP Invoke my-tool without parameters",
874 | api: "http://127.0.0.1:5000/mcp",
875 | enabled: true,
876 | requestHeader: map[string]string{},
877 | requestBody: jsonrpc.JSONRPCRequest{
878 | Jsonrpc: "2.0",
879 | Id: "invoke-without-parameter",
880 | Request: jsonrpc.Request{
881 | Method: "tools/call",
882 | },
883 | Params: map[string]any{
884 | "name": "my-tool",
885 | "arguments": map[string]any{},
886 | },
887 | },
888 | wantStatusCode: http.StatusOK,
889 | wantBody: `{"jsonrpc":"2.0","id":"invoke-without-parameter","error":{"code":-32602,"message":"provided parameters were invalid: parameter \"id\" is required"}}`,
890 | },
891 | {
892 | name: "MCP Invoke my-tool with insufficient parameters",
893 | api: "http://127.0.0.1:5000/mcp",
894 | enabled: true,
895 | requestHeader: map[string]string{},
896 | requestBody: jsonrpc.JSONRPCRequest{
897 | Jsonrpc: "2.0",
898 | Id: "invoke-insufficient-parameter",
899 | Request: jsonrpc.Request{
900 | Method: "tools/call",
901 | },
902 | Params: map[string]any{
903 | "name": "my-tool",
904 | "arguments": map[string]any{"id": 1},
905 | },
906 | },
907 | wantStatusCode: http.StatusOK,
908 | wantBody: `{"jsonrpc":"2.0","id":"invoke-insufficient-parameter","error":{"code":-32602,"message":"provided parameters were invalid: parameter \"name\" is required"}}`,
909 | },
910 | {
911 | name: "MCP Invoke my-auth-required-tool",
912 | api: "http://127.0.0.1:5000/mcp",
913 | enabled: configs.supportSelect1Auth,
914 | requestHeader: map[string]string{"my-google-auth_token": idToken},
915 | requestBody: jsonrpc.JSONRPCRequest{
916 | Jsonrpc: "2.0",
917 | Id: "invoke my-auth-required-tool",
918 | Request: jsonrpc.Request{
919 | Method: "tools/call",
920 | },
921 | Params: map[string]any{
922 | "name": "my-auth-required-tool",
923 | "arguments": map[string]any{},
924 | },
925 | },
926 | wantStatusCode: http.StatusOK,
927 | wantBody: configs.mcpSelect1Want,
928 | },
929 | {
930 | name: "MCP Invoke my-auth-required-tool with invalid auth token",
931 | api: "http://127.0.0.1:5000/mcp",
932 | requestHeader: map[string]string{"my-google-auth_token": "INVALID_TOKEN"},
933 | requestBody: jsonrpc.JSONRPCRequest{
934 | Jsonrpc: "2.0",
935 | Id: "invoke my-auth-required-tool with invalid token",
936 | Request: jsonrpc.Request{
937 | Method: "tools/call",
938 | },
939 | Params: map[string]any{
940 | "name": "my-auth-required-tool",
941 | "arguments": map[string]any{},
942 | },
943 | },
944 | wantStatusCode: http.StatusUnauthorized,
945 | wantBody: "{\"jsonrpc\":\"2.0\",\"id\":\"invoke my-auth-required-tool with invalid token\",\"error\":{\"code\":-32600,\"message\":\"unauthorized Tool call: Please make sure your specify correct auth headers: unauthorized\"}}",
946 | },
947 | {
948 | name: "MCP Invoke my-auth-required-tool without auth token",
949 | api: "http://127.0.0.1:5000/mcp",
950 | requestHeader: map[string]string{},
951 | requestBody: jsonrpc.JSONRPCRequest{
952 | Jsonrpc: "2.0",
953 | Id: "invoke my-auth-required-tool without token",
954 | Request: jsonrpc.Request{
955 | Method: "tools/call",
956 | },
957 | Params: map[string]any{
958 | "name": "my-auth-required-tool",
959 | "arguments": map[string]any{},
960 | },
961 | },
962 | wantStatusCode: http.StatusUnauthorized,
963 | wantBody: "{\"jsonrpc\":\"2.0\",\"id\":\"invoke my-auth-required-tool without token\",\"error\":{\"code\":-32600,\"message\":\"unauthorized Tool call: Please make sure your specify correct auth headers: unauthorized\"}}",
964 | },
965 |
966 | {
967 | name: "MCP Invoke my-client-auth-tool",
968 | enabled: configs.supportClientAuth,
969 | api: "http://127.0.0.1:5000/mcp",
970 | requestHeader: map[string]string{"Authorization": accessToken},
971 | requestBody: jsonrpc.JSONRPCRequest{
972 | Jsonrpc: "2.0",
973 | Id: "invoke my-client-auth-tool",
974 | Request: jsonrpc.Request{
975 | Method: "tools/call",
976 | },
977 | Params: map[string]any{
978 | "name": "my-client-auth-tool",
979 | "arguments": map[string]any{},
980 | },
981 | },
982 | wantStatusCode: http.StatusOK,
983 | wantBody: "{\"jsonrpc\":\"2.0\",\"id\":\"invoke my-client-auth-tool\",\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"{\\\"f0_\\\":1}\"}]}}",
984 | },
985 | {
986 | name: "MCP Invoke my-client-auth-tool without access token",
987 | enabled: configs.supportClientAuth,
988 | api: "http://127.0.0.1:5000/mcp",
989 | requestHeader: map[string]string{},
990 | requestBody: jsonrpc.JSONRPCRequest{
991 | Jsonrpc: "2.0",
992 | Id: "invoke my-client-auth-tool",
993 | Request: jsonrpc.Request{
994 | Method: "tools/call",
995 | },
996 | Params: map[string]any{
997 | "name": "my-client-auth-tool",
998 | "arguments": map[string]any{},
999 | },
1000 | },
1001 | wantStatusCode: http.StatusUnauthorized,
1002 | wantBody: "{\"jsonrpc\":\"2.0\",\"id\":\"invoke my-client-auth-tool\",\"error\":{\"code\":-32600,\"message\":\"missing access token in the 'Authorization' header\"}",
1003 | },
1004 | {
1005 | name: "MCP Invoke my-client-auth-tool with invalid access token",
1006 | enabled: configs.supportClientAuth,
1007 | api: "http://127.0.0.1:5000/mcp",
1008 | requestHeader: map[string]string{"Authorization": "Bearer invalid-token"},
1009 | requestBody: jsonrpc.JSONRPCRequest{
1010 | Jsonrpc: "2.0",
1011 | Id: "invoke my-client-auth-tool",
1012 | Request: jsonrpc.Request{
1013 | Method: "tools/call",
1014 | },
1015 | Params: map[string]any{
1016 | "name": "my-client-auth-tool",
1017 | "arguments": map[string]any{},
1018 | },
1019 | },
1020 | wantStatusCode: http.StatusUnauthorized,
1021 | },
1022 | {
1023 | name: "MCP Invoke my-fail-tool",
1024 | api: "http://127.0.0.1:5000/mcp",
1025 | enabled: true,
1026 | requestHeader: map[string]string{},
1027 | requestBody: jsonrpc.JSONRPCRequest{
1028 | Jsonrpc: "2.0",
1029 | Id: "invoke-fail-tool",
1030 | Request: jsonrpc.Request{
1031 | Method: "tools/call",
1032 | },
1033 | Params: map[string]any{
1034 | "name": "my-fail-tool",
1035 | "arguments": map[string]any{"id": 1},
1036 | },
1037 | },
1038 | wantStatusCode: http.StatusOK,
1039 | wantBody: myFailToolWant,
1040 | },
1041 | }
1042 | for _, tc := range invokeTcs {
1043 | t.Run(tc.name, func(t *testing.T) {
1044 | if !tc.enabled {
1045 | return
1046 | }
1047 | reqMarshal, err := json.Marshal(tc.requestBody)
1048 | if err != nil {
1049 | t.Fatalf("unexpected error during marshaling of request body")
1050 | }
1051 |
1052 | // add headers
1053 | headers := map[string]string{}
1054 | if sessionId != "" {
1055 | headers["Mcp-Session-Id"] = sessionId
1056 | }
1057 | for key, value := range tc.requestHeader {
1058 | headers[key] = value
1059 | }
1060 |
1061 | httpResponse, respBody := RunRequest(t, http.MethodPost, tc.api, bytes.NewBuffer(reqMarshal), headers)
1062 |
1063 | // Check status code
1064 | if httpResponse.StatusCode != tc.wantStatusCode {
1065 | t.Errorf("StatusCode mismatch: got %d, want %d", httpResponse.StatusCode, tc.wantStatusCode)
1066 | }
1067 |
1068 | // Check response body
1069 | got := string(bytes.TrimSpace(respBody))
1070 | if !strings.Contains(got, tc.wantBody) {
1071 | t.Fatalf("Expected substring not found:\ngot: %q\nwant: %q (to be contained within got)", got, tc.wantBody)
1072 | }
1073 | })
1074 | }
1075 | }
1076 |
1077 | func setupPostgresSchemas(t *testing.T, ctx context.Context, pool *pgxpool.Pool, schemaName string) func() {
1078 | createSchemaStmt := fmt.Sprintf("CREATE SCHEMA %s", schemaName)
1079 | _, err := pool.Exec(ctx, createSchemaStmt)
1080 | if err != nil {
1081 | t.Fatalf("failed to create schema: %v", err)
1082 | }
1083 |
1084 | return func() {
1085 | dropSchemaStmt := fmt.Sprintf("DROP SCHEMA %s CASCADE", schemaName)
1086 | _, err := pool.Exec(ctx, dropSchemaStmt)
1087 | if err != nil {
1088 | t.Fatalf("failed to drop schema: %v", err)
1089 | }
1090 | }
1091 | }
1092 |
1093 | func RunPostgresListTablesTest(t *testing.T, tableNameParam, tableNameAuth, user string) {
1094 | // TableNameParam columns to construct want
1095 | paramTableColumns := fmt.Sprintf(`[
1096 | {"data_type": "integer", "column_name": "id", "column_default": "nextval('%s_id_seq'::regclass)", "is_not_nullable": true, "ordinal_position": 1, "column_comment": null},
1097 | {"data_type": "text", "column_name": "name", "column_default": null, "is_not_nullable": false, "ordinal_position": 2, "column_comment": null}
1098 | ]`, tableNameParam)
1099 |
1100 | // TableNameAuth columns to construct want
1101 | authTableColumns := fmt.Sprintf(`[
1102 | {"data_type": "integer", "column_name": "id", "column_default": "nextval('%s_id_seq'::regclass)", "is_not_nullable": true, "ordinal_position": 1, "column_comment": null},
1103 | {"data_type": "text", "column_name": "name", "column_default": null, "is_not_nullable": false, "ordinal_position": 2, "column_comment": null},
1104 | {"data_type": "text", "column_name": "email", "column_default": null, "is_not_nullable": false, "ordinal_position": 3, "column_comment": null}
1105 | ]`, tableNameAuth)
1106 |
1107 | const (
1108 | // Template to construct detailed output want
1109 | detailedObjectTemplate = `{
1110 | "object_name": "%[1]s", "schema_name": "public",
1111 | "object_details": {
1112 | "owner": "%[3]s", "comment": null,
1113 | "indexes": [{"is_primary": true, "is_unique": true, "index_name": "%[1]s_pkey", "index_method": "btree", "index_columns": ["id"], "index_definition": "CREATE UNIQUE INDEX %[1]s_pkey ON public.%[1]s USING btree (id)"}],
1114 | "triggers": [], "columns": %[2]s, "object_name": "%[1]s", "object_type": "TABLE", "schema_name": "public",
1115 | "constraints": [{"constraint_name": "%[1]s_pkey", "constraint_type": "PRIMARY KEY", "constraint_columns": ["id"], "constraint_definition": "PRIMARY KEY (id)", "foreign_key_referenced_table": null, "foreign_key_referenced_columns": null}]
1116 | }
1117 | }`
1118 |
1119 | // Template to construct simple output want
1120 | simpleObjectTemplate = `{"object_name":"%s", "schema_name":"public", "object_details":{"name":"%s"}}`
1121 | )
1122 |
1123 | // Helper to build json for detailed want
1124 | getDetailedWant := func(tableName, columnJSON string) string {
1125 | return fmt.Sprintf(detailedObjectTemplate, tableName, columnJSON, user)
1126 | }
1127 |
1128 | // Helper to build template for simple want
1129 | getSimpleWant := func(tableName string) string {
1130 | return fmt.Sprintf(simpleObjectTemplate, tableName, tableName)
1131 | }
1132 |
1133 | invokeTcs := []struct {
1134 | name string
1135 | api string
1136 | requestBody io.Reader
1137 | wantStatusCode int
1138 | want string
1139 | isAllTables bool
1140 | }{
1141 | {
1142 | name: "invoke list_tables all tables detailed output",
1143 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
1144 | requestBody: bytes.NewBuffer([]byte(`{"table_names": ""}`)),
1145 | wantStatusCode: http.StatusOK,
1146 | want: fmt.Sprintf("[%s,%s]", getDetailedWant(tableNameAuth, authTableColumns), getDetailedWant(tableNameParam, paramTableColumns)),
1147 | isAllTables: true,
1148 | },
1149 | {
1150 | name: "invoke list_tables all tables simple output",
1151 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
1152 | requestBody: bytes.NewBuffer([]byte(`{"table_names": "", "output_format": "simple"}`)),
1153 | wantStatusCode: http.StatusOK,
1154 | want: fmt.Sprintf("[%s,%s]", getSimpleWant(tableNameAuth), getSimpleWant(tableNameParam)),
1155 | isAllTables: true,
1156 | },
1157 | {
1158 | name: "invoke list_tables detailed output",
1159 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
1160 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"table_names": "%s"}`, tableNameAuth))),
1161 | wantStatusCode: http.StatusOK,
1162 | want: fmt.Sprintf("[%s]", getDetailedWant(tableNameAuth, authTableColumns)),
1163 | },
1164 | {
1165 | name: "invoke list_tables simple output",
1166 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
1167 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"table_names": "%s", "output_format": "simple"}`, tableNameAuth))),
1168 | wantStatusCode: http.StatusOK,
1169 | want: fmt.Sprintf("[%s]", getSimpleWant(tableNameAuth)),
1170 | },
1171 | {
1172 | name: "invoke list_tables with invalid output format",
1173 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
1174 | requestBody: bytes.NewBuffer([]byte(`{"table_names": "", "output_format": "abcd"}`)),
1175 | wantStatusCode: http.StatusBadRequest,
1176 | },
1177 | {
1178 | name: "invoke list_tables with malformed table_names parameter",
1179 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
1180 | requestBody: bytes.NewBuffer([]byte(`{"table_names": 12345, "output_format": "detailed"}`)),
1181 | wantStatusCode: http.StatusBadRequest,
1182 | },
1183 | {
1184 | name: "invoke list_tables with multiple table names",
1185 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
1186 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"table_names": "%s,%s"}`, tableNameParam, tableNameAuth))),
1187 | wantStatusCode: http.StatusOK,
1188 | want: fmt.Sprintf("[%s,%s]", getDetailedWant(tableNameAuth, authTableColumns), getDetailedWant(tableNameParam, paramTableColumns)),
1189 | },
1190 | {
1191 | name: "invoke list_tables with non-existent table",
1192 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
1193 | requestBody: bytes.NewBuffer([]byte(`{"table_names": "non_existent_table"}`)),
1194 | wantStatusCode: http.StatusOK,
1195 | want: `[]`,
1196 | },
1197 | {
1198 | name: "invoke list_tables with one existing and one non-existent table",
1199 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
1200 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"table_names": "%s,non_existent_table"}`, tableNameParam))),
1201 | wantStatusCode: http.StatusOK,
1202 | want: fmt.Sprintf("[%s]", getDetailedWant(tableNameParam, paramTableColumns)),
1203 | },
1204 | }
1205 | for _, tc := range invokeTcs {
1206 | t.Run(tc.name, func(t *testing.T) {
1207 | resp, respBytes := RunRequest(t, http.MethodPost, tc.api, tc.requestBody, nil)
1208 | if resp.StatusCode != tc.wantStatusCode {
1209 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBytes))
1210 | }
1211 |
1212 | if tc.wantStatusCode == http.StatusOK {
1213 | var bodyWrapper map[string]json.RawMessage
1214 |
1215 | if err := json.Unmarshal(respBytes, &bodyWrapper); err != nil {
1216 | t.Fatalf("error parsing response wrapper: %s, body: %s", err, string(respBytes))
1217 | }
1218 |
1219 | resultJSON, ok := bodyWrapper["result"]
1220 | if !ok {
1221 | t.Fatal("unable to find 'result' in response body")
1222 | }
1223 |
1224 | var resultString string
1225 | if err := json.Unmarshal(resultJSON, &resultString); err != nil {
1226 | t.Fatalf("'result' is not a JSON-encoded string: %s", err)
1227 | }
1228 |
1229 | var got, want []any
1230 |
1231 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
1232 | t.Fatalf("failed to unmarshal actual result string: %v", err)
1233 | }
1234 | if err := json.Unmarshal([]byte(tc.want), &want); err != nil {
1235 | t.Fatalf("failed to unmarshal expected want string: %v", err)
1236 | }
1237 |
1238 | // Checking only the default public schema where the test tables are created to avoid brittle tests.
1239 | if tc.isAllTables {
1240 | var filteredGot []any
1241 | for _, item := range got {
1242 | if tableMap, ok := item.(map[string]interface{}); ok {
1243 | if schema, ok := tableMap["schema_name"]; ok && schema == "public" {
1244 | filteredGot = append(filteredGot, item)
1245 | }
1246 | }
1247 | }
1248 | got = filteredGot
1249 | }
1250 |
1251 | sort.SliceStable(got, func(i, j int) bool {
1252 | return fmt.Sprintf("%v", got[i]) < fmt.Sprintf("%v", got[j])
1253 | })
1254 | sort.SliceStable(want, func(i, j int) bool {
1255 | return fmt.Sprintf("%v", want[i]) < fmt.Sprintf("%v", want[j])
1256 | })
1257 |
1258 | if !reflect.DeepEqual(got, want) {
1259 | t.Errorf("Unexpected result: got %#v, want: %#v", got, want)
1260 | }
1261 | }
1262 | })
1263 | }
1264 | }
1265 |
1266 | func setUpPostgresViews(t *testing.T, ctx context.Context, pool *pgxpool.Pool, viewName string) func() {
1267 | createView := fmt.Sprintf("CREATE VIEW %s AS SELECT 1 AS col", viewName)
1268 | _, err := pool.Exec(ctx, createView)
1269 | if err != nil {
1270 | t.Fatalf("failed to create view: %v", err)
1271 | }
1272 | return func() {
1273 | dropView := fmt.Sprintf("DROP VIEW %s", viewName)
1274 | _, err := pool.Exec(ctx, dropView)
1275 | if err != nil {
1276 | t.Fatalf("failed to drop view: %v", err)
1277 | }
1278 | }
1279 | }
1280 |
1281 | func RunPostgresListViewsTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
1282 | //adding this line temporarily
1283 | viewName := "test_view_" + strings.ReplaceAll(uuid.New().String(), "-", "")
1284 | dropViewfunc1 := setUpPostgresViews(t, ctx, pool, viewName)
1285 | defer dropViewfunc1()
1286 |
1287 | invokeTcs := []struct {
1288 | name string
1289 | requestBody io.Reader
1290 | wantStatusCode int
1291 | want string
1292 | }{
1293 | {
1294 | name: "invoke list_views with newly created view",
1295 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"view_name": "%s"}`, viewName))),
1296 | wantStatusCode: http.StatusOK,
1297 | want: fmt.Sprintf(`[{"schema_name":"public","view_name":"%s","owner_name":"postgres","definition":" SELECT 1 AS col;"}]`, viewName),
1298 | },
1299 | {
1300 | name: "invoke list_views with non-existent_view",
1301 | requestBody: bytes.NewBuffer([]byte(`{"view_name": "non_existent_view"}`)),
1302 | wantStatusCode: http.StatusOK,
1303 | want: `null`,
1304 | },
1305 | }
1306 | for _, tc := range invokeTcs {
1307 | t.Run(tc.name, func(t *testing.T) {
1308 | const api = "http://127.0.0.1:5000/api/tool/list_views/invoke"
1309 | resp, body := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
1310 |
1311 | if resp.StatusCode != tc.wantStatusCode {
1312 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(body))
1313 | }
1314 | if tc.wantStatusCode != http.StatusOK {
1315 | return
1316 | }
1317 |
1318 | var bodyWrapper struct {
1319 | Result json.RawMessage `json:"result"`
1320 | }
1321 | if err := json.Unmarshal(body, &bodyWrapper); err != nil {
1322 | t.Fatalf("error decoding response wrapper: %v", err)
1323 | }
1324 |
1325 | var resultString string
1326 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
1327 | resultString = string(bodyWrapper.Result)
1328 | }
1329 |
1330 | var got, want any
1331 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
1332 | t.Fatalf("failed to unmarshal nested result string: %v", err)
1333 | }
1334 | if err := json.Unmarshal([]byte(tc.want), &want); err != nil {
1335 | t.Fatalf("failed to unmarshal want string: %v", err)
1336 | }
1337 |
1338 | if diff := cmp.Diff(want, got); diff != "" {
1339 | t.Errorf("Unexpected result (-want +got):\n%s", diff)
1340 | }
1341 | })
1342 | }
1343 | }
1344 |
1345 | func RunPostgresListSchemasTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
1346 | schemaName := "test_schema_" + strings.ReplaceAll(uuid.New().String(), "-", "")
1347 | cleanup := setupPostgresSchemas(t, ctx, pool, schemaName)
1348 | defer cleanup()
1349 |
1350 | wantSchema := map[string]any{"functions": float64(0), "grants": map[string]any{}, "owner": "postgres", "schema_name": schemaName, "tables": float64(0), "views": float64(0)}
1351 |
1352 | invokeTcs := []struct {
1353 | name string
1354 | requestBody io.Reader
1355 | wantStatusCode int
1356 | want []map[string]any
1357 | compareSubset bool
1358 | }{
1359 | {
1360 | name: "invoke list_schemas with schema_name",
1361 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"schema_name": "%s"}`, schemaName))),
1362 | wantStatusCode: http.StatusOK,
1363 | want: []map[string]any{wantSchema},
1364 | },
1365 | {
1366 | name: "invoke list_schemas with owner name",
1367 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"owner": "%s"}`, "postgres"))),
1368 | wantStatusCode: http.StatusOK,
1369 | want: []map[string]any{wantSchema},
1370 | compareSubset: true,
1371 | },
1372 | {
1373 | name: "invoke list_schemas with limit 1",
1374 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"schema_name": "%s","limit": 1}`, schemaName))),
1375 | wantStatusCode: http.StatusOK,
1376 | want: []map[string]any{wantSchema},
1377 | },
1378 | {
1379 | name: "invoke list_schemas with non-existent schema",
1380 | requestBody: bytes.NewBuffer([]byte(`{"schema_name": "non_existent_schema"}`)),
1381 | wantStatusCode: http.StatusOK,
1382 | want: nil,
1383 | },
1384 | }
1385 | for _, tc := range invokeTcs {
1386 | t.Run(tc.name, func(t *testing.T) {
1387 | const api = "http://127.0.0.1:5000/api/tool/list_schemas/invoke"
1388 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
1389 | if resp.StatusCode != tc.wantStatusCode {
1390 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
1391 | }
1392 | if tc.wantStatusCode != http.StatusOK {
1393 | return
1394 | }
1395 |
1396 | var bodyWrapper struct {
1397 | Result json.RawMessage `json:"result"`
1398 | }
1399 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
1400 | t.Fatalf("error decoding response wrapper: %v", err)
1401 | }
1402 |
1403 | var resultString string
1404 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
1405 | resultString = string(bodyWrapper.Result)
1406 | }
1407 |
1408 | var got []map[string]any
1409 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
1410 | t.Fatalf("failed to unmarshal nested result string: %v", err)
1411 | }
1412 |
1413 | if tc.compareSubset {
1414 | // Assert that the 'wantTrigger' is present in the 'got' list.
1415 | found := false
1416 | for _, resultSchema := range got {
1417 | if resultSchema["schema_name"] == wantSchema["schema_name"] {
1418 | found = true
1419 | if diff := cmp.Diff(wantSchema, resultSchema); diff != "" {
1420 | t.Errorf("Mismatch in fields for the expected trigger (-want +got):\n%s", diff)
1421 | }
1422 | break
1423 | }
1424 | }
1425 | if !found {
1426 | t.Errorf("Expected schema '%s' not found in the list of all schemas.", wantSchema)
1427 | }
1428 | } else {
1429 | if diff := cmp.Diff(tc.want, got); diff != "" {
1430 | t.Errorf("Unexpected result (-want +got):\n%s", diff)
1431 | }
1432 | }
1433 | })
1434 | }
1435 | }
1436 |
1437 | func RunPostgresDatabaseOverviewTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
1438 | const api = "http://127.0.0.1:5000/api/tool/database_overview/invoke"
1439 | requestBody := bytes.NewBuffer([]byte(`{}`))
1440 |
1441 | resp, respBody := RunRequest(t, http.MethodPost, api, requestBody, nil)
1442 |
1443 | if resp.StatusCode != http.StatusOK {
1444 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, http.StatusOK, string(respBody))
1445 | }
1446 |
1447 | var bodyWrapper struct {
1448 | Result json.RawMessage `json:"result"`
1449 | }
1450 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
1451 | t.Fatalf("error decoding response wrapper: %v, body: %s", err, string(respBody))
1452 | }
1453 |
1454 | var resultString string
1455 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
1456 | resultString = string(bodyWrapper.Result)
1457 | }
1458 |
1459 | var got []map[string]any
1460 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
1461 | t.Fatalf("failed to unmarshal nested result string: %v, result string: %s", err, resultString)
1462 | }
1463 |
1464 | if len(got) != 1 {
1465 | t.Fatalf("Expected exactly one row in the result, got %d", len(got))
1466 | }
1467 |
1468 | resultRow := got[0]
1469 |
1470 | // Define expected keys based on the SELECT statement
1471 | expectedKeys := []string{
1472 | "pg_version",
1473 | "is_replica",
1474 | "uptime",
1475 | "max_connections",
1476 | "current_connections",
1477 | "active_connections",
1478 | "pct_connections_used",
1479 | }
1480 |
1481 | for _, key := range expectedKeys {
1482 | if _, ok := resultRow[key]; !ok {
1483 | t.Errorf("Missing expected key in result: %s", key)
1484 | }
1485 | }
1486 |
1487 | // Check types of the fields. JSON numbers are unmarshalled into float64.
1488 | if _, ok := resultRow["pg_version"].(string); !ok {
1489 | t.Errorf("Expected 'pg_version' to be a string, got %T", resultRow["pg_version"])
1490 | }
1491 | if _, ok := resultRow["is_replica"].(bool); !ok {
1492 | t.Errorf("Expected 'is_replica' to be a bool, got %T", resultRow["is_replica"])
1493 | }
1494 | if _, ok := resultRow["uptime"].(string); !ok {
1495 | t.Errorf("Expected 'uptime' to be a string, got %T", resultRow["uptime"])
1496 | }
1497 | if _, ok := resultRow["max_connections"].(float64); !ok {
1498 | t.Errorf("Expected 'max_connections' to be a number (float64), got %T", resultRow["max_connections"])
1499 | }
1500 | if _, ok := resultRow["current_connections"].(float64); !ok {
1501 | t.Errorf("Expected 'current_connections' to be a number (float64), got %T", resultRow["current_connections"])
1502 | }
1503 | if _, ok := resultRow["active_connections"].(float64); !ok {
1504 | t.Errorf("Expected 'active_connections' to be a number (float64), got %T", resultRow["active_connections"])
1505 | }
1506 | if _, ok := resultRow["pct_connections_used"].(float64); !ok {
1507 | t.Errorf("Expected 'pct_connections_used' to be a number (float64), got %T", resultRow["pct_connections_used"])
1508 | }
1509 |
1510 | // Basic sanity checks on values
1511 | if maxConn, ok := resultRow["max_connections"].(float64); ok {
1512 | if maxConn <= 0 {
1513 | t.Errorf("Expected 'max_connections' to be positive, got %f", maxConn)
1514 | }
1515 | }
1516 |
1517 | if pctUsed, ok := resultRow["pct_connections_used"].(float64); ok {
1518 | if pctUsed < 0 || pctUsed > 100 {
1519 | t.Errorf("Expected 'pct_connections_used' to be between 0 and 100, got %f", pctUsed)
1520 | }
1521 | }
1522 | }
1523 |
1524 | func setupPostgresTrigger(t *testing.T, ctx context.Context, pool *pgxpool.Pool, schemaName, tableName, functionName, triggerName string) func() {
1525 | t.Helper()
1526 |
1527 | createSchemaStmt := fmt.Sprintf("CREATE SCHEMA %s", schemaName)
1528 | if _, err := pool.Exec(ctx, createSchemaStmt); err != nil {
1529 | t.Fatalf("failed to create schema %s: %v", schemaName, err)
1530 | }
1531 |
1532 | createTableStmt := fmt.Sprintf("CREATE TABLE %s.%s (id SERIAL PRIMARY KEY, name TEXT)", schemaName, tableName)
1533 | if _, err := pool.Exec(ctx, createTableStmt); err != nil {
1534 | t.Fatalf("failed to create table %s.%s: %v", schemaName, tableName, err)
1535 | }
1536 |
1537 | createFunctionStmt := fmt.Sprintf(`
1538 | CREATE OR REPLACE FUNCTION %s.%s() RETURNS TRIGGER AS $$
1539 | BEGIN
1540 | RETURN NEW;
1541 | END;
1542 | $$ LANGUAGE plpgsql;
1543 | `, schemaName, functionName)
1544 | if _, err := pool.Exec(ctx, createFunctionStmt); err != nil {
1545 | t.Fatalf("failed to create function %s.%s: %v", schemaName, functionName, err)
1546 | }
1547 |
1548 | createTriggerStmt := fmt.Sprintf(`
1549 | CREATE TRIGGER %s
1550 | AFTER INSERT ON %s.%s
1551 | FOR EACH ROW
1552 | EXECUTE FUNCTION %s.%s();
1553 | `, triggerName, schemaName, tableName, schemaName, functionName)
1554 | if _, err := pool.Exec(ctx, createTriggerStmt); err != nil {
1555 | t.Fatalf("failed to create trigger %s: %v", triggerName, err)
1556 | }
1557 |
1558 | return func() {
1559 | dropSchemaStmt := fmt.Sprintf("DROP SCHEMA %s CASCADE", schemaName)
1560 | if _, err := pool.Exec(ctx, dropSchemaStmt); err != nil {
1561 | t.Fatalf("failed to drop schema %s: %v", schemaName, err)
1562 | }
1563 | }
1564 | }
1565 |
1566 | func RunPostgresListTriggersTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
1567 | uniqueID := strings.ReplaceAll(uuid.New().String(), "-", "")
1568 | schemaName := "test_schema_" + uniqueID
1569 | tableName := "test_table_" + uniqueID
1570 | functionName := "test_func_" + uniqueID
1571 | triggerName := "test_trigger_" + uniqueID
1572 |
1573 | cleanup := setupPostgresTrigger(t, ctx, pool, schemaName, tableName, functionName, triggerName)
1574 | defer cleanup()
1575 |
1576 | // Definition can vary slightly based on server version/settings, so we fetch it to compare.
1577 | var expectedDef string
1578 | getDefQuery := fmt.Sprintf("SELECT pg_get_triggerdef(oid) FROM pg_trigger WHERE tgname = '%s'", triggerName)
1579 | err := pool.QueryRow(ctx, getDefQuery).Scan(&expectedDef)
1580 | if err != nil {
1581 | t.Fatalf("failed to fetch trigger definition: %v", err)
1582 | }
1583 |
1584 | wantTrigger := map[string]any{
1585 | "trigger_name": triggerName,
1586 | "schema_name": schemaName,
1587 | "table_name": tableName,
1588 | "status": "ENABLED",
1589 | "timing": "AFTER",
1590 | "events": "INSERT",
1591 | "activation_level": "ROW",
1592 | "function_name": functionName,
1593 | "definition": expectedDef,
1594 | }
1595 |
1596 | invokeTcs := []struct {
1597 | name string
1598 | requestBody io.Reader
1599 | wantStatusCode int
1600 | want []map[string]any
1601 | compareSubset bool
1602 | }{
1603 | {
1604 | name: "list all triggers (expecting the one we created)",
1605 | requestBody: bytes.NewBuffer([]byte(`{}`)),
1606 | wantStatusCode: http.StatusOK,
1607 | want: []map[string]any{wantTrigger},
1608 | compareSubset: true, // avoid test flakiness in race condition
1609 | },
1610 | {
1611 | name: "filter by trigger_name",
1612 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"trigger_name": "%s"}`, triggerName))),
1613 | wantStatusCode: http.StatusOK,
1614 | want: []map[string]any{wantTrigger},
1615 | },
1616 | {
1617 | name: "filter by schema_name",
1618 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"schema_name": "%s"}`, schemaName))),
1619 | wantStatusCode: http.StatusOK,
1620 | want: []map[string]any{wantTrigger},
1621 | },
1622 | {
1623 | name: "filter by table_name",
1624 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"table_name": "%s"}`, tableName))),
1625 | wantStatusCode: http.StatusOK,
1626 | want: []map[string]any{wantTrigger},
1627 | },
1628 | {
1629 | name: "filter by non-existent trigger_name",
1630 | requestBody: bytes.NewBuffer([]byte(`{"trigger_name": "non_existent_trigger"}`)),
1631 | wantStatusCode: http.StatusOK,
1632 | want: nil,
1633 | },
1634 | {
1635 | name: "filter by non-existent schema_name",
1636 | requestBody: bytes.NewBuffer([]byte(`{"schema_name": "non_existent_schema"}`)),
1637 | wantStatusCode: http.StatusOK,
1638 | want: nil,
1639 | },
1640 | {
1641 | name: "filter by non-existent table_name",
1642 | requestBody: bytes.NewBuffer([]byte(`{"table_name": "non_existent_table"}`)),
1643 | wantStatusCode: http.StatusOK,
1644 | want: nil,
1645 | },
1646 | }
1647 | for _, tc := range invokeTcs {
1648 | t.Run(tc.name, func(t *testing.T) {
1649 | const api = "http://127.0.0.1:5000/api/tool/list_triggers/invoke"
1650 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
1651 | if resp.StatusCode != tc.wantStatusCode {
1652 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
1653 | }
1654 | if tc.wantStatusCode != http.StatusOK {
1655 | return
1656 | }
1657 |
1658 | var bodyWrapper struct {
1659 | Result json.RawMessage `json:"result"`
1660 | }
1661 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
1662 | t.Fatalf("error decoding response wrapper: %v", err)
1663 | }
1664 |
1665 | var resultString string
1666 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
1667 | resultString = string(bodyWrapper.Result)
1668 | }
1669 |
1670 | var got []map[string]any
1671 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
1672 | t.Fatalf("failed to unmarshal nested result string: %v, content: %s", err, resultString)
1673 | }
1674 |
1675 | if tc.compareSubset {
1676 | // Assert that the 'wantTrigger' is present in the 'got' list.
1677 | found := false
1678 | for _, resultTrigger := range got {
1679 | if resultTrigger["trigger_name"] == wantTrigger["trigger_name"] {
1680 | found = true
1681 | if diff := cmp.Diff(wantTrigger, resultTrigger); diff != "" {
1682 | t.Errorf("Mismatch in fields for the expected trigger (-want +got):\n%s", diff)
1683 | }
1684 | break
1685 | }
1686 | }
1687 | if !found {
1688 | t.Errorf("Expected trigger '%s' not found in the list of all triggers.", triggerName)
1689 | }
1690 | } else {
1691 | if diff := cmp.Diff(tc.want, got); diff != "" {
1692 | t.Errorf("Unexpected result (-want +got):\n%s", diff)
1693 | }
1694 | }
1695 | })
1696 | }
1697 | }
1698 |
1699 | func setupPostgresPublicationTable(t *testing.T, ctx context.Context, pool *pgxpool.Pool, tableName string, pubName string) func(t *testing.T) {
1700 | t.Helper()
1701 | createTableStmt := fmt.Sprintf("CREATE TABLE %s (id SERIAL PRIMARY KEY, name TEXT);", tableName)
1702 | if _, err := pool.Exec(ctx, createTableStmt); err != nil {
1703 | t.Fatalf("unable to create table %s: %v", tableName, err)
1704 | }
1705 |
1706 | createPubStmt := fmt.Sprintf("CREATE PUBLICATION %s FOR TABLE %s;", pubName, tableName)
1707 | if _, err := pool.Exec(ctx, createPubStmt); err != nil {
1708 | if _, dropErr := pool.Exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s;", tableName)); dropErr != nil {
1709 | t.Errorf("unable to drop table after failing to create publication: %v", dropErr)
1710 | }
1711 | t.Fatalf("unable to create publication %s: %v", pubName, err)
1712 | }
1713 |
1714 | return func(t *testing.T) {
1715 | t.Helper()
1716 | if _, err := pool.Exec(ctx, fmt.Sprintf("DROP PUBLICATION IF EXISTS %s;", pubName)); err != nil {
1717 | t.Errorf("unable to drop publication %s: %v", pubName, err)
1718 | }
1719 | if _, err := pool.Exec(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s;", tableName)); err != nil {
1720 | t.Errorf("unable to drop table %s: %v", tableName, err)
1721 | }
1722 | }
1723 | }
1724 |
1725 | func RunPostgresListPublicationTablesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
1726 | table1Name := "pub_table_1"
1727 | pub1Name := "pub_1"
1728 |
1729 | table2Name := "pub_table_2"
1730 | pub2Name := "pub_2"
1731 |
1732 | cleanup := setupPostgresPublicationTable(t, ctx, pool, table1Name, pub1Name)
1733 | defer cleanup(t)
1734 | cleanup2 := setupPostgresPublicationTable(t, ctx, pool, table2Name, pub2Name)
1735 | defer cleanup2(t)
1736 |
1737 | // Fetch the current user to match the publication_owner
1738 | var currentUser string
1739 | err := pool.QueryRow(ctx, "SELECT current_user;").Scan(¤tUser)
1740 | if err != nil {
1741 | t.Fatalf("unable to fetch current user: %v", err)
1742 | }
1743 |
1744 | wantTable1 := map[string]any{
1745 | "publication_name": pub1Name,
1746 | "schema_name": "public",
1747 | "table_name": table1Name,
1748 | "publishes_all_tables": false,
1749 | "publishes_inserts": true,
1750 | "publishes_updates": true,
1751 | "publishes_deletes": true,
1752 | "publishes_truncates": true,
1753 | "publication_owner": currentUser,
1754 | }
1755 |
1756 | wantTable2 := map[string]any{
1757 | "publication_name": pub2Name,
1758 | "schema_name": "public",
1759 | "table_name": table2Name,
1760 | "publishes_all_tables": false,
1761 | "publishes_inserts": true,
1762 | "publishes_updates": true,
1763 | "publishes_deletes": true,
1764 | "publishes_truncates": true,
1765 | "publication_owner": currentUser,
1766 | }
1767 |
1768 | invokeTcs := []struct {
1769 | name string
1770 | requestBody io.Reader
1771 | wantStatusCode int
1772 | want []map[string]any
1773 | }{
1774 | {
1775 | name: "list all publication tables",
1776 | requestBody: bytes.NewBufferString(`{}`),
1777 | wantStatusCode: http.StatusOK,
1778 | want: []map[string]any{wantTable1, wantTable2},
1779 | },
1780 | {
1781 | name: "list all tables for the created publication",
1782 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"publication_names": "%s"}`, pub1Name)),
1783 | wantStatusCode: http.StatusOK,
1784 | want: []map[string]any{wantTable1},
1785 | },
1786 | {
1787 | name: "filter by table_name",
1788 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_names": "%s, %s"}`, table1Name, table2Name)),
1789 | wantStatusCode: http.StatusOK,
1790 | want: []map[string]any{wantTable1, wantTable2},
1791 | },
1792 | {
1793 | name: "filter by schema_name and table_name",
1794 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_names": "public", "table_name": "%s , %s"}`, table1Name, table2Name)),
1795 | wantStatusCode: http.StatusOK,
1796 | want: []map[string]any{wantTable1, wantTable2},
1797 | },
1798 | {
1799 | name: "invoke list_publication_tables with non-existent table",
1800 | requestBody: bytes.NewBufferString(`{"table_names": "non_existent_table"}`),
1801 | wantStatusCode: http.StatusOK,
1802 | want: nil,
1803 | },
1804 | {
1805 | name: "invoke list_publication_tables with non-existent publication",
1806 | requestBody: bytes.NewBufferString(`{"publication_names": "non_existent_pub"}`),
1807 | wantStatusCode: http.StatusOK,
1808 | want: nil,
1809 | },
1810 | }
1811 | for _, tc := range invokeTcs {
1812 | t.Run(tc.name, func(t *testing.T) {
1813 | const api = "http://127.0.0.1:5000/api/tool/list_publication_tables/invoke"
1814 |
1815 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
1816 | if resp.StatusCode != tc.wantStatusCode {
1817 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
1818 | }
1819 | if tc.wantStatusCode != http.StatusOK {
1820 | return
1821 | }
1822 |
1823 | var bodyWrapper struct {
1824 | Result json.RawMessage `json:"result"`
1825 | }
1826 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
1827 | t.Fatalf("error decoding response wrapper: %v", err)
1828 | }
1829 |
1830 | var resultString string
1831 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
1832 | resultString = string(bodyWrapper.Result)
1833 | }
1834 |
1835 | var got []map[string]any
1836 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
1837 | t.Fatalf("failed to unmarshal nested result string: %v, content: %s", err, resultString)
1838 | }
1839 |
1840 | if diff := cmp.Diff(tc.want, got); diff != "" {
1841 | t.Errorf("Unexpected result (-want +got):\n%s", diff)
1842 | }
1843 | })
1844 | }
1845 | }
1846 |
1847 | func RunPostgresListActiveQueriesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
1848 | type queryListDetails struct {
1849 | ProcessId any `json:"pid"`
1850 | User string `json:"user"`
1851 | Datname string `json:"datname"`
1852 | ApplicationName string `json:"application_name"`
1853 | ClientAddress string `json:"client_addr"`
1854 | State string `json:"state"`
1855 | WaitEventType string `json:"wait_event_type"`
1856 | WaitEvent string `json:"wait_event"`
1857 | BackendStart any `json:"backend_start"`
1858 | TransactionStart any `json:"xact_start"`
1859 | QueryStart any `json:"query_start"`
1860 | QueryDuration any `json:"query_duration"`
1861 | Query string `json:"query"`
1862 | }
1863 |
1864 | singleQueryWanted := queryListDetails{
1865 | ProcessId: any(nil),
1866 | User: "",
1867 | Datname: "",
1868 | ApplicationName: "",
1869 | ClientAddress: "",
1870 | State: "",
1871 | WaitEventType: "",
1872 | WaitEvent: "",
1873 | BackendStart: any(nil),
1874 | TransactionStart: any(nil),
1875 | QueryStart: any(nil),
1876 | QueryDuration: any(nil),
1877 | Query: "SELECT pg_sleep(10);",
1878 | }
1879 |
1880 | invokeTcs := []struct {
1881 | name string
1882 | requestBody io.Reader
1883 | clientSleepSecs int
1884 | waitSecsBeforeCheck int
1885 | wantStatusCode int
1886 | want any
1887 | }{
1888 | // exclude background monitoring apps such as "wal_uploader"
1889 | {
1890 | name: "invoke list_active_queries when the system is idle",
1891 | requestBody: bytes.NewBufferString(`{"exclude_application_names": "wal_uploader"}`),
1892 | clientSleepSecs: 0,
1893 | waitSecsBeforeCheck: 0,
1894 | wantStatusCode: http.StatusOK,
1895 | want: []queryListDetails(nil),
1896 | },
1897 | {
1898 | name: "invoke list_active_queries when there is 1 ongoing but lower than the threshold",
1899 | requestBody: bytes.NewBufferString(`{"min_duration": "100 seconds", "exclude_application_names": "wal_uploader"}`),
1900 | clientSleepSecs: 1,
1901 | waitSecsBeforeCheck: 1,
1902 | wantStatusCode: http.StatusOK,
1903 | want: []queryListDetails(nil),
1904 | },
1905 | {
1906 | name: "invoke list_active_queries when 1 ongoing query should show up",
1907 | requestBody: bytes.NewBufferString(`{"min_duration": "1 seconds", "exclude_application_names": "wal_uploader"}`),
1908 | clientSleepSecs: 10,
1909 | waitSecsBeforeCheck: 5,
1910 | wantStatusCode: http.StatusOK,
1911 | want: []queryListDetails{singleQueryWanted},
1912 | },
1913 | }
1914 |
1915 | var wg sync.WaitGroup
1916 | for _, tc := range invokeTcs {
1917 | t.Run(tc.name, func(t *testing.T) {
1918 | if tc.clientSleepSecs > 0 {
1919 | wg.Add(1)
1920 |
1921 | go func() {
1922 | defer wg.Done()
1923 |
1924 | err := pool.Ping(ctx)
1925 | if err != nil {
1926 | t.Errorf("unable to connect to test database: %s", err)
1927 | return
1928 | }
1929 | _, err = pool.Exec(ctx, fmt.Sprintf("SELECT pg_sleep(%d);", tc.clientSleepSecs))
1930 | if err != nil {
1931 | t.Errorf("Executing 'SELECT pg_sleep' failed: %s", err)
1932 | }
1933 | }()
1934 | }
1935 |
1936 | if tc.waitSecsBeforeCheck > 0 {
1937 | time.Sleep(time.Duration(tc.waitSecsBeforeCheck) * time.Second)
1938 | }
1939 |
1940 | const api = "http://127.0.0.1:5000/api/tool/list_active_queries/invoke"
1941 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
1942 | if resp.StatusCode != tc.wantStatusCode {
1943 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
1944 | }
1945 | if tc.wantStatusCode != http.StatusOK {
1946 | return
1947 | }
1948 |
1949 | var bodyWrapper struct {
1950 | Result json.RawMessage `json:"result"`
1951 | }
1952 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
1953 | t.Fatalf("error decoding response wrapper: %v", err)
1954 | }
1955 |
1956 | var resultString string
1957 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
1958 | resultString = string(bodyWrapper.Result)
1959 | }
1960 |
1961 | var got any
1962 | var details []queryListDetails
1963 | if err := json.Unmarshal([]byte(resultString), &details); err != nil {
1964 | t.Fatalf("failed to unmarshal nested ObjectDetails string: %v", err)
1965 | }
1966 | got = details
1967 |
1968 | if diff := cmp.Diff(tc.want, got, cmp.Comparer(func(a, b queryListDetails) bool {
1969 | return a.Query == b.Query
1970 | })); diff != "" {
1971 | t.Errorf("Unexpected result: got %#v, want: %#v", got, tc.want)
1972 | }
1973 | })
1974 | }
1975 | wg.Wait()
1976 | }
1977 |
1978 | func RunPostgresListAvailableExtensionsTest(t *testing.T) {
1979 | invokeTcs := []struct {
1980 | name string
1981 | api string
1982 | requestBody io.Reader
1983 | wantStatusCode int
1984 | }{
1985 | {
1986 | name: "invoke list_available_extensions output",
1987 | api: "http://127.0.0.1:5000/api/tool/list_available_extensions/invoke",
1988 | wantStatusCode: http.StatusOK,
1989 | requestBody: bytes.NewBuffer([]byte(`{}`)),
1990 | },
1991 | }
1992 | for _, tc := range invokeTcs {
1993 | t.Run(tc.name, func(t *testing.T) {
1994 | resp, respBody := RunRequest(t, http.MethodPost, tc.api, tc.requestBody, nil)
1995 | if resp.StatusCode != tc.wantStatusCode {
1996 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBody))
1997 | }
1998 |
1999 | // Intentionally not adding the output check as output depends on the postgres instance used where the the functional test runs.
2000 | // Adding the check will make the test flaky.
2001 | })
2002 | }
2003 | }
2004 |
2005 | func RunPostgresListInstalledExtensionsTest(t *testing.T) {
2006 | invokeTcs := []struct {
2007 | name string
2008 | api string
2009 | requestBody io.Reader
2010 | wantStatusCode int
2011 | }{
2012 | {
2013 | name: "invoke list_installed_extensions output",
2014 | api: "http://127.0.0.1:5000/api/tool/list_installed_extensions/invoke",
2015 | wantStatusCode: http.StatusOK,
2016 | requestBody: bytes.NewBuffer([]byte(`{}`)),
2017 | },
2018 | }
2019 | for _, tc := range invokeTcs {
2020 | t.Run(tc.name, func(t *testing.T) {
2021 | resp, bodyBytes := RunRequest(t, http.MethodPost, tc.api, tc.requestBody, nil)
2022 | if resp.StatusCode != tc.wantStatusCode {
2023 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(bodyBytes))
2024 | }
2025 |
2026 | // Intentionally not adding the output check as output depends on the postgres instance used where the the functional test runs.
2027 | // Adding the check will make the test flaky.
2028 | })
2029 | }
2030 | }
2031 |
2032 | func setupPostgresIndex(t *testing.T, ctx context.Context, pool *pgxpool.Pool, schemaName string, tableName string) func(t *testing.T) {
2033 | t.Helper()
2034 | createSchemaStmt := fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s;", schemaName)
2035 | if _, err := pool.Exec(ctx, createSchemaStmt); err != nil {
2036 | t.Fatalf("unable to create schema %s: %v", schemaName, err)
2037 | }
2038 |
2039 | fullTableName := fmt.Sprintf("%s.%s", schemaName, tableName)
2040 | createTableStmt := fmt.Sprintf("CREATE TABLE %s (id SERIAL PRIMARY KEY, name TEXT, email TEXT);", fullTableName)
2041 | if _, err := pool.Exec(ctx, createTableStmt); err != nil {
2042 | t.Fatalf("unable to create table %s: %v", fullTableName, err)
2043 | }
2044 |
2045 | // Create a unique index on email
2046 | index1Stmt := fmt.Sprintf("CREATE UNIQUE INDEX %s_email_idx ON %s (email);", tableName, fullTableName)
2047 | if _, err := pool.Exec(ctx, index1Stmt); err != nil {
2048 | t.Fatalf("unable to create index %s_email_idx: %v", tableName, err)
2049 | }
2050 |
2051 | // Create a non-unique index on name
2052 | index2Stmt := fmt.Sprintf("CREATE INDEX %s_name_idx ON %s (name);", tableName, fullTableName)
2053 | if _, err := pool.Exec(ctx, index2Stmt); err != nil {
2054 | t.Fatalf("unable to create index %s_name_idx: %v", tableName, err)
2055 | }
2056 |
2057 | return func(t *testing.T) {
2058 | t.Helper()
2059 | if _, err := pool.Exec(ctx, fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE;", schemaName)); err != nil {
2060 | t.Errorf("unable to drop schema: %v", err)
2061 | }
2062 | }
2063 | }
2064 |
2065 | func RunPostgresListIndexesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
2066 | schemaName := "testschema_" + strings.ReplaceAll(uuid.New().String(), "-", "")
2067 | tableName := "table1_" + strings.ReplaceAll(uuid.New().String(), "-", "")
2068 | cleanup := setupPostgresIndex(t, ctx, pool, schemaName, tableName)
2069 | defer cleanup(t)
2070 |
2071 | // Primary key index
2072 | wantIndexPK := map[string]any{
2073 | "schema_name": schemaName,
2074 | "table_name": tableName,
2075 | "index_name": tableName + "_pkey",
2076 | "index_type": "btree",
2077 | "is_unique": true,
2078 | "is_primary": true,
2079 | "is_used": false,
2080 | "index_definition": fmt.Sprintf("CREATE UNIQUE INDEX %s_pkey ON %s.%s USING btree (id)", tableName, schemaName, tableName),
2081 | // Size and scan counts can vary, so omitting them from strict checks or using ranges might be better in real tests.
2082 | }
2083 | // Email unique index
2084 | wantIndexEmail := map[string]any{
2085 | "schema_name": schemaName,
2086 | "table_name": tableName,
2087 | "index_name": tableName + "_email_idx",
2088 | "index_type": "btree",
2089 | "is_unique": true,
2090 | "is_primary": false,
2091 | "is_used": false,
2092 | "index_definition": fmt.Sprintf("CREATE UNIQUE INDEX %s_email_idx ON %s.%s USING btree (email)", tableName, schemaName, tableName),
2093 | }
2094 | // Name non-unique index
2095 | wantIndexName := map[string]any{
2096 | "schema_name": schemaName,
2097 | "table_name": tableName,
2098 | "index_name": tableName + "_name_idx",
2099 | "index_type": "btree",
2100 | "is_unique": false,
2101 | "is_primary": false,
2102 | "is_used": false,
2103 | "index_definition": fmt.Sprintf("CREATE INDEX %s_name_idx ON %s.%s USING btree (name)", tableName, schemaName, tableName),
2104 | }
2105 |
2106 | allWantIndexes := []map[string]any{wantIndexEmail, wantIndexName, wantIndexPK}
2107 |
2108 | invokeTcs := []struct {
2109 | name string
2110 | requestBody io.Reader
2111 | wantStatusCode int
2112 | want []map[string]any
2113 | }{
2114 | // List all indexes is skipped because the output might include indexes for other database tables
2115 | // defined outside of this test, which could make the test flaky.
2116 | {
2117 | name: "list_indexes for a specific schema and table",
2118 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "%s"}`, schemaName, tableName)),
2119 | wantStatusCode: http.StatusOK,
2120 | want: allWantIndexes,
2121 | },
2122 | {
2123 | name: "list_indexes for a specific schema",
2124 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s"}`, schemaName)),
2125 | wantStatusCode: http.StatusOK,
2126 | want: allWantIndexes,
2127 | },
2128 | {
2129 | name: "list_indexes with non-existent schema",
2130 | requestBody: bytes.NewBufferString(`{"schema_name": "non_existent_schema"}`),
2131 | wantStatusCode: http.StatusOK,
2132 | want: nil,
2133 | },
2134 | {
2135 | name: "list_indexes with non-existent table in existing schema",
2136 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "non_existent_table"}`, schemaName)),
2137 | wantStatusCode: http.StatusOK,
2138 | want: nil,
2139 | },
2140 | {
2141 | name: "list_indexes filter by index name",
2142 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "%s", "index_name": "%s"}`, schemaName, tableName, tableName+"_email_idx")),
2143 | wantStatusCode: http.StatusOK,
2144 | want: []map[string]any{wantIndexEmail},
2145 | },
2146 | {
2147 | name: "list_indexes filter by non-existent index name",
2148 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "%s", "index_name": "non_existent_idx"}`, schemaName, tableName)),
2149 | wantStatusCode: http.StatusOK,
2150 | want: nil,
2151 | },
2152 | }
2153 | for _, tc := range invokeTcs {
2154 | t.Run(tc.name, func(t *testing.T) {
2155 | const api = "http://127.0.0.1:5000/api/tool/list_indexes/invoke"
2156 |
2157 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
2158 | if resp.StatusCode != tc.wantStatusCode {
2159 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
2160 | }
2161 | if tc.wantStatusCode != http.StatusOK {
2162 | return
2163 | }
2164 |
2165 | var bodyWrapper struct {
2166 | Result json.RawMessage `json:"result"`
2167 | }
2168 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
2169 | t.Fatalf("error decoding response wrapper: %v", err)
2170 | }
2171 |
2172 | var resultString string
2173 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
2174 | resultString = string(bodyWrapper.Result)
2175 | }
2176 |
2177 | var got []map[string]any
2178 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
2179 | t.Fatalf("failed to unmarshal nested result string: %v, resultString: %s", err, resultString)
2180 | }
2181 | // Normalize got by removing fields that are hard to predict (like size)
2182 | for _, index := range got {
2183 | delete(index, "index_size_bytes")
2184 | delete(index, "index_scans")
2185 | delete(index, "tuples_read")
2186 | delete(index, "tuples_fetched")
2187 | }
2188 |
2189 | if diff := cmp.Diff(tc.want, got); diff != "" {
2190 | t.Errorf("Unexpected result (-want +got):\n%s", diff)
2191 | }
2192 | })
2193 | }
2194 | }
2195 |
2196 | func setupListSequencesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) (string, func(t *testing.T)) {
2197 | sequenceName := "list_sequences_seq1_" + strings.ReplaceAll(uuid.New().String(), "-", "")
2198 | createSequence1Stmt := fmt.Sprintf("CREATE SEQUENCE %s INCREMENT 1 START 1;", sequenceName)
2199 |
2200 | _, err := pool.Exec(ctx, createSequence1Stmt)
2201 | if err != nil {
2202 | t.Fatalf("unable to create sequence %s: %s", sequenceName, err)
2203 | }
2204 | return sequenceName, func(t *testing.T) {
2205 | _, err := pool.Exec(ctx, fmt.Sprintf("DROP SEQUENCE IF EXISTS %s;", sequenceName))
2206 | if err != nil {
2207 | t.Errorf("unable to drop sequences: %v", err)
2208 | }
2209 | }
2210 | }
2211 |
2212 | func RunPostgresListSequencesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
2213 | sequenceName, teardown := setupListSequencesTest(t, ctx, pool)
2214 | defer teardown(t)
2215 |
2216 | wantSequence := map[string]any{
2217 | "sequence_name": sequenceName,
2218 | "schema_name": "public",
2219 | "sequence_owner": "postgres",
2220 | "data_type": "bigint",
2221 | "start_value": float64(1),
2222 | "min_value": float64(1),
2223 | "max_value": float64(9223372036854775807),
2224 | "increment_by": float64(1),
2225 | "last_value": nil,
2226 | }
2227 |
2228 | invokeTcs := []struct {
2229 | name string
2230 | api string
2231 | requestBody io.Reader
2232 | wantStatusCode int
2233 | want []map[string]any
2234 | }{
2235 | {
2236 | name: "invoke list_sequences",
2237 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"sequence_name": "%s"}`, sequenceName)),
2238 | wantStatusCode: http.StatusOK,
2239 | want: []map[string]any{wantSequence},
2240 | },
2241 | {
2242 | name: "invoke list_sequences with non-existent sequence",
2243 | requestBody: bytes.NewBufferString(`{"sequence_name": "non_existent_sequence"}`),
2244 | wantStatusCode: http.StatusOK,
2245 | want: nil,
2246 | },
2247 | }
2248 | for _, tc := range invokeTcs {
2249 | t.Run(tc.name, func(t *testing.T) {
2250 | const api = "http://127.0.0.1:5000/api/tool/list_sequences/invoke"
2251 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
2252 | if resp.StatusCode != tc.wantStatusCode {
2253 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
2254 | }
2255 | if tc.wantStatusCode != http.StatusOK {
2256 | return
2257 | }
2258 |
2259 | var bodyWrapper struct {
2260 | Result json.RawMessage `json:"result"`
2261 | }
2262 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
2263 | t.Fatalf("error decoding response wrapper: %v", err)
2264 | }
2265 |
2266 | var resultString string
2267 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
2268 | resultString = string(bodyWrapper.Result)
2269 | }
2270 |
2271 | var got []map[string]any
2272 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
2273 | t.Fatalf("failed to unmarshal nested result string: %v", err)
2274 | }
2275 |
2276 | if diff := cmp.Diff(tc.want, got); diff != "" {
2277 | t.Errorf("Unexpected result (-want +got):\n%s", diff)
2278 | }
2279 | })
2280 | }
2281 | }
2282 |
2283 | func RunPostgresListTableSpacesTest(t *testing.T) {
2284 | invokeTcs := []struct {
2285 | name string
2286 | api string
2287 | requestBody io.Reader
2288 | wantStatusCode int
2289 | }{
2290 | {
2291 | name: "invoke list_tablespaces output",
2292 | api: "http://127.0.0.1:5000/api/tool/list_tablespaces/invoke",
2293 | wantStatusCode: http.StatusOK,
2294 | requestBody: bytes.NewBuffer([]byte(`{}`)),
2295 | },
2296 | }
2297 | for _, tc := range invokeTcs {
2298 | t.Run(tc.name, func(t *testing.T) {
2299 | resp, respBody := RunRequest(t, http.MethodPost, tc.api, tc.requestBody, nil)
2300 | if resp.StatusCode != tc.wantStatusCode {
2301 | t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBody))
2302 | }
2303 |
2304 | // Intentionally not adding the output check as output depends on the postgres instance used where the the functional test runs.
2305 | // Adding the check will make the test flaky.
2306 | })
2307 | }
2308 | }
2309 |
2310 | func RunPostgresListPgSettingsTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
2311 | targetSetting := "maintenance_work_mem"
2312 | var name, setting, unit, shortDesc, source, contextVal string
2313 |
2314 | // We query the raw pg_settings to get the data needed to reconstruct the logic
2315 | // defined in your listPgSettingQuery.
2316 | err := pool.QueryRow(ctx, `
2317 | SELECT name, setting, unit, short_desc, source, context
2318 | FROM pg_settings
2319 | WHERE name = $1
2320 | `, targetSetting).Scan(&name, &setting, &unit, &shortDesc, &source, &contextVal)
2321 |
2322 | if err != nil {
2323 | t.Fatalf("Setup failed: could not fetch postgres setting '%s': %v", targetSetting, err)
2324 | }
2325 |
2326 | // Replicate the SQL CASE logic for 'requires_restart' field
2327 | requiresRestart := "No"
2328 | switch contextVal {
2329 | case "postmaster":
2330 | requiresRestart = "Yes"
2331 | case "sighup":
2332 | requiresRestart = "No (Reload sufficient)"
2333 | }
2334 |
2335 | expectedObject := map[string]interface{}{
2336 | "name": name,
2337 | "current_value": setting,
2338 | "unit": unit,
2339 | "short_desc": shortDesc,
2340 | "source": source,
2341 | "requires_restart": requiresRestart,
2342 | }
2343 | expectedJSON, _ := json.Marshal([]interface{}{expectedObject})
2344 |
2345 | invokeTcs := []struct {
2346 | name string
2347 | requestBody io.Reader
2348 | wantStatusCode int
2349 | want string
2350 | }{
2351 | {
2352 | name: "invoke list_pg_settings with specific setting",
2353 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"setting_name": "%s"}`, targetSetting))),
2354 | wantStatusCode: http.StatusOK,
2355 | want: string(expectedJSON),
2356 | },
2357 | {
2358 | name: "invoke list_pg_settings with non-existent setting",
2359 | requestBody: bytes.NewBuffer([]byte(`{"setting_name": "non_existent_config_xyz"}`)),
2360 | wantStatusCode: http.StatusOK,
2361 | want: `null`,
2362 | },
2363 | }
2364 |
2365 | for _, tc := range invokeTcs {
2366 | t.Run(tc.name, func(t *testing.T) {
2367 | const api = "http://127.0.0.1:5000/api/tool/list_pg_settings/invoke"
2368 | resp, body := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
2369 |
2370 | if resp.StatusCode != tc.wantStatusCode {
2371 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(body))
2372 | }
2373 | if tc.wantStatusCode != http.StatusOK {
2374 | return
2375 | }
2376 |
2377 | var bodyWrapper struct {
2378 | Result json.RawMessage `json:"result"`
2379 | }
2380 | if err := json.Unmarshal(body, &bodyWrapper); err != nil {
2381 | t.Fatalf("error decoding response wrapper: %v", err)
2382 | }
2383 |
2384 | var resultString string
2385 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
2386 | resultString = string(bodyWrapper.Result)
2387 | }
2388 |
2389 | var got, want any
2390 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
2391 | t.Fatalf("failed to unmarshal nested result string: %v", err)
2392 | }
2393 | if err := json.Unmarshal([]byte(tc.want), &want); err != nil {
2394 | t.Fatalf("failed to unmarshal want string: %v", err)
2395 | }
2396 |
2397 | if diff := cmp.Diff(want, got); diff != "" {
2398 | t.Errorf("Unexpected result (-want +got):\n%s", diff)
2399 | }
2400 | })
2401 | }
2402 | }
2403 |
2404 | // RunPostgresDatabaseStatsTest tests the database_stats tool by comparing API results
2405 | // against a direct query to the database.
2406 | func RunPostgresListDatabaseStatsTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
2407 | dbName1 := "test_db_stats_" + strings.ReplaceAll(uuid.NewString(), "-", "")
2408 | dbOwner1 := "test_user_" + strings.ReplaceAll(uuid.NewString(), "-", "")
2409 | dbName2 := "test_db_stats_" + strings.ReplaceAll(uuid.NewString(), "-", "")
2410 | dbOwner2 := "test_user_" + strings.ReplaceAll(uuid.NewString(), "-", "")
2411 |
2412 | cleanup1 := setUpDatabase(t, ctx, pool, dbName1, dbOwner1)
2413 | defer cleanup1()
2414 | cleanup2 := setUpDatabase(t, ctx, pool, dbName2, dbOwner2)
2415 | defer cleanup2()
2416 |
2417 | requiredKeys := map[string]bool{
2418 | "database_name": true,
2419 | "database_owner": true,
2420 | "default_tablespace": true,
2421 | "is_connectable": true,
2422 | }
2423 |
2424 | db1Want := map[string]interface{}{
2425 | "database_name": dbName1,
2426 | "database_owner": dbOwner1,
2427 | "default_tablespace": "pg_default",
2428 | "is_connectable": true,
2429 | }
2430 |
2431 | db2Want := map[string]interface{}{
2432 | "database_name": dbName2,
2433 | "database_owner": dbOwner2,
2434 | "default_tablespace": "pg_default",
2435 | "is_connectable": true,
2436 | }
2437 |
2438 | invokeTcs := []struct {
2439 | name string
2440 | requestBody io.Reader
2441 | wantStatusCode int
2442 | want []map[string]interface{}
2443 | }{
2444 | {
2445 | name: "invoke database_stats filtering by specific database name",
2446 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"database_name": "%s"}`, dbName1))),
2447 | wantStatusCode: http.StatusOK,
2448 | want: []map[string]interface{}{db1Want},
2449 | },
2450 | {
2451 | name: "invoke database_stats filtering by specific owner",
2452 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"database_owner": "%s"}`, dbOwner2))),
2453 | wantStatusCode: http.StatusOK,
2454 | want: []map[string]interface{}{db2Want},
2455 | },
2456 | {
2457 | name: "filter by tablespace",
2458 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"default_tablespace": "pg_default", "database_name": "%s"}`, dbName1))),
2459 | wantStatusCode: http.StatusOK,
2460 | want: []map[string]interface{}{db1Want},
2461 | },
2462 | {
2463 | name: "sort by size",
2464 | requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"sort_by": "size", "database_name": "%s"}`, dbName2))),
2465 | wantStatusCode: http.StatusOK,
2466 | want: []map[string]interface{}{db2Want},
2467 | },
2468 | }
2469 |
2470 | for _, tc := range invokeTcs {
2471 | t.Run(tc.name, func(t *testing.T) {
2472 | const api = "http://127.0.0.1:5000/api/tool/list_database_stats/invoke"
2473 | resp, body := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
2474 |
2475 | if resp.StatusCode != tc.wantStatusCode {
2476 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(body))
2477 | }
2478 | var bodyWrapper struct {
2479 | Result json.RawMessage `json:"result"`
2480 | }
2481 | if err := json.Unmarshal(body, &bodyWrapper); err != nil {
2482 | t.Fatalf("error decoding response wrapper: %v", err)
2483 | }
2484 |
2485 | var resultString string
2486 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
2487 | resultString = string(bodyWrapper.Result)
2488 | }
2489 |
2490 | var got []map[string]interface{}
2491 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
2492 | t.Fatalf("failed to unmarshal nested result string: %v", err)
2493 | }
2494 |
2495 | // Configuration for comparison
2496 | opts := []cmp.Option{
2497 | // Ensure consistent order based on name for comparison
2498 | cmpopts.SortSlices(func(a, b map[string]interface{}) bool {
2499 | return a["database_name"].(string) < b["database_name"].(string)
2500 | }),
2501 |
2502 | // Ignore Volatile Keys which change in every run and only compare the keys in 'requiredKeys'
2503 | cmpopts.IgnoreMapEntries(func(key string, _ interface{}) bool {
2504 | return !requiredKeys[key]
2505 | }),
2506 |
2507 | // Ignore Irrelevant Databases
2508 | cmpopts.IgnoreSliceElements(func(v map[string]interface{}) bool {
2509 | name, ok := v["database_name"].(string)
2510 | if !ok {
2511 | return true
2512 | }
2513 | return name != dbName1 && name != dbName2
2514 | }),
2515 | }
2516 |
2517 | if diff := cmp.Diff(tc.want, got, opts...); diff != "" {
2518 | t.Errorf("Unexpected result (-want +got):\n%s", diff)
2519 | }
2520 | })
2521 | }
2522 | }
2523 |
2524 | func setUpDatabase(t *testing.T, ctx context.Context, pool *pgxpool.Pool, dbName, dbOwner string) func() {
2525 | _, err := pool.Exec(ctx, fmt.Sprintf("CREATE ROLE %s LOGIN PASSWORD 'password';", dbOwner))
2526 | if err != nil {
2527 | _, _ = pool.Exec(ctx, fmt.Sprintf("DROP ROLE %s;", dbOwner))
2528 | t.Fatalf("failed to create %s: %v", dbOwner, err)
2529 | }
2530 | _, err = pool.Exec(ctx, fmt.Sprintf("GRANT %s TO current_user;", dbOwner))
2531 | if err != nil {
2532 | t.Fatalf("failed to grant %s to current_user: %v", dbOwner, err)
2533 | }
2534 | _, err = pool.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s OWNER %s;", dbName, dbOwner))
2535 | if err != nil {
2536 | t.Fatalf("failed to create %s: %v", dbName, err)
2537 | }
2538 | return func() {
2539 | _, _ = pool.Exec(ctx, fmt.Sprintf("DROP DATABASE IF EXISTS %s;", dbName))
2540 | _, _ = pool.Exec(ctx, fmt.Sprintf("DROP ROLE IF EXISTS %s;", dbOwner))
2541 | }
2542 | }
2543 |
2544 | func setupPostgresRoles(t *testing.T, ctx context.Context, pool *pgxpool.Pool) (string, string, string, func(t *testing.T)) {
2545 | t.Helper()
2546 | suffix := strings.ReplaceAll(uuid.New().String(), "-", "")
2547 |
2548 | adminUser := "test_role_admin_" + suffix
2549 | superUser := "test_role_super_" + suffix
2550 | normalUser := "test_role_normal_" + suffix
2551 |
2552 | createAdminStmt := fmt.Sprintf("CREATE ROLE %s NOLOGIN;", adminUser)
2553 | if _, err := pool.Exec(ctx, createAdminStmt); err != nil {
2554 | t.Fatalf("unable to create role %s: %v", adminUser, err)
2555 | }
2556 |
2557 | createSuperUserStmt := fmt.Sprintf("CREATE ROLE %s LOGIN CREATEDB;", superUser)
2558 | if _, err := pool.Exec(ctx, createSuperUserStmt); err != nil {
2559 | t.Fatalf("unable to create role %s: %v", superUser, err)
2560 | }
2561 |
2562 | createNormalUserStmt := fmt.Sprintf("CREATE ROLE %s LOGIN;", normalUser)
2563 | if _, err := pool.Exec(ctx, createNormalUserStmt); err != nil {
2564 | t.Fatalf("unable to create role %s: %v", normalUser, err)
2565 | }
2566 |
2567 | // Establish Relationships (Admin -> Superuser -> Normal)
2568 | if _, err := pool.Exec(ctx, fmt.Sprintf("GRANT %s TO %s;", adminUser, superUser)); err != nil {
2569 | t.Fatalf("unable to grant %s to %s: %v", adminUser, superUser, err)
2570 | }
2571 | if _, err := pool.Exec(ctx, fmt.Sprintf("GRANT %s TO %s;", superUser, normalUser)); err != nil {
2572 | t.Fatalf("unable to grant %s to %s: %v", superUser, normalUser, err)
2573 | }
2574 |
2575 | return adminUser, superUser, normalUser, func(t *testing.T) {
2576 | t.Helper()
2577 | _, _ = pool.Exec(ctx, fmt.Sprintf("DROP ROLE IF EXISTS %s;", normalUser))
2578 | _, _ = pool.Exec(ctx, fmt.Sprintf("DROP ROLE IF EXISTS %s;", superUser))
2579 | _, _ = pool.Exec(ctx, fmt.Sprintf("DROP ROLE IF EXISTS %s;", adminUser))
2580 | }
2581 | }
2582 |
2583 | func RunPostgresListRolesTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
2584 | adminUser, superUser, normalUser, cleanup := setupPostgresRoles(t, ctx, pool)
2585 | defer cleanup(t)
2586 |
2587 | wantAdmin := map[string]any{
2588 | "role_name": adminUser,
2589 | "connection_limit": float64(-1),
2590 | "is_superuser": false,
2591 | "inherits_privileges": true,
2592 | "can_create_roles": false,
2593 | "can_create_db": false,
2594 | "can_login": false,
2595 | "is_replication_role": false,
2596 | "bypass_rls": false,
2597 | "direct_members": []any{superUser},
2598 | "member_of": []any{},
2599 | }
2600 |
2601 | wantSuperUser := map[string]any{
2602 | "role_name": superUser,
2603 | "connection_limit": float64(-1),
2604 | "is_superuser": false,
2605 | "inherits_privileges": true,
2606 | "can_create_roles": false,
2607 | "can_create_db": true,
2608 | "can_login": true,
2609 | "is_replication_role": false,
2610 | "bypass_rls": false,
2611 | "direct_members": []any{normalUser},
2612 | "member_of": []any{adminUser},
2613 | }
2614 |
2615 | wantNormalUser := map[string]any{
2616 | "role_name": normalUser,
2617 | "connection_limit": float64(-1),
2618 | "is_superuser": false,
2619 | "inherits_privileges": true,
2620 | "can_create_roles": false,
2621 | "can_create_db": false,
2622 | "can_login": true,
2623 | "is_replication_role": false,
2624 | "bypass_rls": false,
2625 | "direct_members": []any{},
2626 | "member_of": []any{superUser},
2627 | }
2628 |
2629 | invokeTcs := []struct {
2630 | name string
2631 | requestBody io.Reader
2632 | wantStatusCode int
2633 | want []map[string]any
2634 | }{
2635 | {
2636 | name: "list_roles with filter for created roles",
2637 | requestBody: bytes.NewBufferString(`{"role_name": "test_role_"}`),
2638 | wantStatusCode: http.StatusOK,
2639 | want: []map[string]any{wantAdmin, wantNormalUser, wantSuperUser},
2640 | },
2641 | {
2642 | name: "list_roles filter specific role",
2643 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"role_name": "%s"}`, superUser)),
2644 | wantStatusCode: http.StatusOK,
2645 | want: []map[string]any{wantSuperUser},
2646 | },
2647 | {
2648 | name: "list_roles non-existent role",
2649 | requestBody: bytes.NewBufferString(`{"role_name": "non_existent_role_xyz"}`),
2650 | wantStatusCode: http.StatusOK,
2651 | want: nil,
2652 | },
2653 | }
2654 |
2655 | for _, tc := range invokeTcs {
2656 | t.Run(tc.name, func(t *testing.T) {
2657 | const api = "http://127.0.0.1:5000/api/tool/list_roles/invoke"
2658 |
2659 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
2660 | if resp.StatusCode != tc.wantStatusCode {
2661 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
2662 | }
2663 | if tc.wantStatusCode != http.StatusOK {
2664 | return
2665 | }
2666 |
2667 | var bodyWrapper struct {
2668 | Result json.RawMessage `json:"result"`
2669 | }
2670 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
2671 | t.Fatalf("error decoding response wrapper: %v", err)
2672 | }
2673 |
2674 | var resultString string
2675 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
2676 | resultString = string(bodyWrapper.Result)
2677 | }
2678 |
2679 | var got []map[string]any
2680 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
2681 | t.Fatalf("failed to unmarshal nested result string: %v, resultString: %s", err, resultString)
2682 | }
2683 |
2684 | gotMap := make(map[string]map[string]any)
2685 | for _, role := range got {
2686 | // Remove fields that change every run
2687 | delete(role, "oid")
2688 | delete(role, "valid_until")
2689 |
2690 | if name, ok := role["role_name"].(string); ok {
2691 | gotMap[name] = role
2692 | }
2693 | }
2694 |
2695 | // Check that every role in 'want' exists in 'got' and matches
2696 | for _, wantRole := range tc.want {
2697 | roleName, _ := wantRole["role_name"].(string)
2698 |
2699 | gotRole, exists := gotMap[roleName]
2700 | if !exists {
2701 | t.Errorf("Expected role %q was not found in the response", roleName)
2702 | continue
2703 | }
2704 |
2705 | if diff := cmp.Diff(wantRole, gotRole); diff != "" {
2706 | t.Errorf("Role %q mismatch (-want +got):\n%s", roleName, diff)
2707 | }
2708 | }
2709 |
2710 | // Verify that if want is nil/empty, got is also empty
2711 | if len(tc.want) == 0 && len(got) != 0 {
2712 | t.Errorf("Expected empty result, but got %d roles", len(got))
2713 | }
2714 | })
2715 | }
2716 | }
2717 |
2718 | // RunMySQLListTablesTest run tests against the mysql-list-tables tool
2719 | func RunMySQLListTablesTest(t *testing.T, databaseName, tableNameParam, tableNameAuth, expectedOwner string) {
2720 | var ownerWant any
2721 | if expectedOwner == "" {
2722 | ownerWant = nil
2723 | } else {
2724 | ownerWant = expectedOwner
2725 | }
2726 |
2727 | type tableInfo struct {
2728 | ObjectName string `json:"object_name"`
2729 | SchemaName string `json:"schema_name"`
2730 | ObjectDetails string `json:"object_details"`
2731 | }
2732 |
2733 | type column struct {
2734 | DataType string `json:"data_type"`
2735 | ColumnName string `json:"column_name"`
2736 | ColumnComment string `json:"column_comment"`
2737 | ColumnDefault any `json:"column_default"`
2738 | IsNotNullable int `json:"is_not_nullable"`
2739 | OrdinalPosition int `json:"ordinal_position"`
2740 | }
2741 |
2742 | type objectDetails struct {
2743 | Owner any `json:"owner"`
2744 | Columns []column `json:"columns"`
2745 | Comment string `json:"comment"`
2746 | Indexes []any `json:"indexes"`
2747 | Triggers []any `json:"triggers"`
2748 | Constraints []any `json:"constraints"`
2749 | ObjectName string `json:"object_name"`
2750 | ObjectType string `json:"object_type"`
2751 | SchemaName string `json:"schema_name"`
2752 | }
2753 |
2754 | paramTableWant := objectDetails{
2755 | ObjectName: tableNameParam,
2756 | SchemaName: databaseName,
2757 | ObjectType: "TABLE",
2758 | Owner: ownerWant,
2759 | Columns: []column{
2760 | {DataType: "int", ColumnName: "id", IsNotNullable: 1, OrdinalPosition: 1},
2761 | {DataType: "varchar(255)", ColumnName: "name", OrdinalPosition: 2},
2762 | },
2763 | Indexes: []any{map[string]any{"index_columns": []any{"id"}, "index_name": "PRIMARY", "is_primary": float64(1), "is_unique": float64(1)}},
2764 | Triggers: []any{},
2765 | Constraints: []any{map[string]any{"constraint_columns": []any{"id"}, "constraint_name": "PRIMARY", "constraint_type": "PRIMARY KEY", "foreign_key_referenced_columns": any(nil), "foreign_key_referenced_table": any(nil), "constraint_definition": ""}},
2766 | }
2767 |
2768 | authTableWant := objectDetails{
2769 | ObjectName: tableNameAuth,
2770 | SchemaName: databaseName,
2771 | ObjectType: "TABLE",
2772 | Owner: ownerWant,
2773 | Columns: []column{
2774 | {DataType: "int", ColumnName: "id", IsNotNullable: 1, OrdinalPosition: 1},
2775 | {DataType: "varchar(255)", ColumnName: "name", OrdinalPosition: 2},
2776 | {DataType: "varchar(255)", ColumnName: "email", OrdinalPosition: 3},
2777 | },
2778 | Indexes: []any{map[string]any{"index_columns": []any{"id"}, "index_name": "PRIMARY", "is_primary": float64(1), "is_unique": float64(1)}},
2779 | Triggers: []any{},
2780 | Constraints: []any{map[string]any{"constraint_columns": []any{"id"}, "constraint_name": "PRIMARY", "constraint_type": "PRIMARY KEY", "foreign_key_referenced_columns": any(nil), "foreign_key_referenced_table": any(nil), "constraint_definition": ""}},
2781 | }
2782 |
2783 | invokeTcs := []struct {
2784 | name string
2785 | requestBody io.Reader
2786 | wantStatusCode int
2787 | want any
2788 | isSimple bool
2789 | isAllTables bool
2790 | }{
2791 | {
2792 | name: "invoke list_tables for all tables detailed output",
2793 | requestBody: bytes.NewBufferString(`{"table_names":""}`),
2794 | wantStatusCode: http.StatusOK,
2795 | want: []objectDetails{authTableWant, paramTableWant},
2796 | isAllTables: true,
2797 | },
2798 | {
2799 | name: "invoke list_tables detailed output",
2800 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_names": "%s"}`, tableNameAuth)),
2801 | wantStatusCode: http.StatusOK,
2802 | want: []objectDetails{authTableWant},
2803 | },
2804 | {
2805 | name: "invoke list_tables simple output",
2806 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_names": "%s", "output_format": "simple"}`, tableNameAuth)),
2807 | wantStatusCode: http.StatusOK,
2808 | want: []map[string]any{{"name": tableNameAuth}},
2809 | isSimple: true,
2810 | },
2811 | {
2812 | name: "invoke list_tables with multiple table names",
2813 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_names": "%s,%s"}`, tableNameParam, tableNameAuth)),
2814 | wantStatusCode: http.StatusOK,
2815 | want: []objectDetails{authTableWant, paramTableWant},
2816 | },
2817 | {
2818 | name: "invoke list_tables with one existing and one non-existent table",
2819 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_names": "%s,non_existent_table"}`, tableNameAuth)),
2820 | wantStatusCode: http.StatusOK,
2821 | want: []objectDetails{authTableWant},
2822 | },
2823 | {
2824 | name: "invoke list_tables with non-existent table",
2825 | requestBody: bytes.NewBufferString(`{"table_names": "non_existent_table"}`),
2826 | wantStatusCode: http.StatusOK,
2827 | want: []objectDetails{},
2828 | },
2829 | }
2830 | for _, tc := range invokeTcs {
2831 | t.Run(tc.name, func(t *testing.T) {
2832 | const api = "http://127.0.0.1:5000/api/tool/list_tables/invoke"
2833 | resp, body := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
2834 | if resp.StatusCode != tc.wantStatusCode {
2835 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(body))
2836 | }
2837 | if tc.wantStatusCode != http.StatusOK {
2838 | return
2839 | }
2840 |
2841 | var bodyWrapper struct {
2842 | Result json.RawMessage `json:"result"`
2843 | }
2844 | if err := json.Unmarshal(body, &bodyWrapper); err != nil {
2845 | t.Fatalf("error decoding response wrapper: %v", err)
2846 | }
2847 |
2848 | var resultString string
2849 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
2850 | resultString = string(bodyWrapper.Result)
2851 | }
2852 |
2853 | var got any
2854 | if tc.isSimple {
2855 | var tables []tableInfo
2856 | if err := json.Unmarshal([]byte(resultString), &tables); err != nil {
2857 | t.Fatalf("failed to unmarshal outer JSON array into []tableInfo: %v", err)
2858 | }
2859 | details := []map[string]any{}
2860 | for _, table := range tables {
2861 | var d map[string]any
2862 | if err := json.Unmarshal([]byte(table.ObjectDetails), &d); err != nil {
2863 | t.Fatalf("failed to unmarshal nested ObjectDetails string: %v", err)
2864 | }
2865 | details = append(details, d)
2866 | }
2867 | got = details
2868 | } else {
2869 | var tables []tableInfo
2870 | if err := json.Unmarshal([]byte(resultString), &tables); err != nil {
2871 | t.Fatalf("failed to unmarshal outer JSON array into []tableInfo: %v", err)
2872 | }
2873 | details := []objectDetails{}
2874 | for _, table := range tables {
2875 | var d objectDetails
2876 | if err := json.Unmarshal([]byte(table.ObjectDetails), &d); err != nil {
2877 | t.Fatalf("failed to unmarshal nested ObjectDetails string: %v", err)
2878 | }
2879 | details = append(details, d)
2880 | }
2881 | got = details
2882 | }
2883 |
2884 | opts := []cmp.Option{
2885 | cmpopts.SortSlices(func(a, b objectDetails) bool { return a.ObjectName < b.ObjectName }),
2886 | cmpopts.SortSlices(func(a, b column) bool { return a.ColumnName < b.ColumnName }),
2887 | cmpopts.SortSlices(func(a, b map[string]any) bool { return a["name"].(string) < b["name"].(string) }),
2888 | }
2889 |
2890 | // Checking only the current database where the test tables are created to avoid brittle tests.
2891 | if tc.isAllTables {
2892 | filteredGot := []objectDetails{}
2893 | if got != nil {
2894 | for _, item := range got.([]objectDetails) {
2895 | if item.SchemaName == databaseName {
2896 | filteredGot = append(filteredGot, item)
2897 | }
2898 | }
2899 | }
2900 | got = filteredGot
2901 | }
2902 |
2903 | if diff := cmp.Diff(tc.want, got, opts...); diff != "" {
2904 | t.Errorf("Unexpected result: got %#v, want: %#v", got, tc.want)
2905 | }
2906 | })
2907 | }
2908 | }
2909 |
2910 | // RunMySQLListActiveQueriesTest run tests against the mysql-list-active-queries tests
2911 | func RunMySQLListActiveQueriesTest(t *testing.T, ctx context.Context, pool *sql.DB) {
2912 | type queryListDetails struct {
2913 | ProcessId any `json:"process_id"`
2914 | Query string `json:"query"`
2915 | TrxStarted any `json:"trx_started"`
2916 | TrxDuration any `json:"trx_duration_seconds"`
2917 | TrxWaitDuration any `json:"trx_wait_duration_seconds"`
2918 | QueryTime any `json:"query_time"`
2919 | TrxState string `json:"trx_state"`
2920 | ProcessState string `json:"process_state"`
2921 | User string `json:"user"`
2922 | TrxRowsLocked any `json:"trx_rows_locked"`
2923 | TrxRowsModified any `json:"trx_rows_modified"`
2924 | Db string `json:"db"`
2925 | }
2926 |
2927 | singleQueryWanted := queryListDetails{
2928 | ProcessId: any(nil),
2929 | Query: "SELECT sleep(10)",
2930 | TrxStarted: any(nil),
2931 | TrxDuration: any(nil),
2932 | TrxWaitDuration: any(nil),
2933 | QueryTime: any(nil),
2934 | TrxState: "",
2935 | ProcessState: "User sleep",
2936 | User: "",
2937 | TrxRowsLocked: any(nil),
2938 | TrxRowsModified: any(nil),
2939 | Db: "",
2940 | }
2941 |
2942 | invokeTcs := []struct {
2943 | name string
2944 | requestBody io.Reader
2945 | clientSleepSecs int
2946 | waitSecsBeforeCheck int
2947 | wantStatusCode int
2948 | want any
2949 | }{
2950 | {
2951 | name: "invoke list_active_queries when the system is idle",
2952 | requestBody: bytes.NewBufferString(`{}`),
2953 | clientSleepSecs: 0,
2954 | waitSecsBeforeCheck: 0,
2955 | wantStatusCode: http.StatusOK,
2956 | want: []queryListDetails(nil),
2957 | },
2958 | {
2959 | name: "invoke list_active_queries when there is 1 ongoing but lower than the threshold",
2960 | requestBody: bytes.NewBufferString(`{"min_duration_secs": 100}`),
2961 | clientSleepSecs: 10,
2962 | waitSecsBeforeCheck: 1,
2963 | wantStatusCode: http.StatusOK,
2964 | want: []queryListDetails(nil),
2965 | },
2966 | {
2967 | name: "invoke list_active_queries when 1 ongoing query should show up",
2968 | requestBody: bytes.NewBufferString(`{"min_duration_secs": 5}`),
2969 | clientSleepSecs: 0,
2970 | waitSecsBeforeCheck: 5,
2971 | wantStatusCode: http.StatusOK,
2972 | want: []queryListDetails{singleQueryWanted},
2973 | },
2974 | {
2975 | name: "invoke list_active_queries when 2 ongoing query should show up",
2976 | requestBody: bytes.NewBufferString(`{"min_duration_secs": 2}`),
2977 | clientSleepSecs: 10,
2978 | waitSecsBeforeCheck: 3,
2979 | wantStatusCode: http.StatusOK,
2980 | want: []queryListDetails{singleQueryWanted, singleQueryWanted},
2981 | },
2982 | }
2983 |
2984 | var wg sync.WaitGroup
2985 | for _, tc := range invokeTcs {
2986 | t.Run(tc.name, func(t *testing.T) {
2987 | if tc.clientSleepSecs > 0 {
2988 | wg.Add(1)
2989 |
2990 | go func() {
2991 | defer wg.Done()
2992 |
2993 | err := pool.PingContext(ctx)
2994 | if err != nil {
2995 | t.Errorf("unable to connect to test database: %s", err)
2996 | return
2997 | }
2998 | _, err = pool.ExecContext(ctx, fmt.Sprintf("SELECT sleep(%d);", tc.clientSleepSecs))
2999 | if err != nil {
3000 | t.Errorf("Executing 'SELECT sleep' failed: %s", err)
3001 | }
3002 | }()
3003 | }
3004 |
3005 | if tc.waitSecsBeforeCheck > 0 {
3006 | time.Sleep(time.Duration(tc.waitSecsBeforeCheck) * time.Second)
3007 | }
3008 |
3009 | const api = "http://127.0.0.1:5000/api/tool/list_active_queries/invoke"
3010 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
3011 | if resp.StatusCode != tc.wantStatusCode {
3012 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
3013 | }
3014 | if tc.wantStatusCode != http.StatusOK {
3015 | return
3016 | }
3017 |
3018 | var bodyWrapper struct {
3019 | Result json.RawMessage `json:"result"`
3020 | }
3021 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
3022 | t.Fatalf("error decoding response wrapper: %v", err)
3023 | }
3024 |
3025 | var resultString string
3026 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
3027 | resultString = string(bodyWrapper.Result)
3028 | }
3029 |
3030 | var got any
3031 | var details []queryListDetails
3032 | if err := json.Unmarshal([]byte(resultString), &details); err != nil {
3033 | t.Fatalf("failed to unmarshal nested ObjectDetails string: %v", err)
3034 | }
3035 | got = details
3036 |
3037 | if diff := cmp.Diff(tc.want, got, cmp.Comparer(func(a, b queryListDetails) bool {
3038 | return a.Query == b.Query && a.ProcessState == b.ProcessState
3039 | })); diff != "" {
3040 | t.Errorf("Unexpected result: got %#v, want: %#v", got, tc.want)
3041 | }
3042 | })
3043 | }
3044 | wg.Wait()
3045 | }
3046 |
3047 | func RunMySQLListTablesMissingUniqueIndexes(t *testing.T, ctx context.Context, pool *sql.DB, databaseName string) {
3048 | type listDetails struct {
3049 | TableSchema string `json:"table_schema"`
3050 | TableName string `json:"table_name"`
3051 | }
3052 |
3053 | // bunch of wanted
3054 | nonUniqueKeyTableName := "t03_non_unqiue_key_table"
3055 | noKeyTableName := "t04_no_key_table"
3056 | nonUniqueKeyTableWant := listDetails{
3057 | TableSchema: databaseName,
3058 | TableName: nonUniqueKeyTableName,
3059 | }
3060 | noKeyTableWant := listDetails{
3061 | TableSchema: databaseName,
3062 | TableName: noKeyTableName,
3063 | }
3064 |
3065 | invokeTcs := []struct {
3066 | name string
3067 | requestBody io.Reader
3068 | newTableName string
3069 | newTablePrimaryKey bool
3070 | newTableUniqueKey bool
3071 | newTableNonUniqueKey bool
3072 | wantStatusCode int
3073 | want any
3074 | }{
3075 | {
3076 | name: "invoke list_tables_missing_unique_indexes when nothing to be found",
3077 | requestBody: bytes.NewBufferString(`{}`),
3078 | newTableName: "",
3079 | newTablePrimaryKey: false,
3080 | newTableUniqueKey: false,
3081 | newTableNonUniqueKey: false,
3082 | wantStatusCode: http.StatusOK,
3083 | want: []listDetails(nil),
3084 | },
3085 | {
3086 | name: "invoke list_tables_missing_unique_indexes pk table will not show",
3087 | requestBody: bytes.NewBufferString(`{}`),
3088 | newTableName: "t01",
3089 | newTablePrimaryKey: true,
3090 | newTableUniqueKey: false,
3091 | newTableNonUniqueKey: false,
3092 | wantStatusCode: http.StatusOK,
3093 | want: []listDetails(nil),
3094 | },
3095 | {
3096 | name: "invoke list_tables_missing_unique_indexes uk table will not show",
3097 | requestBody: bytes.NewBufferString(`{}`),
3098 | newTableName: "t02",
3099 | newTablePrimaryKey: false,
3100 | newTableUniqueKey: true,
3101 | newTableNonUniqueKey: false,
3102 | wantStatusCode: http.StatusOK,
3103 | want: []listDetails(nil),
3104 | },
3105 | {
3106 | name: "invoke list_tables_missing_unique_indexes non-unique key only table will show",
3107 | requestBody: bytes.NewBufferString(`{}`),
3108 | newTableName: nonUniqueKeyTableName,
3109 | newTablePrimaryKey: false,
3110 | newTableUniqueKey: false,
3111 | newTableNonUniqueKey: true,
3112 | wantStatusCode: http.StatusOK,
3113 | want: []listDetails{nonUniqueKeyTableWant},
3114 | },
3115 | {
3116 | name: "invoke list_tables_missing_unique_indexes table with no key at all will show",
3117 | requestBody: bytes.NewBufferString(`{}`),
3118 | newTableName: noKeyTableName,
3119 | newTablePrimaryKey: false,
3120 | newTableUniqueKey: false,
3121 | newTableNonUniqueKey: false,
3122 | wantStatusCode: http.StatusOK,
3123 | want: []listDetails{nonUniqueKeyTableWant, noKeyTableWant},
3124 | },
3125 | {
3126 | name: "invoke list_tables_missing_unique_indexes table w/ both pk & uk will not show",
3127 | requestBody: bytes.NewBufferString(`{}`),
3128 | newTableName: "t05",
3129 | newTablePrimaryKey: true,
3130 | newTableUniqueKey: true,
3131 | newTableNonUniqueKey: false,
3132 | wantStatusCode: http.StatusOK,
3133 | want: []listDetails{nonUniqueKeyTableWant, noKeyTableWant},
3134 | },
3135 | {
3136 | name: "invoke list_tables_missing_unique_indexes table w/ uk & nk will not show",
3137 | requestBody: bytes.NewBufferString(`{}`),
3138 | newTableName: "t06",
3139 | newTablePrimaryKey: false,
3140 | newTableUniqueKey: true,
3141 | newTableNonUniqueKey: true,
3142 | wantStatusCode: http.StatusOK,
3143 | want: []listDetails{nonUniqueKeyTableWant, noKeyTableWant},
3144 | },
3145 | {
3146 | name: "invoke list_tables_missing_unique_indexes table w/ pk & nk will not show",
3147 | requestBody: bytes.NewBufferString(`{}`),
3148 | newTableName: "t07",
3149 | newTablePrimaryKey: true,
3150 | newTableUniqueKey: false,
3151 | newTableNonUniqueKey: true,
3152 | wantStatusCode: http.StatusOK,
3153 | want: []listDetails{nonUniqueKeyTableWant, noKeyTableWant},
3154 | },
3155 | {
3156 | name: "invoke list_tables_missing_unique_indexes with a non-exist database, nothing to show",
3157 | requestBody: bytes.NewBufferString(`{"table_schema": "non-exist-database"}`),
3158 | newTableName: "",
3159 | newTablePrimaryKey: false,
3160 | newTableUniqueKey: false,
3161 | newTableNonUniqueKey: false,
3162 | wantStatusCode: http.StatusOK,
3163 | want: []listDetails(nil),
3164 | },
3165 | {
3166 | name: "invoke list_tables_missing_unique_indexes with the right database, show everything",
3167 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_schema": "%s"}`, databaseName)),
3168 | newTableName: "",
3169 | newTablePrimaryKey: false,
3170 | newTableUniqueKey: false,
3171 | newTableNonUniqueKey: false,
3172 | wantStatusCode: http.StatusOK,
3173 | want: []listDetails{nonUniqueKeyTableWant, noKeyTableWant},
3174 | },
3175 | {
3176 | name: "invoke list_tables_missing_unique_indexes with limited output",
3177 | requestBody: bytes.NewBufferString(`{"limit": 1}`),
3178 | newTableName: "",
3179 | newTablePrimaryKey: false,
3180 | newTableUniqueKey: false,
3181 | newTableNonUniqueKey: false,
3182 | wantStatusCode: http.StatusOK,
3183 | want: []listDetails{nonUniqueKeyTableWant},
3184 | },
3185 | }
3186 |
3187 | createTableHelper := func(t *testing.T, tableName, databaseName string, primaryKey, uniqueKey, nonUniqueKey bool, ctx context.Context, pool *sql.DB) func() {
3188 | var stmt strings.Builder
3189 | stmt.WriteString(fmt.Sprintf("CREATE TABLE %s (", tableName))
3190 | stmt.WriteString("c1 INT")
3191 | if primaryKey {
3192 | stmt.WriteString(" PRIMARY KEY")
3193 | }
3194 | stmt.WriteString(", c2 INT, c3 CHAR(8)")
3195 | if uniqueKey {
3196 | stmt.WriteString(", UNIQUE(c2)")
3197 | }
3198 | if nonUniqueKey {
3199 | stmt.WriteString(", INDEX(c3)")
3200 | }
3201 | stmt.WriteString(")")
3202 |
3203 | t.Logf("Creating table: %s", stmt.String())
3204 | if _, err := pool.ExecContext(ctx, stmt.String()); err != nil {
3205 | t.Fatalf("failed executing %s: %v", stmt.String(), err)
3206 | }
3207 |
3208 | return func() {
3209 | t.Logf("Dropping table: %s", tableName)
3210 | if _, err := pool.ExecContext(ctx, fmt.Sprintf("DROP TABLE %s", tableName)); err != nil {
3211 | t.Errorf("failed to drop table %s: %v", tableName, err)
3212 | }
3213 | }
3214 | }
3215 |
3216 | var cleanups []func()
3217 | defer func() {
3218 | for i := len(cleanups) - 1; i >= 0; i-- {
3219 | cleanups[i]()
3220 | }
3221 | }()
3222 |
3223 | for _, tc := range invokeTcs {
3224 | t.Run(tc.name, func(t *testing.T) {
3225 | if tc.newTableName != "" {
3226 | cleanup := createTableHelper(t, tc.newTableName, databaseName, tc.newTablePrimaryKey, tc.newTableUniqueKey, tc.newTableNonUniqueKey, ctx, pool)
3227 | cleanups = append(cleanups, cleanup)
3228 | }
3229 |
3230 | const api = "http://127.0.0.1:5000/api/tool/list_tables_missing_unique_indexes/invoke"
3231 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
3232 | if resp.StatusCode != tc.wantStatusCode {
3233 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
3234 | }
3235 | if tc.wantStatusCode != http.StatusOK {
3236 | return
3237 | }
3238 |
3239 | var bodyWrapper struct {
3240 | Result json.RawMessage `json:"result"`
3241 | }
3242 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
3243 | t.Fatalf("error decoding response wrapper: %v", err)
3244 | }
3245 |
3246 | var resultString string
3247 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
3248 | resultString = string(bodyWrapper.Result)
3249 | }
3250 |
3251 | var got any
3252 | var details []listDetails
3253 | if err := json.Unmarshal([]byte(resultString), &details); err != nil {
3254 | t.Fatalf("failed to unmarshal nested listDetails string: %v", err)
3255 | }
3256 | got = details
3257 |
3258 | if diff := cmp.Diff(tc.want, got, cmp.Comparer(func(a, b listDetails) bool {
3259 | return a.TableSchema == b.TableSchema && a.TableName == b.TableName
3260 | })); diff != "" {
3261 | t.Errorf("Unexpected result: got %#v, want: %#v", got, tc.want)
3262 | }
3263 | })
3264 | }
3265 | }
3266 |
3267 | func RunMySQLListTableFragmentationTest(t *testing.T, databaseName, tableNameParam, tableNameAuth string) {
3268 | type tableFragmentationDetails struct {
3269 | TableSchema string `json:"table_schema"`
3270 | TableName string `json:"table_name"`
3271 | DataSize any `json:"data_size"`
3272 | IndexSize any `json:"index_size"`
3273 | DataFree any `json:"data_free"`
3274 | FragmentationPercentage any `json:"fragmentation_percentage"`
3275 | }
3276 |
3277 | paramTableEntryWanted := tableFragmentationDetails{
3278 | TableSchema: databaseName,
3279 | TableName: tableNameParam,
3280 | DataSize: any(nil),
3281 | IndexSize: any(nil),
3282 | DataFree: any(nil),
3283 | FragmentationPercentage: any(nil),
3284 | }
3285 | authTableEntryWanted := tableFragmentationDetails{
3286 | TableSchema: databaseName,
3287 | TableName: tableNameAuth,
3288 | DataSize: any(nil),
3289 | IndexSize: any(nil),
3290 | DataFree: any(nil),
3291 | FragmentationPercentage: any(nil),
3292 | }
3293 |
3294 | invokeTcs := []struct {
3295 | name string
3296 | requestBody io.Reader
3297 | wantStatusCode int
3298 | want any
3299 | }{
3300 | {
3301 | name: "invoke list_table_fragmentation on all, no data_free threshold, expected to have 2 results",
3302 | requestBody: bytes.NewBufferString(`{"data_free_threshold_bytes": 0}`),
3303 | wantStatusCode: http.StatusOK,
3304 | want: []tableFragmentationDetails{authTableEntryWanted, paramTableEntryWanted},
3305 | },
3306 | {
3307 | name: "invoke list_table_fragmentation on all, no data_free threshold, limit to 1, expected to have 1 results",
3308 | requestBody: bytes.NewBufferString(`{"data_free_threshold_bytes": 0, "limit": 1}`),
3309 | wantStatusCode: http.StatusOK,
3310 | want: []tableFragmentationDetails{authTableEntryWanted},
3311 | },
3312 | {
3313 | name: "invoke list_table_fragmentation on all databases and 1 specific table name, no data_free threshold, expected to have 1 result",
3314 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_name": "%s","data_free_threshold_bytes": 0}`, tableNameAuth)),
3315 | wantStatusCode: http.StatusOK,
3316 | want: []tableFragmentationDetails{authTableEntryWanted},
3317 | },
3318 | {
3319 | name: "invoke list_table_fragmentation on 1 database and 1 specific table name, no data_free threshold, expected to have 1 result",
3320 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_schema": "%s", "table_name": "%s", "data_free_threshold_bytes": 0}`, databaseName, tableNameParam)),
3321 | wantStatusCode: http.StatusOK,
3322 | want: []tableFragmentationDetails{paramTableEntryWanted},
3323 | },
3324 | {
3325 | name: "invoke list_table_fragmentation on 1 database and 1 specific table name, high data_free threshold, expected to have 0 result",
3326 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_schema": "%s", "table_name": "%s", "data_free_threshold_bytes": 1000000000}`, databaseName, tableNameParam)),
3327 | wantStatusCode: http.StatusOK,
3328 | want: []tableFragmentationDetails(nil),
3329 | },
3330 | {
3331 | name: "invoke list_table_fragmentation on 1 non-exist database, no data_free threshold, expected to have 0 result",
3332 | requestBody: bytes.NewBufferString(`{"table_schema": "non_existent_database", "data_free_threshold_bytes": 0}`),
3333 | wantStatusCode: http.StatusOK,
3334 | want: []tableFragmentationDetails(nil),
3335 | },
3336 | {
3337 | name: "invoke list_table_fragmentation on 1 non-exist table, no data_free threshold, expected to have 0 result",
3338 | requestBody: bytes.NewBufferString(`{"table_name": "non_existent_table", "data_free_threshold_bytes": 0}`),
3339 | wantStatusCode: http.StatusOK,
3340 | want: []tableFragmentationDetails(nil),
3341 | },
3342 | }
3343 | for _, tc := range invokeTcs {
3344 | t.Run(tc.name, func(t *testing.T) {
3345 | const api = "http://127.0.0.1:5000/api/tool/list_table_fragmentation/invoke"
3346 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
3347 | if resp.StatusCode != tc.wantStatusCode {
3348 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
3349 | }
3350 | if tc.wantStatusCode != http.StatusOK {
3351 | return
3352 | }
3353 |
3354 | var bodyWrapper struct {
3355 | Result json.RawMessage `json:"result"`
3356 | }
3357 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
3358 | t.Fatalf("error decoding response wrapper: %v", err)
3359 | }
3360 |
3361 | var resultString string
3362 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
3363 | resultString = string(bodyWrapper.Result)
3364 | }
3365 |
3366 | var got any
3367 | var details []tableFragmentationDetails
3368 | if err := json.Unmarshal([]byte(resultString), &details); err != nil {
3369 | t.Fatalf("failed to unmarshal outer JSON array into []tableInfo: %v", err)
3370 | }
3371 | got = details
3372 |
3373 | if diff := cmp.Diff(tc.want, got, cmp.Comparer(func(a, b tableFragmentationDetails) bool {
3374 | return a.TableSchema == b.TableSchema && a.TableName == b.TableName
3375 | })); diff != "" {
3376 | t.Errorf("Unexpected result: got %#v, want: %#v", got, tc.want)
3377 | }
3378 | })
3379 | }
3380 | }
3381 |
3382 | func RunMySQLGetQueryPlanTest(t *testing.T, ctx context.Context, pool *sql.DB, databaseName, tableNameParam string) {
3383 | // Create a simple query to explain
3384 | query := fmt.Sprintf("SELECT * FROM %s", tableNameParam)
3385 |
3386 | invokeTcs := []struct {
3387 | name string
3388 | requestBody io.Reader
3389 | wantStatusCode int
3390 | checkResult func(t *testing.T, result any)
3391 | }{
3392 | {
3393 | name: "invoke get_query_plan with valid query",
3394 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"sql_statement": "%s"}`, query)),
3395 | wantStatusCode: http.StatusOK,
3396 | checkResult: func(t *testing.T, result any) {
3397 | resultMap, ok := result.(map[string]any)
3398 | if !ok {
3399 | t.Fatalf("result should be a map, got %T", result)
3400 | }
3401 | if _, ok := resultMap["query_block"]; !ok {
3402 | t.Errorf("result should contain 'query_block', got %v", resultMap)
3403 | }
3404 | },
3405 | },
3406 | {
3407 | name: "invoke get_query_plan with invalid query",
3408 | requestBody: bytes.NewBufferString(`{"sql_statement": "SELECT * FROM non_existent_table"}`),
3409 | wantStatusCode: http.StatusBadRequest,
3410 | checkResult: nil,
3411 | },
3412 | }
3413 |
3414 | for _, tc := range invokeTcs {
3415 | t.Run(tc.name, func(t *testing.T) {
3416 | const api = "http://127.0.0.1:5000/api/tool/get_query_plan/invoke"
3417 | resp, respBytes := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
3418 | if resp.StatusCode != tc.wantStatusCode {
3419 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBytes))
3420 | }
3421 | if tc.wantStatusCode != http.StatusOK {
3422 | return
3423 | }
3424 |
3425 | var bodyWrapper map[string]json.RawMessage
3426 |
3427 | if err := json.Unmarshal(respBytes, &bodyWrapper); err != nil {
3428 | t.Fatalf("error parsing response wrapper: %s, body: %s", err, string(respBytes))
3429 | }
3430 |
3431 | resultJSON, ok := bodyWrapper["result"]
3432 | if !ok {
3433 | t.Fatal("unable to find 'result' in response body")
3434 | }
3435 |
3436 | var resultString string
3437 | if err := json.Unmarshal(resultJSON, &resultString); err != nil {
3438 | if string(resultJSON) == "null" {
3439 | resultString = "null"
3440 | } else {
3441 | t.Fatalf("'result' is not a JSON-encoded string: %s", err)
3442 | }
3443 | }
3444 |
3445 | var got map[string]any
3446 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
3447 | t.Fatalf("failed to unmarshal actual result string: %v", err)
3448 | }
3449 |
3450 | if tc.checkResult != nil {
3451 | tc.checkResult(t, got)
3452 | }
3453 | })
3454 | }
3455 | }
3456 |
3457 | // RunMSSQLListTablesTest run tests againsts the mssql-list-tables tools.
3458 | func RunMSSQLListTablesTest(t *testing.T, tableNameParam, tableNameAuth string) {
3459 | // TableNameParam columns to construct want.
3460 | const paramTableColumns = `[
3461 | {"column_name": "id", "data_type": "INT", "column_ordinal_position": 1, "is_not_nullable": true},
3462 | {"column_name": "name", "data_type": "VARCHAR(255)", "column_ordinal_position": 2, "is_not_nullable": false}
3463 | ]`
3464 |
3465 | // TableNameAuth columns to construct want
3466 | const authTableColumns = `[
3467 | {"column_name": "id", "data_type": "INT", "column_ordinal_position": 1, "is_not_nullable": true},
3468 | {"column_name": "name", "data_type": "VARCHAR(255)", "column_ordinal_position": 2, "is_not_nullable": false},
3469 | {"column_name": "email", "data_type": "VARCHAR(255)", "column_ordinal_position": 3, "is_not_nullable": false}
3470 | ]`
3471 |
3472 | const (
3473 | // Template to construct detailed output want.
3474 | detailedObjectTemplate = `{
3475 | "schema_name": "dbo",
3476 | "object_name": "%[1]s",
3477 | "object_details": {
3478 | "owner": "dbo",
3479 | "triggers": [],
3480 | "columns": %[2]s,
3481 | "object_name": "%[1]s",
3482 | "object_type": "TABLE",
3483 | "schema_name": "dbo"
3484 | }
3485 | }`
3486 |
3487 | // Template to construct simple output want
3488 | simpleObjectTemplate = `{"object_name":"%s", "schema_name":"dbo", "object_details":{"name":"%s"}}`
3489 | )
3490 |
3491 | // Helper to build json for detailed want
3492 | getDetailedWant := func(tableName, columnJSON string) string {
3493 | return fmt.Sprintf(detailedObjectTemplate, tableName, columnJSON)
3494 | }
3495 |
3496 | // Helper to build template for simple want
3497 | getSimpleWant := func(tableName string) string {
3498 | return fmt.Sprintf(simpleObjectTemplate, tableName, tableName)
3499 | }
3500 |
3501 | invokeTcs := []struct {
3502 | name string
3503 | api string
3504 | requestBody string
3505 | wantStatusCode int
3506 | want string
3507 | isAllTables bool
3508 | }{
3509 | {
3510 | name: "invoke list_tables for all tables detailed output",
3511 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
3512 | requestBody: `{"table_names": ""}`,
3513 | wantStatusCode: http.StatusOK,
3514 | want: fmt.Sprintf("[%s,%s]", getDetailedWant(tableNameAuth, authTableColumns), getDetailedWant(tableNameParam, paramTableColumns)),
3515 | isAllTables: true,
3516 | },
3517 | {
3518 | name: "invoke list_tables for all tables simple output",
3519 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
3520 | requestBody: `{"table_names": "", "output_format": "simple"}`,
3521 | wantStatusCode: http.StatusOK,
3522 | want: fmt.Sprintf("[%s,%s]", getSimpleWant(tableNameAuth), getSimpleWant(tableNameParam)),
3523 | isAllTables: true,
3524 | },
3525 | {
3526 | name: "invoke list_tables detailed output",
3527 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
3528 | requestBody: fmt.Sprintf(`{"table_names": "%s"}`, tableNameAuth),
3529 | wantStatusCode: http.StatusOK,
3530 | want: fmt.Sprintf("[%s]", getDetailedWant(tableNameAuth, authTableColumns)),
3531 | },
3532 | {
3533 | name: "invoke list_tables simple output",
3534 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
3535 | requestBody: fmt.Sprintf(`{"table_names": "%s", "output_format": "simple"}`, tableNameAuth),
3536 | wantStatusCode: http.StatusOK,
3537 | want: fmt.Sprintf("[%s]", getSimpleWant(tableNameAuth)),
3538 | },
3539 | {
3540 | name: "invoke list_tables with invalid output format",
3541 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
3542 | requestBody: `{"table_names": "", "output_format": "abcd"}`,
3543 | wantStatusCode: http.StatusBadRequest,
3544 | },
3545 | {
3546 | name: "invoke list_tables with malformed table_names parameter",
3547 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
3548 | requestBody: `{"table_names": 12345, "output_format": "detailed"}`,
3549 | wantStatusCode: http.StatusBadRequest,
3550 | },
3551 | {
3552 | name: "invoke list_tables with multiple table names",
3553 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
3554 | requestBody: fmt.Sprintf(`{"table_names": "%s,%s"}`, tableNameParam, tableNameAuth),
3555 | wantStatusCode: http.StatusOK,
3556 | want: fmt.Sprintf("[%s,%s]", getDetailedWant(tableNameAuth, authTableColumns), getDetailedWant(tableNameParam, paramTableColumns)),
3557 | },
3558 | {
3559 | name: "invoke list_tables with non-existent table",
3560 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
3561 | requestBody: `{"table_names": "non_existent_table"}`,
3562 | wantStatusCode: http.StatusOK,
3563 | want: `[]`,
3564 | },
3565 | {
3566 | name: "invoke list_tables with one existing and one non-existent table",
3567 | api: "http://127.0.0.1:5000/api/tool/list_tables/invoke",
3568 | requestBody: fmt.Sprintf(`{"table_names": "%s,non_existent_table"}`, tableNameParam),
3569 | wantStatusCode: http.StatusOK,
3570 | want: fmt.Sprintf("[%s]", getDetailedWant(tableNameParam, paramTableColumns)),
3571 | },
3572 | }
3573 | for _, tc := range invokeTcs {
3574 | t.Run(tc.name, func(t *testing.T) {
3575 | resp, respBytes := RunRequest(t, http.MethodPost, tc.api, bytes.NewBuffer([]byte(tc.requestBody)), nil)
3576 |
3577 | if resp.StatusCode != tc.wantStatusCode {
3578 | t.Fatalf("response status code is not %d, got %d: %s", tc.wantStatusCode, resp.StatusCode, string(respBytes))
3579 | }
3580 |
3581 | if tc.wantStatusCode == http.StatusOK {
3582 | var bodyWrapper map[string]json.RawMessage
3583 |
3584 | if err := json.Unmarshal(respBytes, &bodyWrapper); err != nil {
3585 | t.Fatalf("error parsing response wrapper: %s, body: %s", err, string(respBytes))
3586 | }
3587 |
3588 | resultJSON, ok := bodyWrapper["result"]
3589 | if !ok {
3590 | t.Fatal("unable to find 'result' in response body")
3591 | }
3592 |
3593 | var resultString string
3594 | if err := json.Unmarshal(resultJSON, &resultString); err != nil {
3595 | if string(resultJSON) == "null" {
3596 | resultString = "null"
3597 | } else {
3598 | t.Fatalf("'result' is not a JSON-encoded string: %s", err)
3599 | }
3600 | }
3601 |
3602 | var got, want []any
3603 |
3604 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
3605 | t.Fatalf("failed to unmarshal actual result string: %v", err)
3606 | }
3607 | if err := json.Unmarshal([]byte(tc.want), &want); err != nil {
3608 | t.Fatalf("failed to unmarshal expected want string: %v", err)
3609 | }
3610 |
3611 | for _, item := range got {
3612 | itemMap, ok := item.(map[string]any)
3613 | if !ok {
3614 | continue
3615 | }
3616 |
3617 | detailsStr, ok := itemMap["object_details"].(string)
3618 | if !ok {
3619 | continue
3620 | }
3621 |
3622 | var detailsMap map[string]any
3623 | if err := json.Unmarshal([]byte(detailsStr), &detailsMap); err != nil {
3624 | t.Fatalf("failed to unmarshal nested object_details string: %v", err)
3625 | }
3626 |
3627 | // clean unpredictable fields
3628 | delete(detailsMap, "constraints")
3629 | delete(detailsMap, "indexes")
3630 |
3631 | itemMap["object_details"] = detailsMap
3632 | }
3633 |
3634 | // Checking only the default dbo schema where the test tables are created to avoid brittle tests.
3635 | if tc.isAllTables {
3636 | var filteredGot []any
3637 | for _, item := range got {
3638 | if tableMap, ok := item.(map[string]interface{}); ok {
3639 | if schema, ok := tableMap["schema_name"]; ok && schema == "dbo" {
3640 | filteredGot = append(filteredGot, item)
3641 | }
3642 | }
3643 | }
3644 | got = filteredGot
3645 | }
3646 |
3647 | sort.SliceStable(got, func(i, j int) bool {
3648 | return fmt.Sprintf("%v", got[i]) < fmt.Sprintf("%v", got[j])
3649 | })
3650 | sort.SliceStable(want, func(i, j int) bool {
3651 | return fmt.Sprintf("%v", want[i]) < fmt.Sprintf("%v", want[j])
3652 | })
3653 |
3654 | if !reflect.DeepEqual(got, want) {
3655 | gotJSON, _ := json.MarshalIndent(got, "", " ")
3656 | wantJSON, _ := json.MarshalIndent(want, "", " ")
3657 | t.Errorf("Unexpected result:\ngot:\n%s\n\nwant:\n%s", string(gotJSON), string(wantJSON))
3658 | }
3659 | }
3660 | })
3661 | }
3662 | }
3663 |
3664 | // RunPostgresListLocksTest runs tests for the postgres list-locks tool
3665 | func RunPostgresListLocksTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
3666 | type lockDetails struct {
3667 | Pid any `json:"pid"`
3668 | Usename string `json:"usename"`
3669 | Database string `json:"database"`
3670 | RelName string `json:"relname"`
3671 | LockType string `json:"locktype"`
3672 | Mode string `json:"mode"`
3673 | Granted bool `json:"granted"`
3674 | FastPath bool `json:"fastpath"`
3675 | VirtualXid any `json:"virtualxid"`
3676 | TransactionId any `json:"transactionid"`
3677 | ClassId any `json:"classid"`
3678 | ObjId any `json:"objid"`
3679 | ObjSubId any `json:"objsubid"`
3680 | PageNumber any `json:"page"`
3681 | TupleNumber any `json:"tuple"`
3682 | VirtualBlock any `json:"virtualblock"`
3683 | BlockNumber any `json:"blockno"`
3684 | }
3685 |
3686 | invokeTcs := []struct {
3687 | name string
3688 | requestBody io.Reader
3689 | wantStatusCode int
3690 | expectResults bool
3691 | }{
3692 | {
3693 | name: "invoke list_locks with no arguments",
3694 | requestBody: bytes.NewBuffer([]byte(`{}`)),
3695 | wantStatusCode: http.StatusOK,
3696 | expectResults: false, // locks may or may not exist
3697 | },
3698 | }
3699 | for _, tc := range invokeTcs {
3700 | t.Run(tc.name, func(t *testing.T) {
3701 | const api = "http://127.0.0.1:5000/api/tool/list_locks/invoke"
3702 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
3703 | if resp.StatusCode != tc.wantStatusCode {
3704 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
3705 | }
3706 | if tc.wantStatusCode != http.StatusOK {
3707 | return
3708 | }
3709 |
3710 | var bodyWrapper struct {
3711 | Result json.RawMessage `json:"result"`
3712 | }
3713 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
3714 | t.Fatalf("error decoding response wrapper: %v", err)
3715 | }
3716 |
3717 | var resultString string
3718 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
3719 | resultString = string(bodyWrapper.Result)
3720 | }
3721 |
3722 | var got []lockDetails
3723 | if resultString != "null" {
3724 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
3725 | t.Fatalf("failed to unmarshal result: %v, result string: %s", err, resultString)
3726 | }
3727 | }
3728 |
3729 | // Verify that if results exist, they have the expected structure
3730 | for _, lock := range got {
3731 | if lock.LockType == "" {
3732 | t.Errorf("lock type should not be empty")
3733 | }
3734 | }
3735 | })
3736 | }
3737 | }
3738 |
3739 | // RunPostgresLongRunningTransactionsTest runs tests for the postgres long-running-transactions tool
3740 | func RunPostgresLongRunningTransactionsTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
3741 | type transactionDetails struct {
3742 | Pid any `json:"pid"`
3743 | Usename string `json:"usename"`
3744 | Database string `json:"database"`
3745 | ApplicationName string `json:"application_name"`
3746 | XactStart any `json:"xact_start"`
3747 | XactDurationSecs any `json:"xact_duration_secs"`
3748 | IdleInTransaction string `json:"idle_in_transaction"`
3749 | Query string `json:"query"`
3750 | }
3751 |
3752 | invokeTcs := []struct {
3753 | name string
3754 | requestBody io.Reader
3755 | wantStatusCode int
3756 | }{
3757 | {
3758 | name: "invoke long_running_transactions with default threshold",
3759 | requestBody: bytes.NewBuffer([]byte(`{}`)),
3760 | wantStatusCode: http.StatusOK,
3761 | },
3762 | {
3763 | name: "invoke long_running_transactions with custom threshold",
3764 | requestBody: bytes.NewBuffer([]byte(`{"min_transaction_duration_secs": 3600}`)),
3765 | wantStatusCode: http.StatusOK,
3766 | },
3767 | }
3768 | for _, tc := range invokeTcs {
3769 | t.Run(tc.name, func(t *testing.T) {
3770 | const api = "http://127.0.0.1:5000/api/tool/long_running_transactions/invoke"
3771 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
3772 | if resp.StatusCode != tc.wantStatusCode {
3773 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
3774 | }
3775 | if tc.wantStatusCode != http.StatusOK {
3776 | return
3777 | }
3778 |
3779 | var bodyWrapper struct {
3780 | Result json.RawMessage `json:"result"`
3781 | }
3782 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
3783 | t.Fatalf("error decoding response wrapper: %v", err)
3784 | }
3785 |
3786 | var resultString string
3787 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
3788 | resultString = string(bodyWrapper.Result)
3789 | }
3790 |
3791 | var got []transactionDetails
3792 | if resultString != "null" {
3793 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
3794 | t.Fatalf("failed to unmarshal result: %v, result string: %s", err, resultString)
3795 | }
3796 | }
3797 |
3798 | // Verify that if results exist, they have the expected structure
3799 | for _, tx := range got {
3800 | if tx.XactDurationSecs == nil {
3801 | t.Errorf("transaction duration should not be null for long-running transactions")
3802 | }
3803 | }
3804 | })
3805 | }
3806 | }
3807 |
3808 | // RunPostgresReplicationStatsTest runs tests for the postgres replication-stats tool
3809 | func RunPostgresReplicationStatsTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
3810 | type replicationStats struct {
3811 | ClientAddr string `json:"client_addr"`
3812 | Username string `json:"usename"`
3813 | ApplicationName string `json:"application_name"`
3814 | ClientHostname string `json:"client_hostname"`
3815 | BackendStart any `json:"backend_start"`
3816 | State string `json:"state"`
3817 | SyncState string `json:"sync_state"`
3818 | ReplyTime any `json:"reply_time"`
3819 | FlushLsn string `json:"flush_lsn"`
3820 | ReplayLsn string `json:"replay_lsn"`
3821 | WriteLag any `json:"write_lag"`
3822 | FlushLag any `json:"flush_lag"`
3823 | ReplayLag any `json:"replay_lag"`
3824 | SyncPriority any `json:"sync_priority"`
3825 | ReplicationSlotName any `json:"slot_name"`
3826 | IsStreaming bool `json:"is_streaming"`
3827 | }
3828 |
3829 | invokeTcs := []struct {
3830 | name string
3831 | requestBody io.Reader
3832 | wantStatusCode int
3833 | }{
3834 | {
3835 | name: "invoke replication_stats with no arguments",
3836 | requestBody: bytes.NewBuffer([]byte(`{}`)),
3837 | wantStatusCode: http.StatusOK,
3838 | },
3839 | }
3840 | for _, tc := range invokeTcs {
3841 | t.Run(tc.name, func(t *testing.T) {
3842 | const api = "http://127.0.0.1:5000/api/tool/replication_stats/invoke"
3843 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
3844 | if resp.StatusCode != tc.wantStatusCode {
3845 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
3846 | }
3847 | if tc.wantStatusCode != http.StatusOK {
3848 | return
3849 | }
3850 |
3851 | var bodyWrapper struct {
3852 | Result json.RawMessage `json:"result"`
3853 | }
3854 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
3855 | t.Fatalf("error decoding response wrapper: %v", err)
3856 | }
3857 |
3858 | var resultString string
3859 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
3860 | resultString = string(bodyWrapper.Result)
3861 | }
3862 |
3863 | var got []replicationStats
3864 | if resultString != "null" {
3865 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
3866 | t.Fatalf("failed to unmarshal result: %v, result string: %s", err, resultString)
3867 | }
3868 | }
3869 |
3870 | // Verify that if results exist, they have the expected structure
3871 | for _, stat := range got {
3872 | if stat.State == "" {
3873 | t.Errorf("replication state should not be empty")
3874 | }
3875 | }
3876 | })
3877 | }
3878 | }
3879 |
3880 | func RunPostgresGetColumnCardinalityTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
3881 | schemaName := "testschema_" + strings.ReplaceAll(uuid.New().String(), "-", "")
3882 | tableName := "table1_" + strings.ReplaceAll(uuid.New().String(), "-", "")
3883 | cleanup := setupPostgresSchemas(t, ctx, pool, schemaName)
3884 | defer cleanup()
3885 |
3886 | // Create table with multiple columns
3887 | createTableStmt := fmt.Sprintf(`
3888 | CREATE TABLE %s.%s (
3889 | id SERIAL PRIMARY KEY,
3890 | email VARCHAR(100) UNIQUE,
3891 | name VARCHAR(50),
3892 | status VARCHAR(20),
3893 | created_at TIMESTAMP
3894 | )
3895 | `, schemaName, tableName)
3896 |
3897 | if _, err := pool.Exec(ctx, createTableStmt); err != nil {
3898 | t.Fatalf("unable to create table: %s", err)
3899 | }
3900 |
3901 | // Insert larger sample data to ensure statistics are collected
3902 | insertStmt := fmt.Sprintf(`
3903 | INSERT INTO %s.%s (email, name, status, created_at) VALUES
3904 | ('[email protected]', 'Alice', 'active', NOW()),
3905 | ('[email protected]', 'Bob', 'inactive', NOW()),
3906 | ('[email protected]', 'Charlie', 'active', NOW()),
3907 | ('[email protected]', 'David', 'active', NOW()),
3908 | ('[email protected]', 'Eve', 'inactive', NOW()),
3909 | ('[email protected]', 'Frank', 'active', NOW()),
3910 | ('[email protected]', 'Grace', 'inactive', NOW()),
3911 | ('[email protected]', 'Henry', 'active', NOW()),
3912 | ('[email protected]', 'Ivy', 'active', NOW()),
3913 | ('[email protected]', 'Jack', 'inactive', NOW())
3914 | `, schemaName, tableName)
3915 |
3916 | if _, err := pool.Exec(ctx, insertStmt); err != nil {
3917 | t.Fatalf("unable to insert data: %s", err)
3918 | }
3919 |
3920 | // Run ANALYZE to update statistics
3921 | analyzeStmt := fmt.Sprintf(`ANALYZE %s.%s`, schemaName, tableName)
3922 | if _, err := pool.Exec(ctx, analyzeStmt); err != nil {
3923 | t.Fatalf("unable to run ANALYZE: %s", err)
3924 | }
3925 |
3926 | invokeTcs := []struct {
3927 | name string
3928 | requestBody io.Reader
3929 | wantStatusCode int
3930 | shouldHaveData bool // Whether we expect data in the response
3931 | }{
3932 | {
3933 | name: "get cardinality for a specific column",
3934 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "%s", "column_name": "email"}`, schemaName, tableName)),
3935 | wantStatusCode: http.StatusOK,
3936 | shouldHaveData: true,
3937 | },
3938 | {
3939 | name: "get cardinality for all columns",
3940 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "%s"}`, schemaName, tableName)),
3941 | wantStatusCode: http.StatusOK,
3942 | shouldHaveData: true,
3943 | },
3944 | {
3945 | name: "get cardinality with non-existent column",
3946 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "table_name": "%s", "column_name": "non_existent"}`, schemaName, tableName)),
3947 | wantStatusCode: http.StatusOK,
3948 | shouldHaveData: false,
3949 | },
3950 | {
3951 | name: "get cardinality with non-existent schema",
3952 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "non_existent_schema", "table_name": "%s"}`, tableName)),
3953 | wantStatusCode: http.StatusOK,
3954 | shouldHaveData: false,
3955 | },
3956 | }
3957 |
3958 | for _, tc := range invokeTcs {
3959 | t.Run(tc.name, func(t *testing.T) {
3960 | const api = "http://127.0.0.1:5000/api/tool/get_column_cardinality/invoke"
3961 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
3962 | if resp.StatusCode != tc.wantStatusCode {
3963 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
3964 | }
3965 | if tc.wantStatusCode != http.StatusOK {
3966 | return
3967 | }
3968 |
3969 | var bodyWrapper struct {
3970 | Result json.RawMessage `json:"result"`
3971 | }
3972 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
3973 | t.Fatalf("error decoding response wrapper: %v", err)
3974 | }
3975 |
3976 | var resultString string
3977 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
3978 | resultString = string(bodyWrapper.Result)
3979 | }
3980 |
3981 | var got []map[string]any
3982 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
3983 | t.Fatalf("failed to unmarshal nested result string: %v", err)
3984 | }
3985 |
3986 | // Verify that we got the expected data presence
3987 | if tc.shouldHaveData {
3988 | if len(got) == 0 {
3989 | t.Logf("warning: expected data but got empty result. This can happen if pg_stats is not populated yet.")
3990 | return
3991 | }
3992 |
3993 | // Verify column names and cardinality values
3994 | for _, row := range got {
3995 | columnName, ok := row["column_name"].(string)
3996 | if !ok {
3997 | t.Fatalf("column_name is not a string: %v", row["column_name"])
3998 | }
3999 |
4000 | // Check that estimated_cardinality is present and is a number
4001 | cardinality, ok := row["estimated_cardinality"]
4002 | if !ok {
4003 | t.Fatalf("estimated_cardinality is missing for column %s", columnName)
4004 | }
4005 |
4006 | // Convert to float64 for numeric checks
4007 | cardinalityFloat, ok := cardinality.(float64)
4008 | if !ok {
4009 | t.Fatalf("estimated_cardinality is not a number: %v", cardinality)
4010 | }
4011 |
4012 | // Cardinality should be >= 0
4013 | if cardinalityFloat < 0 {
4014 | t.Errorf("cardinality for column %s is negative: %v", columnName, cardinalityFloat)
4015 | }
4016 | }
4017 | } else {
4018 | if len(got) != 0 {
4019 | t.Errorf("expected no data but got: %v", got)
4020 | }
4021 | }
4022 | })
4023 | }
4024 | }
4025 |
4026 | func createPostgresExtension(t *testing.T, ctx context.Context, pool *pgxpool.Pool, extensionName string) func() {
4027 | createExtensionCmd := fmt.Sprintf("CREATE EXTENSION IF NOT EXISTS %s", extensionName)
4028 | _, err := pool.Exec(ctx, createExtensionCmd)
4029 | if err != nil {
4030 | t.Fatalf("failed to create extension: %v", err)
4031 | }
4032 | return func() {
4033 | dropExtensionCmd := fmt.Sprintf("DROP EXTENSION IF EXISTS %s", extensionName)
4034 | _, err := pool.Exec(ctx, dropExtensionCmd)
4035 | if err != nil {
4036 | t.Fatalf("failed to drop extension: %v", err)
4037 | }
4038 | }
4039 | }
4040 |
4041 | func RunPostgresListQueryStatsTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
4042 | // Insert a simple query by running a SELECT statement
4043 | // This will record statistics in pg_stat_statements
4044 | selectStmt := "SELECT 1 as test_query"
4045 | if _, err := pool.Exec(ctx, selectStmt); err != nil {
4046 | t.Logf("warning: unable to execute test query: %s", err)
4047 | }
4048 |
4049 | dropExtensionFunc := createPostgresExtension(t, ctx, pool, "pg_stat_statements")
4050 | defer dropExtensionFunc()
4051 |
4052 | type queryStatDetails struct {
4053 | Datname string `json:"datname"`
4054 | Query string `json:"query"`
4055 | Calls any `json:"calls"`
4056 | TotalExecTime any `json:"total_exec_time"`
4057 | MinExecTime any `json:"min_exec_time"`
4058 | MaxExecTime any `json:"max_exec_time"`
4059 | MeanExecTime any `json:"mean_exec_time"`
4060 | Rows any `json:"rows"`
4061 | SharedBlksHit any `json:"shared_blks_hit"`
4062 | SharedBlksRead any `json:"shared_blks_read"`
4063 | }
4064 |
4065 | invokeTcs := []struct {
4066 | name string
4067 | requestBody io.Reader
4068 | wantStatusCode int
4069 | }{
4070 | {
4071 | name: "list query stats with default limit",
4072 | requestBody: bytes.NewBufferString(`{}`),
4073 | wantStatusCode: http.StatusOK,
4074 | },
4075 | {
4076 | name: "list query stats with custom limit",
4077 | requestBody: bytes.NewBufferString(`{"limit": 10}`),
4078 | wantStatusCode: http.StatusOK,
4079 | },
4080 | {
4081 | name: "list query stats for specific database",
4082 | requestBody: bytes.NewBufferString(`{"database_name": "postgres"}`),
4083 | wantStatusCode: http.StatusOK,
4084 | },
4085 | {
4086 | name: "list query stats with non-existent database name",
4087 | requestBody: bytes.NewBufferString(`{"database_name": "non_existent_db_xyz"}`),
4088 | wantStatusCode: http.StatusOK,
4089 | },
4090 | }
4091 |
4092 | for _, tc := range invokeTcs {
4093 | t.Run(tc.name, func(t *testing.T) {
4094 | const api = "http://127.0.0.1:5000/api/tool/list_query_stats/invoke"
4095 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
4096 | if resp.StatusCode != tc.wantStatusCode {
4097 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
4098 | }
4099 | if tc.wantStatusCode != http.StatusOK {
4100 | return
4101 | }
4102 |
4103 | var bodyWrapper struct {
4104 | Result json.RawMessage `json:"result"`
4105 | }
4106 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
4107 | t.Fatalf("error decoding response wrapper: %v", err)
4108 | }
4109 |
4110 | var resultString string
4111 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
4112 | resultString = string(bodyWrapper.Result)
4113 | }
4114 |
4115 | var got []map[string]any
4116 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
4117 | t.Fatalf("failed to unmarshal nested result string: %v, resultString: %s", err, resultString)
4118 | }
4119 |
4120 | // For databases with pg_stat_statements available, verify response structure
4121 | if len(got) > 0 {
4122 | // Verify the response has the expected fields
4123 | requiredFields := []string{"datname", "query", "calls", "total_exec_time", "min_exec_time", "max_exec_time", "mean_exec_time", "rows", "shared_blks_hit", "shared_blks_read"}
4124 | for _, field := range requiredFields {
4125 | if _, ok := got[0][field]; !ok {
4126 | t.Errorf("missing expected field: %s in result: %v", field, got[0])
4127 | }
4128 | }
4129 |
4130 | // Verify data types
4131 | var stat queryStatDetails
4132 | statData, _ := json.Marshal(got[0])
4133 | if err := json.Unmarshal(statData, &stat); err != nil {
4134 | t.Logf("warning: unable to unmarshal query stat: %v", err)
4135 | }
4136 |
4137 | // Verify that results are ordered by total_exec_time (descending)
4138 | if len(got) > 1 {
4139 | for i := 0; i < len(got)-1; i++ {
4140 | currentTime, ok1 := got[i]["total_exec_time"].(float64)
4141 | nextTime, ok2 := got[i+1]["total_exec_time"].(float64)
4142 | if ok1 && ok2 && currentTime < nextTime {
4143 | t.Logf("warning: results may not be ordered by total_exec_time descending: %f vs %f", currentTime, nextTime)
4144 | }
4145 | }
4146 | }
4147 | }
4148 | })
4149 | }
4150 | }
4151 |
4152 | // RunPostgresListTableStatsTest runs tests for the postgres list-table-stats tool
4153 | func RunPostgresListTableStatsTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
4154 | type tableStatsDetails struct {
4155 | SchemaName string `json:"schema_name"`
4156 | TableName string `json:"table_name"`
4157 | Owner string `json:"owner"`
4158 | TotalSizeBytes any `json:"total_size_bytes"`
4159 | SeqScan any `json:"seq_scan"`
4160 | IdxScan any `json:"idx_scan"`
4161 | IdxScanRatioPercent float64 `json:"idx_scan_ratio_percent"`
4162 | LiveRows any `json:"live_rows"`
4163 | DeadRows any `json:"dead_rows"`
4164 | DeadRowRatioPercent float64 `json:"dead_row_ratio_percent"`
4165 | NTupIns any `json:"n_tup_ins"`
4166 | NTupUpd any `json:"n_tup_upd"`
4167 | NTupDel any `json:"n_tup_del"`
4168 | LastVacuum any `json:"last_vacuum"`
4169 | LastAutovacuum any `json:"last_autovacuum"`
4170 | LastAutoanalyze any `json:"last_autoanalyze"`
4171 | }
4172 |
4173 | // Create a test table to generate statistics
4174 | testTableName := "test_list_table_stats_" + strings.ReplaceAll(uuid.New().String(), "-", "")
4175 | createTableStmt := fmt.Sprintf(`
4176 | CREATE TABLE %s (
4177 | id SERIAL PRIMARY KEY,
4178 | name VARCHAR(100),
4179 | email VARCHAR(100)
4180 | )
4181 | `, testTableName)
4182 |
4183 | if _, err := pool.Exec(ctx, createTableStmt); err != nil {
4184 | t.Fatalf("unable to create test table: %s", err)
4185 | }
4186 | defer func() {
4187 | dropTableStmt := fmt.Sprintf("DROP TABLE IF EXISTS %s", testTableName)
4188 | if _, err := pool.Exec(ctx, dropTableStmt); err != nil {
4189 | t.Logf("warning: unable to drop test table: %v", err)
4190 | }
4191 | }()
4192 |
4193 | // Insert some data to generate statistics
4194 | insertStmt := fmt.Sprintf(`
4195 | INSERT INTO %s (name, email) VALUES
4196 | ('Alice', '[email protected]'),
4197 | ('Bob', '[email protected]'),
4198 | ('Charlie', '[email protected]'),
4199 | ('David', '[email protected]'),
4200 | ('Eve', '[email protected]')
4201 | `, testTableName)
4202 |
4203 | if _, err := pool.Exec(ctx, insertStmt); err != nil {
4204 | t.Fatalf("unable to insert test data: %s", err)
4205 | }
4206 |
4207 | // Run some sequential scans to generate statistics
4208 | for i := 0; i < 3; i++ {
4209 | selectStmt := fmt.Sprintf("SELECT * FROM %s WHERE name = 'Alice'", testTableName)
4210 | if _, err := pool.Exec(ctx, selectStmt); err != nil {
4211 | t.Logf("warning: unable to execute select: %v", err)
4212 | }
4213 | }
4214 |
4215 | // Run ANALYZE to update statistics
4216 | analyzeStmt := fmt.Sprintf("ANALYZE %s", testTableName)
4217 | if _, err := pool.Exec(ctx, analyzeStmt); err != nil {
4218 | t.Logf("warning: unable to run ANALYZE: %v", err)
4219 | }
4220 |
4221 | invokeTcs := []struct {
4222 | name string
4223 | requestBody io.Reader
4224 | wantStatusCode int
4225 | shouldHaveData bool
4226 | filterTable bool
4227 | }{
4228 | {
4229 | name: "list table stats with no arguments (default limit)",
4230 | requestBody: bytes.NewBufferString(`{}`),
4231 | wantStatusCode: http.StatusOK,
4232 | shouldHaveData: false, // may or may not have data depending on what's in the database
4233 | },
4234 | {
4235 | name: "list table stats with default limit",
4236 | requestBody: bytes.NewBufferString(`{"schema_name": "public"}`),
4237 | wantStatusCode: http.StatusOK,
4238 | shouldHaveData: false,
4239 | },
4240 | {
4241 | name: "list table stats filtering by specific table",
4242 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"table_name": "%s"}`, testTableName)),
4243 | wantStatusCode: http.StatusOK,
4244 | shouldHaveData: true,
4245 | filterTable: true,
4246 | },
4247 | {
4248 | name: "list table stats with custom limit",
4249 | requestBody: bytes.NewBufferString(`{"limit": 10}`),
4250 | wantStatusCode: http.StatusOK,
4251 | shouldHaveData: false,
4252 | },
4253 | {
4254 | name: "list table stats sorted by size",
4255 | requestBody: bytes.NewBufferString(`{"sort_by": "size", "limit": 5}`),
4256 | wantStatusCode: http.StatusOK,
4257 | shouldHaveData: false,
4258 | },
4259 | {
4260 | name: "list table stats sorted by seq_scan",
4261 | requestBody: bytes.NewBufferString(`{"sort_by": "seq_scan", "limit": 5}`),
4262 | wantStatusCode: http.StatusOK,
4263 | shouldHaveData: false,
4264 | },
4265 | {
4266 | name: "list table stats sorted by idx_scan",
4267 | requestBody: bytes.NewBufferString(`{"sort_by": "idx_scan", "limit": 5}`),
4268 | wantStatusCode: http.StatusOK,
4269 | shouldHaveData: false,
4270 | },
4271 | {
4272 | name: "list table stats sorted by dead_rows",
4273 | requestBody: bytes.NewBufferString(`{"sort_by": "dead_rows", "limit": 5}`),
4274 | wantStatusCode: http.StatusOK,
4275 | shouldHaveData: false,
4276 | },
4277 | {
4278 | name: "list table stats with non-existent table filter",
4279 | requestBody: bytes.NewBufferString(`{"table_name": "non_existent_table_xyz"}`),
4280 | wantStatusCode: http.StatusOK,
4281 | shouldHaveData: false,
4282 | },
4283 | {
4284 | name: "list table stats with non-existent schema filter",
4285 | requestBody: bytes.NewBufferString(`{"schema_name": "non_existent_schema_xyz"}`),
4286 | wantStatusCode: http.StatusOK,
4287 | shouldHaveData: false,
4288 | },
4289 | {
4290 | name: "list table stats with owner filter",
4291 | requestBody: bytes.NewBufferString(`{"owner": "postgres"}`),
4292 | wantStatusCode: http.StatusOK,
4293 | shouldHaveData: false,
4294 | },
4295 | }
4296 |
4297 | for _, tc := range invokeTcs {
4298 | t.Run(tc.name, func(t *testing.T) {
4299 | const api = "http://127.0.0.1:5000/api/tool/list_table_stats/invoke"
4300 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
4301 | if resp.StatusCode != tc.wantStatusCode {
4302 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
4303 | }
4304 | if tc.wantStatusCode != http.StatusOK {
4305 | return
4306 | }
4307 |
4308 | var bodyWrapper struct {
4309 | Result json.RawMessage `json:"result"`
4310 | }
4311 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
4312 | t.Fatalf("error decoding response wrapper: %v", err)
4313 | }
4314 |
4315 | var resultString string
4316 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
4317 | resultString = string(bodyWrapper.Result)
4318 | }
4319 |
4320 | var got []tableStatsDetails
4321 | if resultString != "null" {
4322 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
4323 | t.Fatalf("failed to unmarshal result: %v, result string: %s", err, resultString)
4324 | }
4325 | }
4326 |
4327 | // Verify expected data presence
4328 | if tc.shouldHaveData {
4329 | if len(got) == 0 {
4330 | t.Fatalf("expected data but got empty result")
4331 | }
4332 |
4333 | // Verify the test table is in results
4334 | found := false
4335 | for _, row := range got {
4336 | if row.TableName == testTableName {
4337 | found = true
4338 | // Verify expected fields are present
4339 | if row.SchemaName == "" {
4340 | t.Errorf("schema_name should not be empty")
4341 | }
4342 | if row.Owner == "" {
4343 | t.Errorf("owner should not be empty")
4344 | }
4345 | if row.TotalSizeBytes == nil {
4346 | t.Errorf("total_size_bytes should not be null")
4347 | }
4348 | if row.LiveRows == nil {
4349 | t.Errorf("live_rows should not be null")
4350 | }
4351 | break
4352 | }
4353 | }
4354 |
4355 | if !found {
4356 | t.Errorf("test table %s not found in results", testTableName)
4357 | }
4358 | } else if tc.filterTable {
4359 | // For filtered queries that shouldn't find anything
4360 | if len(got) != 0 {
4361 | t.Logf("warning: expected no data but got: %v", len(got))
4362 | }
4363 | }
4364 |
4365 | // Verify result structure and data types
4366 | for _, stat := range got {
4367 | // Verify schema_name and table_name are strings
4368 | if stat.SchemaName == "" && stat.TableName != "" {
4369 | t.Errorf("schema_name is empty for table %s", stat.TableName)
4370 | }
4371 |
4372 | // Verify numeric fields are valid
4373 | if stat.IdxScanRatioPercent < 0 || stat.IdxScanRatioPercent > 100 {
4374 | t.Errorf("idx_scan_ratio_percent should be between 0 and 100, got %f", stat.IdxScanRatioPercent)
4375 | }
4376 |
4377 | if stat.DeadRowRatioPercent < 0 || stat.DeadRowRatioPercent > 100 {
4378 | t.Errorf("dead_row_ratio_percent should be between 0 and 100, got %f", stat.DeadRowRatioPercent)
4379 | }
4380 | }
4381 |
4382 | // Verify sorting for specific sort_by options
4383 | if tc.name == "list table stats sorted by size" && len(got) > 1 {
4384 | for i := 0; i < len(got)-1; i++ {
4385 | current, ok1 := got[i].TotalSizeBytes.(float64)
4386 | next, ok2 := got[i+1].TotalSizeBytes.(float64)
4387 | if ok1 && ok2 && current < next {
4388 | t.Logf("warning: results may not be sorted by total_size_bytes descending")
4389 | }
4390 | }
4391 | }
4392 | })
4393 | }
4394 | }
4395 |
4396 | // RunPostgresListStoredProcedureTest runs tests for the postgres list-stored-procedure tool
4397 | func RunPostgresListStoredProcedureTest(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
4398 | type storedProcedureDetails struct {
4399 | SchemaName string `json:"schema_name"`
4400 | Name string `json:"name"`
4401 | Owner string `json:"owner"`
4402 | Language string `json:"language"`
4403 | Definition string `json:"definition"`
4404 | Description any `json:"description"`
4405 | }
4406 |
4407 | // Create test schema
4408 | testSchemaName := "test_proc_schema_" + strings.ReplaceAll(uuid.New().String(), "-", "")
4409 | createSchemaStmt := fmt.Sprintf("CREATE SCHEMA %s", testSchemaName)
4410 | if _, err := pool.Exec(ctx, createSchemaStmt); err != nil {
4411 | t.Fatalf("unable to create test schema: %v", err)
4412 | }
4413 | defer func() {
4414 | dropSchemaStmt := fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE", testSchemaName)
4415 | if _, err := pool.Exec(ctx, dropSchemaStmt); err != nil {
4416 | t.Logf("warning: unable to drop test schema: %v", err)
4417 | }
4418 | }()
4419 |
4420 | // Create test procedures
4421 | proc1Name := "test_proc_1_" + strings.ReplaceAll(uuid.New().String(), "-", "")
4422 | createProc1Stmt := fmt.Sprintf(`
4423 | CREATE PROCEDURE %s.%s(p_count INT)
4424 | LANGUAGE plpgsql
4425 | AS $$
4426 | BEGIN
4427 | INSERT INTO test_table VALUES (p_count);
4428 | COMMIT;
4429 | END;
4430 | $$
4431 | `, testSchemaName, proc1Name)
4432 |
4433 | if _, err := pool.Exec(ctx, createProc1Stmt); err != nil {
4434 | t.Fatalf("unable to create test procedure 1: %v", err)
4435 | }
4436 |
4437 | // Add a comment/description to the procedure
4438 | commentStmt := fmt.Sprintf("COMMENT ON PROCEDURE %s.%s(INT) IS 'Test procedure that inserts a record'", testSchemaName, proc1Name)
4439 | if _, err := pool.Exec(ctx, commentStmt); err != nil {
4440 | t.Logf("warning: unable to add comment to procedure: %v", err)
4441 | }
4442 |
4443 | // Create a second test procedure
4444 | proc2Name := "test_proc_2_" + strings.ReplaceAll(uuid.New().String(), "-", "")
4445 | createProc2Stmt := fmt.Sprintf(`
4446 | CREATE PROCEDURE %s.%s()
4447 | LANGUAGE plpgsql
4448 | AS $$
4449 | DECLARE
4450 | v_count INT;
4451 | BEGIN
4452 | SELECT COUNT(*) INTO v_count FROM test_table;
4453 | RAISE NOTICE 'Total records: %%', v_count;
4454 | END;
4455 | $$
4456 | `, testSchemaName, proc2Name)
4457 |
4458 | if _, err := pool.Exec(ctx, createProc2Stmt); err != nil {
4459 | t.Fatalf("unable to create test procedure 2: %v", err)
4460 | }
4461 |
4462 | invokeTcs := []struct {
4463 | name string
4464 | requestBody io.Reader
4465 | wantStatusCode int
4466 | shouldHaveData bool
4467 | expectedCount int
4468 | filterByRole string
4469 | filterBySchema string
4470 | }{
4471 | {
4472 | name: "list stored procedures with no arguments (default limit 20)",
4473 | requestBody: bytes.NewBufferString(`{}`),
4474 | wantStatusCode: http.StatusOK,
4475 | shouldHaveData: false, // may or may not have data depending on what's in the database
4476 | },
4477 | {
4478 | name: "list stored procedures filtering by specific schema",
4479 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s"}`, testSchemaName)),
4480 | wantStatusCode: http.StatusOK,
4481 | shouldHaveData: true,
4482 | expectedCount: 2,
4483 | filterBySchema: testSchemaName,
4484 | },
4485 | {
4486 | name: "list stored procedures filtering by procedure owner (postgres)",
4487 | requestBody: bytes.NewBufferString(`{"role_name": "postgres"}`),
4488 | wantStatusCode: http.StatusOK,
4489 | shouldHaveData: false, // might have procedures owned by postgres
4490 | },
4491 | {
4492 | name: "list stored procedures with custom limit",
4493 | requestBody: bytes.NewBufferString(`{"limit": 5}`),
4494 | wantStatusCode: http.StatusOK,
4495 | shouldHaveData: false,
4496 | },
4497 | {
4498 | name: "list stored procedures filtering by schema and role",
4499 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "role_name": "postgres"}`, testSchemaName)),
4500 | wantStatusCode: http.StatusOK,
4501 | shouldHaveData: true,
4502 | expectedCount: 2,
4503 | filterBySchema: testSchemaName,
4504 | filterByRole: "postgres",
4505 | },
4506 | {
4507 | name: "list stored procedures with non-existent schema",
4508 | requestBody: bytes.NewBufferString(`{"schema_name": "non_existent_schema_xyz"}`),
4509 | wantStatusCode: http.StatusOK,
4510 | shouldHaveData: false,
4511 | },
4512 | {
4513 | name: "list stored procedures with non-existent role",
4514 | requestBody: bytes.NewBufferString(`{"role_name": "non_existent_role_xyz"}`),
4515 | wantStatusCode: http.StatusOK,
4516 | shouldHaveData: false,
4517 | },
4518 | {
4519 | name: "list stored procedures with partial schema name match",
4520 | requestBody: bytes.NewBufferString(`{"schema_name": "test_proc"}`),
4521 | wantStatusCode: http.StatusOK,
4522 | shouldHaveData: true,
4523 | expectedCount: 2,
4524 | },
4525 | {
4526 | name: "list stored procedures with limit 1",
4527 | requestBody: bytes.NewBufferString(fmt.Sprintf(`{"schema_name": "%s", "limit": 1}`, testSchemaName)),
4528 | wantStatusCode: http.StatusOK,
4529 | shouldHaveData: true,
4530 | expectedCount: 1,
4531 | filterBySchema: testSchemaName,
4532 | },
4533 | }
4534 |
4535 | for _, tc := range invokeTcs {
4536 | t.Run(tc.name, func(t *testing.T) {
4537 | const api = "http://127.0.0.1:5000/api/tool/list_stored_procedure/invoke"
4538 | resp, respBody := RunRequest(t, http.MethodPost, api, tc.requestBody, nil)
4539 | if resp.StatusCode != tc.wantStatusCode {
4540 | t.Fatalf("wrong status code: got %d, want %d, body: %s", resp.StatusCode, tc.wantStatusCode, string(respBody))
4541 | }
4542 | if tc.wantStatusCode != http.StatusOK {
4543 | return
4544 | }
4545 |
4546 | var bodyWrapper struct {
4547 | Result json.RawMessage `json:"result"`
4548 | }
4549 | if err := json.Unmarshal(respBody, &bodyWrapper); err != nil {
4550 | t.Fatalf("error decoding response wrapper: %v", err)
4551 | }
4552 |
4553 | var resultString string
4554 | if err := json.Unmarshal(bodyWrapper.Result, &resultString); err != nil {
4555 | resultString = string(bodyWrapper.Result)
4556 | }
4557 |
4558 | var got []storedProcedureDetails
4559 | if resultString != "null" {
4560 | if err := json.Unmarshal([]byte(resultString), &got); err != nil {
4561 | t.Fatalf("failed to unmarshal result: %v, result string: %s", err, resultString)
4562 | }
4563 | }
4564 |
4565 | // Verify expected data presence
4566 | if tc.shouldHaveData {
4567 | if len(got) == 0 {
4568 | t.Fatalf("expected data but got empty result")
4569 | }
4570 |
4571 | // If filtering by schema, verify all results are from that schema
4572 | if tc.filterBySchema != "" {
4573 | for _, proc := range got {
4574 | if proc.SchemaName != tc.filterBySchema && !strings.Contains(proc.SchemaName, tc.filterBySchema) {
4575 | t.Errorf("procedure schema %s does not match filter %s", proc.SchemaName, tc.filterBySchema)
4576 | }
4577 | }
4578 | }
4579 |
4580 | // If filtering by role, verify all results are owned by that role
4581 | if tc.filterByRole != "" {
4582 | for _, proc := range got {
4583 | if proc.Owner != tc.filterByRole {
4584 | t.Errorf("procedure owner %s does not match filter %s", proc.Owner, tc.filterByRole)
4585 | }
4586 | }
4587 | }
4588 |
4589 | // Verify expected count if specified
4590 | if tc.expectedCount > 0 && len(got) != tc.expectedCount {
4591 | t.Errorf("expected %d procedures but got %d", tc.expectedCount, len(got))
4592 | }
4593 | }
4594 |
4595 | // Verify result structure and data types
4596 | for _, proc := range got {
4597 | // Verify all required fields are present and non-empty
4598 | if proc.SchemaName == "" {
4599 | t.Errorf("schema_name should not be empty")
4600 | }
4601 | if proc.Name == "" {
4602 | t.Errorf("procedure name should not be empty")
4603 | }
4604 | if proc.Owner == "" {
4605 | t.Errorf("owner should not be empty")
4606 | }
4607 | if proc.Language == "" {
4608 | t.Errorf("language should not be empty")
4609 | }
4610 | if proc.Definition == "" {
4611 | t.Errorf("definition should not be empty")
4612 | }
4613 |
4614 | // Verify definition contains CREATE PROCEDURE
4615 | if !strings.Contains(proc.Definition, "CREATE PROCEDURE") {
4616 | t.Logf("warning: definition may not be a valid CREATE PROCEDURE statement: %s", proc.Definition)
4617 | }
4618 |
4619 | // Verify language is a valid PostgreSQL language
4620 | validLanguages := []string{"plpgsql", "sql", "c", "internal", "plperl", "pltcl", "plpython"}
4621 | found := false
4622 | for _, lang := range validLanguages {
4623 | if proc.Language == lang {
4624 | found = true
4625 | break
4626 | }
4627 | }
4628 | if !found {
4629 | t.Logf("warning: language %s may not be a standard PostgreSQL language", proc.Language)
4630 | }
4631 | }
4632 |
4633 | // Verify results are sorted by schema_name and name
4634 | if len(got) > 1 {
4635 | for i := 0; i < len(got)-1; i++ {
4636 | currentKey := fmt.Sprintf("%s.%s", got[i].SchemaName, got[i].Name)
4637 | nextKey := fmt.Sprintf("%s.%s", got[i+1].SchemaName, got[i+1].Name)
4638 | if currentKey > nextKey {
4639 | t.Logf("warning: results may not be sorted by schema_name and name")
4640 | }
4641 | }
4642 | }
4643 | })
4644 | }
4645 | }
4646 |
4647 | // RunRequest is a helper function to send HTTP requests and return the response
4648 | func RunRequest(t *testing.T, method, url string, body io.Reader, headers map[string]string) (*http.Response, []byte) {
4649 | // Send request
4650 | req, err := http.NewRequest(method, url, body)
4651 | if err != nil {
4652 | t.Fatalf("unable to create request: %s", err)
4653 | }
4654 |
4655 | req.Header.Set("Content-type", "application/json")
4656 |
4657 | for k, v := range headers {
4658 | req.Header.Set(k, v)
4659 | }
4660 |
4661 | resp, err := http.DefaultClient.Do(req)
4662 | if err != nil {
4663 | t.Fatalf("unable to send request: %s", err)
4664 | }
4665 | respBody, err := io.ReadAll(resp.Body)
4666 | if err != nil {
4667 | t.Fatalf("unable to read request body: %s", err)
4668 | }
4669 |
4670 | defer resp.Body.Close()
4671 | return resp, respBody
4672 | }
4673 |
```