Merge pull request #309 from bmeg/feature/grip-client-context

Feature/grip client context
bmeg · Apr 15, 2024 · 09f1a55 · 09f1a55
2 parents 87975df + e350317
commit 09f1a55
Show file tree

Hide file tree

Showing 7 changed files with 17 additions and 12 deletions.
diff --git 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
@@ -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
@@ -1,6 +1,7 @@
 package query
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/bmeg/grip/gripql"
@@ -47,7 +48,7 @@ Example:
 		}
 
 		log.Debugf("Query: %s\n", query.String())
-		res, err := conn.Traversal(query)
+		res, err := conn.Traversal(context.Background(), query)
 		if err != nil {
 			log.Errorf("Traversal error: %s", err)
 			return err

diff --git a/gripql/client.go b/gripql/client.go
@@ -189,9 +189,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
@@ -1,6 +1,8 @@
 package schema
 
 import (
+	"context"
+
 	"github.com/bmeg/grip/gripql"
 	"github.com/bmeg/grip/log"
 	"github.com/bmeg/grip/util"
@@ -40,7 +42,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()

diff --git a/server/metagraphs.go b/server/metagraphs.go
@@ -30,15 +30,15 @@ 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)
 	}
 	vertices := []*gripql.Vertex{}
 	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)
 	}

diff --git a/test/server/auth_test.go b/test/server/auth_test.go
@@ -61,7 +61,7 @@ func TestBasicAuthFail(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	_, err = cli.Traversal(&gripql.GraphQuery{Graph: "test", Query: gripql.NewQuery().V().Statements})
+	_, err = cli.Traversal(context.Background(), &gripql.GraphQuery{Graph: "test", Query: gripql.NewQuery().V().Statements})
 	if err == nil || !strings.Contains(err.Error(), "PermissionDenied") {
 		t.Errorf("expected PermissionDenied error; got: %v", err)
 	}
@@ -131,7 +131,7 @@ func TestBasicAuth(t *testing.T) {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
-	_, err = cli.Traversal(&gripql.GraphQuery{Graph: "test", Query: gripql.NewQuery().V().Statements})
+	_, err = cli.Traversal(context.Background(), &gripql.GraphQuery{Graph: "test", Query: gripql.NewQuery().V().Statements})
 	if err != nil {
 		t.Errorf("unexpected error: %v", err)
 	}