From 057d25197f4d9e36f1122117601ea4eaa803dbaa Mon Sep 17 00:00:00 2001 From: Nester Marchenko Date: Wed, 22 Oct 2025 15:50:09 -0400 Subject: [PATCH 1/3] feat(neo4j): add dry_run parameter to validate Cypher queries without execution --- .../tools/neo4j/neo4j-execute-cypher.md | 5 +- .../neo4jexecutecypher/neo4jexecutecypher.go | 57 ++++++++++++++++++- tests/neo4j/neo4j_integration_test.go | 42 ++++++++++++++ 3 files changed, 100 insertions(+), 4 deletions(-) diff --git a/docs/en/resources/tools/neo4j/neo4j-execute-cypher.md b/docs/en/resources/tools/neo4j/neo4j-execute-cypher.md index c03b843761d7..f544e1c25de2 100644 --- a/docs/en/resources/tools/neo4j/neo4j-execute-cypher.md +++ b/docs/en/resources/tools/neo4j/neo4j-execute-cypher.md @@ -27,8 +27,9 @@ Cypher](https://neo4j.com/docs/cypher-manual/current/queries/) syntax and supports all Cypher features, including pattern matching, filtering, and aggregation. -`neo4j-execute-cypher` takes one input parameter `cypher` and run the cypher -query against the `source`. +`neo4j-execute-cypher` takes a required input parameter `cypher` and run the cypher +query against the `source`. It also supports an optional `dry_run` +parameter to validate a query without executing it. > **Note:** This tool is intended for developer assistant workflows with > human-in-the-loop and shouldn't be used for production agents. diff --git a/internal/tools/neo4j/neo4jexecutecypher/neo4jexecutecypher.go b/internal/tools/neo4j/neo4jexecutecypher/neo4jexecutecypher.go index ee2633085468..1a461aa58da9 100644 --- a/internal/tools/neo4j/neo4jexecutecypher/neo4jexecutecypher.go +++ b/internal/tools/neo4j/neo4jexecutecypher/neo4jexecutecypher.go @@ -84,7 +84,13 @@ func (cfg Config) Initialize(srcs map[string]sources.Source) (tools.Tool, error) } cypherParameter := tools.NewStringParameter("cypher", "The cypher to execute.") - parameters := tools.Parameters{cypherParameter} + dryRunParameter := tools.NewBooleanParameterWithDefault( + "dry_run", + false, + "If set to true, the query will be validated and information about the execution "+ + "will be returned without running the query. Defaults to false.", + ) + parameters := tools.Parameters{cypherParameter, dryRunParameter} mcpManifest := tools.GetMcpManifest(cfg.Name, cfg.Description, cfg.AuthRequired, parameters) @@ -124,13 +130,18 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken paramsMap := params.AsMap() cypherStr, ok := paramsMap["cypher"].(string) if !ok { - return nil, fmt.Errorf("unable to get cast %s", paramsMap["cypher"]) + return nil, fmt.Errorf("unable to cast cypher parameter %s", paramsMap["cypher"]) } if cypherStr == "" { return nil, fmt.Errorf("parameter 'cypher' must be a non-empty string") } + dryRun, ok := paramsMap["dry_run"].(bool) + if !ok { + return nil, fmt.Errorf("unable to cast dry_run parameter %s", paramsMap["dry_run"]) + } + // validate the cypher query before executing cf := t.classifier.Classify(cypherStr) if cf.Error != nil { @@ -141,6 +152,11 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken return nil, fmt.Errorf("this tool is read-only and cannot execute write queries") } + if dryRun { + // Add EXPLAIN to the beginning of the query to validate it without executing + cypherStr = "EXPLAIN " + cypherStr + } + config := neo4j.ExecuteQueryWithDatabase(t.Database) results, err := neo4j.ExecuteQuery(ctx, t.Driver, cypherStr, nil, neo4j.EagerResultTransformer, config) @@ -148,9 +164,28 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken return nil, fmt.Errorf("unable to execute query: %w", err) } + // If dry run, return the summary information only + if dryRun { + summary := results.Summary + plan := summary.Plan() + execPlan := map[string]any{ + "queryType": cf.Type.String(), + "statementType": summary.StatementType(), + "operator": plan.Operator(), + "arguments": plan.Arguments(), + "identifiers": plan.Identifiers(), + "childrenCount": len(plan.Children()), + } + if len(plan.Children()) > 0 { + execPlan["children"] = t.addPlanChildren(plan) + } + return []map[string]any{execPlan}, nil + } + var out []any keys := results.Keys records := results.Records + for _, record := range records { vMap := make(map[string]any) for col, value := range record.Values { @@ -181,3 +216,21 @@ func (t Tool) Authorized(verifiedAuthServices []string) bool { func (t Tool) RequiresClientAuthorization() bool { return false } + +// Recursive function to add plan children +func (t Tool) addPlanChildren(p neo4j.Plan) []map[string]any { + var children []map[string]any + for _, child := range p.Children() { + childMap := map[string]any{ + "operator": child.Operator(), + "arguments": child.Arguments(), + "identifiers": child.Identifiers(), + "children_count": len(child.Children()), + } + if len(child.Children()) > 0 { + childMap["children"] = t.addPlanChildren(child) + } + children = append(children, childMap) + } + return children +} diff --git a/tests/neo4j/neo4j_integration_test.go b/tests/neo4j/neo4j_integration_test.go index 7d60083c66cd..a7742ee1f24d 100644 --- a/tests/neo4j/neo4j_integration_test.go +++ b/tests/neo4j/neo4j_integration_test.go @@ -160,6 +160,13 @@ func TestNeo4jToolEndpoints(t *testing.T) { "description": "The cypher to execute.", "authSources": []any{}, }, + map[string]any{ + "name": "dry_run", + "type": "boolean", + "required": false, + "description": "If set to true, the query will be validated and information about the execution will be returned without running the query. Defaults to false.", + "authSources": []any{}, + }, }, "authRequired": []any{}, }, @@ -240,6 +247,34 @@ func TestNeo4jToolEndpoints(t *testing.T) { want: "[{\"a\":1}]", wantStatus: http.StatusOK, }, + { + name: "invoke my-simple-execute-cypher-tool with dry_run", + api: "http://127.0.0.1:5000/api/tool/my-simple-execute-cypher-tool/invoke", + requestBody: bytes.NewBuffer([]byte(`{"cypher": "MATCH (n:Test) RETURN n", "dry_run": true}`)), + wantStatus: http.StatusOK, + validateFunc: func(t *testing.T, body string) { + var result []map[string]any + if err := json.Unmarshal([]byte(body), &result); err != nil { + t.Fatalf("failed to unmarshal dry_run result: %v", err) + } + if len(result) == 0 { + t.Fatalf("expected a query plan, but got an empty result") + } + if _, ok := result[0]["operator"]; !ok { + t.Errorf("expected key 'Operator' not found in dry_run response: %s", body) + } + if _, ok := result[0]["childrenCount"]; !ok { + t.Errorf("expected key 'ChildrenCount' not found in dry_run response: %s", body) + } + }, + }, + { + name: "invoke my-simple-execute-cypher-tool with dry_run and invalid syntax", + api: "http://127.0.0.1:5000/api/tool/my-simple-execute-cypher-tool/invoke", + requestBody: bytes.NewBuffer([]byte(`{"cypher": "RTN 1", "dry_run": true}`)), + wantStatus: http.StatusBadRequest, + wantErrorSubstring: "unable to execute query", + }, { name: "invoke readonly tool with write query", api: "http://127.0.0.1:5000/api/tool/my-readonly-execute-cypher-tool/invoke", @@ -247,6 +282,13 @@ func TestNeo4jToolEndpoints(t *testing.T) { wantStatus: http.StatusBadRequest, wantErrorSubstring: "this tool is read-only and cannot execute write queries", }, + { + name: "invoke readonly tool with write query and dry_run", + api: "http://127.0.0.1:5000/api/tool/my-readonly-execute-cypher-tool/invoke", + requestBody: bytes.NewBuffer([]byte(`{"cypher": "CREATE (n:TestNode)", "dry_run": true}`)), + wantStatus: http.StatusBadRequest, + wantErrorSubstring: "this tool is read-only and cannot execute write queries", + }, { name: "invoke my-schema-tool", api: "http://127.0.0.1:5000/api/tool/my-schema-tool/invoke", From 86cad9a6df855da18c20da9299d4c55b01428667 Mon Sep 17 00:00:00 2001 From: Nester Marchenko Date: Wed, 22 Oct 2025 17:00:12 -0400 Subject: [PATCH 2/3] feat(neo4j): set custom User-Agent for Neo4j driver connection --- internal/sources/neo4j/neo4j.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/internal/sources/neo4j/neo4j.go b/internal/sources/neo4j/neo4j.go index 9c1b9c76f6c3..51ec85b01d6e 100644 --- a/internal/sources/neo4j/neo4j.go +++ b/internal/sources/neo4j/neo4j.go @@ -21,10 +21,12 @@ import ( "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" "github.com/neo4j/neo4j-go-driver/v5/neo4j" + neo4jconf "github.com/neo4j/neo4j-go-driver/v5/neo4j/config" "go.opentelemetry.io/otel/trace" ) const SourceKind string = "neo4j" +const UserAgent string = "genai-toolbox/neo4j-source" // validate interface var _ sources.SourceConfig = Config{} @@ -106,7 +108,9 @@ func initNeo4jDriver(ctx context.Context, tracer trace.Tracer, uri, user, passwo defer span.End() auth := neo4j.BasicAuth(user, password, "") - driver, err := neo4j.NewDriverWithContext(uri, auth) + driver, err := neo4j.NewDriverWithContext(uri, auth, func(config *neo4jconf.Config) { + config.UserAgent = UserAgent + }) if err != nil { return nil, fmt.Errorf("unable to create connection driver: %w", err) } From c05243166fffce1d9e2a9f50e0b5aa7e9fef4de7 Mon Sep 17 00:00:00 2001 From: Nester Marchenko Date: Thu, 23 Oct 2025 12:36:34 -0400 Subject: [PATCH 3/3] feat(neo4j): use dynamic User-Agent from context and refactor plan children handling --- internal/sources/neo4j/neo4j.go | 8 ++++-- .../neo4jexecutecypher/neo4jexecutecypher.go | 6 ++--- tests/neo4j/neo4j_integration_test.go | 25 ++++++++++++++++--- 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/internal/sources/neo4j/neo4j.go b/internal/sources/neo4j/neo4j.go index 51ec85b01d6e..2262a54bc9e2 100644 --- a/internal/sources/neo4j/neo4j.go +++ b/internal/sources/neo4j/neo4j.go @@ -20,13 +20,13 @@ import ( "github.com/goccy/go-yaml" "github.com/googleapis/genai-toolbox/internal/sources" + "github.com/googleapis/genai-toolbox/internal/util" "github.com/neo4j/neo4j-go-driver/v5/neo4j" neo4jconf "github.com/neo4j/neo4j-go-driver/v5/neo4j/config" "go.opentelemetry.io/otel/trace" ) const SourceKind string = "neo4j" -const UserAgent string = "genai-toolbox/neo4j-source" // validate interface var _ sources.SourceConfig = Config{} @@ -108,8 +108,12 @@ func initNeo4jDriver(ctx context.Context, tracer trace.Tracer, uri, user, passwo defer span.End() auth := neo4j.BasicAuth(user, password, "") + userAgent, err := util.UserAgentFromContext(ctx) + if err != nil { + return nil, err + } driver, err := neo4j.NewDriverWithContext(uri, auth, func(config *neo4jconf.Config) { - config.UserAgent = UserAgent + config.UserAgent = userAgent }) if err != nil { return nil, fmt.Errorf("unable to create connection driver: %w", err) diff --git a/internal/tools/neo4j/neo4jexecutecypher/neo4jexecutecypher.go b/internal/tools/neo4j/neo4jexecutecypher/neo4jexecutecypher.go index 1a461aa58da9..33c474d5091a 100644 --- a/internal/tools/neo4j/neo4jexecutecypher/neo4jexecutecypher.go +++ b/internal/tools/neo4j/neo4jexecutecypher/neo4jexecutecypher.go @@ -177,7 +177,7 @@ func (t Tool) Invoke(ctx context.Context, params tools.ParamValues, accessToken "childrenCount": len(plan.Children()), } if len(plan.Children()) > 0 { - execPlan["children"] = t.addPlanChildren(plan) + execPlan["children"] = addPlanChildren(plan) } return []map[string]any{execPlan}, nil } @@ -218,7 +218,7 @@ func (t Tool) RequiresClientAuthorization() bool { } // Recursive function to add plan children -func (t Tool) addPlanChildren(p neo4j.Plan) []map[string]any { +func addPlanChildren(p neo4j.Plan) []map[string]any { var children []map[string]any for _, child := range p.Children() { childMap := map[string]any{ @@ -228,7 +228,7 @@ func (t Tool) addPlanChildren(p neo4j.Plan) []map[string]any { "children_count": len(child.Children()), } if len(child.Children()) > 0 { - childMap["children"] = t.addPlanChildren(child) + childMap["children"] = addPlanChildren(child) } children = append(children, childMap) } diff --git a/tests/neo4j/neo4j_integration_test.go b/tests/neo4j/neo4j_integration_test.go index a7742ee1f24d..816edba6e9ef 100644 --- a/tests/neo4j/neo4j_integration_test.go +++ b/tests/neo4j/neo4j_integration_test.go @@ -260,11 +260,28 @@ func TestNeo4jToolEndpoints(t *testing.T) { if len(result) == 0 { t.Fatalf("expected a query plan, but got an empty result") } - if _, ok := result[0]["operator"]; !ok { - t.Errorf("expected key 'Operator' not found in dry_run response: %s", body) + + operatorValue, ok := result[0]["operator"] + if !ok { + t.Fatalf("expected key 'Operator' not found in dry_run response: %s", body) + } + + operatorStr, ok := operatorValue.(string) + if !ok { + t.Fatalf("expected 'Operator' to be a string, but got %T", operatorValue) } - if _, ok := result[0]["childrenCount"]; !ok { - t.Errorf("expected key 'ChildrenCount' not found in dry_run response: %s", body) + + if operatorStr != "ProduceResults@neo4j" { + t.Errorf("unexpected operator: got %q, want %q", operatorStr, "ProduceResults@neo4j") + } + + childrenCount, ok := result[0]["childrenCount"] + if !ok { + t.Fatalf("expected key 'ChildrenCount' not found in dry_run response: %s", body) + } + + if childrenCount.(float64) != 1 { + t.Errorf("unexpected children count: got %v, want %d", childrenCount, 1) } }, },