wip

Author: Sovietaced

datacatalog/pkg/manager/impl/artifact_manager.go

@@ -223,6 +223,11 @@ func (m *artifactManager) findArtifact(ctx context.Context, datasetID *datacatal
 		artifactModel = tag.Artifact
 	}
 
+	// If the artifact is expired consider this tag expired too
+	if artifactModel.ExpiresAt != nil && artifactModel.ExpiresAt.Before(m.clock.Now()) {
+		return models.Artifact{}, errors.NewDataCatalogErrorf(codes.NotFound, "entry not found")
+	}
+
 	if len(artifactModel.ArtifactData) == 0 {
 		return models.Artifact{}, errors.NewDataCatalogErrorf(codes.Internal, "artifact [%+v] with key %v does not have artifact data associated", artifactModel, key)
 	}

datacatalog/pkg/manager/impl/artifact_manager_test.go

@@ -201,7 +201,9 @@ func getExpectedArtifactModel(ctx context.Context, t *testing.T, datastore *stor
 		Tags: []models.Tag{
 			{TagKey: models.TagKey{TagName: "test-tag"}, DatasetUUID: expectedDataset.GetUUID(), ArtifactID: artifact.GetId()},
 		},
-		CreatedAt: getTestTimestamp(),
+		BaseModel: models.BaseModel{
+			CreatedAt: getTestTimestamp(),
+		},
 	}
 }
 

datacatalog/pkg/repositories/gormimpl/artifact.go

@@ -3,13 +3,11 @@ package gormimpl
 import (
 	"context"
 
-	"google.golang.org/grpc/codes"
 	"gorm.io/gorm"
 	"gorm.io/gorm/clause"
 	"k8s.io/utils/clock"
 
 	"github.com/flyteorg/flyte/datacatalog/pkg/common"
-	catalogErrors "github.com/flyteorg/flyte/datacatalog/pkg/errors"
 	"github.com/flyteorg/flyte/datacatalog/pkg/repositories/errors"
 	"github.com/flyteorg/flyte/datacatalog/pkg/repositories/interfaces"
 	"github.com/flyteorg/flyte/datacatalog/pkg/repositories/models"
@@ -43,38 +41,21 @@ func (h *artifactRepo) Create(ctx context.Context, artifact models.Artifact) err
 	timer := h.repoMetrics.CreateDuration.Start(ctx)
 	defer timer.Stop()
 
-	err := h.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
-		var existing models.Artifact
-		getResult := tx.
-			Clauses(clause.Locking{Strength: "UPDATE"}).
-			Order("artifacts.created_at DESC"). // Always pick the most recent
-			First(
-				&existing,
-				&models.Artifact{ArtifactKey: artifact.ArtifactKey},
-			)
-
-		if getResult.Error == nil {
-			if existing.ExpiresAt != nil && h.clock.Now().Before(*existing.ExpiresAt) {
-				// If the previous artifact is not expired return already exists
-				return catalogErrors.NewDataCatalogErrorf(codes.AlreadyExists, "value already exists")
-			}
-		} else {
-			if getResult.Error.Error() != gorm.ErrRecordNotFound.Error() {
-				return h.errorTransformer.ToDataCatalogError(getResult.Error)
-			}
-			// Not found is desirable
-		}
+	tx := h.db.WithContext(ctx).Begin()
 
-		createResult := tx.Create(&artifact)
+	tx = tx.Create(&artifact)
 
-		if createResult.Error != nil {
-			return h.errorTransformer.ToDataCatalogError(createResult.Error)
-		}
+	if tx.Error != nil {
+		tx.Rollback()
+		return h.errorTransformer.ToDataCatalogError(tx.Error)
+	}
 
-		return nil
-	})
+	tx = tx.Commit()
+	if tx.Error != nil {
+		return h.errorTransformer.ToDataCatalogError(tx.Error)
+	}
 
-	return err
+	return nil
 }
 
 func (h *artifactRepo) GetAndFilterExpired(ctx context.Context, in models.ArtifactKey) (models.Artifact, error) {

datacatalog/pkg/repositories/gormimpl/tag.go

@@ -2,7 +2,6 @@ package gormimpl
 
 import (
 	"context"
-
 	"gorm.io/gorm"
 
 	"github.com/flyteorg/flyte/datacatalog/pkg/repositories/errors"

datacatalog/pkg/repositories/models/artifact.go

@@ -11,10 +11,7 @@ type ArtifactKey struct {
 }
 
 type Artifact struct {
-	CreatedAt time.Time `gorm:"primary_key"`
-	UpdatedAt time.Time
-	DeletedAt *time.Time `sql:"index"`
-
+	BaseModel
 	ArtifactKey
 	DatasetUUID        string         `gorm:"type:uuid;index:artifacts_dataset_uuid_idx"`
 	Dataset            Dataset        `gorm:"association_autocreate:false"`

datacatalog/pkg/repositories/models/tag.go

@@ -11,7 +11,7 @@ type TagKey struct {
 type Tag struct {
 	BaseModel
 	TagKey
-	ArtifactID  string
+	ArtifactID  string   `gorm:"primary_key"`
 	DatasetUUID string   `gorm:"type:uuid;index:tags_dataset_uuid_idx"`
 	Artifact    Artifact `gorm:"references:DatasetProject,DatasetName,DatasetDomain,DatasetVersion,ArtifactID;foreignkey:DatasetProject,DatasetName,DatasetDomain,DatasetVersion,ArtifactID"`
 }

datacatalog/pkg/repositories/transformers/artifact_test.go

@@ -123,7 +123,9 @@ func TestFromArtifactModel(t *testing.T) {
 		SerializedMetadata: []byte{},
 		Partitions:         getTestPartitions(),
 		Tags:               getTestTags(),
-		CreatedAt:          createdAt,
+		BaseModel: models.BaseModel{
+			CreatedAt: createdAt,
+		},
 	}
 
 	actual, err := FromArtifactModel(artifactModel)