From 91038d90fe05232c49ea09aaf234b2653388bee2 Mon Sep 17 00:00:00 2001
From: Katrina Rogan <katroganGH@gmail.com>
Date: Mon, 18 Mar 2024 13:42:06 -0700
Subject: [PATCH] SVR-323: Add task org to plugins log link templating (#165)

Signed-off-by: Katrina Rogan <katroganGH@gmail.com>
---
 .../go/tasks/pluginmachinery/internal/webapi/launcher.go    | 2 +-
 flyteplugins/go/tasks/pluginmachinery/tasklog/template.go   | 6 ++++++
 .../go/tasks/pluginmachinery/tasklog/template_test.go       | 6 ++++--
 flyteplugins/go/tasks/plugins/webapi/agent/plugin.go        | 2 +-
 4 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/flyteplugins/go/tasks/pluginmachinery/internal/webapi/launcher.go b/flyteplugins/go/tasks/pluginmachinery/internal/webapi/launcher.go
index 99a3ccdf7a..9482b3df95 100644
--- a/flyteplugins/go/tasks/pluginmachinery/internal/webapi/launcher.go
+++ b/flyteplugins/go/tasks/pluginmachinery/internal/webapi/launcher.go
@@ -2,9 +2,9 @@ package webapi
 
 import (
 	"context"
-	pluginErrors "github.com/flyteorg/flyte/flyteplugins/go/tasks/errors"
 	"time"
 
+	pluginErrors "github.com/flyteorg/flyte/flyteplugins/go/tasks/errors"
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core"
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/webapi"
 	"github.com/flyteorg/flyte/flytestdlib/cache"
diff --git a/flyteplugins/go/tasks/pluginmachinery/tasklog/template.go b/flyteplugins/go/tasks/pluginmachinery/tasklog/template.go
index 1ca37525f3..d755594985 100644
--- a/flyteplugins/go/tasks/pluginmachinery/tasklog/template.go
+++ b/flyteplugins/go/tasks/pluginmachinery/tasklog/template.go
@@ -33,6 +33,7 @@ type templateRegexes struct {
 	PodUnixFinishTime    *regexp.Regexp
 	TaskID               *regexp.Regexp
 	TaskVersion          *regexp.Regexp
+	TaskOrg              *regexp.Regexp
 	TaskProject          *regexp.Regexp
 	TaskDomain           *regexp.Regexp
 	TaskRetryAttempt     *regexp.Regexp
@@ -59,6 +60,7 @@ func initDefaultRegexes() templateRegexes {
 		MustCreateRegex("podUnixFinishTime"),
 		MustCreateRegex("taskID"),
 		MustCreateRegex("taskVersion"),
+		MustCreateRegex("taskOrg"),
 		MustCreateRegex("taskProject"),
 		MustCreateRegex("taskDomain"),
 		MustCreateRegex("taskRetryAttempt"),
@@ -136,6 +138,10 @@ func (input Input) templateVars() []TemplateVar {
 					defaultRegexes.TaskVersion,
 					taskExecutionIdentifier.TaskId.Version,
 				},
+				TemplateVar{
+					defaultRegexes.TaskOrg,
+					taskExecutionIdentifier.TaskId.Org,
+				},
 				TemplateVar{
 					defaultRegexes.TaskProject,
 					taskExecutionIdentifier.TaskId.Project,
diff --git a/flyteplugins/go/tasks/pluginmachinery/tasklog/template_test.go b/flyteplugins/go/tasks/pluginmachinery/tasklog/template_test.go
index b32d451574..56d549e2dc 100644
--- a/flyteplugins/go/tasks/pluginmachinery/tasklog/template_test.go
+++ b/flyteplugins/go/tasks/pluginmachinery/tasklog/template_test.go
@@ -29,6 +29,7 @@ func dummyTaskExecID() pluginCore.TaskExecutionID {
 			Project:      "my-task-project",
 			Domain:       "my-task-domain",
 			Version:      "1",
+			Org:          "my-task-org",
 		},
 		NodeExecutionId: &core.NodeExecutionIdentifier{
 			ExecutionId: &core.WorkflowExecutionIdentifier{
@@ -145,6 +146,7 @@ func Test_Input_templateVars(t *testing.T) {
 				{defaultRegexes.TaskRetryAttempt, "1"},
 				{defaultRegexes.TaskID, "my-task-name"},
 				{defaultRegexes.TaskVersion, "1"},
+				{defaultRegexes.TaskOrg, "my-task-org"},
 				{defaultRegexes.TaskProject, "my-task-project"},
 				{defaultRegexes.TaskDomain, "my-task-domain"},
 				{defaultRegexes.ExecutionName, "my-execution-name"},
@@ -420,7 +422,7 @@ func TestTemplateLogPlugin(t *testing.T) {
 		{
 			"task-with-task-execution-identifier-with-org",
 			TemplateLogPlugin{
-				TemplateURIs:  []TemplateURI{"https://flyte.corp.net/console/org/{{ .executionOrg }}/projects/{{ .executionProject }}/domains/{{ .executionDomain }}/executions/{{ .executionName }}/nodeId/{{ .nodeID }}/taskId/{{ .taskID }}/attempt/{{ .taskRetryAttempt }}/view/logs"},
+				TemplateURIs:  []TemplateURI{"https://flyte.corp.net/console/org/{{ .executionOrg }}/projects/{{ .executionProject }}/domains/{{ .executionDomain }}/executions/{{ .executionName }}/nodeId/{{ .nodeID }}/taskId/{{ .taskID }}/taskOrg/{{ .taskOrg }}/attempt/{{ .taskRetryAttempt }}/view/logs"},
 				MessageFormat: core.TaskLog_JSON,
 			},
 			args{
@@ -441,7 +443,7 @@ func TestTemplateLogPlugin(t *testing.T) {
 			Output{
 				TaskLogs: []*core.TaskLog{
 					{
-						Uri:           "https://flyte.corp.net/console/org/my-execution-org/projects/my-execution-project/domains/my-execution-domain/executions/my-execution-name/nodeId/n0-0-n0/taskId/my-task-name/attempt/1/view/logs",
+						Uri:           "https://flyte.corp.net/console/org/my-execution-org/projects/my-execution-project/domains/my-execution-domain/executions/my-execution-name/nodeId/n0-0-n0/taskId/my-task-name/taskOrg/my-task-org/attempt/1/view/logs",
 						MessageFormat: core.TaskLog_JSON,
 						Name:          "main_logs",
 					},
diff --git a/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go b/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go
index ea96980b78..f31dcd91c1 100644
--- a/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go
+++ b/flyteplugins/go/tasks/plugins/webapi/agent/plugin.go
@@ -4,13 +4,13 @@ import (
 	"context"
 	"encoding/gob"
 	"fmt"
-	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service"
 	"time"
 
 	"golang.org/x/exp/maps"
 
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin"
 	flyteIdl "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
+	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service"
 	pluginErrors "github.com/flyteorg/flyte/flyteplugins/go/tasks/errors"
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery"
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core"