Make DatasetID optional for Get Artifact (#26)

* Allow getting artifact with just the artifactID

* Validate dataset only if it was passed

datacatalog/pkg/manager/impl/artifact_manager.go

@@ -128,20 +128,20 @@ func (m *artifactManager) GetArtifact(ctx context.Context, request datacatalog.G
 	timer := m.systemMetrics.getResponseTime.Start(ctx)
 	defer timer.Stop()
 
-	datasetID := request.Dataset
 	err := validators.ValidateGetArtifactRequest(request)
 	if err != nil {
 		logger.Warningf(ctx, "Invalid get artifact request %v, err: %v", request, err)
 		m.systemMetrics.validationErrorCounter.Inc(ctx)
 		return nil, err
 	}
 
-	ctx = contextutils.WithProjectDomain(ctx, datasetID.Project, datasetID.Domain)
+	datasetID := request.Dataset
+
 	var artifactModel models.Artifact
 	switch request.QueryHandle.(type) {
 	case *datacatalog.GetArtifactRequest_ArtifactId:
 		logger.Debugf(ctx, "Get artifact by id %v", request.GetArtifactId())
-		artifactKey := transformers.ToArtifactKey(*datasetID, request.GetArtifactId())
+		artifactKey := transformers.ToArtifactKey(datasetID, request.GetArtifactId())
 		artifactModel, err = m.repo.ArtifactRepo().Get(ctx, artifactKey)
 
 		if err != nil {

datacatalog/pkg/manager/impl/tag_manager.go

@@ -46,10 +46,10 @@ func (m *tagManager) AddTag(ctx context.Context, request datacatalog.AddTagReque
 	}
 
 	// verify the artifact and dataset exists before adding a tag to it
-	datasetID := *request.Tag.Dataset
+	datasetID := request.Tag.Dataset
 	ctx = contextutils.WithProjectDomain(ctx, datasetID.Project, datasetID.Domain)
 
-	datasetKey := transformers.FromDatasetID(datasetID)
+	datasetKey := transformers.FromDatasetID(*datasetID)
 	dataset, err := m.repo.DatasetRepo().Get(ctx, datasetKey)
 	if err != nil {
 		m.systemMetrics.addTagFailureCounter.Inc(ctx)
@@ -63,7 +63,7 @@ func (m *tagManager) AddTag(ctx context.Context, request datacatalog.AddTagReque
 		return nil, err
 	}
 
-	tagKey := transformers.ToTagKey(datasetID, request.Tag.Name)
+	tagKey := transformers.ToTagKey(*datasetID, request.Tag.Name)
 	err = m.repo.TagRepo().Create(ctx, models.Tag{
 		TagKey:      tagKey,
 		ArtifactID:  request.Tag.ArtifactId,

datacatalog/pkg/manager/impl/validators/artifact_validator.go

@@ -14,20 +14,27 @@ const (
 )
 
 func ValidateGetArtifactRequest(request datacatalog.GetArtifactRequest) error {
-	if err := ValidateDatasetID(request.Dataset); err != nil {
-		return err
-	}
-
 	if request.QueryHandle == nil {
 		return NewMissingArgumentError(fmt.Sprintf("one of %s/%s", artifactID, tagName))
 	}
 
 	switch request.QueryHandle.(type) {
 	case *datacatalog.GetArtifactRequest_ArtifactId:
+		if request.Dataset != nil {
+			err := ValidateDatasetID(request.Dataset)
+			if err != nil {
+				return err
+			}
+		}
+
 		if err := ValidateEmptyStringField(request.GetArtifactId(), artifactID); err != nil {
 			return err
 		}
 	case *datacatalog.GetArtifactRequest_TagName:
+		if err := ValidateDatasetID(request.Dataset); err != nil {
+			return err
+		}
+
 		if err := ValidateEmptyStringField(request.GetTagName(), tagName); err != nil {
 			return err
 		}

datacatalog/pkg/repositories/gormimpl/artifact_test.go

@@ -204,6 +204,36 @@ func TestGetArtifact(t *testing.T) {
 	assert.EqualValues(t, 1, len(response.Tags))
 }
 
+func TestGetArtifactByID(t *testing.T) {
+	artifact := getTestArtifact()
+
+	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"."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{
+		ArtifactID: artifact.ArtifactID,
+	}
+
+	artifactRepo := NewArtifactRepo(utils.GetDbForTest(t), errors.NewPostgresErrorTransformer(), promutils.NewTestScope())
+	response, err := artifactRepo.Get(context.Background(), getInput)
+	assert.NoError(t, err)
+	assert.Equal(t, artifact.ArtifactID, response.ArtifactID)
+}
+
 func TestGetArtifactDoesNotExist(t *testing.T) {
 	artifact := getTestArtifact()
 

datacatalog/pkg/repositories/transformers/artifact.go

@@ -86,12 +86,17 @@ func FromArtifactModels(artifacts []models.Artifact) ([]*datacatalog.Artifact, e
 	return retArtifacts, nil
 }
 
-func ToArtifactKey(datasetID datacatalog.DatasetID, artifactID string) models.ArtifactKey {
-	return models.ArtifactKey{
-		DatasetProject: datasetID.Project,
-		DatasetDomain:  datasetID.Domain,
-		DatasetName:    datasetID.Name,
-		DatasetVersion: datasetID.Version,
-		ArtifactID:     artifactID,
+// Transforms datasetID and artifact combination into an ArtifactKey
+// The DatasetID is optional since artifactIDs are unique per Artifact
+func ToArtifactKey(datasetID *datacatalog.DatasetID, artifactID string) models.ArtifactKey {
+	artifactKey := models.ArtifactKey{
+		ArtifactID: artifactID,
 	}
+	if datasetID != nil {
+		artifactKey.DatasetProject = datasetID.Project
+		artifactKey.DatasetDomain = datasetID.Domain
+		artifactKey.DatasetName = datasetID.Name
+		artifactKey.DatasetVersion = datasetID.Version
+	}
+	return artifactKey
 }

datacatalog/pkg/repositories/transformers/artifact_test.go

@@ -133,10 +133,19 @@ func TestFromArtifactModel(t *testing.T) {
 }
 
 func TestToArtifactKey(t *testing.T) {
-	artifactKey := ToArtifactKey(datasetID, "artifactID-1")
+	artifactKey := ToArtifactKey(&datasetID, "artifactID-1")
 	assert.Equal(t, datasetID.Project, artifactKey.DatasetProject)
 	assert.Equal(t, datasetID.Domain, artifactKey.DatasetDomain)
 	assert.Equal(t, datasetID.Name, artifactKey.DatasetName)
 	assert.Equal(t, datasetID.Version, artifactKey.DatasetVersion)
 	assert.Equal(t, artifactKey.ArtifactID, "artifactID-1")
 }
+
+func TestToArtifactKeyNoDataset(t *testing.T) {
+	artifactKey := ToArtifactKey(nil, "artifactID-1")
+	assert.Equal(t, artifactKey.DatasetProject, "")
+	assert.Equal(t, artifactKey.DatasetDomain, "")
+	assert.Equal(t, artifactKey.DatasetName, "")
+	assert.Equal(t, artifactKey.DatasetVersion, "")
+	assert.Equal(t, artifactKey.ArtifactID, "artifactID-1")
+}