Skip to content

Commit a587fff

Browse files
sudo-suhassudo-suhas
authored
fix: Patch asset's lineage (#195)
- For the UpsertPatchAsset RPC, when no upstreams and downstreams are specified in the request, skip overwriting the asset's lineage unless OverwriteLineage flag is set to true. - Bump up the proton commit in Makefile to pull in changes from raystack/proton#238 + generate protos.
1 parent 7df7acd commit a587fff

File tree

9 files changed

+1467
-1211
lines changed

9 files changed

+1467
-1211
lines changed

‎Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
NAME="github.com/odpf/compass"
22
VERSION=$(shell git describe --always --tags 2>/dev/null)
33
COVERFILE="/tmp/compass.coverprofile"
4-
PROTON_COMMIT := "838f2a8c9ddc8fa6dfbd6f3ebe6201e76e2368f2"
4+
PROTON_COMMIT := "c7639b42da0679b2340a52155d2fe577b9d45aa2"
55
.PHONY: all build test clean install proto
66

77
all: build

‎core/asset/service.go‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -39,21 +39,27 @@ func (s *Service) GetAllAssets(ctx context.Context, flt Filter, withTotal bool)
3939
}
4040

4141
func (s *Service) UpsertAsset(ctx context.Context, ast *Asset, upstreams, downstreams []string) (string, error) {
42-
var assetID string
43-
var err error
42+
assetID, err := s.UpsertAssetWithoutLineage(ctx, ast)
43+
if err != nil {
44+
return "", err
45+
}
46+
47+
if err := s.lineageRepository.Upsert(ctx, ast.URN, upstreams, downstreams); err != nil {
48+
return "", err
49+
}
4450

45-
assetID, err = s.assetRepository.Upsert(ctx, ast)
51+
return assetID, nil
52+
}
53+
54+
func (s *Service) UpsertAssetWithoutLineage(ctx context.Context, ast *Asset) (string, error) {
55+
assetID, err := s.assetRepository.Upsert(ctx, ast)
4656
if err != nil {
47-
return assetID, err
57+
return "", err
4858
}
4959

5060
ast.ID = assetID
5161
if err := s.discoveryRepository.Upsert(ctx, *ast); err != nil {
52-
return assetID, err
53-
}
54-
55-
if err := s.lineageRepository.Upsert(ctx, ast.URN, upstreams, downstreams); err != nil {
56-
return assetID, err
62+
return "", err
5763
}
5864

5965
return assetID, nil

‎core/asset/service_test.go‎

Lines changed: 62 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -233,12 +233,70 @@ func TestService_UpsertAsset(t *testing.T) {
233233

234234
svc := asset.NewService(mockAssetRepo, mockDiscoveryRepo, mockLineageRepo)
235235
rid, err := svc.UpsertAsset(ctx, tc.Asset, tc.Upstreams, tc.Downstreams)
236-
if err != nil && errors.Is(tc.Err, err) {
237-
t.Fatalf("got error %v, expected error was %v", err, tc.Err)
236+
if tc.Err != nil {
237+
assert.EqualError(t, err, tc.Err.Error())
238+
return
239+
}
240+
assert.NoError(t, err)
241+
assert.Equal(t, tc.ReturnedID, rid)
242+
})
243+
}
244+
}
245+
246+
func TestService_UpsertAssetWithoutLineage(t *testing.T) {
247+
sampleAsset := &asset.Asset{ID: "some-id", URN: "some-urn", Type: asset.TypeDashboard, Service: "some-service"}
248+
var testCases = []struct {
249+
Description string
250+
Asset *asset.Asset
251+
Err error
252+
ReturnedID string
253+
Setup func(context.Context, *mocks.AssetRepository, *mocks.DiscoveryRepository)
254+
}{
255+
{
256+
Description: `should return error if asset repository upsert return error`,
257+
Asset: sampleAsset,
258+
Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository) {
259+
ar.EXPECT().Upsert(ctx, sampleAsset).Return("", errors.New("unknown error"))
260+
},
261+
Err: errors.New("unknown error"),
262+
},
263+
{
264+
Description: `should return error if discovery repository upsert return error`,
265+
Asset: sampleAsset,
266+
Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository) {
267+
ar.EXPECT().Upsert(ctx, sampleAsset).Return(sampleAsset.ID, nil)
268+
dr.EXPECT().Upsert(ctx, *sampleAsset).Return(errors.New("unknown error"))
269+
},
270+
Err: errors.New("unknown error"),
271+
},
272+
{
273+
Description: `should return no error if all repositories upsert return no error`,
274+
Asset: sampleAsset,
275+
Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository) {
276+
ar.EXPECT().Upsert(ctx, sampleAsset).Return(sampleAsset.ID, nil)
277+
dr.EXPECT().Upsert(ctx, *sampleAsset).Return(nil)
278+
},
279+
ReturnedID: sampleAsset.ID,
280+
},
281+
}
282+
for _, tc := range testCases {
283+
t.Run(tc.Description, func(t *testing.T) {
284+
ctx := context.Background()
285+
286+
mockAssetRepo := mocks.NewAssetRepository(t)
287+
mockDiscoveryRepo := mocks.NewDiscoveryRepository(t)
288+
if tc.Setup != nil {
289+
tc.Setup(ctx, mockAssetRepo, mockDiscoveryRepo)
238290
}
239-
if tc.ReturnedID != rid {
240-
t.Fatalf("got returned id %v, expected returned id was %v", rid, tc.ReturnedID)
291+
292+
svc := asset.NewService(mockAssetRepo, mockDiscoveryRepo, mocks.NewLineageRepository(t))
293+
rid, err := svc.UpsertAssetWithoutLineage(ctx, tc.Asset)
294+
if tc.Err != nil {
295+
assert.EqualError(t, err, tc.Err.Error())
296+
return
241297
}
298+
assert.NoError(t, err)
299+
assert.Equal(t, tc.ReturnedID, rid)
242300
})
243301
}
244302
}

‎internal/server/v1beta1/asset.go‎

Lines changed: 39 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ type AssetService interface {
3030
GetAssetByVersion(ctx context.Context, id string, version string) (asset.Asset, error)
3131
GetAssetVersionHistory(ctx context.Context, flt asset.Filter, id string) ([]asset.Asset, error)
3232
UpsertAsset(ctx context.Context, ast *asset.Asset, upstreams, downstreams []string) (string, error)
33+
UpsertAssetWithoutLineage(ctx context.Context, ast *asset.Asset) (string, error)
3334
DeleteAsset(ctx context.Context, id string) error
3435

3536
GetLineage(ctx context.Context, urn string, query asset.LineageQuery) (asset.Lineage, error)
@@ -268,13 +269,14 @@ func (server *APIServer) UpsertPatchAsset(ctx context.Context, req *compassv1bet
268269
ast.Patch(patchAssetMap)
269270
ast.UpdatedBy.ID = userID
270271

271-
assetID, err := server.upsertAsset(
272-
ctx,
273-
ast,
274-
"asset_upsert_patch",
275-
req.GetUpstreams(),
276-
req.GetDownstreams(),
277-
)
272+
var assetID string
273+
if len(req.Upstreams) != 0 || len(req.Downstreams) != 0 || req.OverwriteLineage {
274+
assetID, err = server.upsertAsset(
275+
ctx, ast, "asset_upsert_patch", req.GetUpstreams(), req.GetDownstreams(),
276+
)
277+
} else {
278+
assetID, err = server.upsertAssetWithoutLineage(ctx, ast)
279+
}
278280
if err != nil {
279281
return nil, err
280282
}
@@ -382,6 +384,36 @@ func (server *APIServer) upsertAsset(
382384
return
383385
}
384386

387+
func (server *APIServer) upsertAssetWithoutLineage(ctx context.Context, ast asset.Asset) (string, error) {
388+
const mode = "asset_upsert_patch_without_lineage"
389+
390+
if err := server.validateAsset(ast); err != nil {
391+
return "", status.Error(codes.InvalidArgument, err.Error())
392+
}
393+
394+
assetID, err := server.assetService.UpsertAssetWithoutLineage(ctx, &ast)
395+
if err != nil {
396+
switch {
397+
case errors.As(err, new(asset.InvalidError)):
398+
return "", status.Error(codes.InvalidArgument, err.Error())
399+
400+
case errors.As(err, new(asset.DiscoveryError)):
401+
server.sendStatsDCounterMetric("discovery_error",
402+
map[string]string{
403+
"method": mode,
404+
})
405+
}
406+
407+
return "", internalServerError(server.logger, err.Error())
408+
}
409+
410+
server.sendStatsDCounterMetric(mode, map[string]string{
411+
"type": ast.Type.String(),
412+
"service": ast.Service,
413+
})
414+
return assetID, nil
415+
}
416+
385417
func (server *APIServer) buildAsset(baseAsset *compassv1beta1.UpsertAssetRequest_Asset) asset.Asset {
386418
ast := asset.Asset{
387419
URN: baseAsset.GetUrn(),

‎internal/server/v1beta1/asset_test.go‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -678,6 +678,97 @@ func TestUpsertPatchAsset(t *testing.T) {
678678

679679
},
680680
},
681+
{
682+
Description: "without explicit overwrite_lineage, should upsert asset without lineage",
683+
Setup: func(ctx context.Context, as *mocks.AssetService) {
684+
patchedAsset := asset.Asset{
685+
URN: "test dagger",
686+
Type: asset.TypeTable,
687+
Name: "new-name",
688+
Service: "kafka",
689+
UpdatedBy: user.User{ID: userID},
690+
Data: map[string]interface{}{},
691+
Owners: []user.User{{ID: "id", UUID: "", Email: "[email protected]", Provider: "provider"}},
692+
}
693+
694+
assetWithID := patchedAsset
695+
assetWithID.ID = assetID
696+
697+
as.EXPECT().GetAssetByID(ctx, "test dagger").Return(currentAsset, nil)
698+
as.EXPECT().UpsertAssetWithoutLineage(ctx, &patchedAsset).
699+
Return(assetWithID.ID, nil).
700+
Run(func(ctx context.Context, ast *asset.Asset) {
701+
patchedAsset.ID = assetWithID.ID
702+
})
703+
},
704+
Request: &compassv1beta1.UpsertPatchAssetRequest{
705+
Asset: &compassv1beta1.UpsertPatchAssetRequest_Asset{
706+
Urn: "test dagger",
707+
Type: "table",
708+
Name: wrapperspb.String("new-name"),
709+
Service: "kafka",
710+
Data: &structpb.Struct{},
711+
Owners: []*compassv1beta1.User{{Id: "id", Uuid: "", Email: "[email protected]", Provider: "provider"}},
712+
},
713+
},
714+
ExpectStatus: codes.OK,
715+
PostCheck: func(resp *compassv1beta1.UpsertPatchAssetResponse) error {
716+
expected := &compassv1beta1.UpsertPatchAssetResponse{
717+
Id: assetID,
718+
}
719+
if diff := cmp.Diff(resp, expected, protocmp.Transform()); diff != "" {
720+
return fmt.Errorf("expected response to be %+v, was %+v", expected, resp)
721+
}
722+
return nil
723+
724+
},
725+
},
726+
{
727+
Description: "with explicit overwrite_lineage, should upsert asset when lineage is not in the request",
728+
Setup: func(ctx context.Context, as *mocks.AssetService) {
729+
patchedAsset := asset.Asset{
730+
URN: "test dagger",
731+
Type: asset.TypeTable,
732+
Name: "new-name",
733+
Service: "kafka",
734+
UpdatedBy: user.User{ID: userID},
735+
Data: map[string]interface{}{},
736+
Owners: []user.User{{ID: "id", UUID: "", Email: "[email protected]", Provider: "provider"}},
737+
}
738+
739+
assetWithID := patchedAsset
740+
assetWithID.ID = assetID
741+
742+
as.EXPECT().GetAssetByID(ctx, "test dagger").Return(currentAsset, nil)
743+
as.EXPECT().UpsertAsset(ctx, &patchedAsset, []string{}, []string{}).
744+
Return(assetWithID.ID, nil).
745+
Run(func(ctx context.Context, ast *asset.Asset, _, _ []string) {
746+
patchedAsset.ID = assetWithID.ID
747+
})
748+
},
749+
Request: &compassv1beta1.UpsertPatchAssetRequest{
750+
Asset: &compassv1beta1.UpsertPatchAssetRequest_Asset{
751+
Urn: "test dagger",
752+
Type: "table",
753+
Name: wrapperspb.String("new-name"),
754+
Service: "kafka",
755+
Data: &structpb.Struct{},
756+
Owners: []*compassv1beta1.User{{Id: "id", Uuid: "", Email: "[email protected]", Provider: "provider"}},
757+
},
758+
OverwriteLineage: true,
759+
},
760+
ExpectStatus: codes.OK,
761+
PostCheck: func(resp *compassv1beta1.UpsertPatchAssetResponse) error {
762+
expected := &compassv1beta1.UpsertPatchAssetResponse{
763+
Id: assetID,
764+
}
765+
if diff := cmp.Diff(resp, expected, protocmp.Transform()); diff != "" {
766+
return fmt.Errorf("expected response to be %+v, was %+v", expected, resp)
767+
}
768+
return nil
769+
770+
},
771+
},
681772
}
682773
for _, tc := range testCases {
683774
t.Run(tc.Description, func(t *testing.T) {

‎internal/server/v1beta1/mocks/asset_service.go‎

Lines changed: 45 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎proto/compass.swagger.yaml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2278,6 +2278,13 @@ definitions:
22782278
type: array
22792279
items:
22802280
$ref: '#/definitions/LineageNode'
2281+
overwrite_lineage:
2282+
type: boolean
2283+
description: |-
2284+
overwrite_lineage determines whether the asset's lineage should be
2285+
overwritten with the upstreams and downstreams specified in the request.
2286+
Currently, it is only applicable when both upstreams and downstreams are
2287+
empty/not specified.
22812288
upstreams:
22822289
type: array
22832290
items:

0 commit comments

Comments
 (0)