Use a map of TemplateLogPlugin in log config

Signed-off-by: Eduardo Apolinario <[email protected]>
flyteorg · Jan 29, 2024 · 28aaa05 · 28aaa05
1 parent 5c3b813
commit 28aaa05
Show file tree

Hide file tree

Showing 10 changed files with 43 additions and 111 deletions.
diff --git a/flyteplugins/go/tasks/logs/config.go b/flyteplugins/go/tasks/logs/config.go
@@ -28,7 +28,7 @@ type LogConfig struct {
 	StackdriverLogResourceName string              `json:"stackdriver-logresourcename" pflag:",Name of the logresource in stackdriver"`
 	StackDriverTemplateURI     tasklog.TemplateURI `json:"stackdriver-template-uri" pflag:",Template Uri to use when building stackdriver log links"`
 
-	DynamicLogLinks map[string]tasklog.TemplateURI `json:"dynamic-log-links" pflag:",Map of dynamic log links"`
+	DynamicLogLinks map[string]tasklog.TemplateLogPlugin `json:"dynamic-log-links" pflag:"-,Map of dynamic log links"`
 
 	Templates []tasklog.TemplateLogPlugin `json:"templates" pflag:"-,"`
 }

diff --git a/flyteplugins/go/tasks/logs/logconfig_flags.go b/flyteplugins/go/tasks/logs/logconfig_flags.go
diff --git a/flyteplugins/go/tasks/logs/logconfig_flags_test.go b/flyteplugins/go/tasks/logs/logconfig_flags_test.go
diff --git a/flyteplugins/go/tasks/logs/logging_utils.go b/flyteplugins/go/tasks/logs/logging_utils.go
@@ -115,12 +115,10 @@ func InitializeLogPlugins(cfg *LogConfig) (tasklog.Plugin, error) {
 		dynamicPlugins = append(
 			dynamicPlugins,
 			tasklog.TemplateLogPlugin{
-				Name:        logLinkType,
-				DisplayName: logLinkType,
-				DynamicTemplateURIs: []tasklog.TemplateURI{
-					dynamicLogLink,
-				},
-				MessageFormat: core.TaskLog_JSON,
+				Name:                logLinkType,
+				DisplayName:         dynamicLogLink.DisplayName,
+				DynamicTemplateURIs: dynamicLogLink.TemplateURIs,
+				MessageFormat:       core.TaskLog_JSON,
 			})
 	}
 

diff --git a/flyteplugins/go/tasks/logs/logging_utils_test.go b/flyteplugins/go/tasks/logs/logging_utils_test.go
@@ -360,8 +360,13 @@ func TestGetLogsForContainerInPod_Flyteinteractive(t *testing.T) {
 		{
 			"Flyteinteractive enabled but no task template",
 			&LogConfig{
-				DynamicLogLinks: map[string]tasklog.TemplateURI{
-					"vscode": "https://flyteinteractive.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+				DynamicLogLinks: map[string]tasklog.TemplateLogPlugin{
+					"vscode": tasklog.TemplateLogPlugin{
+						DisplayName: "vscode link",
+						TemplateURIs: []tasklog.TemplateURI{
+							"https://flyteinteractive.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+						},
+					},
 				},
 			},
 			nil,
@@ -370,8 +375,13 @@ func TestGetLogsForContainerInPod_Flyteinteractive(t *testing.T) {
 		{
 			"Flyteinteractive enabled but config not found in task template",
 			&LogConfig{
-				DynamicLogLinks: map[string]tasklog.TemplateURI{
-					"vscode": "https://flyteinteractive.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+				DynamicLogLinks: map[string]tasklog.TemplateLogPlugin{
+					"vscode": tasklog.TemplateLogPlugin{
+						DisplayName: "vscode link",
+						TemplateURIs: []tasklog.TemplateURI{
+							"https://flyteinteractive.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+						},
+					},
 				},
 			},
 			&core.TaskTemplate{},
@@ -391,8 +401,13 @@ func TestGetLogsForContainerInPod_Flyteinteractive(t *testing.T) {
 		{
 			"Flyteinteractive - multiple dynamic options",
 			&LogConfig{
-				DynamicLogLinks: map[string]tasklog.TemplateURI{
-					"vscode": "https://abc.com:{{ .taskConfig.port }}/{{ .taskConfig.route }}",
+				DynamicLogLinks: map[string]tasklog.TemplateLogPlugin{
+					"vscode": tasklog.TemplateLogPlugin{
+						DisplayName: "vscode link",
+						TemplateURIs: []tasklog.TemplateURI{
+							"https://abc.com:{{ .taskConfig.port }}/{{ .taskConfig.route }}",
+						},
+					},
 				},
 			},
 			&core.TaskTemplate{
@@ -406,15 +421,20 @@ func TestGetLogsForContainerInPod_Flyteinteractive(t *testing.T) {
 				{
 					Uri:           "https://abc.com:65535/a-route",
 					MessageFormat: core.TaskLog_JSON,
-					Name:          "vscode my-Suffix",
+					Name:          "vscode link my-Suffix",
 				},
 			},
 		},
 		{
 			"Flyteinteractive - multiple uses of the template (invalid use of ports in a URI)",
 			&LogConfig{
-				DynamicLogLinks: map[string]tasklog.TemplateURI{
-					"vscode": "https://abc.com:{{ .taskConfig.port }}:{{ .taskConfig.port}}",
+				DynamicLogLinks: map[string]tasklog.TemplateLogPlugin{
+					"vscode": tasklog.TemplateLogPlugin{
+						DisplayName: "vscode link",
+						TemplateURIs: []tasklog.TemplateURI{
+							"https://abc.com:{{ .taskConfig.port }}:{{ .taskConfig.port}}",
+						},
+					},
 				},
 			},
 			&core.TaskTemplate{
@@ -427,7 +447,7 @@ func TestGetLogsForContainerInPod_Flyteinteractive(t *testing.T) {
 				{
 					Uri:           "https://abc.com:65535:65535",
 					MessageFormat: core.TaskLog_JSON,
-					Name:          "vscode my-Suffix",
+					Name:          "vscode link my-Suffix",
 				},
 			},
 		},
@@ -456,8 +476,13 @@ func TestGetLogsForContainerInPod_Flyteinteractive(t *testing.T) {
 			&LogConfig{
 				IsKubernetesEnabled:   true,
 				KubernetesTemplateURI: "https://k8s.com/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
-				DynamicLogLinks: map[string]tasklog.TemplateURI{
-					"vscode": "https://flyteinteractive.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+				DynamicLogLinks: map[string]tasklog.TemplateLogPlugin{
+					"vscode": tasklog.TemplateLogPlugin{
+						DisplayName: "vscode link",
+						TemplateURIs: []tasklog.TemplateURI{
+							"https://flyteinteractive.mydomain.com:{{ .taskConfig.port }}/{{ .namespace }}/{{ .podName }}/{{ .containerName }}/{{ .containerId }}",
+						},
+					},
 				},
 			},
 			&core.TaskTemplate{
@@ -475,7 +500,7 @@ func TestGetLogsForContainerInPod_Flyteinteractive(t *testing.T) {
 				{
 					Uri:           "https://flyteinteractive.mydomain.com:65535/my-namespace/my-pod/ContainerName/ContainerID",
 					MessageFormat: core.TaskLog_JSON,
-					Name:          "vscode my-Suffix",
+					Name:          "vscode link my-Suffix",
 				},
 			},
 		},

diff --git a/flyteplugins/go/tasks/pluginmachinery/tasklog/plugin.go b/flyteplugins/go/tasks/pluginmachinery/tasklog/plugin.go
@@ -46,7 +46,6 @@ type Plugin interface {
 }
 
 type TemplateLogPlugin struct {
-	// TODO: these don't need pflags anymore?
 	Name                string                     `json:"name" pflag:",Name of the plugin."`
 	DisplayName         string                     `json:"displayName" pflag:",Display name for the generated log when displayed in the console."`
 	TemplateURIs        []TemplateURI              `json:"templateUris" pflag:",URI Templates for generating task log links."`

diff --git a/flyteplugins/go/tasks/plugins/array/k8s/config_flags.go b/flyteplugins/go/tasks/plugins/array/k8s/config_flags.go
diff --git a/flyteplugins/go/tasks/plugins/array/k8s/config_flags_test.go b/flyteplugins/go/tasks/plugins/array/k8s/config_flags_test.go
diff --git a/flyteplugins/go/tasks/plugins/k8s/spark/config_flags.go b/flyteplugins/go/tasks/plugins/k8s/spark/config_flags.go
diff --git a/flyteplugins/go/tasks/plugins/k8s/spark/config_flags_test.go b/flyteplugins/go/tasks/plugins/k8s/spark/config_flags_test.go