From ec9201e933123728861bc201111370f23d3a47ad Mon Sep 17 00:00:00 2001 From: Suhas Karanth Date: Mon, 21 Aug 2023 12:32:11 +0530 Subject: [PATCH] perf: skip probes fetch during upsert patch asset (#56) --- core/asset/service.go | 39 ++++--- core/asset/service_test.go | 109 ++++++++++++++++++ internal/server/v1beta1/asset.go | 3 +- internal/server/v1beta1/asset_test.go | 14 +-- .../server/v1beta1/mocks/asset_service.go | 53 +++++++++ 5 files changed, 197 insertions(+), 21 deletions(-) diff --git a/core/asset/service.go b/core/asset/service.go index 6e0fa528..9786e206 100644 --- a/core/asset/service.go +++ b/core/asset/service.go @@ -121,32 +121,45 @@ func (s *Service) DeleteAsset(ctx context.Context, id string) (err error) { return s.lineageRepository.DeleteByURN(ctx, urn) } -func (s *Service) GetAssetByID(ctx context.Context, id string) (ast Asset, err error) { +func (s *Service) GetAssetByID(ctx context.Context, id string) (Asset, error) { + ast, err := s.assetByIDWithoutProbes(ctx, "GetAssetByID", id) + if err != nil { + return Asset{}, err + } + + probes, err := s.assetRepository.GetProbes(ctx, ast.URN) + if err != nil { + return Asset{}, fmt.Errorf("get probes: %w", err) + } + + ast.Probes = probes + + return ast, nil +} + +func (s *Service) GetAssetByIDWithoutProbes(ctx context.Context, id string) (Asset, error) { + return s.assetByIDWithoutProbes(ctx, "GetAssetByIDWithoutProbes", id) +} + +func (s *Service) assetByIDWithoutProbes(ctx context.Context, op, id string) (ast Asset, err error) { defer func() { - s.instrumentAssetOp(ctx, "GetAssetByID", id, err) + s.instrumentAssetOp(ctx, op, id, err) }() if isValidUUID(id) { - var err error ast, err = s.assetRepository.GetByID(ctx, id) if err != nil { return Asset{}, fmt.Errorf("get asset by id: %w", err) } - } else { - var err error - ast, err = s.assetRepository.GetByURN(ctx, id) - if err != nil { - return Asset{}, fmt.Errorf("get asset by urn: %w", err) - } + + return ast, nil } - probes, err := s.assetRepository.GetProbes(ctx, ast.URN) + ast, err = s.assetRepository.GetByURN(ctx, id) if err != nil { - return Asset{}, fmt.Errorf("get probes: %w", err) + return Asset{}, fmt.Errorf("get asset by urn: %w", err) } - ast.Probes = probes - return ast, nil } diff --git a/core/asset/service_test.go b/core/asset/service_test.go index 94861e4a..35219586 100644 --- a/core/asset/service_test.go +++ b/core/asset/service_test.go @@ -571,6 +571,115 @@ func TestService_GetAssetByID(t *testing.T) { } } +func TestService_GetAssetByIDWithoutProbes(t *testing.T) { + assetID := "f742aa61-1100-445c-8d72-355a42e2fb59" + urn := "my-test-urn" + now := time.Now().UTC() + ast := asset.Asset{ + ID: assetID, + } + + testCases := []struct { + Description string + ID string + Expected *asset.Asset + ExpectedErr error + Setup func(context.Context, *mocks.AssetRepository) + }{ + { + Description: `should return error if the repository return error without id`, + ID: assetID, + Setup: func(ctx context.Context, ar *mocks.AssetRepository) { + ar.EXPECT().GetByID(ctx, assetID).Return(asset.Asset{}, asset.NotFoundError{}) + }, + ExpectedErr: asset.NotFoundError{}, + }, + { + Description: `should return error if the repository return error, with id`, + ID: assetID, + Setup: func(ctx context.Context, ar *mocks.AssetRepository) { + ar.EXPECT().GetByID(ctx, assetID).Return(asset.Asset{}, asset.NotFoundError{AssetID: ast.ID}) + }, + ExpectedErr: asset.NotFoundError{AssetID: ast.ID}, + }, + { + Description: `should return error if the repository return error, with invalid id`, + ID: assetID, + Setup: func(ctx context.Context, ar *mocks.AssetRepository) { + ar.EXPECT().GetByID(ctx, assetID).Return(asset.Asset{}, asset.InvalidError{AssetID: ast.ID}) + }, + ExpectedErr: asset.InvalidError{AssetID: ast.ID}, + }, + { + Description: `with URN, should return error from repository`, + ID: urn, + Setup: func(ctx context.Context, ar *mocks.AssetRepository) { + ar.EXPECT().GetByURN(ctx, urn).Return(asset.Asset{}, errors.New("the world exploded")) + }, + ExpectedErr: errors.New("the world exploded"), + }, + { + Description: `with ID, should return no error if asset is found`, + ID: assetID, + Expected: &asset.Asset{ + ID: assetID, + URN: urn, + CreatedAt: now, + }, + Setup: func(ctx context.Context, ar *mocks.AssetRepository) { + ar.EXPECT().GetByID(ctx, assetID).Return(asset.Asset{ + ID: assetID, + URN: urn, + CreatedAt: now, + }, nil) + }, + ExpectedErr: nil, + }, + { + Description: `with URN, should return no error if asset is found`, + ID: urn, + Expected: &asset.Asset{ + ID: assetID, + URN: urn, + CreatedAt: now, + }, + Setup: func(ctx context.Context, ar *mocks.AssetRepository) { + ar.EXPECT().GetByURN(ctx, urn).Return(asset.Asset{ + ID: assetID, + URN: urn, + CreatedAt: now, + }, nil) + }, + ExpectedErr: nil, + }, + } + for _, tc := range testCases { + t.Run(tc.Description, func(t *testing.T) { + ctx := context.Background() + + mockAssetRepo := mocks.NewAssetRepository(t) + if tc.Setup != nil { + tc.Setup(ctx, mockAssetRepo) + } + + svc := asset.NewService(asset.ServiceDeps{ + AssetRepo: mockAssetRepo, + DiscoveryRepo: mocks.NewDiscoveryRepository(t), + LineageRepo: mocks.NewLineageRepository(t), + }) + + actual, err := svc.GetAssetByIDWithoutProbes(ctx, tc.ID) + if tc.Expected != nil { + assert.Equal(t, *tc.Expected, actual) + } + if tc.ExpectedErr != nil { + assert.ErrorContains(t, err, tc.ExpectedErr.Error()) + assert.ErrorAs(t, err, &tc.ExpectedErr) + } + }) + } +} + func TestService_GetAssetByVersion(t *testing.T) { assetID := "f742aa61-1100-445c-8d72-355a42e2fb59" urn := "my-test-urn" diff --git a/internal/server/v1beta1/asset.go b/internal/server/v1beta1/asset.go index 1104fade..07fb88f3 100644 --- a/internal/server/v1beta1/asset.go +++ b/internal/server/v1beta1/asset.go @@ -30,6 +30,7 @@ type StatsDClient interface { type AssetService interface { GetAllAssets(ctx context.Context, flt asset.Filter, withTotal bool) ([]asset.Asset, uint32, error) GetAssetByID(ctx context.Context, id string) (asset.Asset, error) + GetAssetByIDWithoutProbes(ctx context.Context, id string) (asset.Asset, error) GetAssetByVersion(ctx context.Context, id, version string) (asset.Asset, error) GetAssetVersionHistory(ctx context.Context, flt asset.Filter, id string) ([]asset.Asset, error) UpsertAsset(ctx context.Context, ast *asset.Asset, upstreams, downstreams []string) (string, error) @@ -264,7 +265,7 @@ func (server *APIServer) UpsertPatchAsset(ctx context.Context, req *compassv1bet return nil, status.Error(codes.InvalidArgument, err.Error()) } - ast, err := server.assetService.GetAssetByID(ctx, urn) + ast, err := server.assetService.GetAssetByIDWithoutProbes(ctx, urn) if err != nil && !errors.As(err, &asset.NotFoundError{}) { return nil, internalServerError(server.logger, err.Error()) } diff --git a/internal/server/v1beta1/asset_test.go b/internal/server/v1beta1/asset_test.go index 7508d560..79ced36b 100644 --- a/internal/server/v1beta1/asset_test.go +++ b/internal/server/v1beta1/asset_test.go @@ -715,7 +715,7 @@ func TestUpsertPatchAsset(t *testing.T) { Description: "should return internal server error when finding asset failed", Setup: func(ctx context.Context, as *mocks.AssetService, _ *mocks.UserService) { expectedErr := errors.New("unknown error") - as.EXPECT().GetAssetByID(ctx, "test dagger").Return(currentAsset, expectedErr) + as.EXPECT().GetAssetByIDWithoutProbes(ctx, "test dagger").Return(currentAsset, expectedErr) }, Request: validPayload, ExpectStatus: codes.Internal, @@ -724,7 +724,7 @@ func TestUpsertPatchAsset(t *testing.T) { Description: "should return internal server error when upserting asset service failed", Setup: func(ctx context.Context, as *mocks.AssetService, _ *mocks.UserService) { expectedErr := errors.New("unknown error") - as.EXPECT().GetAssetByID(ctx, "test dagger").Return(currentAsset, nil) + as.EXPECT().GetAssetByIDWithoutProbes(ctx, "test dagger").Return(currentAsset, nil) as.EXPECT().UpsertAsset(ctx, mock.AnythingOfType("*asset.Asset"), mock.AnythingOfType("[]string"), mock.AnythingOfType("[]string")).Return("1234-5678", expectedErr) }, Request: validPayload, @@ -733,7 +733,7 @@ func TestUpsertPatchAsset(t *testing.T) { { Description: "should return invalid argument error when upserting asset without lineage failed, with invalid error", Setup: func(ctx context.Context, as *mocks.AssetService, _ *mocks.UserService) { - as.EXPECT().GetAssetByID(ctx, "test dagger").Return(currentAssetNew, nil) + as.EXPECT().GetAssetByIDWithoutProbes(ctx, "test dagger").Return(currentAssetNew, nil) as.EXPECT().UpsertAssetWithoutLineage(ctx, ¤tAssetNew).Return("", asset.InvalidError{}) }, Request: validPayloadWithoutStreams, @@ -742,7 +742,7 @@ func TestUpsertPatchAsset(t *testing.T) { { Description: "should return internal server error when upserting asset without asset failed, with discovery error ", Setup: func(ctx context.Context, as *mocks.AssetService, _ *mocks.UserService) { - as.EXPECT().GetAssetByID(ctx, "test dagger").Return(currentAssetNew, nil) + as.EXPECT().GetAssetByIDWithoutProbes(ctx, "test dagger").Return(currentAssetNew, nil) as.EXPECT().UpsertAssetWithoutLineage(ctx, ¤tAssetNew).Return("", asset.DiscoveryError{Err: errors.New("discovery error")}) }, Request: validPayloadWithoutStreams, @@ -767,7 +767,7 @@ func TestUpsertPatchAsset(t *testing.T) { assetWithID := patchedAsset assetWithID.ID = assetID - as.EXPECT().GetAssetByID(ctx, "test dagger").Return(currentAsset, nil) + as.EXPECT().GetAssetByIDWithoutProbes(ctx, "test dagger").Return(currentAsset, nil) as.EXPECT().UpsertAsset(ctx, &patchedAsset, upstreams, downstreams).Return(assetWithID.ID, nil).Run(func(ctx context.Context, ast *asset.Asset, upstreams, downstreams []string) { patchedAsset.ID = assetWithID.ID }) @@ -801,7 +801,7 @@ func TestUpsertPatchAsset(t *testing.T) { assetWithID := patchedAsset assetWithID.ID = assetID - as.EXPECT().GetAssetByID(ctx, "test dagger").Return(currentAsset, nil) + as.EXPECT().GetAssetByIDWithoutProbes(ctx, "test dagger").Return(currentAsset, nil) as.EXPECT().UpsertAssetWithoutLineage(ctx, &patchedAsset). Return(assetWithID.ID, nil). Run(func(ctx context.Context, ast *asset.Asset) { @@ -846,7 +846,7 @@ func TestUpsertPatchAsset(t *testing.T) { assetWithID := patchedAsset assetWithID.ID = assetID - as.EXPECT().GetAssetByID(ctx, "test dagger").Return(currentAsset, nil) + as.EXPECT().GetAssetByIDWithoutProbes(ctx, "test dagger").Return(currentAsset, nil) as.EXPECT().UpsertAsset(ctx, &patchedAsset, []string{}, []string{}). Return(assetWithID.ID, nil). Run(func(ctx context.Context, ast *asset.Asset, _, _ []string) { diff --git a/internal/server/v1beta1/mocks/asset_service.go b/internal/server/v1beta1/mocks/asset_service.go index 7834952e..d80073fe 100644 --- a/internal/server/v1beta1/mocks/asset_service.go +++ b/internal/server/v1beta1/mocks/asset_service.go @@ -226,6 +226,59 @@ func (_c *AssetService_GetAssetByID_Call) RunAndReturn(run func(context.Context, return _c } +// GetAssetByIDWithoutProbes provides a mock function with given fields: ctx, id +func (_m *AssetService) GetAssetByIDWithoutProbes(ctx context.Context, id string) (asset.Asset, error) { + ret := _m.Called(ctx, id) + + var r0 asset.Asset + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, string) (asset.Asset, error)); ok { + return rf(ctx, id) + } + if rf, ok := ret.Get(0).(func(context.Context, string) asset.Asset); ok { + r0 = rf(ctx, id) + } else { + r0 = ret.Get(0).(asset.Asset) + } + + if rf, ok := ret.Get(1).(func(context.Context, string) error); ok { + r1 = rf(ctx, id) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// AssetService_GetAssetByIDWithoutProbes_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetAssetByIDWithoutProbes' +type AssetService_GetAssetByIDWithoutProbes_Call struct { + *mock.Call +} + +// GetAssetByIDWithoutProbes is a helper method to define mock.On call +// - ctx context.Context +// - id string +func (_e *AssetService_Expecter) GetAssetByIDWithoutProbes(ctx interface{}, id interface{}) *AssetService_GetAssetByIDWithoutProbes_Call { + return &AssetService_GetAssetByIDWithoutProbes_Call{Call: _e.mock.On("GetAssetByIDWithoutProbes", ctx, id)} +} + +func (_c *AssetService_GetAssetByIDWithoutProbes_Call) Run(run func(ctx context.Context, id string)) *AssetService_GetAssetByIDWithoutProbes_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string)) + }) + return _c +} + +func (_c *AssetService_GetAssetByIDWithoutProbes_Call) Return(_a0 asset.Asset, _a1 error) *AssetService_GetAssetByIDWithoutProbes_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *AssetService_GetAssetByIDWithoutProbes_Call) RunAndReturn(run func(context.Context, string) (asset.Asset, error)) *AssetService_GetAssetByIDWithoutProbes_Call { + _c.Call.Return(run) + return _c +} + // GetAssetByVersion provides a mock function with given fields: ctx, id, version func (_m *AssetService) GetAssetByVersion(ctx context.Context, id string, version string) (asset.Asset, error) { ret := _m.Called(ctx, id, version)