feat(aws): add AWS_REGION to aws_prompt_info (#10062)

plugins/aws/README.md

@@ -1,7 +1,7 @@
 # aws
 
 This plugin provides completion support for [awscli](https://docs.aws.amazon.com/cli/latest/reference/index.html)
-and a few utilities to manage AWS profiles and display them in the prompt.
+and a few utilities to manage AWS profiles/regions and display them in the prompt.
 
 To use it, add `aws` to the plugins array in your zshrc file.
 
@@ -16,6 +16,9 @@ plugins=(... aws)
   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. 
 
+* `asr [<region>]`: sets `$AWS_REGION` and `$AWS_DEFAULT_REGION` (legacy) to `<region>`.
+  Run `asr` without arguments to clear the profile.
+
 * `acp [<profile>] [<mfa_token>]`: in addition to `asp` functionality, it actually changes
    the profile by assuming the role specified in the `<profile>` configuration. It supports
    MFA and sets `$AWS_ACCESS_KEY_ID`, `$AWS_SECRET_ACCESS_KEY` and `$AWS_SESSION_TOKEN`, if
@@ -25,25 +28,34 @@ plugins=(... aws)
 
 * `agp`: gets the current value of `$AWS_PROFILE`.
 
+* `agr`: gets the current value of `$AWS_REGION`.
+
 * `aws_change_access_key`: changes the AWS access key of a profile.
 
 * `aws_profiles`: lists the available profiles in the  `$AWS_CONFIG_FILE` (default: `~/.aws/config`).
   Used to provide completion for the `asp` function.
 
+* `aws_regions`: lists the available regions.
+  Used to provide completion for the `asr` function.
+
 ## Plugin options
 
 * Set `SHOW_AWS_PROMPT=false` in your zshrc file if you want to prevent the plugin from modifying your RPROMPT.
   Some themes might overwrite the value of RPROMPT instead of appending to it, so they'll need to be fixed to
-  see the AWS profile prompt.
+  see the AWS profile/region prompt.
 
 ## Theme
 
 The plugin creates an `aws_prompt_info` function that you can use in your theme, which displays
-the current `$AWS_PROFILE`. It uses two variables to control how that is shown:
+the current `$AWS_PROFILE` and `$AWS_REGION`. It uses four variables to control how that is shown:
+
+* ZSH_THEME_AWS_PROFILE_PREFIX: sets the prefix of the AWS_PROFILE. Defaults to `<aws:`.
+
+* ZSH_THEME_AWS_PROFILE_SUFFIX: sets the suffix of the AWS_PROFILE. Defaults to `>`.
 
-* ZSH_THEME_AWS_PREFIX: sets the prefix of the AWS_PROFILE. Defaults to `<aws:`.
+* ZSH_THEME_AWS_REGION_PREFIX: sets the prefix of the AWS_REGION. Defaults to `<region:`.
 
-* ZSH_THEME_AWS_SUFFIX: sets the suffix of the AWS_PROFILE. Defaults to `>`.
+* ZSH_THEME_AWS_REGION_SUFFIX: sets the suffix of the AWS_REGION. Defaults to `>`.
 
 ## Configuration
 

plugins/aws/aws.plugin.zsh

@@ -2,6 +2,10 @@ function agp() {
   echo $AWS_PROFILE
 }
 
+function agr() {
+  echo $AWS_REGION
+}
+
 # AWS profile selection
 function asp() {
   if [[ -z "$1" ]]; then
@@ -27,6 +31,25 @@ function asp() {
   fi
 }
 
+# AWS region selection
+function asr() {
+  if [[ -z "$1" ]]; then
+    unset AWS_DEFAULT_REGION AWS_REGION
+    echo AWS region cleared.
+    return
+  fi
+
+  local -a available_regions
+  available_regions=($(aws_regions))
+  if [[ -z "${available_regions[(r)$1]}" ]]; then
+    echo "${fg[red]}Available regions: \n$(aws_regions)"
+    return 1
+  fi
+
+  export AWS_REGION=$1
+  export AWS_DEFAULT_REGION=$1
+}
+
 # AWS profile switch
 function acp() {
   if [[ -z "$1" ]]; then
@@ -145,12 +168,25 @@ function aws_change_access_key() {
   AWS_PAGER="" aws iam list-access-keys
 }
 
+function aws_regions() {
+  if [[ $AWS_DEFAULT_PROFILE || $AWS_PROFILE ]];then
+    aws ec2 describe-regions |grep RegionName | awk -F ':' '{gsub(/"/, "", $2);gsub(/,/, "", $2);gsub(/ /, "", $2);  print $2}'
+  else
+    echo "You must specify a AWS profile."
+  fi
+}
+
 function aws_profiles() {
   aws --no-cli-pager configure list-profiles 2> /dev/null && return
   [[ -r "${AWS_CONFIG_FILE:-$HOME/.aws/config}" ]] || return 1
   grep --color=never -Eo '\[.*\]' "${AWS_CONFIG_FILE:-$HOME/.aws/config}" | sed -E 's/^[[:space:]]*\[(profile)?[[:space:]]*([^[:space:]]+)\][[:space:]]*$/\2/g'
 }
 
+function _aws_regions() {
+  reply=($(aws_regions))
+}
+compctl -K _aws_regions asr
+
 function _aws_profiles() {
   reply=($(aws_profiles))
 }
@@ -158,8 +194,8 @@ compctl -K _aws_profiles asp acp aws_change_access_key
 
 # AWS prompt
 function aws_prompt_info() {
-  [[ -n "$AWS_PROFILE" ]] || return
-  echo "${ZSH_THEME_AWS_PREFIX=<aws:}${AWS_PROFILE:gs/%/%%}${ZSH_THEME_AWS_SUFFIX=>}"
+  if [[ -z $AWS_REGION && -z $AWS_PROFILE ]];then return; fi
+  echo "${ZSH_THEME_AWS_PROFILE_PREFIX:=<aws:}${AWS_PROFILE}${ZSH_THEME_AWS_PROFILE_SUFFIX:=>} ${ZSH_THEME_AWS_REGION_PREFIX:=<region:}${AWS_REGION}${ZSH_THEME_AWS_REGION_SUFFIX:=>}"
 }
 
 if [[ "$SHOW_AWS_PROMPT" != false && "$RPROMPT" != *'$(aws_prompt_info)'* ]]; then
@@ -211,3 +247,4 @@ else
   [[ -r $_aws_zsh_completer_path ]] && source $_aws_zsh_completer_path
   unset _aws_zsh_completer_path _brew_prefix
 fi
+