feat: implement optional attributes

goto · Jul 25, 2024 · c78f89c · c78f89c
1 parent 626ce47
commit c78f89c
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 11 deletions.
diff --git a/core/asset/lineage.go b/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

diff --git a/core/asset/service.go b/core/asset/service.go
@@ -190,6 +190,13 @@ 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,
+			NodeAttrs: make(map[string]NodeAttributes),
+		}, nil
+	}
+
 	urns := newUniqueStrings(len(edges))
 	urns.add(urn)
 	for _, edge := range edges {

diff --git a/core/asset/service_test.go b/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},
@@ -881,8 +894,11 @@ func TestService_GetLineage(t *testing.T) {
 		{
 			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 +958,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 {

diff --git a/internal/server/v1beta1/lineage.go b/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())

diff --git a/internal/server/v1beta1/lineage_test.go b/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})