perf: skip probes fetch during upsert patch asset (#56)

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
 }
 

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"

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())
 	}

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, &currentAssetNew).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, &currentAssetNew).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) {