Skip to content

Commit 8318c34

Browse files
irainiairainia
authored
feat: optional attributes when getting graph (#76)
* chore: regenerate proto * feat: implement optional attributes * chore: update proton commit * test: add unit test for no attributes selection
1 parent 3143032 commit 8318c34

File tree

9 files changed

+1518
-1446
lines changed

9 files changed

+1518
-1446
lines changed

‎Makefile

+1-1
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ help: ##@help show this help
2626
NAME="github.com/goto/compass"
2727
VERSION=$(shell git describe --always --tags 2>/dev/null)
2828
COVERFILE="/tmp/compass.coverprofile"
29-
PROTON_COMMIT := "eaca9798d1c1d7b3101ec1259c7e5fb949afba28"
29+
PROTON_COMMIT := "fe99cc96e060085d6052096e9ba59b4038c691c6"
3030

3131
TOOLS_MOD_DIR = ./tools
3232
TOOLS_DIR = $(abspath ./.tools)

‎core/asset/lineage.go

+3-2
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,9 @@ const (
2121
)
2222

2323
type LineageQuery struct {
24-
Level int
25-
Direction LineageDirection
24+
Level int
25+
Direction LineageDirection
26+
WithAttributes bool
2627
}
2728

2829
//go:generate mockery --name=LineageRepository -r --case underscore --with-expecter --structname=LineageRepository --filename=lineage_repository.go --output=./mocks

‎core/asset/service.go

+6
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,12 @@ func (s *Service) GetLineage(ctx context.Context, urn string, query LineageQuery
190190
return Lineage{}, fmt.Errorf("get lineage: get graph edges: %w", err)
191191
}
192192

193+
if !query.WithAttributes {
194+
return Lineage{
195+
Edges: edges,
196+
}, nil
197+
}
198+
193199
urns := newUniqueStrings(len(edges))
194200
urns.add(urn)
195201
for _, edge := range edges {

‎core/asset/service_test.go

+44-6
Original file line numberDiff line numberDiff line change
@@ -807,6 +807,7 @@ func TestService_GetLineage(t *testing.T) {
807807
type testCase struct {
808808
Description string
809809
ID string
810+
Query asset.LineageQuery
810811
Setup func(context.Context, *mocks.AssetRepository, *mocks.DiscoveryRepository, *mocks.LineageRepository)
811812
Expected asset.Lineage
812813
Err error
@@ -816,8 +817,11 @@ func TestService_GetLineage(t *testing.T) {
816817
{
817818
Description: `should return error if the GetGraph function return error`,
818819
ID: assetID,
820+
Query: asset.LineageQuery{
821+
WithAttributes: true,
822+
},
819823
Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
820-
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{}).
824+
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: true}).
821825
Return(asset.LineageGraph{}, errors.New("error fetching graph"))
822826
},
823827
Expected: asset.Lineage{},
@@ -826,8 +830,11 @@ func TestService_GetLineage(t *testing.T) {
826830
{
827831
Description: `should return no error if graph with 0 edges are returned`,
828832
ID: assetID,
833+
Query: asset.LineageQuery{
834+
WithAttributes: true,
835+
},
829836
Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
830-
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{}).
837+
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: true}).
831838
Return(asset.LineageGraph{}, nil)
832839
ar.EXPECT().GetProbesWithFilter(ctx, asset.ProbesFilter{
833840
AssetURNs: []string{"urn-source-1"},
@@ -840,8 +847,11 @@ func TestService_GetLineage(t *testing.T) {
840847
{
841848
Description: `should return an error if GetProbesWithFilter function returns error`,
842849
ID: assetID,
850+
Query: asset.LineageQuery{
851+
WithAttributes: true,
852+
},
843853
Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
844-
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{}).Return(asset.LineageGraph{
854+
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: true}).Return(asset.LineageGraph{
845855
{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
846856
{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
847857
{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
@@ -857,8 +867,11 @@ func TestService_GetLineage(t *testing.T) {
857867
{
858868
Description: `should return no error if GetProbesWithFilter function returns 0 probes`,
859869
ID: assetID,
870+
Query: asset.LineageQuery{
871+
WithAttributes: true,
872+
},
860873
Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
861-
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{}).Return(asset.LineageGraph{
874+
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: true}).Return(asset.LineageGraph{
862875
{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
863876
{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
864877
{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
@@ -878,11 +891,36 @@ func TestService_GetLineage(t *testing.T) {
878891
},
879892
Err: nil,
880893
},
894+
{
895+
Description: `should return lineage with edges and without node attributes`,
896+
ID: assetID,
897+
Query: asset.LineageQuery{
898+
WithAttributes: false,
899+
},
900+
Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
901+
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: false}).Return(asset.LineageGraph{
902+
{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
903+
{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
904+
{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
905+
}, nil)
906+
},
907+
Expected: asset.Lineage{
908+
Edges: []asset.LineageEdge{
909+
{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
910+
{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
911+
{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
912+
},
913+
},
914+
Err: nil,
915+
},
881916
{
882917
Description: `should return lineage with edges and node attributes`,
883918
ID: assetID,
919+
Query: asset.LineageQuery{
920+
WithAttributes: true,
921+
},
884922
Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
885-
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{}).Return(asset.LineageGraph{
923+
lr.EXPECT().GetGraph(ctx, "urn-source-1", asset.LineageQuery{WithAttributes: true}).Return(asset.LineageGraph{
886924
{Source: "urn-source-1", Target: "urn-target-1", Prop: nil},
887925
{Source: "urn-source-1", Target: "urn-target-2", Prop: nil},
888926
{Source: "urn-target-2", Target: "urn-target-3", Prop: nil},
@@ -942,7 +980,7 @@ func TestService_GetLineage(t *testing.T) {
942980
LineageRepo: mockLineageRepo,
943981
})
944982

945-
actual, err := svc.GetLineage(ctx, "urn-source-1", asset.LineageQuery{})
983+
actual, err := svc.GetLineage(ctx, "urn-source-1", tc.Query)
946984
if tc.Err == nil {
947985
assert.NoError(t, err)
948986
} else {

‎internal/server/v1beta1/lineage.go

+8-2
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,15 @@ func (server *APIServer) GetGraph(ctx context.Context, req *compassv1beta1.GetGr
2222
return nil, status.Error(codes.InvalidArgument, "invalid direction value")
2323
}
2424

25+
withAttributes := true
26+
if req != nil && req.WithAttributes != nil {
27+
withAttributes = *req.WithAttributes
28+
}
29+
2530
lineage, err := server.assetService.GetLineage(ctx, req.GetUrn(), asset.LineageQuery{
26-
Level: int(req.GetLevel()),
27-
Direction: direction,
31+
Level: int(req.GetLevel()),
32+
Direction: direction,
33+
WithAttributes: withAttributes,
2834
})
2935
if err != nil {
3036
return nil, internalServerError(server.logger, err.Error())

‎internal/server/v1beta1/lineage_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ func TestGetLineageGraph(t *testing.T) {
5858
defer mockUserSvc.AssertExpectations(t)
5959
defer mockSvc.AssertExpectations(t)
6060

61-
mockSvc.EXPECT().GetLineage(ctx, nodeURN, asset.LineageQuery{Level: level, Direction: direction}).Return(lineage, nil)
61+
mockSvc.EXPECT().GetLineage(ctx, nodeURN, asset.LineageQuery{Level: level, Direction: direction, WithAttributes: true}).Return(lineage, nil)
6262
mockUserSvc.EXPECT().ValidateUser(ctx, userUUID, "").Return(userID, nil)
6363

6464
handler := NewAPIServer(APIServerDeps{AssetSvc: mockSvc, UserSvc: mockUserSvc, Logger: logger})

‎proto/compass.swagger.yaml

+4
Original file line numberDiff line numberDiff line change
@@ -994,6 +994,10 @@ paths:
994994
in: query
995995
required: false
996996
type: string
997+
- name: with_attributes
998+
in: query
999+
required: false
1000+
type: boolean
9971001
tags:
9981002
- Lineage
9991003
- Asset

0 commit comments

Comments
 (0)