From 804b1a5df53c41774c310fbe72309ec59e2013d8 Mon Sep 17 00:00:00 2001 From: Kyle Ellrott Date: Thu, 11 Apr 2024 16:18:18 -0700 Subject: [PATCH] Adding context variable to go client traversal --- cmd/dump/main.go | 5 +++-- cmd/info/main.go | 5 +++-- cmd/query/main.go | 3 ++- gripql/client.go | 4 ++-- gripql/schema/scan.go | 5 +++-- server/metagraphs.go | 4 ++-- 6 files changed, 15 insertions(+), 11 deletions(-) diff --git a/cmd/dump/main.go b/cmd/dump/main.go index dbaa6cf2..8c9fc6ae 100644 --- a/cmd/dump/main.go +++ b/cmd/dump/main.go @@ -1,6 +1,7 @@ package dump import ( + "context" "fmt" "github.com/bmeg/grip/gripql" @@ -29,7 +30,7 @@ var Cmd = &cobra.Command{ if vertexDump { q := gripql.V() - elems, err := conn.Traversal(&gripql.GraphQuery{Graph: graph, Query: q.Statements}) + elems, err := conn.Traversal(context.Background(), &gripql.GraphQuery{Graph: graph, Query: q.Statements}) if err != nil { return err } @@ -44,7 +45,7 @@ var Cmd = &cobra.Command{ if edgeDump { q := gripql.E() - elems, err := conn.Traversal(&gripql.GraphQuery{Graph: graph, Query: q.Statements}) + elems, err := conn.Traversal(context.Background(), &gripql.GraphQuery{Graph: graph, Query: q.Statements}) if err != nil { return err } diff --git a/cmd/info/main.go b/cmd/info/main.go index 6dccf082..d1beafe3 100644 --- a/cmd/info/main.go +++ b/cmd/info/main.go @@ -1,6 +1,7 @@ package info import ( + "context" "fmt" "github.com/bmeg/grip/gripql" @@ -27,7 +28,7 @@ var Cmd = &cobra.Command{ fmt.Printf("Graph: %s\n", graph) q := gripql.V().Count() - res, err := conn.Traversal(&gripql.GraphQuery{Graph: graph, Query: q.Statements}) + res, err := conn.Traversal(context.Background(), &gripql.GraphQuery{Graph: graph, Query: q.Statements}) if err != nil { return err } @@ -36,7 +37,7 @@ var Cmd = &cobra.Command{ } q = gripql.E().Count() - res, err = conn.Traversal(&gripql.GraphQuery{Graph: graph, Query: q.Statements}) + res, err = conn.Traversal(context.Background(), &gripql.GraphQuery{Graph: graph, Query: q.Statements}) if err != nil { return err } diff --git a/cmd/query/main.go b/cmd/query/main.go index f2474d7d..8a24978c 100644 --- a/cmd/query/main.go +++ b/cmd/query/main.go @@ -1,6 +1,7 @@ package query import ( + "context" "encoding/json" "fmt" @@ -67,7 +68,7 @@ Example: return err } - res, err := conn.Traversal(&query) + res, err := conn.Traversal(context.Background(), &query) if err != nil { return err } diff --git a/gripql/client.go b/gripql/client.go index 7fd6c880..1fbb6ecd 100644 --- a/gripql/client.go +++ b/gripql/client.go @@ -186,9 +186,9 @@ func (client Client) GetVertex(graph string, id string) (*Vertex, error) { } // Traversal runs a graph traversal query -func (client Client) Traversal(query *GraphQuery) (chan *QueryResult, error) { +func (client Client) Traversal(ctx context.Context, query *GraphQuery) (chan *QueryResult, error) { out := make(chan *QueryResult, 100) - tclient, err := client.QueryC.Traversal(context.Background(), query) + tclient, err := client.QueryC.Traversal(ctx, query) if err != nil { return nil, err } diff --git a/gripql/schema/scan.go b/gripql/schema/scan.go index 85e2b257..be66dd05 100644 --- a/gripql/schema/scan.go +++ b/gripql/schema/scan.go @@ -1,6 +1,7 @@ package schema import ( + "context" "fmt" "github.com/bmeg/grip/gripql" @@ -42,7 +43,7 @@ func ScanSchema(conn gripql.Client, graph string, sampleCount uint32, exclude [] schema := map[string]interface{}{} log.Infof("Scanning %s\n", label) nodeQuery := gripql.V().HasLabel(label).Limit(sampleCount) - nodeRes, err := conn.Traversal(&gripql.GraphQuery{Graph: graph, Query: nodeQuery.Statements}) + nodeRes, err := conn.Traversal(context.Background(), &gripql.GraphQuery{Graph: graph, Query: nodeQuery.Statements}) if err == nil { for row := range nodeRes { v := row.GetVertex() @@ -67,7 +68,7 @@ func ScanSchema(conn gripql.Client, graph string, sampleCount uint32, exclude [] } log.Infof("Scanning edge %s\n", elabel) edgeQuery := gripql.E().HasLabel(elabel).Limit(sampleCount).As("edge").Out().Fields().As("to").Select("edge").In().Fields().As("from").Select("edge", "from", "to") - edgeRes, err := conn.Traversal(&gripql.GraphQuery{Graph: graph, Query: edgeQuery.Statements}) + edgeRes, err := conn.Traversal(context.Background(), &gripql.GraphQuery{Graph: graph, Query: edgeQuery.Statements}) if err == nil { labelSchema := edgeMap{} for row := range edgeRes { diff --git a/server/metagraphs.go b/server/metagraphs.go index f0cdc44b..856189cf 100644 --- a/server/metagraphs.go +++ b/server/metagraphs.go @@ -30,7 +30,7 @@ func (server *GripServer) getGraph(graph string) (*gripql.Graph, error) { if err != nil { return nil, fmt.Errorf("failed to load existing schema: %v", err) } - res, err := conn.Traversal(&gripql.GraphQuery{Graph: graph, Query: gripql.NewQuery().V().Statements}) + res, err := conn.Traversal(context.Background(), &gripql.GraphQuery{Graph: graph, Query: gripql.NewQuery().V().Statements}) if err != nil { return nil, fmt.Errorf("failed to load existing schema: %v", err) } @@ -38,7 +38,7 @@ func (server *GripServer) getGraph(graph string) (*gripql.Graph, error) { for row := range res { vertices = append(vertices, row.GetVertex()) } - res, err = conn.Traversal(&gripql.GraphQuery{Graph: graph, Query: gripql.NewQuery().E().Statements}) + res, err = conn.Traversal(context.Background(), &gripql.GraphQuery{Graph: graph, Query: gripql.NewQuery().E().Statements}) if err != nil { return nil, fmt.Errorf("failed to load existing schema: %v", err) }