Removed omit id on FirstOrCreate (flyteorg#357)

eapolinario · Mar 2, 2022 · b68ba58 · b68ba58
1 parent 1d8c521
commit b68ba58
Show file tree

Hide file tree

Showing 4 changed files with 187 additions and 6 deletions.
diff --git a/flyteadmin/pkg/repositories/gormimpl/resource_repo.go b/flyteadmin/pkg/repositories/gormimpl/resource_repo.go
@@ -59,7 +59,7 @@ func (r *ResourceRepo) CreateOrUpdate(ctx context.Context, input models.Resource
 	}
 	timer := r.metrics.GetDuration.Start()
 	var record models.Resource
-	tx := r.db.Omit("id").FirstOrCreate(&record, models.Resource{
+	tx := r.db.FirstOrCreate(&record, models.Resource{
 		Project:      input.Project,
 		Domain:       input.Domain,
 		Workflow:     input.Workflow,

diff --git a/flyteadmin/pkg/repositories/gormimpl/resource_repo_test.go b/flyteadmin/pkg/repositories/gormimpl/resource_repo_test.go
@@ -36,6 +36,48 @@ func TestCreateWorkflowAttributes(t *testing.T) {
 	assert.True(t, query.Triggered)
 }
 
+func getMockResourceResponseFromDb(expected models.Resource) map[string]interface{} {
+	metadata := make(map[string]interface{})
+	metadata["resource_type"] = expected.ResourceType
+	metadata["project"] = expected.Project
+	metadata["domain"] = expected.Domain
+	metadata["priority"] = 2
+	return metadata
+}
+
+func TestUpdateWorkflowAttributes_WithExisting(t *testing.T) {
+	resourceRepo := NewResourceRepo(GetDbForTest(t), errors.NewTestErrorTransformer(), mockScope.NewTestScope())
+	results := make([]map[string]interface{}, 0)
+	metadata := getMockResourceResponseFromDb(models.Resource{
+		ResourceType: resourceType.String(),
+		Project:      project,
+		Domain:       domain,
+		Priority:     2,
+	})
+	results = append(results, metadata)
+
+	GlobalMock := mocket.Catcher.Reset()
+	GlobalMock.Logging = true
+
+	mockSelectQuery := GlobalMock.NewMock()
+	mockSelectQuery.WithQuery(
+		`SELECT * FROM "resources" WHERE "resources"."project" = $1 AND "resources"."domain" = $2 AND "resources"."resource_type" = $3 AND "resources"."priority" = $4 ORDER BY "resources"."id" LIMIT 1`).WithReply(results)
+
+	mockSaveQuery := GlobalMock.NewMock()
+	mockSaveQuery.WithQuery(
+		`INSERT INTO "resources" ("created_at","updated_at","deleted_at","project","domain","workflow","launch_plan","resource_type","priority","attributes") VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10)`)
+
+	err := resourceRepo.CreateOrUpdate(context.Background(), models.Resource{
+		ResourceType: resourceType.String(),
+		Project:      project,
+		Domain:       domain,
+		Priority:     2,
+	})
+	assert.NoError(t, err)
+	assert.True(t, mockSelectQuery.Triggered)
+	assert.True(t, mockSaveQuery.Triggered)
+}
+
 func TestGetWorkflowAttributes(t *testing.T) {
 	resourceRepo := NewResourceRepo(GetDbForTest(t), errors.NewTestErrorTransformer(), mockScope.NewTestScope())
 	GlobalMock := mocket.Catcher.Reset()

diff --git a/flyteadmin/pkg/repositories/models/resource.go b/flyteadmin/pkg/repositories/models/resource.go
@@ -18,11 +18,11 @@ type Resource struct {
 	CreatedAt    time.Time
 	UpdatedAt    time.Time
 	DeletedAt    *time.Time `sql:"index"`
-	Project      string     `gorm:"unique_index:resource_idx" valid:"length(0|255)"`
-	Domain       string     `gorm:"unique_index:resource_idx" valid:"length(0|255)"`
-	Workflow     string     `gorm:"unique_index:resource_idx" valid:"length(0|255)"`
-	LaunchPlan   string     `gorm:"unique_index:resource_idx" valid:"length(0|255)"`
-	ResourceType string     `gorm:"unique_index:resource_idx" valid:"length(0|255)"`
+	Project      string     `gorm:"uniqueIndex:resource_idx" valid:"length(0|255)"`
+	Domain       string     `gorm:"uniqueIndex:resource_idx" valid:"length(0|255)"`
+	Workflow     string     `gorm:"uniqueIndex:resource_idx" valid:"length(0|255)"`
+	LaunchPlan   string     `gorm:"uniqueIndex:resource_idx" valid:"length(0|255)"`
+	ResourceType string     `gorm:"uniqueIndex:resource_idx" valid:"length(0|255)"`
 	Priority     ResourcePriority
 	// Serialized flyteidl.admin.MatchingAttributes.
 	Attributes []byte

diff --git a/flyteadmin/tests/attributes_test.go b/flyteadmin/tests/attributes_test.go
@@ -1,3 +1,4 @@
+//go:build integration
 // +build integration
 
 package tests
@@ -26,6 +27,108 @@ var matchingAttributes = &admin.MatchingAttributes{
 	},
 }
 
+func TestUpdateClusterResourceAttributes(t *testing.T) {
+	matchingAttributes = &admin.MatchingAttributes{
+		Target: &admin.MatchingAttributes_ClusterResourceAttributes{
+			ClusterResourceAttributes: &admin.ClusterResourceAttributes{
+				Attributes: map[string]string{
+					"key": "value",
+				},
+			},
+		},
+	}
+
+	ctx := context.Background()
+	client, conn := GetTestAdminServiceClient()
+	defer conn.Close()
+	db, err := repositories.GetDB(ctx, getDbConfig(), getLoggerConfig())
+	assert.Nil(t, err)
+	truncateTableForTesting(db, "resources")
+	sqlDB, err := db.DB()
+	assert.Nil(t, err)
+	err = sqlDB.Close()
+	assert.Nil(t, err)
+
+	req := admin.ProjectDomainAttributesUpdateRequest{
+		Attributes: &admin.ProjectDomainAttributes{
+			Project:            "admintests",
+			Domain:             "development",
+			MatchingAttributes: matchingAttributes,
+		},
+	}
+
+	_, err = client.UpdateProjectDomainAttributes(ctx, &req)
+	fmt.Println(err)
+	assert.Nil(t, err)
+
+
+	listResp, err := client.ListMatchableAttributes(ctx, &admin.ListMatchableAttributesRequest{
+		ResourceType: admin.MatchableResource_CLUSTER_RESOURCE,
+	})
+	fmt.Println(err)
+	assert.Nil(t, err)
+	assert.Equal(t, 1, len(listResp.GetConfigurations()))
+
+	response, err := client.GetProjectDomainAttributes(ctx, &admin.ProjectDomainAttributesGetRequest{
+		Project:      "admintests",
+		Domain:       "development",
+		ResourceType: admin.MatchableResource_CLUSTER_RESOURCE,
+	})
+	fmt.Println(err)
+	assert.Nil(t, err)
+	assert.True(t, proto.Equal(&admin.ProjectDomainAttributesGetResponse{
+		Attributes: &admin.ProjectDomainAttributes{
+			Project:            "admintests",
+			Domain:             "development",
+			MatchingAttributes: matchingAttributes,
+		},
+	}, response))
+
+	var updatedMatchingAttributes = &admin.MatchingAttributes{
+		Target: &admin.MatchingAttributes_ClusterResourceAttributes{
+			ClusterResourceAttributes: &admin.ClusterResourceAttributes{
+				Attributes: map[string]string{
+					"key": "value2",
+				},
+			},
+		},
+	}
+	req = admin.ProjectDomainAttributesUpdateRequest{
+		Attributes: &admin.ProjectDomainAttributes{
+			Project:            "admintests",
+			Domain:             "development",
+			MatchingAttributes: updatedMatchingAttributes,
+		},
+	}
+
+	_, err = client.UpdateProjectDomainAttributes(ctx, &req)
+	fmt.Println(err)
+	assert.Nil(t, err)
+
+	listResp, err = client.ListMatchableAttributes(ctx, &admin.ListMatchableAttributesRequest{
+		ResourceType: admin.MatchableResource_CLUSTER_RESOURCE,
+	})
+	fmt.Println(err)
+	assert.Nil(t, err)
+	assert.Equal(t, 1, len(listResp.GetConfigurations()))
+
+	response, err = client.GetProjectDomainAttributes(ctx, &admin.ProjectDomainAttributesGetRequest{
+		Project:      "admintests",
+		Domain:       "development",
+		ResourceType: admin.MatchableResource_CLUSTER_RESOURCE,
+	})
+	fmt.Println(err)
+	assert.Nil(t, err)
+	assert.True(t, proto.Equal(&admin.ProjectDomainAttributesGetResponse{
+		Attributes: &admin.ProjectDomainAttributes{
+			Project:            "admintests",
+			Domain:             "development",
+			MatchingAttributes: updatedMatchingAttributes,
+		},
+	}, response))
+
+}
+
 func TestUpdateProjectDomainAttributes(t *testing.T) {
 	ctx := context.Background()
 	client, conn := GetTestAdminServiceClient()
@@ -65,6 +168,42 @@ func TestUpdateProjectDomainAttributes(t *testing.T) {
 		},
 	}, response))
 
+	var updatedMatchingAttributes = &admin.MatchingAttributes{
+		Target: &admin.MatchingAttributes_TaskResourceAttributes{
+			TaskResourceAttributes: &admin.TaskResourceAttributes{
+				Defaults: &admin.TaskResourceSpec{
+					Cpu: "1",
+				},
+			},
+		},
+	}
+	req = admin.ProjectDomainAttributesUpdateRequest{
+		Attributes: &admin.ProjectDomainAttributes{
+			Project:            "admintests",
+			Domain:             "development",
+			MatchingAttributes: updatedMatchingAttributes,
+		},
+	}
+
+	_, err = client.UpdateProjectDomainAttributes(ctx, &req)
+	fmt.Println(err)
+	assert.Nil(t, err)
+
+	response, err = client.GetProjectDomainAttributes(ctx, &admin.ProjectDomainAttributesGetRequest{
+		Project:      "admintests",
+		Domain:       "development",
+		ResourceType: admin.MatchableResource_TASK_RESOURCE,
+	})
+	fmt.Println(err)
+	assert.Nil(t, err)
+	assert.True(t, proto.Equal(&admin.ProjectDomainAttributesGetResponse{
+		Attributes: &admin.ProjectDomainAttributes{
+			Project:            "admintests",
+			Domain:             "development",
+			MatchingAttributes: updatedMatchingAttributes,
+		},
+	}, response))
+
 	workflowResponse, err := client.GetWorkflowAttributes(ctx, &admin.WorkflowAttributesGetRequest{
 		Project:      "admintests",
 		Domain:       "development",