From 02fd8717c1c8a49abe3e4dac3756634460928ed9 Mon Sep 17 00:00:00 2001
From: Katrina Rogan <katroganGH@gmail.com>
Date: Tue, 12 Apr 2022 12:01:09 -0700
Subject: [PATCH] Allow archiving of task named entities (#402)

---
 flyteadmin/boilerplate/flyte/end2end/run-tests.py     |  4 ++--
 .../manager/impl/validation/named_entity_validator.go |  9 ++++++---
 .../impl/validation/named_entity_validator_test.go    | 11 +++++++++++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/flyteadmin/boilerplate/flyte/end2end/run-tests.py b/flyteadmin/boilerplate/flyte/end2end/run-tests.py
index 6645c60fde..6e172f70a4 100755
--- a/flyteadmin/boilerplate/flyte/end2end/run-tests.py
+++ b/flyteadmin/boilerplate/flyte/end2end/run-tests.py
@@ -36,9 +36,9 @@
         ("core.control_flow.subworkflows.parent_wf", {"a": 3}),
         ("core.control_flow.subworkflows.nested_parent_wf", {"a": 3}),
         ("core.flyte_basics.basic_workflow.my_wf", {"a": 50, "b": "hello"}),
-        # Getting a 403 for the wikipedia image
+        # TODO: enable new files and folders workflows
         # ("core.flyte_basics.files.rotate_one_workflow", {"in_image": "https://upload.wikimedia.org/wikipedia/commons/d/d2/Julia_set_%28C_%3D_0.285%2C_0.01%29.jpg"}),
-        ("core.flyte_basics.folders.download_and_rotate", {}),
+        # ("core.flyte_basics.folders.download_and_rotate", {}),
         ("core.flyte_basics.hello_world.my_wf", {}),
         ("core.flyte_basics.lp.my_wf", {"val": 4}),
         ("core.flyte_basics.lp.go_greet", {"day_of_week": "5", "number": 3, "am": True}),
diff --git a/flyteadmin/pkg/manager/impl/validation/named_entity_validator.go b/flyteadmin/pkg/manager/impl/validation/named_entity_validator.go
index 04e9603254..a23d609bb3 100644
--- a/flyteadmin/pkg/manager/impl/validation/named_entity_validator.go
+++ b/flyteadmin/pkg/manager/impl/validation/named_entity_validator.go
@@ -6,8 +6,11 @@ import (
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/admin"
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
 	"google.golang.org/grpc/codes"
+	"k8s.io/apimachinery/pkg/util/sets"
 )
 
+var archivableResourceTypes = sets.NewInt32(int32(core.ResourceType_WORKFLOW), int32(core.ResourceType_TASK))
+
 func ValidateNamedEntityGetRequest(request admin.NamedEntityGetRequest) error {
 	if err := ValidateResourceType(request.ResourceType); err != nil {
 		return err
@@ -29,11 +32,11 @@ func ValidateNamedEntityUpdateRequest(request admin.NamedEntityUpdateRequest) er
 		return shared.GetMissingArgumentError(shared.Metadata)
 	}
 
-	// Anything but the default state is only permitted for workflow resources.
+	// Only tasks and workflow resources can be modified from the default state.
 	if request.Metadata.State != admin.NamedEntityState_NAMED_ENTITY_ACTIVE &&
-		request.ResourceType != core.ResourceType_WORKFLOW {
+		!archivableResourceTypes.Has(int32(request.ResourceType)) {
 		return errors.NewFlyteAdminErrorf(codes.InvalidArgument,
-			"Only workflow name entities can have their state updated")
+			"Resource [%s] cannot have its state updated", request.ResourceType.String())
 	}
 	return nil
 }
diff --git a/flyteadmin/pkg/manager/impl/validation/named_entity_validator_test.go b/flyteadmin/pkg/manager/impl/validation/named_entity_validator_test.go
index 55e461cf4a..f69d6fa683 100644
--- a/flyteadmin/pkg/manager/impl/validation/named_entity_validator_test.go
+++ b/flyteadmin/pkg/manager/impl/validation/named_entity_validator_test.go
@@ -130,6 +130,17 @@ func TestValidateNamedEntityUpdateRequest(t *testing.T) {
 			State: admin.NamedEntityState_NAMED_ENTITY_ARCHIVED,
 		},
 	}))
+	assert.Nil(t, ValidateNamedEntityUpdateRequest(admin.NamedEntityUpdateRequest{
+		ResourceType: core.ResourceType_TASK,
+		Id: &admin.NamedEntityIdentifier{
+			Project: "project",
+			Domain:  "domain",
+			Name:    "name",
+		},
+		Metadata: &admin.NamedEntityMetadata{
+			State: admin.NamedEntityState_NAMED_ENTITY_ARCHIVED,
+		},
+	}))
 }
 
 func TestValidateNamedEntityListRequest(t *testing.T) {