stuff

Author: Sovietaced

datacatalog/pkg/repositories/errors/postgres.go

@@ -24,10 +24,10 @@ type postgresErrorTransformer struct {
 }
 
 const (
-	unexpectedType            = "unexpected error type for: %v"
-	uniqueConstraintViolation = "value with matching already exists (%s)"
-	defaultPgError            = "failed database operation with code [%s] and msg [%s]"
-	unsupportedTableOperation = "cannot query with specified table attributes: %s"
+	unexpectedType             = "unexpected error type for: %v"
+	uUniqueConstraintViolation = "value with matching already exists (%s)"
+	defaultPgError             = "failed database operation with code [%s] and msg [%s]"
+	unsupportedTableOperation  = "cannot query with specified table attributes: %s"
 )
 
 func (p *postgresErrorTransformer) fromGormError(err error) error {
@@ -40,9 +40,14 @@ func (p *postgresErrorTransformer) fromGormError(err error) error {
 }
 
 func (p *postgresErrorTransformer) ToDataCatalogError(err error) error {
+	var dce catalogErrors.DataCatalogError
+	if errors.As(err, &dce) {
+		return dce // already a data catalog error
+	}
+
 	// First try the stdlib error handling
 	if database.IsPgErrorWithCode(err, uniqueConstraintViolationCode) {
-		return catalogErrors.NewDataCatalogErrorf(codes.AlreadyExists, uniqueConstraintViolation, err.Error())
+		return catalogErrors.NewDataCatalogErrorf(codes.AlreadyExists, uUniqueConstraintViolation, err.Error())
 	}
 
 	if unwrappedErr := errors.Unwrap(err); unwrappedErr != nil {
@@ -58,7 +63,7 @@ func (p *postgresErrorTransformer) ToDataCatalogError(err error) error {
 
 	switch pqError.Code {
 	case uniqueConstraintViolationCode:
-		return catalogErrors.NewDataCatalogErrorf(codes.AlreadyExists, uniqueConstraintViolation, pqError.Message)
+		return catalogErrors.NewDataCatalogErrorf(codes.AlreadyExists, uUniqueConstraintViolation, pqError.Message)
 	case undefinedTable:
 		return catalogErrors.NewDataCatalogErrorf(codes.InvalidArgument, unsupportedTableOperation, pqError.Message)
 	default:

datacatalog/pkg/repositories/gormimpl/artifact.go

@@ -3,11 +3,14 @@ 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"
@@ -41,18 +44,28 @@ func (h *artifactRepo) Create(ctx context.Context, artifact models.Artifact) err
 	timer := h.repoMetrics.CreateDuration.Start(ctx)
 	defer timer.Stop()
 
-	tx := h.db.WithContext(ctx).Begin()
+	err := h.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		tx = tx.
+			Where("artifacts.expires_at is null or artifacts.expires_at < ?", h.clock.Now().UTC()).
+			Order("artifacts.created_at DESC"). // Always pick the most recent
+			Find(
+				&models.Artifact{ArtifactKey: artifact.ArtifactKey},
+			)
 
-	tx = tx.Create(&artifact)
+		if tx.Error != nil {
+			return tx.Error
+		}
 
-	if tx.Error != nil {
-		tx.Rollback()
-		return h.errorTransformer.ToDataCatalogError(tx.Error)
-	}
+		if tx.RowsAffected > 0 {
+			return catalogErrors.NewDataCatalogErrorf(codes.AlreadyExists, "artifact already exists")
+		}
 
-	tx = tx.Commit()
-	if tx.Error != nil {
-		return h.errorTransformer.ToDataCatalogError(tx.Error)
+		tx = tx.Create(&artifact)
+		return tx.Error
+	})
+
+	if err != nil {
+		return h.errorTransformer.ToDataCatalogError(err)
 	}
 
 	return nil

datacatalog/pkg/repositories/gormimpl/artifact_test.go

@@ -124,15 +124,23 @@ func getDBTagResponse(artifact models.Artifact) []map[string]interface{} {
 }
 
 func TestCreateArtifact(t *testing.T) {
+	testClock := testclock.NewFakeClock(time.Unix(0, 0))
 	artifact := getTestArtifact()
 
+	existingChecked := false
 	artifactCreated := false
 	GlobalMock := mocket.Catcher.Reset()
 	GlobalMock.Logging = true
 
 	numArtifactDataCreated := 0
 	numPartitionsCreated := 0
 
+	GlobalMock.NewMock().WithQuery(
+		`SELECT * FROM "artifacts" WHERE (artifacts.expires_at is null or artifacts.expires_at < $1) AND "artifacts"."dataset_project" = $2 AND "artifacts"."dataset_name" = $3 AND "artifacts"."dataset_domain" = $4 AND "artifacts"."dataset_version" = $5 AND "artifacts"."artifact_id" = $6 ORDER BY artifacts.created_at DESC%!!(string=123)!(string=testVersion)!(string=testDomain)!(string=testName)!(string=testProject)(EXTRA time.Time=1970-01-01 00:00:00 +0000 UTC)`).WithCallback(
+		func(s string, values []driver.NamedValue) {
+			existingChecked = true
+		})
+
 	// Only match on queries that append expected filters
 	GlobalMock.NewMock().WithQuery(
 		`INSERT INTO "artifacts" ("created_at","updated_at","deleted_at","dataset_project","dataset_name","dataset_domain","dataset_version","artifact_id","dataset_uuid","serialized_metadata","expires_at") VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)`).WithCallback(
@@ -173,12 +181,38 @@ func TestCreateArtifact(t *testing.T) {
 
 	artifact.Partitions = partitions
 
-	artifactRepo := NewArtifactRepo(utils.GetDbForTest(t), errors.NewPostgresErrorTransformer(), promutils.NewTestScope(), clock.RealClock{})
+	artifactRepo := NewArtifactRepo(utils.GetDbForTest(t), errors.NewPostgresErrorTransformer(), promutils.NewTestScope(), testClock)
 	err := artifactRepo.Create(context.Background(), artifact)
 	assert.NoError(t, err)
 	assert.True(t, artifactCreated)
 	assert.Equal(t, 2, numArtifactDataCreated)
 	assert.Equal(t, 1, numPartitionsCreated)
+	assert.True(t, existingChecked)
+}
+
+func TestCreateArtifactAlreadyExists(t *testing.T) {
+	testClock := testclock.NewFakeClock(time.Unix(0, 0))
+
+	artifact := getTestArtifactWithExpiration(testClock.Now().Add(time.Second))
+	expectedArtifactResponse := getDBArtifactResponse(artifact)
+
+	existingChecked := false
+	GlobalMock := mocket.Catcher.Reset()
+	GlobalMock.Logging = true
+
+	GlobalMock.NewMock().WithQuery(
+		`SELECT * FROM "artifacts" WHERE (artifacts.expires_at is null or artifacts.expires_at < $1) AND "artifacts"."dataset_project" = $2 AND "artifacts"."dataset_name" = $3 AND "artifacts"."dataset_domain" = $4 AND "artifacts"."dataset_version" = $5 AND "artifacts"."artifact_id" = $6 ORDER BY artifacts.created_at DESC%!!(string=123)!(string=testVersion)!(string=testDomain)!(string=testName)!(string=testProject)(EXTRA time.Time=1970-01-01 00:00:00 +0000 UTC)`).WithCallback(
+		func(s string, values []driver.NamedValue) {
+			existingChecked = true
+		}).WithReply(expectedArtifactResponse)
+
+	artifactRepo := NewArtifactRepo(utils.GetDbForTest(t), errors.NewPostgresErrorTransformer(), promutils.NewTestScope(), testClock)
+	err := artifactRepo.Create(context.Background(), artifact)
+	assert.Error(t, err)
+	dcErr, ok := err.(apiErrors.DataCatalogError)
+	assert.True(t, ok)
+	assert.Equal(t, codes.AlreadyExists.String(), dcErr.Code().String())
+	assert.True(t, existingChecked)
 }
 
 func TestGetArtifactNotExpired(t *testing.T) {