2022-10-06 14:10:53 -06:00
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
# frozen_string_literal: true
2024-09-25 15:25:54 -04:00
require 'aws-sdk-lambda'
require 'aws-sdk-s3'
require 'aws-sdk-iam'
require 'aws-sdk-cloudwatchlogs'
require 'logger'
require 'json'
require 'zip'
2022-10-06 14:10:53 -06:00
# snippet-start:[ruby.example_code.ruby.LambdaWrapper.full]
# snippet-start:[ruby.example_code.ruby.LambdaWrapper.decl]
class LambdaWrapper
2024-09-25 15:25:54 -04:00
attr_accessor :lambda_client , :cloudwatch_client , :iam_client
2022-10-06 14:10:53 -06:00
def initialize
@lambda_client = Aws :: Lambda :: Client . new
2024-09-25 15:25:54 -04:00
@cloudwatch_client = Aws :: CloudWatchLogs :: Client . new ( region : 'us-east-1' )
@iam_client = Aws :: IAM :: Client . new ( region : 'us-east-1' )
2022-10-06 14:10:53 -06:00
@logger = Logger . new ( $stdout )
@logger . level = Logger :: WARN
end
# snippet-end:[ruby.example_code.ruby.LambdaWrapper.decl]
# snippet-start:[ruby.example_code.lambda.setup_iam]
# Get an AWS Identity and Access Management (IAM) role.
#
# @param iam_role_name: The name of the role to retrieve.
# @param action: Whether to create or destroy the IAM apparatus.
# @return: The IAM role.
def manage_iam ( iam_role_name , action )
2024-09-25 15:25:54 -04:00
case action
when 'create'
create_iam_role ( iam_role_name )
when 'destroy'
destroy_iam_role ( iam_role_name )
else
raise " Incorrect action provided. Must provide 'create' or 'destroy' "
end
end
private
def create_iam_role ( iam_role_name )
2022-10-06 14:10:53 -06:00
role_policy = {
2024-09-25 15:25:54 -04:00
'Version' : '2012-10-17' ,
2022-10-06 14:10:53 -06:00
'Statement' : [
{
2024-09-25 15:25:54 -04:00
'Effect' : 'Allow' ,
'Principal' : { 'Service' : 'lambda.amazonaws.com' } ,
'Action' : 'sts:AssumeRole'
2022-10-06 14:10:53 -06:00
}
]
}
2024-09-25 15:25:54 -04:00
role = @iam_client . create_role (
role_name : iam_role_name ,
assume_role_policy_document : role_policy . to_json
)
@iam_client . attach_role_policy (
{
policy_arn : 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole' ,
2022-10-06 14:10:53 -06:00
role_name : iam_role_name
2024-09-25 15:25:54 -04:00
}
)
wait_for_role_to_exist ( iam_role_name )
@logger . debug ( " Successfully created IAM role: #{ role [ 'role' ] [ 'arn' ] } " )
sleep ( 10 )
[ role , role_policy . to_json ]
end
def destroy_iam_role ( iam_role_name )
@iam_client . detach_role_policy (
{
policy_arn : 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole' ,
role_name : iam_role_name
}
)
@iam_client . delete_role ( role_name : iam_role_name )
@logger . debug ( " Detached policy & deleted IAM role: #{ iam_role_name } " )
end
def wait_for_role_to_exist ( iam_role_name )
@iam_client . wait_until ( :role_exists , { role_name : iam_role_name } ) do | w |
w . max_attempts = 5
w . delay = 5
2022-10-06 14:10:53 -06:00
end
end
# snippet-end:[ruby.example_code.lambda.setup_iam]
# snippet-start:[ruby.example_code.lambda.create_deployment_package]
# Creates a Lambda deployment package in .zip format.
#
# @param source_file: The name of the object, without suffix, for the Lambda file and zip.
# @return: The deployment package.
def create_deployment_package ( source_file )
Dir . chdir ( File . dirname ( __FILE__ ) )
2024-09-25 15:25:54 -04:00
if File . exist? ( 'lambda_function.zip' )
File . delete ( 'lambda_function.zip' )
@logger . debug ( 'Deleting old zip: lambda_function.zip' )
end
Zip :: File . open ( 'lambda_function.zip' , create : true ) do | zipfile |
zipfile . add ( 'lambda_function.rb' , " #{ source_file } .rb " )
2022-10-06 14:10:53 -06:00
end
@logger . debug ( " Zipping #{ source_file } .rb into: lambda_function.zip. " )
2024-09-25 15:25:54 -04:00
File . read ( 'lambda_function.zip' ) . to_s
2022-10-06 14:10:53 -06:00
rescue StandardError = > e
@logger . error ( " There was an error creating deployment package: \n #{ e . message } " )
end
# snippet-end:[ruby.example_code.lambda.create_deployment_package]
# snippet-start:[ruby.example_code.lambda.GetFunction]
# Gets data about a Lambda function.
#
# @param function_name: The name of the function.
# @return response: The function data, or nil if no such function exists.
def get_function ( function_name )
@lambda_client . get_function (
{
2023-06-08 16:39:36 -06:00
function_name : function_name
2022-10-06 14:10:53 -06:00
}
)
rescue Aws :: Lambda :: Errors :: ResourceNotFoundException = > e
@logger . debug ( " Could not find function: #{ function_name } : \n #{ e . message } " )
nil
end
# snippet-end:[ruby.example_code.lambda.GetFunction]
# snippet-start:[ruby.example_code.lambda.CreateFunction]
# Deploys a Lambda function.
#
# @param function_name: The name of the Lambda function.
2024-09-25 15:25:54 -04:00
# @param handler_name: The fully qualified name of the handler function.
2022-10-06 14:10:53 -06:00
# @param role_arn: The IAM role to use for the function.
2024-09-25 15:25:54 -04:00
# @param deployment_package: The deployment package that contains the function code in .zip format.
2022-10-06 14:10:53 -06:00
# @return: The Amazon Resource Name (ARN) of the newly created function.
def create_function ( function_name , handler_name , role_arn , deployment_package )
response = @lambda_client . create_function ( {
role : role_arn . to_s ,
2023-06-08 16:39:36 -06:00
function_name : function_name ,
2022-10-06 14:10:53 -06:00
handler : handler_name ,
2024-09-25 15:25:54 -04:00
runtime : 'ruby2.7' ,
2022-10-06 14:10:53 -06:00
code : {
zip_file : deployment_package
} ,
environment : {
variables : {
2024-09-25 15:25:54 -04:00
'LOG_LEVEL' = > 'info'
2022-10-06 14:10:53 -06:00
}
}
} )
2024-09-25 15:25:54 -04:00
@lambda_client . wait_until ( :function_active_v2 , { function_name : function_name } ) do | w |
2022-10-06 14:10:53 -06:00
w . max_attempts = 5
w . delay = 5
end
response
rescue Aws :: Lambda :: Errors :: ServiceException = > e
@logger . error ( " There was an error creating #{ function_name } : \n #{ e . message } " )
rescue Aws :: Waiters :: Errors :: WaiterFailed = > e
@logger . error ( " Failed waiting for #{ function_name } to activate: \n #{ e . message } " )
end
# snippet-end:[ruby.example_code.lambda.CreateFunction]
# snippet-start:[ruby.example_code.lambda.Invoke]
# Invokes a Lambda function.
# @param function_name [String] The name of the function to invoke.
# @param payload [nil] Payload containing runtime parameters.
# @return [Object] The response from the function invocation.
def invoke_function ( function_name , payload = nil )
2024-09-25 15:25:54 -04:00
params = { function_name : function_name }
2022-10-06 14:10:53 -06:00
params [ :payload ] = payload unless payload . nil?
@lambda_client . invoke ( params )
rescue Aws :: Lambda :: Errors :: ServiceException = > e
@logger . error ( " There was an error executing #{ function_name } : \n #{ e . message } " )
end
# snippet-end:[ruby.example_code.lambda.Invoke]
2024-09-25 15:25:54 -04:00
# snippet-start:[ruby.example_code.lambda.GetLogs]
2022-10-06 14:10:53 -06:00
# Get function logs from the latest Amazon CloudWatch log stream.
# @param function_name [String] The name of the function.
# @param string_match [String] A string to look for in the logs.
# @return all_logs [Array] An array of all the log messages found for that stream.
def get_cloudwatch_logs ( function_name , string_match )
2024-09-25 15:25:54 -04:00
@logger . debug ( 'Enforcing a 10 second sleep to allow CloudWatch logs to appear.' )
2022-10-06 14:10:53 -06:00
sleep ( 10 )
2024-09-25 15:25:54 -04:00
streams = @cloudwatch_client . describe_log_streams ( { log_group_name : " /aws/lambda/ #{ function_name } " } )
streams [ 'log_streams' ] . each do | log_stream |
resp = @cloudwatch_client . get_log_events ( {
2022-10-06 14:10:53 -06:00
log_group_name : " /aws/lambda/ #{ function_name } " ,
2024-09-25 15:25:54 -04:00
log_stream_name : log_stream [ 'log_stream_name' ]
2022-10-06 14:10:53 -06:00
} )
2024-09-25 15:25:54 -04:00
resp . events . each do | event |
next unless " / #{ event . message } / " . match ( string_match )
msg = event . message . split ( ' -- : ' ) [ 1 ] . green
puts " CloudWatch log stream: #{ msg } "
return msg
2022-10-06 14:10:53 -06:00
end
end
rescue Aws :: Lambda :: Errors :: ServiceException = > e
@logger . error ( " There was an error fetching CloudWatch logs: \n #{ e . message } " )
end
2024-09-25 15:25:54 -04:00
# snippet-end:[ruby.example_code.lambda.GetLogs]
2022-10-06 14:10:53 -06:00
def invoke_and_verify ( function_name , log_statement , payload = nil )
response = invoke_function ( function_name , payload )
print " Done! \n " . green
JSON . pretty_generate ( response ) . green
get_cloudwatch_logs ( function_name , log_statement )
rescue Aws :: Lambda :: Errors :: ServiceException = > e
@logger . error ( " There was an error executing #{ function_name } : \n #{ e . message } " )
end
# snippet-start:[ruby.example_code.lambda.UpdateFunctionConfiguration]
# Updates the environment variables for a Lambda function.
# @param function_name: The name of the function to update.
# @param log_level: The log level of the function.
# @return: Data about the update, including the status.
def update_function_configuration ( function_name , log_level )
@lambda_client . update_function_configuration ( {
2023-06-08 16:39:36 -06:00
function_name : function_name ,
2022-10-06 14:10:53 -06:00
environment : {
variables : {
2024-09-25 15:25:54 -04:00
'LOG_LEVEL' = > log_level
2022-10-06 14:10:53 -06:00
}
}
} )
2024-09-25 15:25:54 -04:00
@lambda_client . wait_until ( :function_updated_v2 , { function_name : function_name } ) do | w |
2022-10-06 14:10:53 -06:00
w . max_attempts = 5
w . delay = 5
end
rescue Aws :: Lambda :: Errors :: ServiceException = > e
@logger . error ( " There was an error updating configurations for #{ function_name } : \n #{ e . message } " )
rescue Aws :: Waiters :: Errors :: WaiterFailed = > e
@logger . error ( " Failed waiting for #{ function_name } to activate: \n #{ e . message } " )
end
# snippet-end:[ruby.example_code.lambda.UpdateFunctionConfiguration]
# snippet-start:[ruby.example_code.lambda.UpdateFunctionCode]
# Updates the code for a Lambda function by submitting a .zip archive that contains
# the code for the function.
2024-09-25 15:25:54 -04:00
#
2022-10-06 14:10:53 -06:00
# @param function_name: The name of the function to update.
# @param deployment_package: The function code to update, packaged as bytes in
# .zip format.
# @return: Data about the update, including the status.
def update_function_code ( function_name , deployment_package )
@lambda_client . update_function_code (
2023-06-08 16:39:36 -06:00
function_name : function_name ,
2022-10-06 14:10:53 -06:00
zip_file : deployment_package
)
2024-09-25 15:25:54 -04:00
@lambda_client . wait_until ( :function_updated_v2 , { function_name : function_name } ) do | w |
2022-10-06 14:10:53 -06:00
w . max_attempts = 5
w . delay = 5
end
rescue Aws :: Lambda :: Errors :: ServiceException = > e
@logger . error ( " There was an error updating function code for: #{ function_name } : \n #{ e . message } " )
nil
rescue Aws :: Waiters :: Errors :: WaiterFailed = > e
@logger . error ( " Failed waiting for #{ function_name } to update: \n #{ e . message } " )
end
# snippet-end:[ruby.example_code.lambda.UpdateFunctionCode]
# snippet-start:[ruby.example_code.lambda.ListFunctions]
# Lists the Lambda functions for the current account.
def list_functions
functions = [ ]
@lambda_client . list_functions . each do | response |
2024-09-25 15:25:54 -04:00
response [ 'functions' ] . each do | function |
functions . append ( function [ 'function_name' ] )
2022-10-06 14:10:53 -06:00
end
end
functions
rescue Aws :: Lambda :: Errors :: ServiceException = > e
2024-09-25 15:25:54 -04:00
@logger . error ( " There was an error listing functions: \n #{ e . message } " )
2022-10-06 14:10:53 -06:00
end
# snippet-end:[ruby.example_code.lambda.ListFunctions]
# snippet-start:[ruby.example_code.lambda.DeleteFunction]
# Deletes a Lambda function.
# @param function_name: The name of the function to delete.
def delete_function ( function_name )
print " Deleting function: #{ function_name } ... "
@lambda_client . delete_function (
2023-06-08 16:39:36 -06:00
function_name : function_name
2022-10-06 14:10:53 -06:00
)
2024-09-25 15:25:54 -04:00
print 'Done!' . green
2022-10-06 14:10:53 -06:00
rescue Aws :: Lambda :: Errors :: ServiceException = > e
@logger . error ( " There was an error deleting #{ function_name } : \n #{ e . message } " )
end
# snippet-end:[ruby.example_code.lambda.DeleteFunction]
end
2024-09-25 15:25:54 -04:00
# snippet-end:[ruby.example_code.ruby.LambdaWrapper.full]
