feat: optional attributes when getting graph

Makefile

@@ -26,7 +26,7 @@ help:           ##@help show this help
 NAME="github.com/goto/compass"
 VERSION=$(shell git describe --always --tags 2>/dev/null)
 COVERFILE="/tmp/compass.coverprofile"
-PROTON_COMMIT := "eaca9798d1c1d7b3101ec1259c7e5fb949afba28"
+PROTON_COMMIT := "fe99cc96e060085d6052096e9ba59b4038c691c6"
 
 TOOLS_MOD_DIR = ./tools
 TOOLS_DIR = $(abspath ./.tools)

core/asset/lineage.go

@@ -21,8 +21,9 @@ const (
 )
 
 type LineageQuery struct {
-	Level     int
-	Direction LineageDirection
+	Level          int
+	Direction      LineageDirection
+	WithAttributes bool
 }
 
 //go:generate mockery --name=LineageRepository -r --case underscore --with-expecter --structname=LineageRepository --filename=lineage_repository.go --output=./mocks

core/asset/service.go

@@ -190,6 +190,12 @@ func (s *Service) GetLineage(ctx context.Context, urn string, query LineageQuery
 		return Lineage{}, fmt.Errorf("get lineage: get graph edges: %w", err)
 	}
 
+	if !query.WithAttributes {
+		return Lineage{
+			Edges: edges,
+		}, nil
+	}
+
 	urns := newUniqueStrings(len(edges))
 	urns.add(urn)
 	for _, edge := range edges {

core/asset/service_test.go

@@ -807,6 +807,7 @@ func TestService_GetLineage(t *testing.T) {
 	type testCase struct {
 		Description string
 		ID          string
+		Query       asset.LineageQuery
 		Setup       func(context.Context, *mocks.AssetRepository, *mocks.DiscoveryRepository, *mocks.LineageRepository)
 		Expected    asset.Lineage
 		Err         error
@@ -816,8 +817,11 @@ func TestService_GetLineage(t *testing.T) {
 		{
 			Description: `should return error if the GetGraph function return error`,
 			ID:          assetID,
+			Query: asset.LineageQuery{
+				WithAttributes: true,
+			},
 			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
-				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{}).
+				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: true}).
 					Return(asset.LineageGraph{}, errors.New("error fetching graph"))
 			},
 			Expected: asset.Lineage{},
@@ -826,8 +830,11 @@ func TestService_GetLineage(t *testing.T) {
 		{
 			Description: `should return no error if graph with 0 edges are returned`,
 			ID:          assetID,
+			Query: asset.LineageQuery{
+				WithAttributes: true,
+			},
 			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
-				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{}).
+				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: true}).
 					Return(asset.LineageGraph{}, nil)
 				ar.EXPECT().GetProbesWithFilter(ctx, asset.ProbesFilter{
 					AssetURNs: []string{"urn-source-1"},
@@ -840,8 +847,11 @@ func TestService_GetLineage(t *testing.T) {
 		{
 			Description: `should return an error if GetProbesWithFilter function returns error`,
 			ID:          assetID,
+			Query: asset.LineageQuery{
+				WithAttributes: true,
+			},
 			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
-				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{}).Return(asset.LineageGraph{
+				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: true}).Return(asset.LineageGraph{
 					{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
 					{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
 					{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
@@ -857,8 +867,11 @@ func TestService_GetLineage(t *testing.T) {
 		{
 			Description: `should return no error if GetProbesWithFilter function returns 0 probes`,
 			ID:          assetID,
+			Query: asset.LineageQuery{
+				WithAttributes: true,
+			},
 			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
-				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{}).Return(asset.LineageGraph{
+				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: true}).Return(asset.LineageGraph{
 					{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
 					{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
 					{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
@@ -878,11 +891,36 @@ func TestService_GetLineage(t *testing.T) {
 			},
 			Err: nil,
 		},
+		{
+			Description: `should return lineage with edges and without node attributes`,
+			ID:          assetID,
+			Query: asset.LineageQuery{
+				WithAttributes: false,
+			},
+			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
+				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: false}).Return(asset.LineageGraph{
+					{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
+					{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
+					{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
+				}, nil)
+			},
+			Expected: asset.Lineage{
+				Edges: []asset.LineageEdge{
+					{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
+					{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
+					{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
+				},
+			},
+			Err: nil,
+		},
 		{
 			Description: `should return lineage with edges and node attributes`,
 			ID:          assetID,
+			Query: asset.LineageQuery{
+				WithAttributes: true,
+			},
 			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
-				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{}).Return(asset.LineageGraph{
+				lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: true}).Return(asset.LineageGraph{
 					{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
 					{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
 					{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
@@ -942,7 +980,7 @@ func TestService_GetLineage(t *testing.T) {
 				LineageRepo:   mockLineageRepo,
 			})
 
-			actual, err := svc.GetLineage(ctx, "urn-source-1", asset.LineageQuery{})
+			actual, err := svc.GetLineage(ctx, "urn-source-1", tc.Query)
 			if tc.Err == nil {
 				assert.NoError(t, err)
 			} else {

internal/server/v1beta1/lineage.go

@@ -22,9 +22,15 @@ func (server *APIServer) GetGraph(ctx context.Context, req *compassv1beta1.GetGr
 		return nil, status.Error(codes.InvalidArgument, "invalid direction value")
 	}
 
+	withAttributes := true
+	if req != nil && req.WithAttributes != nil {
+		withAttributes = *req.WithAttributes
+	}
+
 	lineage, err := server.assetService.GetLineage(ctx, req.GetUrn(), asset.LineageQuery{
-		Level:     int(req.GetLevel()),
-		Direction: direction,
+		Level:          int(req.GetLevel()),
+		Direction:      direction,
+		WithAttributes: withAttributes,
 	})
 	if err != nil {
 		return nil, internalServerError(server.logger, err.Error())

internal/server/v1beta1/lineage_test.go

@@ -58,7 +58,7 @@ func TestGetLineageGraph(t *testing.T) {
 			defer mockUserSvc.AssertExpectations(t)
 			defer mockSvc.AssertExpectations(t)
 
-			mockSvc.EXPECT().GetLineage(ctx, nodeURN, asset.LineageQuery{Level: level, Direction: direction}).Return(lineage, nil)
+			mockSvc.EXPECT().GetLineage(ctx, nodeURN, asset.LineageQuery{Level: level, Direction: direction, WithAttributes: true}).Return(lineage, nil)
 			mockUserSvc.EXPECT().ValidateUser(ctx, userUUID, "").Return(userID, nil)
 
 			handler := NewAPIServer(APIServerDeps{AssetSvc: mockSvc, UserSvc: mockUserSvc, Logger: logger})

proto/compass.swagger.yaml

@@ -994,6 +994,10 @@ paths:
           in: query
           required: false
           type: string
+        - name: with_attributes
+          in: query
+          required: false
+          type: boolean
       tags:
         - Lineage
         - Asset