adding consoleUrl parameterization based on partition (#4375)

* adding consoleUrl parameterization based on partition Signed-off-by: Laura Lin <[email protected]> * Update plugin.go Signed-off-by: Laura Lin <[email protected]> * Update plugin_test.go Signed-off-by: Laura Lin <[email protected]> * lint Signed-off-by: Laura Lin <[email protected]> --------- Signed-off-by: Laura Lin <[email protected]> Co-authored-by: Kevin Su <[email protected]>
flyteorg · Nov 11, 2023 · 31748fe · 31748fe
1 parent 21b92f4
commit 31748fe
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/flyteplugins/go/tasks/plugins/webapi/athena/plugin.go b/flyteplugins/go/tasks/plugins/webapi/athena/plugin.go
@@ -3,6 +3,7 @@ package athena
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	awsSdk "github.com/aws/aws-sdk-go-v2/aws"
@@ -177,12 +178,19 @@ func (p Plugin) Status(ctx context.Context, tCtx webapi.StatusContext) (phase co
 
 func createTaskInfo(queryID string, cfg awsSdk.Config) *core.TaskInfo {
 	timeNow := time.Now()
+	var consoleURL string
+	if strings.Contains(cfg.Region, "gov") {
+		consoleURL = "console.amazonaws-us-gov.com"
+	} else {
+		consoleURL = "console.aws.amazon.com"
+	}
 	return &core.TaskInfo{
 		OccurredAt: &timeNow,
 		Logs: []*idlCore.TaskLog{
 			{
-				Uri: fmt.Sprintf("https://%v.console.aws.amazon.com/athena/home?force&region=%v#query/history/%v",
+				Uri: fmt.Sprintf("https://%v.%v/athena/home?force&region=%v#query/history/%v",
 					cfg.Region,
+					consoleURL,
 					cfg.Region,
 					queryID),
 				Name: "Athena Query Console",

diff --git a/flyteplugins/go/tasks/plugins/webapi/athena/plugin_test.go b/flyteplugins/go/tasks/plugins/webapi/athena/plugin_test.go
@@ -22,3 +22,18 @@ func TestCreateTaskInfo(t *testing.T) {
 	assert.Len(t, taskInfo.ExternalResources, 1)
 	assert.Equal(t, taskInfo.ExternalResources[0].ExternalID, "query_id")
 }
+
+
+func TestCreateTaskInfoGovAWS(t *testing.T) {
+	taskInfo := createTaskInfo("query_id", awsSdk.Config{
+		Region: "us-gov-east-1",
+	})
+	assert.EqualValues(t, []*idlCore.TaskLog{
+		{
+			Uri:  "https://us-gov-east-1.console.amazonaws-us-gov.com/athena/home?force&region=us-gov-east-1#query/history/query_id",
+			Name: "Athena Query Console",
+		},
+	}, taskInfo.Logs)
+	assert.Len(t, taskInfo.ExternalResources, 1)
+	assert.Equal(t, taskInfo.ExternalResources[0].ExternalID, "query_id")
+}