From 3e92d251ba7a950d050234d8780ce75085b1b1d4 Mon Sep 17 00:00:00 2001
From: "B. Hamonangan" <70834922+hamonangann@users.noreply.github.com>
Date: Wed, 27 Dec 2023 17:09:23 +0700
Subject: [PATCH] feat(aws): add support for sso sessions login (#12090)

Co-authored-by: Carlo Sala <carlosalag@protonmail.com>
---
 plugins/aws/README.md      | 1 +
 plugins/aws/aws.plugin.zsh | 6 +++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/plugins/aws/README.md b/plugins/aws/README.md
index 4850e223e..dbf8dcfeb 100644
--- a/plugins/aws/README.md
+++ b/plugins/aws/README.md
@@ -16,6 +16,7 @@ plugins=(... aws)
   It also sets `$AWS_EB_PROFILE` to `<profile>` for the Elastic Beanstalk CLI. It sets `$AWS_PROFILE_REGION` for display in `aws_prompt_info`.
   Run `asp` without arguments to clear the profile.
 * `asp [<profile>] login`: If AWS SSO has been configured in your aws profile, it will run the `aws sso login` command following profile selection.
+* `asp [<profile>] login [<sso_session>]`: In addition to `asp [<profile>] login`, if SSO session has been configured in your aws profile, it will run the `aws sso login --sso-session <sso_session>` command following profile selection.
 * `asp [<profile>] logout`: If AWS SSO has been configured in your aws profile, it will run the `aws sso logout` command following profile selection.
 
 * `asr [<region>]`: sets `$AWS_REGION` and `$AWS_DEFAULT_REGION` (legacy) to `<region>`.
diff --git a/plugins/aws/aws.plugin.zsh b/plugins/aws/aws.plugin.zsh
index 0d7040f92..fd00714e7 100644
--- a/plugins/aws/aws.plugin.zsh
+++ b/plugins/aws/aws.plugin.zsh
@@ -29,7 +29,11 @@ function asp() {
   export AWS_PROFILE_REGION=$(aws configure get region)
 
   if [[ "$2" == "login" ]]; then
-    aws sso login
+    if [[ -n "$3" ]]; then
+      aws sso login --sso-session $3
+    else
+      aws sso login
+    fi
   elif [[ "$2" == "logout" ]]; then
     aws sso logout
   fi