tag stealing in tx

flyteorg · Dec 10, 2019 · 3893716 · 3893716
1 parent 3b6f8b8
commit 3893716
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 26 deletions.
diff --git a/pkg/repositories/gormimpl/tag.go b/pkg/repositories/gormimpl/tag.go
@@ -50,15 +50,15 @@ func (h *tagRepo) Create(ctx context.Context, tag models.Tag) error {
 
 	filters = append(filters, NewGormValueFilter(common.Artifact, common.Equal, "tag_name", tag.TagName))
 
-	listTaggedArtifacts := models.ListModelsInput{
+	listTaggedInput := models.ListModelsInput{
 		JoinEntityToConditionMap: map[common.Entity]models.ModelJoinCondition{
 			common.Tag:       NewGormJoinCondition(common.Artifact, common.Tag),
 			common.Partition: NewGormJoinCondition(common.Artifact, common.Partition),
 		},
 		Filters: filters,
 	}
 
-	tx, err := applyListModelsInput(tx, common.Artifact, listTaggedArtifacts)
+	tx, err := applyListModelsInput(tx, common.Artifact, listTaggedInput)
 	if err != nil {
 		tx.Rollback()
 		return err
@@ -72,28 +72,30 @@ func (h *tagRepo) Create(ctx context.Context, tag models.Tag) error {
 		return h.errorTransformer.ToDataCatalogError(tx.Error)
 	}
 
-	if len(artifacts) != 0 {
-		// Soft-delete the existing tags on the artifacts that are tagged by this tag in the partition
-		oldTags := make([]models.Tag, 0, len(artifacts))
-		for _, artifact := range artifacts {
-			oldTags = append(oldTags, models.Tag{
-				TagKey:     models.TagKey{TagName: tag.TagName},
-				ArtifactID: artifact.ArtifactID,
-			})
-		}
-		tx = tx.Delete(&models.Tag{}, oldTags)
-	}
-
-	// Check if the artifact was ever previously tagged with this tag, if so undelete the record
-	var previouslyTagged *models.Artifact
-	tx.Unscoped().Find(previouslyTagged, tag)
-	if previouslyTagged != nil {
-		previouslyTagged.DeletedAt = nil
-		tx = tx.Update(previouslyTagged)
-	} else {
-		// Tag the new artifact
-		tx = tx.Create(&tag)
-	}
+	// if len(artifacts) != 0 {
+	// 	// Soft-delete the existing tags on the artifacts that are tagged by this tag in the partition
+	// 	for _, artifact := range artifacts {
+	// 		oldTag := models.Tag{
+	// 			TagKey:      models.TagKey{TagName: tag.TagName},
+	// 			ArtifactID:  artifact.ArtifactID,
+	// 			DatasetUUID: artifact.DatasetUUID,
+	// 		}
+	// 		tx = tx.Where(oldTag).Delete(&models.Tag{})
+	// 	}
+	// }
+
+	// If the artifact was ever previously tagged with this tag, we need to
+	// undelete the record because we cannot tag the artifact again since
+	// the primary keys are the same.
+	// var previouslyTagged *models.Artifact
+	// tx = tx.Unscoped().Find(previouslyTagged, tag) // unscope will ignore deletedAt
+	// if previouslyTagged != nil {
+	// 	previouslyTagged.DeletedAt = nil
+	// 	tx = tx.Update(previouslyTagged)
+	// } else {
+	// 	// Tag the new artifact
+	// 	tx = tx.Create(&tag)
+	// }
 
 	tx = tx.Commit()
 	if tx.Error != nil {

diff --git a/pkg/repositories/gormimpl/tag_test.go b/pkg/repositories/gormimpl/tag_test.go
@@ -49,12 +49,17 @@ func TestCreateTag(t *testing.T) {
 	GlobalMock := mocket.Catcher.Reset()
 	GlobalMock.Logging = true
 
+	oldArtifact := getTestArtifact()
+
 	// Only match on queries that append expected filters
 	GlobalMock.NewMock().WithQuery(
-		`SELECT * FROM "artifacts"  WHERE "artifacts"."deleted_at" IS NULL AND (("artifacts"."artifact_id" = 123))`).WithReply(getDBArtifactResponse(getTestArtifact()))
+		`SELECT * FROM "artifacts"  WHERE "artifacts"."deleted_at" IS NULL AND (("artifacts"."artifact_id" = 123))`).WithReply(getDBArtifactResponse(oldArtifact))
+
+	GlobalMock.NewMock().WithQuery(
+		`SELECT * FROM "partitions"  WHERE "partitions"."deleted_at" IS NULL AND (("artifact_id" IN (123)))`).WithReply(getDBPartitionResponse(oldArtifact))
 
 	GlobalMock.NewMock().WithQuery(
-		`SELECT * FROM "partitions"  WHERE "partitions"."deleted_at" IS NULL AND (("artifact_id" IN (123)))`).WithReply(getDBArtifactResponse(getTestArtifact()))
+		`SELECT "artifacts".* FROM "artifacts" JOIN tags ON artifacts.artifact_id = tags.artifact_id JOIN partitions ON artifacts.artifact_id = partitions.artifact_id WHERE "artifacts"."deleted_at" IS NULL AND ((partitions.key = region) AND (partitions.value = SEA) AND (artifacts.tag_name = test-tagname)) LIMIT 0 OFFSET 0`).WithReply(getDBArtifactResponse(oldArtifact))
 
 	GlobalMock.NewMock().WithQuery(
 		`INSERT  INTO "tags" ("created_at","updated_at","deleted_at","dataset_project","dataset_name","dataset_domain","dataset_version","tag_name","artifact_id","dataset_uuid") VALUES (?,?,?,?,?,?,?,?,?,?)`).WithCallback(