diff --git a/flyteplugins/go/tasks/plugins/webapi/athena/plugin.go b/flyteplugins/go/tasks/plugins/webapi/athena/plugin.go index a1f16163fe..826a12e45f 100644 --- 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®ion=%v#query/history/%v", + Uri: fmt.Sprintf("https://%v.%v/athena/home?force®ion=%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 index 6e3238d813..e19829447e 100644 --- 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®ion=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") +}