From 028fa393988c9f3475c35de40b3bcd754400c190 Mon Sep 17 00:00:00 2001
From: Jesse Tatasciore <jessetatasciore7@gmail.com>
Date: Tue, 1 Feb 2022 15:30:07 -0500
Subject: [PATCH] fix: Set bazel workspace root for the bazel query verbs

---
 cmd/aspect/aquery/BUILD.bazel |  1 +
 cmd/aspect/aquery/aquery.go   | 17 ++++++++++++++---
 cmd/aspect/cquery/BUILD.bazel |  1 +
 cmd/aspect/cquery/cquery.go   | 17 ++++++++++++++---
 cmd/aspect/query/BUILD.bazel  |  1 +
 cmd/aspect/query/query.go     | 17 ++++++++++++++---
 6 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/cmd/aspect/aquery/BUILD.bazel b/cmd/aspect/aquery/BUILD.bazel
index f01c9713f..26791a82e 100644
--- a/cmd/aspect/aquery/BUILD.bazel
+++ b/cmd/aspect/aquery/BUILD.bazel
@@ -8,6 +8,7 @@ go_library(
     deps = [
         "//pkg/aspect/aquery",
         "//pkg/bazel",
+        "//pkg/interceptors",
         "//pkg/ioutils",
         "@com_github_spf13_cobra//:cobra",
     ],
diff --git a/cmd/aspect/aquery/aquery.go b/cmd/aspect/aquery/aquery.go
index 0da3b0758..f7cfd1f70 100644
--- a/cmd/aspect/aquery/aquery.go
+++ b/cmd/aspect/aquery/aquery.go
@@ -7,10 +7,13 @@ Not licensed for re-use
 package aquery
 
 import (
+	"context"
+
 	"github.com/spf13/cobra"
 
 	"aspect.build/cli/pkg/aspect/aquery"
 	"aspect.build/cli/pkg/bazel"
+	"aspect.build/cli/pkg/interceptors"
 	"aspect.build/cli/pkg/ioutils"
 )
 
@@ -19,13 +22,21 @@ func NewDefaultAQueryCmd() *cobra.Command {
 }
 
 func NewAQueryCommand(streams ioutils.Streams, bzl bazel.Bazel) *cobra.Command {
-	q := aquery.New(streams, bzl, true)
-
 	cmd := &cobra.Command{
 		Use:   "aquery",
 		Short: "Executes an aquery.",
 		Long:  "Executes a query language expression over a specified subgraph of the build dependency graph using aquery.",
-		RunE:  q.Run,
+		RunE: interceptors.Run(
+			[]interceptors.Interceptor{
+				interceptors.WorkspaceRootInterceptor(),
+			},
+			func(ctx context.Context, cmd *cobra.Command, args []string) (exitErr error) {
+				workspaceRoot := ctx.Value(interceptors.WorkspaceRootKey).(string)
+				bzl.SetWorkspaceRoot(workspaceRoot)
+				q := aquery.New(streams, bzl, true)
+				return q.Run(cmd, args)
+			},
+		),
 	}
 
 	return cmd
diff --git a/cmd/aspect/cquery/BUILD.bazel b/cmd/aspect/cquery/BUILD.bazel
index 9c17d88ee..b26bee663 100644
--- a/cmd/aspect/cquery/BUILD.bazel
+++ b/cmd/aspect/cquery/BUILD.bazel
@@ -8,6 +8,7 @@ go_library(
     deps = [
         "//pkg/aspect/cquery",
         "//pkg/bazel",
+        "//pkg/interceptors",
         "//pkg/ioutils",
         "@com_github_spf13_cobra//:cobra",
     ],
diff --git a/cmd/aspect/cquery/cquery.go b/cmd/aspect/cquery/cquery.go
index 204683098..cdf03a608 100644
--- a/cmd/aspect/cquery/cquery.go
+++ b/cmd/aspect/cquery/cquery.go
@@ -7,10 +7,13 @@ Not licensed for re-use
 package cquery
 
 import (
+	"context"
+
 	"github.com/spf13/cobra"
 
 	"aspect.build/cli/pkg/aspect/cquery"
 	"aspect.build/cli/pkg/bazel"
+	"aspect.build/cli/pkg/interceptors"
 	"aspect.build/cli/pkg/ioutils"
 )
 
@@ -19,13 +22,21 @@ func NewDefaultCQueryCmd() *cobra.Command {
 }
 
 func NewCQueryCommand(streams ioutils.Streams, bzl bazel.Bazel) *cobra.Command {
-	q := cquery.New(streams, bzl, true)
-
 	cmd := &cobra.Command{
 		Use:   "cquery",
 		Short: "Executes a cquery.",
 		Long:  "Executes a query language expression over a specified subgraph of the build dependency graph using cquery.",
-		RunE:  q.Run,
+		RunE: interceptors.Run(
+			[]interceptors.Interceptor{
+				interceptors.WorkspaceRootInterceptor(),
+			},
+			func(ctx context.Context, cmd *cobra.Command, args []string) (exitErr error) {
+				workspaceRoot := ctx.Value(interceptors.WorkspaceRootKey).(string)
+				bzl.SetWorkspaceRoot(workspaceRoot)
+				q := cquery.New(streams, bzl, true)
+				return q.Run(cmd, args)
+			},
+		),
 	}
 
 	return cmd
diff --git a/cmd/aspect/query/BUILD.bazel b/cmd/aspect/query/BUILD.bazel
index e09b9669a..5c5d27133 100644
--- a/cmd/aspect/query/BUILD.bazel
+++ b/cmd/aspect/query/BUILD.bazel
@@ -8,6 +8,7 @@ go_library(
     deps = [
         "//pkg/aspect/query",
         "//pkg/bazel",
+        "//pkg/interceptors",
         "//pkg/ioutils",
         "@com_github_spf13_cobra//:cobra",
     ],
diff --git a/cmd/aspect/query/query.go b/cmd/aspect/query/query.go
index 472972c26..a15d8be98 100644
--- a/cmd/aspect/query/query.go
+++ b/cmd/aspect/query/query.go
@@ -7,10 +7,13 @@ Not licensed for re-use
 package query
 
 import (
+	"context"
+
 	"github.com/spf13/cobra"
 
 	"aspect.build/cli/pkg/aspect/query"
 	"aspect.build/cli/pkg/bazel"
+	"aspect.build/cli/pkg/interceptors"
 	"aspect.build/cli/pkg/ioutils"
 )
 
@@ -19,13 +22,21 @@ func NewDefaultQueryCmd() *cobra.Command {
 }
 
 func NewQueryCommand(streams ioutils.Streams, bzl bazel.Bazel) *cobra.Command {
-	q := query.New(streams, bzl, true)
-
 	cmd := &cobra.Command{
 		Use:   "query",
 		Short: "Executes a dependency graph query.",
 		Long:  "Executes a query language expression over a specified subgraph of the build dependency graph.",
-		RunE:  q.Run,
+		RunE: interceptors.Run(
+			[]interceptors.Interceptor{
+				interceptors.WorkspaceRootInterceptor(),
+			},
+			func(ctx context.Context, cmd *cobra.Command, args []string) (exitErr error) {
+				workspaceRoot := ctx.Value(interceptors.WorkspaceRootKey).(string)
+				bzl.SetWorkspaceRoot(workspaceRoot)
+				q := query.New(streams, bzl, true)
+				return q.Run(cmd, args)
+			},
+		),
 	}
 
 	return cmd