Add parent task action label to TaskActions

actions/k8s/client.go

@@ -56,6 +56,7 @@ type ActionUpdate struct {
 
 const (
 	labelTerminalStatusRecorded = "flyte.org/terminal-status-recorded"
+	labelParentTaskAction       = "flyte.org/parent-task-action"
 )
 
 // ActionsClient handles all etcd/K8s TaskAction CR operations for the Actions service.
@@ -1007,7 +1008,8 @@ func (c *ActionsClient) newTaskActionCR(actionID *common.ActionIdentifier, actio
 }
 
 // setParentOwnership resolves the parent TaskAction and wires an OwnerReference
-// on the child so K8s cascades deletion. Returns the parent CR.
+// and a parent-name label on the child so K8s cascades deletion and children can
+// be selected by their immediate parent. Returns the parent CR.
 func (c *ActionsClient) setParentOwnership(ctx context.Context, child *executorv1.TaskAction, runID *common.RunIdentifier, parentActionName string) (*executorv1.TaskAction, error) {
 	parentID := &common.ActionIdentifier{
 		Run:  runID,
@@ -1021,6 +1023,10 @@ func (c *ActionsClient) setParentOwnership(ctx context.Context, child *executorv
 	}
 
 	blockOwnerDeletion := true
+	if child.Labels == nil {
+		child.Labels = make(map[string]string)
+	}
+	child.Labels[labelParentTaskAction] = parent.Name
 	child.OwnerReferences = []metav1.OwnerReference{
 		{
 			APIVersion:         "flyte.org/v1",

actions/k8s/client_task_test.go

@@ -68,6 +68,44 @@ func newChildTaskAction() *actions.Action {
 	}
 }
 
+func TestEnqueue_ParentTaskActionLabel(t *testing.T) {
+	for _, actionType := range []string{"task", "condition"} {
+		t.Run(actionType, func(t *testing.T) {
+			ctx := context.Background()
+			c := newTaskClient(t, nil)
+			newAction := func() *actions.Action {
+				if actionType == "condition" {
+					return newConditionAction(core.SimpleType_BOOLEAN)
+				}
+				return newChildTaskAction()
+			}
+
+			root := newAction()
+			root.ActionId.Name = "root"
+			root.ParentActionName = nil
+			require.NoError(t, c.Enqueue(ctx, root, nil))
+			rootCR, err := c.GetTaskAction(ctx, root.ActionId)
+			require.NoError(t, err)
+			assert.NotContains(t, rootCR.Labels, "flyte.org/parent-task-action")
+
+			child := newAction()
+			child.ParentActionName = &root.ActionId.Name
+			require.NoError(t, c.Enqueue(ctx, child, nil))
+			childCR, err := c.GetTaskAction(ctx, child.ActionId)
+			require.NoError(t, err)
+			assert.Equal(t, rootCR.Name, childCR.Labels["flyte.org/parent-task-action"])
+
+			grandchild := newAction()
+			grandchild.ActionId.Name = "grandchild"
+			grandchild.ParentActionName = &child.ActionId.Name
+			require.NoError(t, c.Enqueue(ctx, grandchild, nil))
+			grandchildCR, err := c.GetTaskAction(ctx, grandchild.ActionId)
+			require.NoError(t, err)
+			assert.Equal(t, childCR.Name, grandchildCR.Labels["flyte.org/parent-task-action"])
+		})
+	}
+}
+
 // A restarted parent re-enqueues children with the same deterministic action
 // IDs; the second Enqueue must return AlreadyExists and leave the existing CR
 // untouched.