From c78f89c0188edaa555f9b9f8e183bfc0b7d93cd5 Mon Sep 17 00:00:00 2001 From: Anwar Hidayat <15167551+irainia@users.noreply.github.com> Date: Thu, 25 Jul 2024 15:33:46 +0700 Subject: [PATCH] feat: implement optional attributes --- core/asset/lineage.go | 5 +++-- core/asset/service.go | 7 +++++++ core/asset/service_test.go | 28 +++++++++++++++++++------ internal/server/v1beta1/lineage.go | 10 +++++++-- internal/server/v1beta1/lineage_test.go | 2 +- 5 files changed, 41 insertions(+), 11 deletions(-) diff --git a/core/asset/lineage.go b/core/asset/lineage.go index b6ceaf88..61acded7 100644 --- 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 index 189ee719..c4052a0c 100644 --- 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 index 35219586..16ca2d7f 100644 --- 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 index a571dcae..d6c889c1 100644 --- 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 index d1cca1bd..ef01be6b 100644 --- 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})