Include Tags in Artifact Responses (#24)

* Include tag in list artifacts

* Add tests

* Attach tags and partitions to getArtifact call

* Order by for getByTag

* Fix tests

datacatalog/pkg/manager/impl/artifact_manager_test.go

@@ -49,16 +49,16 @@ func getTestStringLiteral() *core.Literal {
 }
 
 func getTestArtifact() *datacatalog.Artifact {
-
+	datasetID := &datacatalog.DatasetID{
+		Project: "test-project",
+		Domain:  "test-domain",
+		Name:    "test-name",
+		Version: "test-version",
+		UUID:    "test-uuid",
+	}
 	return &datacatalog.Artifact{
-		Id: "test-id",
-		Dataset: &datacatalog.DatasetID{
-			Project: "test-project",
-			Domain:  "test-domain",
-			Name:    "test-name",
-			Version: "test-version",
-			UUID:    "test-uuid",
-		},
+		Id:      "test-id",
+		Dataset: datasetID,
 		Metadata: &datacatalog.Metadata{
 			KeyMap: map[string]string{"key1": "value1"},
 		},
@@ -72,6 +72,9 @@ func getTestArtifact() *datacatalog.Artifact {
 			{Key: "key1", Value: "value1"},
 			{Key: "key2", Value: "value2"},
 		},
+		Tags: []*datacatalog.Tag{
+			{Name: "test-tag", Dataset: datasetID, ArtifactId: "test-id"},
+		},
 	}
 }
 
@@ -128,6 +131,9 @@ func getExpectedArtifactModel(ctx context.Context, t *testing.T, datastore *stor
 			{Key: "key1", Value: "value1"},
 			{Key: "key2", Value: "value2"},
 		},
+		Tags: []models.Tag{
+			{TagKey: models.TagKey{TagName: "test-tag"}, DatasetUUID: expectedDataset.UUID, ArtifactID: artifact.Id},
+		},
 	}
 }
 

datacatalog/pkg/repositories/gormimpl/artifact.go

@@ -53,9 +53,12 @@ func (h *artifactRepo) Get(ctx context.Context, in models.ArtifactKey) (models.A
 	defer timer.Stop()
 
 	var artifact models.Artifact
-	result := h.db.Preload("ArtifactData").Preload("Partitions").First(&artifact, &models.Artifact{
-		ArtifactKey: in,
-	})
+	result := h.db.Preload("ArtifactData").Preload("Partitions").Preload("Tags").
+		Order("artifacts.created_at DESC").
+		First(
+			&artifact,
+			&models.Artifact{ArtifactKey: in},
+		)
 
 	if result.Error != nil {
 		return models.Artifact{}, h.errorTransformer.ToDataCatalogError(result.Error)
@@ -99,10 +102,9 @@ func (h *artifactRepo) List(ctx context.Context, datasetKey models.DatasetKey, i
 		return []models.Artifact{}, h.errorTransformer.ToDataCatalogError(tx.Error)
 	}
 
-	tx = tx.Preload("ArtifactData").Preload("Partitions").Find(&artifacts)
+	tx = tx.Preload("ArtifactData").Preload("Partitions").Preload("Tags").Find(&artifacts)
 	if tx.Error != nil {
 		return []models.Artifact{}, h.errorTransformer.ToDataCatalogError(tx.Error)
 	}
-
 	return artifacts, nil
 }

datacatalog/pkg/repositories/gormimpl/artifact_test.go

@@ -90,6 +90,21 @@ func getDBPartitionResponse(artifact models.Artifact) []map[string]interface{} {
 	return expectedPartitionResponse
 }
 
+// Raw db response to return on raw queries for tags
+func getDBTagResponse(artifact models.Artifact) []map[string]interface{} {
+	expectedTagResponse := make([]map[string]interface{}, 0)
+	sampleTag := make(map[string]interface{})
+	sampleTag["tag_name"] = "test-tag"
+	sampleTag["artifact_id"] = artifact.ArtifactID
+	sampleTag["dataset_uuid"] = "test-uuid"
+	sampleTag["dataset_project"] = artifact.DatasetProject
+	sampleTag["dataset_domain"] = artifact.DatasetDomain
+	sampleTag["dataset_name"] = artifact.DatasetName
+	sampleTag["dataset_version"] = artifact.DatasetVersion
+	expectedTagResponse = append(expectedTagResponse, sampleTag)
+	return expectedTagResponse
+}
+
 func TestCreateArtifact(t *testing.T) {
 	artifact := getTestArtifact()
 
@@ -153,17 +168,20 @@ func TestGetArtifact(t *testing.T) {
 	expectedArtifactDataResponse := getDBArtifactDataResponse(artifact)
 	expectedArtifactResponse := getDBArtifactResponse(artifact)
 	expectedPartitionResponse := getDBPartitionResponse(artifact)
+	expectedTagResponse := getDBTagResponse(artifact)
 
 	GlobalMock := mocket.Catcher.Reset()
 	GlobalMock.Logging = true
 
 	// Only match on queries that append expected filters
 	GlobalMock.NewMock().WithQuery(
-		`SELECT * FROM "artifacts"  WHERE "artifacts"."deleted_at" IS NULL AND (("artifacts"."dataset_project" = testProject) AND ("artifacts"."dataset_name" = testName) AND ("artifacts"."dataset_domain" = testDomain) AND ("artifacts"."dataset_version" = testVersion) AND ("artifacts"."artifact_id" = 123)) ORDER BY "artifacts"."dataset_project" ASC LIMIT 1`).WithReply(expectedArtifactResponse)
+		`SELECT * FROM "artifacts"  WHERE "artifacts"."deleted_at" IS NULL AND (("artifacts"."dataset_project" = testProject) AND ("artifacts"."dataset_name" = testName) AND ("artifacts"."dataset_domain" = testDomain) AND ("artifacts"."dataset_version" = testVersion) AND ("artifacts"."artifact_id" = 123)) ORDER BY artifacts.created_at DESC,"artifacts"."dataset_project" ASC LIMIT 1`).WithReply(expectedArtifactResponse)
 	GlobalMock.NewMock().WithQuery(
 		`SELECT * FROM "artifact_data"  WHERE "artifact_data"."deleted_at" IS NULL AND ((("dataset_project","dataset_name","dataset_domain","dataset_version","artifact_id") IN ((testProject,testName,testDomain,testVersion,123)))) ORDER BY "artifact_data"."dataset_project" ASC`).WithReply(expectedArtifactDataResponse)
 	GlobalMock.NewMock().WithQuery(
 		`SELECT * FROM "partitions"  WHERE "partitions"."deleted_at" IS NULL AND (("artifact_id" IN (123))) ORDER BY "partitions"."dataset_uuid" ASC`).WithReply(expectedPartitionResponse)
+	GlobalMock.NewMock().WithQuery(
+		`SELECT * FROM "tags"  WHERE "tags"."deleted_at" IS NULL AND ((("artifact_id","dataset_uuid") IN ((123,test-uuid)))) ORDER BY "tags"."dataset_project" ASC`).WithReply(expectedTagResponse)
 	getInput := models.ArtifactKey{
 		DatasetProject: artifact.DatasetProject,
 		DatasetDomain:  artifact.DatasetDomain,
@@ -183,6 +201,7 @@ func TestGetArtifact(t *testing.T) {
 
 	assert.Equal(t, 1, len(response.ArtifactData))
 	assert.Equal(t, 1, len(response.Partitions))
+	assert.EqualValues(t, 1, len(response.Tags))
 }
 
 func TestGetArtifactDoesNotExist(t *testing.T) {
@@ -238,13 +257,15 @@ func TestListArtifactsWithPartition(t *testing.T) {
 	expectedArtifactDataResponse := getDBArtifactDataResponse(artifact)
 	expectedArtifactResponse := getDBArtifactResponse(artifact)
 	expectedPartitionResponse := getDBPartitionResponse(artifact)
-
+	expectedTagResponse := getDBTagResponse(artifact)
 	GlobalMock.NewMock().WithQuery(
 		`SELECT "artifacts".* FROM "artifacts" JOIN partitions partitions0 ON artifacts.artifact_id = partitions0.artifact_id WHERE "artifacts"."deleted_at" IS NULL AND ((partitions0.key = val1) AND (partitions0.val = val2) AND (artifacts.dataset_uuid = test-uuid)) ORDER BY artifacts.created_at desc LIMIT 10 OFFSET 10`).WithReply(expectedArtifactResponse)
 	GlobalMock.NewMock().WithQuery(
 		`SELECT * FROM "artifact_data"  WHERE "artifact_data"."deleted_at" IS NULL AND ((("dataset_project","dataset_name","dataset_domain","dataset_version","artifact_id") IN ((testProject,testName,testDomain,testVersion,123))))`).WithReply(expectedArtifactDataResponse)
 	GlobalMock.NewMock().WithQuery(
 		`SELECT * FROM "partitions"  WHERE "partitions"."deleted_at" IS NULL AND (("artifact_id" IN (123)))`).WithReply(expectedPartitionResponse)
+	GlobalMock.NewMock().WithQuery(
+		`SELECT * FROM "tags"  WHERE "tags"."deleted_at" IS NULL AND ((("artifact_id","dataset_uuid") IN ((123,test-uuid))))`).WithReply(expectedTagResponse)
 
 	artifactRepo := NewArtifactRepo(utils.GetDbForTest(t), errors.NewPostgresErrorTransformer(), promutils.NewTestScope())
 	listInput := models.ListModelsInput{
@@ -267,6 +288,7 @@ func TestListArtifactsWithPartition(t *testing.T) {
 	assert.Equal(t, artifacts[0].ArtifactID, artifact.ArtifactID)
 	assert.Len(t, artifacts[0].ArtifactData, 1)
 	assert.Len(t, artifacts[0].Partitions, 1)
+	assert.Len(t, artifacts[0].Tags, 1)
 }
 
 func TestListArtifactsNoPartitions(t *testing.T) {

datacatalog/pkg/repositories/gormimpl/tag.go

@@ -42,9 +42,12 @@ func (h *tagRepo) Get(ctx context.Context, in models.TagKey) (models.Tag, error)
 	defer timer.Stop()
 
 	var tag models.Tag
-	result := h.db.Preload("Artifact").Preload("Artifact.ArtifactData").Find(&tag, &models.Tag{
-		TagKey: in,
-	})
+	result := h.db.Preload("Artifact").Preload("Artifact.ArtifactData").
+		Preload("Artifact.Partitions").Preload("Artifact.Tags").
+		Order("tags.created_at DESC").
+		First(&tag, &models.Tag{
+			TagKey: in,
+		})
 
 	if result.Error != nil {
 		return models.Tag{}, h.errorTransformer.ToDataCatalogError(result.Error)

datacatalog/pkg/repositories/gormimpl/tag_test.go

@@ -66,51 +66,20 @@ func TestCreateTag(t *testing.T) {
 func TestGetTag(t *testing.T) {
 	artifact := getTestArtifact()
 
-	expectedArtifactDataResponse := make([]map[string]interface{}, 0)
-	sampleArtifactData := make(map[string]interface{})
-	sampleArtifactData["dataset_project"] = artifact.DatasetProject
-	sampleArtifactData["dataset_domain"] = artifact.DatasetDomain
-	sampleArtifactData["dataset_name"] = artifact.DatasetName
-	sampleArtifactData["dataset_version"] = artifact.DatasetVersion
-	sampleArtifactData["artifact_id"] = artifact.ArtifactID
-	sampleArtifactData["name"] = "test-dataloc-name"
-	sampleArtifactData["location"] = "test-dataloc-location"
-	sampleArtifactData["dataset_uuid"] = artifact.DatasetUUID
-
-	expectedArtifactDataResponse = append(expectedArtifactDataResponse, sampleArtifactData)
-
-	expectedArtifactResponse := make([]map[string]interface{}, 0)
-	sampleArtifact := make(map[string]interface{})
-	sampleArtifact["dataset_project"] = artifact.DatasetProject
-	sampleArtifact["dataset_domain"] = artifact.DatasetDomain
-	sampleArtifact["dataset_name"] = artifact.DatasetName
-	sampleArtifact["dataset_version"] = artifact.DatasetVersion
-	sampleArtifact["artifact_id"] = artifact.ArtifactID
-	sampleArtifact["dataset_uuid"] = artifact.DatasetUUID
-	expectedArtifactResponse = append(expectedArtifactResponse, sampleArtifact)
-
-	expectedTagResponse := make([]map[string]interface{}, 0)
-	sampleTag := make(map[string]interface{})
-	sampleTag["dataset_project"] = artifact.DatasetProject
-	sampleTag["dataset_domain"] = artifact.DatasetDomain
-	sampleTag["dataset_name"] = artifact.DatasetName
-	sampleTag["dataset_version"] = artifact.DatasetVersion
-	sampleTag["artifact_id"] = artifact.ArtifactID
-	sampleTag["name"] = "test-tag"
-	sampleTag["dataset_uuid"] = artifact.DatasetUUID
-	expectedTagResponse = append(expectedTagResponse, sampleTag)
-
 	GlobalMock := mocket.Catcher.Reset()
 	GlobalMock.Logging = true
 
 	// Only match on queries that append expected filters
 	GlobalMock.NewMock().WithQuery(
-		`SELECT * FROM "tags"  WHERE "tags"."deleted_at" IS NULL AND (("tags"."dataset_project" = testProject) AND ("tags"."dataset_name" = testName) AND ("tags"."dataset_domain" = testDomain) AND ("tags"."dataset_version" = testVersion) AND ("tags"."tag_name" = test-tag))`).WithReply(expectedTagResponse)
+		`SELECT * FROM "tags"  WHERE "tags"."deleted_at" IS NULL AND (("tags"."dataset_project" = testProject) AND ("tags"."dataset_name" = testName) AND ("tags"."dataset_domain" = testDomain) AND ("tags"."dataset_version" = testVersion) AND ("tags"."tag_name" = test-tag)) ORDER BY tags.created_at DESC,"tags"."dataset_project" ASC LIMIT 1`).WithReply(getDBTagResponse(artifact))
 	GlobalMock.NewMock().WithQuery(
-		`SELECT * FROM "artifacts"  WHERE "artifacts"."deleted_at" IS NULL AND ((("dataset_project","dataset_name","dataset_domain","dataset_version","artifact_id") IN ((testProject,testName,testDomain,testVersion,123))))`).WithReply(expectedArtifactResponse)
+		`SELECT * FROM "artifacts"  WHERE "artifacts"."deleted_at" IS NULL AND ((("dataset_project","dataset_name","dataset_domain","dataset_version","artifact_id") IN ((testProject,testName,testDomain,testVersion,123))))`).WithReply(getDBArtifactResponse(artifact))
 	GlobalMock.NewMock().WithQuery(
-		`SELECT * FROM "artifact_data"  WHERE "artifact_data"."deleted_at" IS NULL AND ((("dataset_project","dataset_name","dataset_domain","dataset_version","artifact_id") IN ((testProject,testName,testDomain,testVersion,123))))`).WithReply(expectedArtifactDataResponse)
-
+		`SELECT * FROM "artifact_data"  WHERE "artifact_data"."deleted_at" IS NULL AND ((("dataset_project","dataset_name","dataset_domain","dataset_version","artifact_id") IN ((testProject,testName,testDomain,testVersion,123))))`).WithReply(getDBArtifactDataResponse(artifact))
+	GlobalMock.NewMock().WithQuery(
+		`SELECT * FROM "partitions"  WHERE "partitions"."deleted_at" IS NULL AND (("artifact_id" IN (123)))`).WithReply(getDBPartitionResponse(artifact))
+	GlobalMock.NewMock().WithQuery(
+		`SELECT * FROM "tags"  WHERE "tags"."deleted_at" IS NULL AND ((("artifact_id","dataset_uuid") IN ((123,test-uuid))))`).WithReply(getDBTagResponse(artifact))
 	getInput := models.TagKey{
 		DatasetProject: artifact.DatasetProject,
 		DatasetDomain:  artifact.DatasetDomain,
@@ -125,6 +94,8 @@ func TestGetTag(t *testing.T) {
 	assert.Equal(t, artifact.ArtifactID, response.ArtifactID)
 	assert.Equal(t, artifact.ArtifactID, response.Artifact.ArtifactID)
 	assert.Len(t, response.Artifact.ArtifactData, 1)
+	assert.Len(t, response.Artifact.Partitions, 1)
+	assert.Len(t, response.Artifact.Tags, 1)
 }
 
 func TestTagAlreadyExists(t *testing.T) {

datacatalog/pkg/repositories/models/artifact.go

@@ -15,6 +15,7 @@ type Artifact struct {
 	Dataset            Dataset        `gorm:"association_autocreate:false"`
 	ArtifactData       []ArtifactData `gorm:"association_foreignkey:DatasetProject,DatasetName,DatasetDomain,DatasetVersion,ArtifactID;foreignkey:DatasetProject,DatasetName,DatasetDomain,DatasetVersion,ArtifactID"`
 	Partitions         []Partition    `gorm:"association_foreignkey:ArtifactID;foreignkey:ArtifactID"`
+	Tags               []Tag          `gorm:"association_foreignkey:ArtifactID,DatasetUUID;foreignkey:ArtifactID,DatasetUUID"`
 	SerializedMetadata []byte
 }
 

datacatalog/pkg/repositories/transformers/artifact.go

@@ -38,6 +38,14 @@ func CreateArtifactModel(request datacatalog.CreateArtifactRequest, artifactData
 }
 
 func FromArtifactModel(artifact models.Artifact) (datacatalog.Artifact, error) {
+	datasetID := datacatalog.DatasetID{
+		Project: artifact.DatasetProject,
+		Domain:  artifact.DatasetDomain,
+		Name:    artifact.DatasetName,
+		Version: artifact.DatasetVersion,
+		UUID:    artifact.DatasetUUID,
+	}
+
 	metadata, err := unmarshalMetadata(artifact.SerializedMetadata)
 	if err != nil {
 		return datacatalog.Artifact{}, err
@@ -51,17 +59,16 @@ func FromArtifactModel(artifact models.Artifact) (datacatalog.Artifact, error) {
 		}
 	}
 
+	tags := make([]*datacatalog.Tag, len(artifact.Tags))
+	for i, tag := range artifact.Tags {
+		tags[i] = FromTagModel(datasetID, tag)
+	}
 	return datacatalog.Artifact{
-		Id: artifact.ArtifactID,
-		Dataset: &datacatalog.DatasetID{
-			Project: artifact.DatasetProject,
-			Domain:  artifact.DatasetDomain,
-			Name:    artifact.DatasetName,
-			Version: artifact.DatasetVersion,
-			UUID:    artifact.DatasetUUID,
-		},
+		Id:         artifact.ArtifactID,
+		Dataset:    &datasetID,
 		Metadata:   metadata,
 		Partitions: partitions,
+		Tags:       tags,
 	}, nil
 }
 

datacatalog/pkg/repositories/transformers/artifact_test.go

@@ -32,6 +32,12 @@ func getTestPartitions() []models.Partition {
 	}
 }
 
+func getTestTags() []models.Tag {
+	return []models.Tag{
+		{TagKey: models.TagKey{TagName: "test"}},
+	}
+}
+
 func getDatasetModel() models.Dataset {
 	return models.Dataset{
 		DatasetKey: models.DatasetKey{
@@ -104,10 +110,8 @@ func TestFromArtifactModel(t *testing.T) {
 			ArtifactID:     "id1",
 		},
 		SerializedMetadata: []byte{},
-		Partitions: []models.Partition{
-			{DatasetUUID: "test-uuid", Key: "key1", Value: "value1"},
-			{DatasetUUID: "test-uuid", Key: "key2", Value: "value2"},
-		},
+		Partitions:         getTestPartitions(),
+		Tags:               getTestTags(),
 	}
 
 	actual, err := FromArtifactModel(artifactModel)
@@ -119,7 +123,13 @@ func TestFromArtifactModel(t *testing.T) {
 	assert.Equal(t, artifactModel.DatasetVersion, actual.Dataset.Version)
 
 	assert.Len(t, actual.Partitions, 2)
-	assert.EqualValues(t, artifactModel.Partitions, getTestPartitions())
+	assert.EqualValues(t, artifactModel.Partitions[0].Key, actual.Partitions[0].Key)
+	assert.EqualValues(t, artifactModel.Partitions[0].Value, actual.Partitions[0].Value)
+	assert.EqualValues(t, artifactModel.Partitions[1].Value, actual.Partitions[1].Value)
+	assert.EqualValues(t, artifactModel.Partitions[1].Value, actual.Partitions[1].Value)
+
+	assert.Len(t, actual.Tags, 1)
+	assert.EqualValues(t, artifactModel.Tags[0].TagName, actual.Tags[0].Name)
 }
 
 func TestToArtifactKey(t *testing.T) {

datacatalog/pkg/repositories/transformers/tag.go

@@ -14,3 +14,11 @@ func ToTagKey(datasetID datacatalog.DatasetID, tagName string) models.TagKey {
 		TagName:        tagName,
 	}
 }
+
+func FromTagModel(datasetID datacatalog.DatasetID, tag models.Tag) *datacatalog.Tag {
+	return &datacatalog.Tag{
+		Name:       tag.TagName,
+		ArtifactId: tag.ArtifactID,
+		Dataset:    &datasetID,
+	}
+}

datacatalog/pkg/repositories/transformers/tag_test.go

@@ -3,6 +3,7 @@ package transformers
 import (
 	"testing"
 
+	"github.com/lyft/datacatalog/pkg/repositories/models"
 	datacatalog "github.com/lyft/datacatalog/protos/gen"
 	"github.com/stretchr/testify/assert"
 )
@@ -25,3 +26,29 @@ func TestToTagKey(t *testing.T) {
 	assert.Equal(t, datasetID.Name, tagKey.DatasetName)
 	assert.Equal(t, datasetID.Version, tagKey.DatasetVersion)
 }
+
+func TestFromTagModel(t *testing.T) {
+	datasetID := datacatalog.DatasetID{
+		Project: "testProj",
+		Domain:  "testDomain",
+		Name:    "testName",
+		Version: "testVersion",
+		UUID:    "test-uuid",
+	}
+
+	tagModel := models.Tag{
+		TagKey: models.TagKey{
+			TagName: "test-tag",
+		},
+		DatasetUUID: "dataset-uuid",
+	}
+
+	tag := FromTagModel(datasetID, tagModel)
+
+	assert.Equal(t, tag.Name, tagModel.TagName)
+	assert.Equal(t, datasetID.Project, tag.Dataset.Project)
+	assert.Equal(t, datasetID.Domain, tag.Dataset.Domain)
+	assert.Equal(t, datasetID.Name, tag.Dataset.Name)
+	assert.Equal(t, datasetID.Version, tag.Dataset.Version)
+	assert.Equal(t, datasetID.UUID, tag.Dataset.UUID)
+}