Tool to manage deployment of a lambda function.

aws-lambda.sh

#!/bin/bash

function fatal {
  echo "$@" >&2
  exit 1
}

# --- Parse arguments ---

usage="""\
Usage: $(basename $0) <command>

Commands:
  deploy             create or update lambda function
  delete-role        delete lambda function iam role
  invoke             invoke lambda function
  logs [since]       view lambda function logs. Since: '5m', '1h', etc.
  delete-function    delete lambda function
  get-arn            get lambda function arn
"""

COMMAND="$1"
case "$COMMAND" in
  deploy|delete-role|invoke|logs|delete-function|get-arn) ;; # valid
  *)
    fatal "\
Unrecognized command: $COMMAND

$usage" ;;
esac

# --- Required tools ---

command -v aws > /dev/null || { fatal "aws is not installed"; }
command -v shasum > /dev/null || { fatal "shasum is not installed"; }
command -v nu > /dev/null || { fatal "nu is not installed"; }
command -v uv > /dev/null || { fatal "uv is not installed"; }

# check if aws cli is configured
if ! aws sts get-caller-identity >/dev/null ; then
  fatal "please configure aws cli"
fi

###############
#  Utilities  #
###############

if [ -z "$LAMBDAX_ROOT" ]; then
  fatal "Please specify a configuration root. E.g. \`export LAMBDAX_ROOT=lambda/\`"
fi
LAMBDAX_ROOT="${LAMBDAX_ROOT%/}" # strip trailing '/'

export CONFIG_FILE="$LAMBDAX_ROOT/config.toml"

if [ ! -f $CONFIG_FILE ]; then
  fatal "File not found: $CONFIG_FILE"
fi

# usage: role_name=$(get-config-value role_name) || exit 1
function get-config-value {
  local output
  output=$(
    nu -c "open \$env.CONFIG_FILE | try { get $1 } catch { |err| print '\"$1\"'; exit 1 } | to text"
  ) || {
    echo "Key not found in config.toml: $output" >&2
    return 1
  }

  echo "$output"
}

# usage: jpath "<json>" 'key1.key2'
function jpath {
  echo $1 | nu --stdin -c "from json | get $2 | to text"
}

function get-function {
  function_info=$(aws lambda get-function --function-name $function_name 2>/dev/null)
  if [ $? = 254 ]; then
    return 1 # function doesn't exist
  else
    return 0
  fi
}

###############

# --- Get config values ---

role_name=$(get-config-value role_name) || exit 1
trust_policy_file=$(get-config-value role_trust_policy_file) || exit 1
role_policy_file=$(get-config-value role_permission_policy_file 2>/dev/null)

case "$COMMAND" in
  deploy|invoke|delete-function|get-arn)
    function_name=$(get-config-value function_name) || exit 1
    ;;
esac

if [ "$COMMAND" = "deploy" ]; then
  python_version=$(get-config-value python_version) || exit 1
fi

# --- Setup cache folder ---

CACHE=$LAMBDAX_ROOT/.cache
mkdir -p $CACHE

##########################
#  Command: delete-role  #
##########################

if [ "$COMMAND" = "delete-role" ]; then

  # Delete attached role policies (trust policy)
  role_policies=$(aws iam list-attached-role-policies --role-name $role_name 2>/dev/null)
  if [ $? != 254 ]; then
    policy_arns=$(jpath "$role_policies" 'AttachedPolicies.PolicyArn')
    for policy_arn in $policy_arns; do
      echo "Detaching policy: $policy_arn" >&2
      aws iam detach-role-policy --role-name $role_name --policy-arn $policy_arn
    done
  fi

  # Delete role policies (permissions policy)
  role_policies=$(aws iam list-role-policies --role-name test-lambda-function-role 2>/dev/null)
  if [ $? != 254 ]; then
    policy_names=$(jpath "$role_policies" 'PolicyNames')
    for policy_name in $policy_names; do
      echo "Deleting policy: $policy_name" >&2
      aws iam delete-role-policy --role-name $role_name --policy-name $policy_name
    done
  fi

  echo "Deleting role: $role_name" >&2
  aws iam delete-role --role-name $role_name 2>/dev/null
  if [ $? = 254 ]; then
    fatal "Role not found" 
  fi

  rm -f $CACHE/$trust_policy_file.sha

  exit
fi

#####################
#  Command: invoke  #
#####################

if [ "$COMMAND" = "invoke" ]; then
  aws lambda invoke --function-name $function_name /dev/null
  exit
fi

###################
#  Command: logs  #
###################

if [ "$COMMAND" = "logs" ]; then

  if [ ! -f $CACHE/lambda_function_log_group ]; then
    fatal "File doesn't exist: $CACHE/lambda_function_log_group" 
  fi

  SINCE="${@:2}"

  uv run --with zmunk-awslogs -m awslogs $(cat $CACHE/lambda_function_log_group) $SINCE

  exit
fi

##############################
#  Command: delete-function  #
##############################

if [ "$COMMAND" = "delete-function" ]; then
  aws lambda delete-function --function-name $function_name 1>/dev/null
  exit
fi

######################
#  Command: get-arn  #
######################

if [ "$COMMAND" = "get-arn" ]; then
  if [ ! -f $CACHE/lambda_function_arn ]; then
    echo "ARN not cached, retrieving..." >&2
    if ! get-function; then
      fatal "Function doesn't exist."
    fi
    arn=$(jpath "$function_info" 'Configuration.FunctionArn')
    echo "$arn" > $CACHE/lambda_function_arn
  fi
  cat $CACHE/lambda_function_arn
  exit
fi

#####################
#  Command: deploy  #
#####################

if [ "$COMMAND" != "deploy" ]; then
  fatal "Unhandled command: $COMMAND" 
fi

###  IAM Role  ###

function get-role {
  role_info=$(aws iam get-role --role-name $role_name 2>/dev/null)
  if [ $? = 254 ]; then
    return 1 # role doesn't exist
  else
    return 0
  fi
}

function update-trust-policy-sha {
  shasum < $LAMBDAX_ROOT/$trust_policy_file > $CACHE/$trust_policy_file.sha
}

function update-role-policy-sha {
  shasum < $LAMBDAX_ROOT/$role_policy_file > $CACHE/$role_policy_file.sha
}

role_policy_modified=false
if [ -n "$role_policy_file" ]; then
  if [ ! -f $CACHE/$role_policy_file.sha ]; then
    role_policy_modified=true # sha file doesn't exist
  elif [ "$(shasum < $LAMBDAX_ROOT/$role_policy_file)" != "$(cat $CACHE/$role_policy_file.sha)" ]; then
    role_policy_modified=true # sha doesn't match
  fi
fi

trust_policy_modified=false
if [ ! -f $CACHE/$trust_policy_file.sha ]; then
  trust_policy_modified=true # sha file doesn't exist
elif [ "$(shasum < $LAMBDAX_ROOT/$trust_policy_file)" != "$(cat $CACHE/$trust_policy_file.sha)" ]; then
  trust_policy_modified=true # sha doesn't match
fi

if ! get-role; then
  echo "Role doesn't exist. Creating..." >&2

  role_info=$(
    aws iam create-role \
      --role-name $role_name \
      --assume-role-policy-document file://$LAMBDAX_ROOT/$trust_policy_file
  )

  aws iam wait role-exists --role-name $role_name

  echo "Role created" >&2

  update-trust-policy-sha

  output=$(
    aws iam attach-role-policy \
    --role-name $role_name \
    --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole \
    2>/dev/null
  )
  if [ $? != 0 ]; then
    echo "Command: aws iam attach-role-policy" >&2
    echo "$output" >&2
    exit 1
  fi
fi

if $trust_policy_modified; then
  echo "Updating trust policy..." >&2

  aws iam update-assume-role-policy \
    --role-name $role_name \
    --policy-document file://$LAMBDAX_ROOT/$trust_policy_file || exit 1

  update-trust-policy-sha
fi

if $role_policy_modified; then

  echo "Updating role policy..." >&2

  # Update permissions policy
  aws iam put-role-policy \
    --role-name $role_name \
    --policy-name lambda-permissions-policy \
    --policy-document file://$LAMBDAX_ROOT/$role_policy_file || exit 1

fi

role_arn=$(jpath "$role_info" 'Role.Arn')

### Dependency layers ###

function_config_args=()

function_config_args+=(--function-name $function_name)

# --- Create lambda layer script ---

# See https://gist.github.com/zmunk/2bc71b940d4a821b59ab073299de5f55

gist_id=2bc71b940d4a821b59ab073299de5f55
gist_version=0d6f830fc63026e2630ccd4d938ed5c9b8b3e587
gist_url=https://gist.githubusercontent.com/zmunk/$gist_id/raw/$gist_version/create-lambda-layer.sh

export python_version # allow nushell to access it

output=$(nu -c $'
  let ver = ($env.python_version | str replace "python" "" | str replace "." "")
  let layers = (open $env.CONFIG_FILE | get -o layers)
  if ($layers | is-not-empty) {
      $layers | transpose layer_name layer_data | each { |row|
        let libraries = ($row.layer_data.libraries | str join " ")
        let layer_name = $"($row.layer_name)-py($ver)-lambda-layer"
        mut script_args = $"($env.python_version) ($libraries) ($layer_name)"

        let pip_args = ($row.layer_data.pip_args?)
        if ($pip_args | is-not-empty) {
          $script_args = $"($script_args) -- ($pip_args | str join \' \')"
        }

        print $script_args
      }
  }
  | ignore
')

if [ -n "$output" ]; then

  layer_arns=()

  while IFS= read -r args; do

    layer_build_cmd="curl $gist_url | bash -s $args"

    # SHA of command to build layer
    # We will use it as a key to save and look up the layer ARN
    layer_sha=$(echo "$layer_build_cmd" | shasum)
    layer_sha="${layer_sha%% *}" # remove '  -' from end

    if [ -f $CACHE/$layer_sha ]; then
      layer_arn="$(cat $CACHE/$layer_sha)"
    else

      layer_name="${args% -- *}"     # remove ' -- ' and everything after it (extra pip args)
      layer_name="${layer_name##* }" # get last word

      echo "Building new lambda layer: $layer_name" >&2

      # execute layer build command
      layer_arn=$(eval "$layer_build_cmd" 2>/dev/null) || {
        fatal "Failed to build layer. Command: '$layer_build_cmd'" 
      }

      echo "$layer_arn" > $CACHE/$layer_sha
    fi

    layer_arns+=($layer_arn)

  done <<< "$output"

  function_config_args+=(--layers "${layer_arns[@]}")
fi

# --- Compile environment variables ---

# NOTE: env vars in config.local.toml take precedence
export SECRETS_CONFIG=$LAMBDAX_ROOT/config.local.toml

output=$(
  nu -c '
    let env_vars = (open $env.CONFIG_FILE | get env? | default {})
    let env_secrets = if ($env.SECRETS_CONFIG | path exists) { open $env.SECRETS_CONFIG | get env? | default {} } else { {} }
    $env_vars
    | merge $env_secrets
    | transpose key value
    | each { |row|
      $"($row.key)=($row.value)"
    }
  | str join ","
  '
)

if [ -n "$output" ]; then
  function_config_args+=(--environment "Variables={$output}")
fi

# --- Check if update-function-config arguments have been modified ---
func_config_sha="$(echo "${function_config_args[@]}" | shasum)"
func_config_modified=false
if [ ! -f $CACHE/update-function-configuration-args.sha ]; then
  func_config_modified=true # sha file doesn't exist
elif [ "$func_config_sha" != "$(cat $CACHE/update-function-configuration-args.sha)" ]; then
  func_config_modified=true # sha doesn't match
fi

###  Lambda Function  ###

# --- Compute package file list ---

# Pick up files specified in package.include while excluding files specified in
# package.exclude, respecting globs
function get-package-files {
  nu -c $'
  # e.g. "expand-patterns $.package.include"
  def expand-patterns [key: cell-path] {
    let patterns = (open $env.CONFIG_FILE | get -o $key)
    if ($patterns | is-empty) { return [] }

    $patterns | each { |pat|
      if ($pat | str contains "*") or ($pat | str contains "?") {
        if ($pat | str starts-with "**/") {
          fd --base-directory $env.LAMBDAX_ROOT --glob $pat | lines
        } else {
          fd --base-directory $env.LAMBDAX_ROOT --max-depth 1 --glob $pat | lines
        }
      } else {
        [$pat]
      }
    } | flatten | uniq
  }
  let included = (expand-patterns $.package.include)
  let excluded = (expand-patterns $.package.exclude)
  $included | where { |f| $f not-in $excluded } | to text
  '
}

# Usage:
# 
#     for file in $package_files; do
#       echo "file: $file"
#     done
#
package_files=$(get-package-files)

if [[ -z "$package_files" ]]; then
  fatal "No package files specified"
fi

# Compute aggregate sha of all package files, or any configuration
# related to `create-function` command
function get-package-sha {
  local contents=""
  for file in $package_files; do
    contents+="$(cat $LAMBDAX_ROOT/$file)"$'\n'
  done
  contents+="$function_name"$'\n'
  contents+="$python_version"$'\n'
  echo "$contents" | shasum
}

# -- Check if any files have been modified

if get-function; then
  function_exists=true
else
  function_exists=false
fi

package_sha="$(get-package-sha)"

lambda_files_modified=false
if [ ! -f $CACHE/lambda-package.sha ]; then
  lambda_files_modified=true # sha file doesn't exist
elif [ "$package_sha" != "$(cat $CACHE/lambda-package.sha)" ]; then
  lambda_files_modified=true # sha doesn't match
elif ! $function_exists; then
  lambda_files_modified=true # function not yet created
fi

if ! $lambda_files_modified && ! $function_config_modified; then
  fatal "
No files or configuration has been modified.
If this is a mistake, remove the sha file(s) and rerun the deployment
  - SHA file for package files: $CACHE/lambda-package.sha
  - SHA file for 'update-function-config' function arguments: $CACHE/update-function-configuration-args.sha
"
fi

function_newly_created=false

if $lambda_files_modified; then

  # --- Clear lambda package CACHE folder

  rm -rf $CACHE/$LAMBDAX_ROOT/
  mkdir -p $CACHE/$LAMBDAX_ROOT/

  # --- Copy files into lambda package folder ---

  for file in $package_files; do
    mkdir -p "$CACHE/$LAMBDAX_ROOT/$(dirname "$file")"
    cp $LAMBDAX_ROOT/$file $CACHE/$LAMBDAX_ROOT/$file || exit 1
  done

  # --- Create zip file ---

  (
    cd $CACHE/$LAMBDAX_ROOT/
    zip -r ../function.zip . >/dev/null
  )

  if ! $function_exists; then
    echo "Function doesn't exist. Creating..." >&2

    function_newly_created=true

    success=false
    for attempt in 1 2; do
      output=$(
        aws lambda create-function \
          --function-name $function_name \
          --runtime $python_version \
          --role $role_arn \
          --handler lambda_function.lambda_handler \
          --zip-file fileb://$CACHE/function.zip 2>&1
        )
      exit_code=$?
      case $exit_code in
        0)
          function_info="$output";
          success=true
          break
          ;;
        254)
          if [[ "$output" = *"cannot be assumed"* ]]; then
            sleep 2 # Retry after 2 seconds
            continue
          else
            echo "Error: $output" >&2
            exit 1
          fi
          ;;
        *)
          echo "Error: $output" >&2
          echo "Exit code: $exit_code" >&2
          exit 1
          ;;
      esac
    done

    if ! $success; then
      echo "Attempted to create function twice." >&2
      echo "Error: $output" >&2
      exit 1
    fi

    echo "Function created." >&2

  else

    echo "Updating function code..." >&2

    function_info=$(
      aws lambda update-function-code \
        --function-name $function_name \
        --zip-file fileb://$CACHE/function.zip
      )
  fi

  # Update lambda package sha
  echo "$package_sha" > $CACHE/lambda-package.sha

  jpath "$function_info" 'LoggingConfig.LogGroup' > $CACHE/lambda_function_log_group

  jpath "$function_info" 'FunctionArn' > $CACHE/lambda_function_arn
fi

# --- Update function configuration ---

# Note: for environment variables see '--environment' in 'aws lambda update-function-configuration help'

if $func_config_modified || $function_newly_created; then

  update_config_cmd=(aws lambda update-function-configuration "${function_config_args[@]}")

  echo "Updating configuration..." >&2

  # Need to check if function is active before updating configuration
  aws lambda wait function-active --function-name $function_name

  # Execute command
  error=$("${update_config_cmd[@]}" 2>&1 1>/dev/null) || {
    echo "Failed to update function configuration." >&2
    echo "Command: '${update_config_cmd[@]}'" >&2
    echo "Error: $error" >&2
    exit 1
  }

  # Set sha
  echo "$func_config_sha" > $CACHE/update-function-configuration-args.sha

fi

# ---

echo "Deployment complete." >&2