From a0f44d34ea3f044dd08501be616f70ddfd63ab45 Mon Sep 17 00:00:00 2001 From: gRedHeadphone Date: Tue, 25 Nov 2025 21:39:25 +0530 Subject: [PATCH] feat(tools/spanner-list-graph): tool impl + docs + tests (#1923) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Spanner List Graphs tool, similar to list tables it can be used to get all/specific graph details ## PR Checklist - [x] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/genai-toolbox/blob/main/CONTRIBUTING.md) - [x] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/genai-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary) - [x] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #1916 --------- Co-authored-by: Averi Kitsch --- cmd/root.go | 1 + docs/en/resources/sources/spanner.md | 3 + .../tools/spanner/spanner-list-graphs.md | 276 ++++++++++++++++++ .../spannerlistgraphs/spannerlistgraphs.go | 255 ++++++++++++++++ .../spannerlistgraphs_test.go | 112 +++++++ tests/spanner/spanner_integration_test.go | 197 ++++++++++++- 6 files changed, 842 insertions(+), 2 deletions(-) create mode 100644 docs/en/resources/tools/spanner/spanner-list-graphs.md create mode 100644 internal/tools/spanner/spannerlistgraphs/spannerlistgraphs.go create mode 100644 internal/tools/spanner/spannerlistgraphs/spannerlistgraphs_test.go diff --git a/cmd/root.go b/cmd/root.go index 6f850309e0..04ad06c3ab 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -200,6 +200,7 @@ import ( _ "github.com/googleapis/genai-toolbox/internal/tools/singlestore/singlestoreexecutesql" _ "github.com/googleapis/genai-toolbox/internal/tools/singlestore/singlestoresql" _ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerexecutesql" + _ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerlistgraphs" _ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerlisttables" _ "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannersql" _ "github.com/googleapis/genai-toolbox/internal/tools/sqlite/sqliteexecutesql" diff --git a/docs/en/resources/sources/spanner.md b/docs/en/resources/sources/spanner.md index 2137332b33..373ba75446 100644 --- a/docs/en/resources/sources/spanner.md +++ b/docs/en/resources/sources/spanner.md @@ -34,6 +34,9 @@ the Google Cloud console][spanner-quickstart]. - [`spanner-list-tables`](../tools/spanner/spanner-list-tables.md) Retrieve schema information about tables in a Spanner database. +- [`spanner-list-graphs`](../tools/spanner/spanner-list-graphs.md) + Retrieve schema information about graphs in a Spanner database. + ### Pre-built Configurations - [Spanner using MCP](https://googleapis.github.io/genai-toolbox/how-to/connect-ide/spanner_mcp/) diff --git a/docs/en/resources/tools/spanner/spanner-list-graphs.md b/docs/en/resources/tools/spanner/spanner-list-graphs.md new file mode 100644 index 0000000000..e846fdfe77 --- /dev/null +++ b/docs/en/resources/tools/spanner/spanner-list-graphs.md @@ -0,0 +1,276 @@ +--- +title: "spanner-list-graphs" +type: docs +weight: 3 +description: > + A "spanner-list-graphs" tool retrieves schema information about graphs in a + Google Cloud Spanner database. +--- + +## About + +A `spanner-list-graphs` tool retrieves comprehensive schema information about +graphs in a Cloud Spanner database. It returns detailed metadata including +node tables, edge tables, labels and property declarations. It's compatible with: + +- [spanner](../../sources/spanner.md) + +This tool is read-only and executes pre-defined SQL queries against the +`INFORMATION_SCHEMA` tables to gather metadata. +{{< notice warning >}} +The tool only works for the GoogleSQL +source dialect, as Spanner Graph isn't available in the PostgreSQL dialect. +{{< /notice >}} + +## Features + +- **Comprehensive Schema Information**: Returns node tables, edge tables, labels + and property declarations +- **Flexible Filtering**: Can list all graphs or filter by specific graph names +- **Output Format Options**: Choose between simple (graph names only) or detailed + (full schema information) output + +## Example + +### Basic Usage - List All Graphs + +```yaml +sources: + my-spanner-db: + kind: spanner + project: ${SPANNER_PROJECT} + instance: ${SPANNER_INSTANCE} + database: ${SPANNER_DATABASE} + dialect: googlesql # wont work for postgresql + +tools: + list_all_graphs: + kind: spanner-list-graphs + source: my-spanner-db + description: Lists all graphs with their complete schema information +``` + +### List Specific Graphs + +```yaml +tools: + list_specific_graphs: + kind: spanner-list-graphs + source: my-spanner-db + description: | + Lists schema information for specific graphs. + Example usage: + { + "graph_names": "FinGraph,SocialGraph", + "output_format": "detailed" + } +``` + +## Parameters + +The tool accepts two optional parameters: + +| **parameter** | **type** | **default** | **description** | +|---------------|:--------:|:-----------:|------------------------------------------------------------------------------------------------------| +| graph_names | string | "" | Comma-separated list of graph names to filter. If empty, lists all graphs in user-accessible schemas | +| output_format | string | "detailed" | Output format: "simple" returns only graph names, "detailed" returns full schema information | + +## Output Format + +### Simple Format + +When `output_format` is set to "simple", the tool returns a minimal JSON structure: + +```json +[ + { + "object_details": { + "name": "FinGraph" + }, + "object_name": "FinGraph", + "schema_name": "" + }, + { + "object_details": { + "name": "SocialGraph" + }, + "object_name": "SocialGraph", + "schema_name": "" + } +] +``` + +### Detailed Format + +When `output_format` is set to "detailed" (default), the tool returns +comprehensive schema information: + +```json +[ + { + "object_details": { + "catalog": "", + "edge_tables": [ + { + "baseCatalogName": "", + "baseSchemaName": "", + "baseTableName": "Knows", + "destinationNodeTable": { + "edgeTableColumns": [ + "DstId" + ], + "nodeTableColumns": [ + "Id" + ], + "nodeTableName": "Person" + }, + "keyColumns": [ + "SrcId", + "DstId" + ], + "kind": "EDGE", + "labelNames": [ + "Knows" + ], + "name": "Knows", + "propertyDefinitions": [ + { + "propertyDeclarationName": "DstId", + "valueExpressionSql": "DstId" + }, + { + "propertyDeclarationName": "SrcId", + "valueExpressionSql": "SrcId" + } + ], + "sourceNodeTable": { + "edgeTableColumns": [ + "SrcId" + ], + "nodeTableColumns": [ + "Id" + ], + "nodeTableName": "Person" + } + } + ], + "labels": [ + { + "name": "Knows", + "propertyDeclarationNames": [ + "DstId", + "SrcId" + ] + }, + { + "name": "Person", + "propertyDeclarationNames": [ + "Id", + "Name" + ] + } + ], + "node_tables": [ + { + "baseCatalogName": "", + "baseSchemaName": "", + "baseTableName": "Person", + "keyColumns": [ + "Id" + ], + "kind": "NODE", + "labelNames": [ + "Person" + ], + "name": "Person", + "propertyDefinitions": [ + { + "propertyDeclarationName": "Id", + "valueExpressionSql": "Id" + }, + { + "propertyDeclarationName": "Name", + "valueExpressionSql": "Name" + } + ] + } + ], + "object_name": "SocialGraph", + "property_declarations": [ + { + "name": "DstId", + "type": "INT64" + }, + { + "name": "Id", + "type": "INT64" + }, + { + "name": "Name", + "type": "STRING" + }, + { + "name": "SrcId", + "type": "INT64" + } + ], + "schema_name": "" + }, + "object_name": "SocialGraph", + "schema_name": "" + } +] +``` + +## Use Cases + +1. **Database Documentation**: Generate comprehensive documentation of your + database schema +2. **Schema Validation**: Verify that expected graphs, node and edge exist +3. **Migration Planning**: Understand the current schema before making changes +4. **Development Tools**: Build tools that need to understand database structure +5. **Audit and Compliance**: Track schema changes and ensure compliance with + data governance policies + +## Example with Agent Integration + +```yaml +sources: + spanner-db: + kind: spanner + project: my-project + instance: my-instance + database: my-database + dialect: googlesql + +tools: + schema_inspector: + kind: spanner-list-graphs + source: spanner-db + description: | + Use this tool to inspect database schema information. + You can: + - List all graphs by leaving graph_names empty + - Get specific graph schemas by providing comma-separated graph names + - Choose between simple (names only) or detailed (full schema) output + + Examples: + 1. List all graphs with details: {"output_format": "detailed"} + 2. Get specific graphs: {"graph_names": "FinGraph,SocialGraph", "output_format": "detailed"} + 3. Just get graph names: {"output_format": "simple"} +``` + +## Reference + +| **field** | **type** | **required** | **description** | +|--------------|:--------:|:------------:|-----------------------------------------------------------------| +| kind | string | true | Must be "spanner-list-graphs" | +| source | string | true | Name of the Spanner source to query (dialect must be GoogleSQL) | +| description | string | false | Description of the tool that is passed to the LLM | +| authRequired | string[] | false | List of auth services required to invoke this tool | + +## Notes + +- This tool is read-only and does not modify any data +- The tool only works for the GoogleSQL source dialect +- Large databases with many graphs may take longer to query diff --git a/internal/tools/spanner/spannerlistgraphs/spannerlistgraphs.go b/internal/tools/spanner/spannerlistgraphs/spannerlistgraphs.go new file mode 100644 index 0000000000..6125872a8d --- /dev/null +++ b/internal/tools/spanner/spannerlistgraphs/spannerlistgraphs.go @@ -0,0 +1,255 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package spannerlistgraphs + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "cloud.google.com/go/spanner" + yaml "github.com/goccy/go-yaml" + "github.com/googleapis/genai-toolbox/internal/sources" + spannerdb "github.com/googleapis/genai-toolbox/internal/sources/spanner" + "github.com/googleapis/genai-toolbox/internal/tools" + "github.com/googleapis/genai-toolbox/internal/util/parameters" + "google.golang.org/api/iterator" +) + +const kind string = "spanner-list-graphs" + +func init() { + if !tools.Register(kind, newConfig) { + panic(fmt.Sprintf("tool kind %q already registered", kind)) + } +} + +func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) { + actual := Config{Name: name} + if err := decoder.DecodeContext(ctx, &actual); err != nil { + return nil, err + } + return actual, nil +} + +type compatibleSource interface { + SpannerClient() *spanner.Client + DatabaseDialect() string +} + +// validate compatible sources are still compatible +var _ compatibleSource = &spannerdb.Source{} + +var compatibleSources = [...]string{spannerdb.SourceKind} + +type Config struct { + Name string `yaml:"name" validate:"required"` + Kind string `yaml:"kind" validate:"required"` + Source string `yaml:"source" validate:"required"` + Description string `yaml:"description"` + AuthRequired []string `yaml:"authRequired"` +} + +// validate interface +var _ tools.ToolConfig = Config{} + +func (cfg Config) ToolConfigKind() string { + return kind +} + +func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) { + // verify source exists + rawS, ok := srcs[cfg.Source] + if !ok { + return nil, fmt.Errorf("no source named %q configured", cfg.Source) + } + + // verify the source is compatible + s, ok := rawS.(compatibleSource) + if !ok { + return nil, fmt.Errorf("invalid source for %q tool: source kind must be one of %q", kind, compatibleSources) + } + + // verify the dialect is GoogleSQL + if strings.ToLower(s.DatabaseDialect()) != "googlesql" { + return nil, fmt.Errorf("invalid source dialect for %q tool: source dialect must be GoogleSQL", kind) + } + + // Define parameters for the tool + allParameters := parameters.Parameters{ + parameters.NewStringParameterWithDefault( + "graph_names", + "", + "Optional: A comma-separated list of graph names. If empty, details for all graphs in user-accessible schemas will be listed.", + ), + parameters.NewStringParameterWithDefault( + "output_format", + "detailed", + "Optional: Use 'simple' to return graph names only or use 'detailed' to return the full information schema.", + ), + } + + description := cfg.Description + if description == "" { + description = "Lists detailed graph schema information (node tables, edge tables, labels and property declarations) as JSON for user-created graphs. Filters by a comma-separated list of names. If names are omitted, lists all graphs in user schemas." + } + mcpManifest := tools.GetMcpManifest(cfg.Name, description, cfg.AuthRequired, allParameters) + + // finish tool setup + t := Tool{ + Config: cfg, + AllParams: allParameters, + Client: s.SpannerClient(), + dialect: s.DatabaseDialect(), + manifest: tools.Manifest{Description: description, Parameters: allParameters.Manifest(), AuthRequired: cfg.AuthRequired}, + mcpManifest: mcpManifest, + } + return t, nil +} + +// validate interface +var _ tools.Tool = Tool{} + +type Tool struct { + Config + AllParams parameters.Parameters `yaml:"allParams"` + Client *spanner.Client + dialect string + manifest tools.Manifest + mcpManifest tools.McpManifest +} + +// processRows iterates over the spanner.RowIterator and converts each row to a map[string]any. +func processRows(iter *spanner.RowIterator) ([]any, error) { + var out []any + defer iter.Stop() + + for { + row, err := iter.Next() + if err == iterator.Done { + break + } + if err != nil { + return nil, fmt.Errorf("unable to parse row: %w", err) + } + + vMap := make(map[string]any) + cols := row.ColumnNames() + for i, c := range cols { + if c == "object_details" { + jsonString := row.ColumnValue(i).AsInterface().(string) + var details map[string]interface{} + if err := json.Unmarshal([]byte(jsonString), &details); err != nil { + return nil, fmt.Errorf("unable to unmarshal JSON: %w", err) + } + vMap[c] = details + } else { + vMap[c] = row.ColumnValue(i) + } + } + out = append(out, vMap) + } + return out, nil +} + +func (t Tool) Invoke(ctx context.Context, params parameters.ParamValues, accessToken tools.AccessToken) (any, error) { + paramsMap := params.AsMap() + + graphNames, _ := paramsMap["graph_names"].(string) + outputFormat, _ := paramsMap["output_format"].(string) + if outputFormat == "" { + outputFormat = "detailed" + } + + stmtParams := map[string]interface{}{ + "graph_names": graphNames, + "output_format": outputFormat, + } + + stmt := spanner.Statement{ + SQL: googleSQLStatement, + Params: stmtParams, + } + + // Execute the query (read-only) + iter := t.Client.Single().Query(ctx, stmt) + results, err := processRows(iter) + if err != nil { + return nil, fmt.Errorf("unable to execute query: %w", err) + } + + return results, nil +} + +func (t Tool) ParseParams(data map[string]any, claims map[string]map[string]any) (parameters.ParamValues, error) { + return parameters.ParseParams(t.AllParams, data, claims) +} + +func (t Tool) Manifest() tools.Manifest { + return t.manifest +} + +func (t Tool) McpManifest() tools.McpManifest { + return t.mcpManifest +} + +func (t Tool) Authorized(verifiedAuthServices []string) bool { + return tools.IsAuthorized(t.AuthRequired, verifiedAuthServices) +} + +func (t Tool) RequiresClientAuthorization() bool { + return false +} + +func (t Tool) ToConfig() tools.ToolConfig { + return t.Config +} + +func (t Tool) GetAuthTokenHeaderName() string { + return "Authorization" +} + +// GoogleSQL statement for listing graphs +const googleSQLStatement = ` +WITH FilterGraphNames AS ( + SELECT DISTINCT TRIM(name) AS GRAPH_NAME + FROM UNNEST(IF(@graph_names = '' OR @graph_names IS NULL, ['%'], SPLIT(@graph_names, ','))) AS name +) + +SELECT + PG.PROPERTY_GRAPH_SCHEMA AS schema_name, + PG.PROPERTY_GRAPH_NAME AS object_name, + CASE + WHEN @output_format = 'simple' THEN + -- IF format is 'simple', return basic JSON + CONCAT('{"name":"', IFNULL(REPLACE(PG.PROPERTY_GRAPH_NAME, '"', '\"'), ''), '"}') + ELSE + CONCAT( + '{', + '"schema_name":"', IFNULL(PG.PROPERTY_GRAPH_SCHEMA, ''), '",', + '"object_name":"', IFNULL(PG.PROPERTY_GRAPH_NAME, ''), '",', + '"catalog":"', IFNULL(JSON_VALUE(PG.PROPERTY_GRAPH_METADATA_JSON,"$.catalog"), ''), '",', + '"node_tables":', TO_JSON_STRING(PG.PROPERTY_GRAPH_METADATA_JSON.nodeTables), ',', + '"edge_tables":', TO_JSON_STRING(PG.PROPERTY_GRAPH_METADATA_JSON.edgeTables), ',', + '"labels":', TO_JSON_STRING(PG.PROPERTY_GRAPH_METADATA_JSON.labels), ',', + '"property_declarations":', TO_JSON_STRING(PG.PROPERTY_GRAPH_METADATA_JSON.propertyDeclarations), + '}' + ) + END AS object_details +FROM INFORMATION_SCHEMA.PROPERTY_GRAPHS PG +WHERE + EXISTS (SELECT 1 FROM FilterGraphNames WHERE FilterGraphNames.GRAPH_NAME = '%') OR PG.PROPERTY_GRAPH_NAME IN (SELECT GRAPH_NAME FROM FilterGraphNames) +` diff --git a/internal/tools/spanner/spannerlistgraphs/spannerlistgraphs_test.go b/internal/tools/spanner/spannerlistgraphs/spannerlistgraphs_test.go new file mode 100644 index 0000000000..eb4cca9e95 --- /dev/null +++ b/internal/tools/spanner/spannerlistgraphs/spannerlistgraphs_test.go @@ -0,0 +1,112 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package spannerlistgraphs_test + +import ( + "testing" + + yaml "github.com/goccy/go-yaml" + "github.com/google/go-cmp/cmp" + "github.com/googleapis/genai-toolbox/internal/server" + "github.com/googleapis/genai-toolbox/internal/testutils" + "github.com/googleapis/genai-toolbox/internal/tools/spanner/spannerlistgraphs" +) + +func TestParseFromYamlListGraphs(t *testing.T) { + ctx, err := testutils.ContextWithNewLogger() + if err != nil { + t.Fatalf("unexpected error: %s", err) + } + tcs := []struct { + desc string + in string + want server.ToolConfigs + }{ + { + desc: "basic example", + in: ` + tools: + example_tool: + kind: spanner-list-graphs + source: my-spanner-instance + description: Lists graphs in the database + `, + want: server.ToolConfigs{ + "example_tool": spannerlistgraphs.Config{ + Name: "example_tool", + Kind: "spanner-list-graphs", + Source: "my-spanner-instance", + Description: "Lists graphs in the database", + AuthRequired: []string{}, + }, + }, + }, + { + desc: "with auth required", + in: ` + tools: + example_tool: + kind: spanner-list-graphs + source: my-spanner-instance + description: Lists graphs in the database + authRequired: + - auth1 + - auth2 + `, + want: server.ToolConfigs{ + "example_tool": spannerlistgraphs.Config{ + Name: "example_tool", + Kind: "spanner-list-graphs", + Source: "my-spanner-instance", + Description: "Lists graphs in the database", + AuthRequired: []string{"auth1", "auth2"}, + }, + }, + }, + { + desc: "minimal config", + in: ` + tools: + example_tool: + kind: spanner-list-graphs + source: my-spanner-instance + `, + want: server.ToolConfigs{ + "example_tool": spannerlistgraphs.Config{ + Name: "example_tool", + Kind: "spanner-list-graphs", + Source: "my-spanner-instance", + Description: "", + AuthRequired: []string{}, + }, + }, + }, + } + for _, tc := range tcs { + t.Run(tc.desc, func(t *testing.T) { + got := struct { + Tools server.ToolConfigs `yaml:"tools"` + }{} + // Parse contents + err := yaml.UnmarshalContext(ctx, testutils.FormatYaml(tc.in), &got) + if err != nil { + t.Fatalf("unable to unmarshal: %s", err) + } + if diff := cmp.Diff(tc.want, got.Tools); diff != "" { + t.Fatalf("incorrect parse: diff %v", diff) + } + }) + } +} diff --git a/tests/spanner/spanner_integration_test.go b/tests/spanner/spanner_integration_test.go index 4cdbd379f8..324738f6cb 100644 --- a/tests/spanner/spanner_integration_test.go +++ b/tests/spanner/spanner_integration_test.go @@ -128,12 +128,47 @@ func TestSpannerToolEndpoints(t *testing.T) { teardownTableTmpl := setupSpannerTable(t, ctx, adminClient, dataClient, createStatementTmpl, "", tableNameTemplateParam, dbString, nil) defer teardownTableTmpl(t) + // set up for graph tool + nodeTableName := "node_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") + createNodeStatementTmpl := fmt.Sprintf("CREATE TABLE %s (id INT64 NOT NULL) PRIMARY KEY (id)", nodeTableName) + teardownNodeTableTmpl := setupSpannerTable(t, ctx, adminClient, dataClient, createNodeStatementTmpl, "", nodeTableName, dbString, nil) + defer teardownNodeTableTmpl(t) + + edgeTableName := "edge_table_" + strings.ReplaceAll(uuid.New().String(), "-", "") + createEdgeStatementTmpl := fmt.Sprintf(` + CREATE TABLE %[1]s ( + id INT64 NOT NULL, + target_id INT64 NOT NULL, + FOREIGN KEY (target_id) REFERENCES %[2]s (id) + ) PRIMARY KEY (id, target_id), + INTERLEAVE IN PARENT %[2]s ON DELETE CASCADE + `, edgeTableName, nodeTableName) + teardownEdgeTableTmpl := setupSpannerTable(t, ctx, adminClient, dataClient, createEdgeStatementTmpl, "", edgeTableName, dbString, nil) + defer teardownEdgeTableTmpl(t) + + graphName := "graph_" + strings.ReplaceAll(uuid.New().String(), "-", "") + createGraphStmt := fmt.Sprintf(` + CREATE PROPERTY GRAPH %[3]s + NODE TABLES ( + %[1]s + ) + EDGE TABLES ( + %[2]s + SOURCE KEY (id) REFERENCES %[1]s + DESTINATION KEY (target_id) REFERENCES %[1]s + LABEL EDGE + ) + `, nodeTableName, edgeTableName, graphName) + teardownGraph := setupSpannerGraph(t, ctx, adminClient, createGraphStmt, graphName, dbString) + defer teardownGraph(t) + // Write config into a file and pass it to command toolsFile := tests.GetToolsConfig(sourceConfig, SpannerToolKind, paramToolStmt, idParamToolStmt, nameParamToolStmt, arrayToolStmt, authToolStmt) toolsFile = addSpannerExecuteSqlConfig(t, toolsFile) toolsFile = addSpannerReadOnlyConfig(t, toolsFile) toolsFile = addTemplateParamConfig(t, toolsFile) toolsFile = addSpannerListTablesConfig(t, toolsFile) + toolsFile = addSpannerListGraphsConfig(t, toolsFile) cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...) if err != nil { @@ -175,8 +210,9 @@ func TestSpannerToolEndpoints(t *testing.T) { tests.DisableDdlTest(), ) runSpannerSchemaToolInvokeTest(t, accessSchemaWant) - runSpannerExecuteSqlToolInvokeTest(t, select1Want, invokeParamWant, tableNameParam, tableNameAuth) + runSpannerExecuteSqlToolInvokeTest(t, select1Want, invokeParamWant, tableNameParam) runSpannerListTablesTest(t, tableNameParam, tableNameAuth, tableNameTemplateParam) + runSpannerListGraphsTest(t, graphName) } // getSpannerToolInfo returns statements and param for my-tool for spanner-sql kind @@ -255,6 +291,39 @@ func setupSpannerTable(t *testing.T, ctx context.Context, adminClient *database. } } +// setupSpannerGraph creates a graph and inserts data into it. +func setupSpannerGraph(t *testing.T, ctx context.Context, adminClient *database.DatabaseAdminClient, createStatement, graphName, dbString string) func(*testing.T) { + // Create graph + op, err := adminClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{ + Database: dbString, + Statements: []string{createStatement}, + }) + if err != nil { + t.Fatalf("unable to start create graph operation %s: %s", graphName, err) + } + err = op.Wait(ctx) + if err != nil { + t.Fatalf("unable to create test graph %s: %s", graphName, err) + } + + return func(t *testing.T) { + // tear down test + op, err = adminClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{ + Database: dbString, + Statements: []string{fmt.Sprintf("DROP PROPERTY GRAPH %s", graphName)}, + }) + if err != nil { + t.Errorf("unable to start drop %s operation: %s", graphName, err) + return + } + + opErr := op.Wait(ctx) + if opErr != nil { + t.Errorf("Teardown failed: %s", opErr) + } + } +} + // addSpannerExecuteSqlConfig gets the tools config for `spanner-execute-sql` func addSpannerExecuteSqlConfig(t *testing.T, config map[string]any) map[string]any { tools, ok := config["tools"].(map[string]any) @@ -324,6 +393,24 @@ func addSpannerListTablesConfig(t *testing.T, config map[string]any) map[string] return config } +// addSpannerListGraphsConfig adds the spanner-list-graphs tool configuration +func addSpannerListGraphsConfig(t *testing.T, config map[string]any) map[string]any { + tools, ok := config["tools"].(map[string]any) + if !ok { + t.Fatalf("unable to get tools from config") + } + + // Add spanner-list-graphs tool + tools["list-graphs-tool"] = map[string]any{ + "kind": "spanner-list-graphs", + "source": "my-instance", + "description": "Lists graphs with their schema information", + } + + config["tools"] = tools + return config +} + func addTemplateParamConfig(t *testing.T, config map[string]any) map[string]any { toolsMap, ok := config["tools"].(map[string]any) if !ok { @@ -384,7 +471,7 @@ func addTemplateParamConfig(t *testing.T, config map[string]any) map[string]any return config } -func runSpannerExecuteSqlToolInvokeTest(t *testing.T, select1Want, invokeParamWant, tableNameParam, tableNameAuth string) { +func runSpannerExecuteSqlToolInvokeTest(t *testing.T, select1Want, invokeParamWant, tableNameParam string) { // Get ID token idToken, err := tests.GetGoogleIdToken(tests.ClientId) if err != nil { @@ -655,6 +742,112 @@ func runSpannerListTablesTest(t *testing.T, tableNameParam, tableNameAuth, table } } +// Helper function to verify graph list results +func verifyGraphListResult(t *testing.T, body map[string]interface{}, expectedGraphs []string, expectedSimpleFormat bool) { + // Parse the result + result, ok := body["result"].(string) + if !ok { + t.Fatalf("unable to find result in response body") + } + + var graphs []interface{} + err := json.Unmarshal([]byte(result), &graphs) + if err != nil { + t.Fatalf("unable to parse result as JSON array: %s", err) + } + + // If we expect specific graphs, verify they exist + if len(expectedGraphs) > 0 { + graphNames := make(map[string]bool) + requiredKeys := []string{"schema_name", "object_name", "catalog", "node_tables", "edge_tables", "labels", "property_declarations"} + if expectedSimpleFormat { + requiredKeys = []string{"name"} + } + + for _, graph := range graphs { + graphMap, ok := graph.(map[string]interface{}) + if !ok { + continue + } + + objectDetails, ok := graphMap["object_details"].(map[string]interface{}) + if !ok { + t.Fatalf("object_details is not of type map[string]interface{}, got: %T", graphMap["object_details"]) + } + for _, reqKey := range requiredKeys { + if _, hasKey := objectDetails[reqKey]; !hasKey { + t.Errorf("missing required key '%s', for object_details: %v", reqKey, objectDetails) + } + } + + if name, ok := graphMap["object_name"].(string); ok { + graphNames[name] = true + } + } + + for _, expected := range expectedGraphs { + if !graphNames[expected] { + t.Errorf("expected graph %s not found in results", expected) + } + } + } +} + +// runSpannerListGraphsTest tests the spanner-list-graphs tool +func runSpannerListGraphsTest(t *testing.T, graphName string) { + invokeTcs := []struct { + name string + requestBody io.Reader + expectedGraphs []string // empty means don't check specific graphs + useSimpleFormat bool + }{ + { + name: "list all graphs with detailed format", + requestBody: bytes.NewBuffer([]byte(`{}`)), + expectedGraphs: []string{graphName}, + }, + { + name: "list graphs with simple format", + requestBody: bytes.NewBuffer([]byte(`{"output_format": "simple"}`)), + expectedGraphs: []string{graphName}, + useSimpleFormat: true, + }, + { + name: "list specific graphs", + requestBody: bytes.NewBuffer([]byte(fmt.Sprintf(`{"graph_names": "%s"}`, graphName))), + expectedGraphs: []string{graphName}, + }, + { + name: "list non-existent graph", + requestBody: bytes.NewBuffer([]byte(`{"graph_names": "non_existent_graph_xyz"}`)), + expectedGraphs: []string{}, + }, + } + + for _, tc := range invokeTcs { + t.Run(tc.name, func(t *testing.T) { + // Use RunRequest helper function from tests package + url := "http://127.0.0.1:5000/api/tool/list-graphs-tool/invoke" + headers := map[string]string{} + + resp, respBody := tests.RunRequest(t, http.MethodPost, url, tc.requestBody, headers) + + if resp.StatusCode != http.StatusOK { + t.Fatalf("response status code is not 200, got %d: %s", resp.StatusCode, string(respBody)) + } + + // Check response body + var body map[string]interface{} + err := json.Unmarshal(respBody, &body) + if err != nil { + t.Fatalf("error parsing response body: %s", err) + } + + verifyGraphListResult(t, body, tc.expectedGraphs, tc.useSimpleFormat) + }) + } +} + func runSpannerSchemaToolInvokeTest(t *testing.T, accessSchemaWant string) { invokeTcs := []struct { name string