From b50cf4fd0b96549fc864d0ec0f5db2fd78fe43ea Mon Sep 17 00:00:00 2001 From: Preshen Goobiah Date: Fri, 5 May 2023 11:21:07 +0200 Subject: [PATCH 1/3] Use redshift_connector module --- .../lambda_function.py | 206 +++++++++++++----- 1 file changed, 150 insertions(+), 56 deletions(-) diff --git a/SecretsManagerRedshiftRotationSingleUser/lambda_function.py b/SecretsManagerRedshiftRotationSingleUser/lambda_function.py index 78ed4037..28553f25 100644 --- a/SecretsManagerRedshiftRotationSingleUser/lambda_function.py +++ b/SecretsManagerRedshiftRotationSingleUser/lambda_function.py @@ -5,8 +5,7 @@ import json import logging import os -import pg -import pgdb +import redshift_connector logger = logging.getLogger() logger.setLevel(logging.INFO) @@ -45,28 +44,41 @@ def lambda_handler(event, context): KeyError: If the secret json does not contain the expected keys """ - arn = event['SecretId'] - token = event['ClientRequestToken'] - step = event['Step'] + logger.info(event) + arn = event["SecretId"] + token = event["ClientRequestToken"] + step = event["Step"] # Setup the client - service_client = boto3.client('secretsmanager', endpoint_url=os.environ['SECRETS_MANAGER_ENDPOINT']) + service_client = boto3.client("secretsmanager") # Make sure the version is staged correctly metadata = service_client.describe_secret(SecretId=arn) - if "RotationEnabled" in metadata and not metadata['RotationEnabled']: + if "RotationEnabled" in metadata and not metadata["RotationEnabled"]: logger.error("Secret %s is not enabled for rotation" % arn) raise ValueError("Secret %s is not enabled for rotation" % arn) - versions = metadata['VersionIdsToStages'] + versions = metadata["VersionIdsToStages"] if token not in versions: - logger.error("Secret version %s has no stage for rotation of secret %s." % (token, arn)) - raise ValueError("Secret version %s has no stage for rotation of secret %s." % (token, arn)) + logger.error( + "Secret version %s has no stage for rotation of secret %s." % (token, arn) + ) + raise ValueError( + "Secret version %s has no stage for rotation of secret %s." % (token, arn) + ) if "AWSCURRENT" in versions[token]: - logger.info("Secret version %s already set as AWSCURRENT for secret %s." % (token, arn)) + logger.info( + "Secret version %s already set as AWSCURRENT for secret %s." % (token, arn) + ) return elif "AWSPENDING" not in versions[token]: - logger.error("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token, arn)) - raise ValueError("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token, arn)) + logger.error( + "Secret version %s not set as AWSPENDING for rotation of secret %s." + % (token, arn) + ) + raise ValueError( + "Secret version %s not set as AWSPENDING for rotation of secret %s." + % (token, arn) + ) # Call the appropriate step if step == "createSecret": @@ -82,7 +94,9 @@ def lambda_handler(event, context): finish_secret(service_client, arn, token) else: - logger.error("lambda_handler: Invalid step parameter %s for secret %s" % (step, arn)) + logger.error( + "lambda_handler: Invalid step parameter %s for secret %s" % (step, arn) + ) raise ValueError("Invalid step parameter %s for secret %s" % (step, arn)) @@ -114,14 +128,28 @@ def create_secret(service_client, arn, token): logger.info("createSecret: Successfully retrieved secret for %s." % arn) except service_client.exceptions.ResourceNotFoundException: # Get exclude characters from environment variable - exclude_characters = os.environ['EXCLUDE_CHARACTERS'] if 'EXCLUDE_CHARACTERS' in os.environ else '/@"\'\\:' + exclude_characters = ( + os.environ["EXCLUDE_CHARACTERS"] + if "EXCLUDE_CHARACTERS" in os.environ + else "/@\"'\\:" + ) # Generate a random password - passwd = service_client.get_random_password(ExcludeCharacters=exclude_characters) - current_dict['password'] = passwd['RandomPassword'] + passwd = service_client.get_random_password( + ExcludeCharacters=exclude_characters + ) + current_dict["password"] = passwd["RandomPassword"] # Put the secret - service_client.put_secret_value(SecretId=arn, ClientRequestToken=token, SecretString=json.dumps(current_dict), VersionStages=['AWSPENDING']) - logger.info("createSecret: Successfully put secret for ARN %s and version %s." % (arn, token)) + service_client.put_secret_value( + SecretId=arn, + ClientRequestToken=token, + SecretString=json.dumps(current_dict), + VersionStages=["AWSPENDING"], + ) + logger.info( + "createSecret: Successfully put secret for ARN %s and version %s." + % (arn, token) + ) def set_secret(service_client, arn, token): @@ -157,18 +185,33 @@ def set_secret(service_client, arn, token): conn = get_connection(pending_dict) if conn: conn.close() - logger.info("setSecret: AWSPENDING secret is already set as password in Redshift DB for secret arn %s." % arn) + logger.info( + "setSecret: AWSPENDING secret is already set as password in Redshift DB for secret arn %s." + % arn + ) return # Make sure the user from current and pending match - if current_dict['username'] != pending_dict['username']: - logger.error("setSecret: Attempting to modify user %s other than current user %s" % (pending_dict['username'], current_dict['username'])) - raise ValueError("Attempting to modify user %s other than current user %s" % (pending_dict['username'], current_dict['username'])) + if current_dict["username"] != pending_dict["username"]: + logger.error( + "setSecret: Attempting to modify user %s other than current user %s" + % (pending_dict["username"], current_dict["username"]) + ) + raise ValueError( + "Attempting to modify user %s other than current user %s" + % (pending_dict["username"], current_dict["username"]) + ) # Make sure the host from current and pending match - if current_dict['host'] != pending_dict['host']: - logger.error("setSecret: Attempting to modify user for host %s other than current host %s" % (pending_dict['host'], current_dict['host'])) - raise ValueError("Attempting to modify user for host %s other than current host %s" % (pending_dict['host'], current_dict['host'])) + if current_dict["host"] != pending_dict["host"]: + logger.error( + "setSecret: Attempting to modify user for host %s other than current host %s" + % (pending_dict["host"], current_dict["host"]) + ) + raise ValueError( + "Attempting to modify user for host %s other than current host %s" + % (pending_dict["host"], current_dict["host"]) + ) # Now try the current password conn = get_connection(current_dict) @@ -177,29 +220,52 @@ def set_secret(service_client, arn, token): conn = get_connection(previous_dict) # Make sure the user/host from current and pending match - if previous_dict['username'] != pending_dict['username']: - logger.error("setSecret: Attempting to modify user %s other than previous valid user %s" % (pending_dict['username'], previous_dict['username'])) - raise ValueError("Attempting to modify user %s other than previous valid user %s" % (pending_dict['username'], previous_dict['username'])) - if previous_dict['host'] != pending_dict['host']: - logger.error("setSecret: Attempting to modify user for host %s other than previous host %s" % (pending_dict['host'], previous_dict['host'])) - raise ValueError("Attempting to modify user for host %s other than previous host %s" % (pending_dict['host'], previous_dict['host'])) + if previous_dict["username"] != pending_dict["username"]: + logger.error( + "setSecret: Attempting to modify user %s other than previous valid user %s" + % (pending_dict["username"], previous_dict["username"]) + ) + raise ValueError( + "Attempting to modify user %s other than previous valid user %s" + % (pending_dict["username"], previous_dict["username"]) + ) + if previous_dict["host"] != pending_dict["host"]: + logger.error( + "setSecret: Attempting to modify user for host %s other than previous host %s" + % (pending_dict["host"], previous_dict["host"]) + ) + raise ValueError( + "Attempting to modify user for host %s other than previous host %s" + % (pending_dict["host"], previous_dict["host"]) + ) # If we still don't have a connection, raise a ValueError if not conn: - logger.error("setSecret: Unable to log into database with previous, current, or pending secret of secret arn %s" % arn) - raise ValueError("Unable to log into database with previous, current, or pending secret of secret arn %s" % arn) + logger.error( + "setSecret: Unable to log into database with previous, current, or pending secret of secret arn %s" + % arn + ) + raise ValueError( + "Unable to log into database with previous, current, or pending secret of secret arn %s" + % arn + ) # Now set the password to the pending password try: with conn.cursor() as cur: + # Get escaped username via quote_ident - cur.execute("SELECT quote_ident(%s)", (pending_dict['username'],)) + cur.execute("SELECT quote_ident(%s)", (pending_dict["username"],)) escaped_username = cur.fetchone()[0] + password = pending_dict["password"] - alter_role = "ALTER USER %s" % escaped_username - cur.execute(alter_role + " WITH PASSWORD %s", (pending_dict['password'],)) + cur.execute(f"alter user {escaped_username} password '{password}'") conn.commit() - logger.info("setSecret: Successfully set password for user %s in Redshift DB for secret arn %s." % (pending_dict['username'], arn)) + + logger.info( + "setSecret: Successfully set password for user %s in Redshift DB for secret arn %s." + % (pending_dict["username"], arn) + ) finally: conn.close() @@ -237,11 +303,19 @@ def test_secret(service_client, arn, token): finally: conn.close() - logger.info("testSecret: Successfully signed into Redshift DB with AWSPENDING secret in %s." % arn) + logger.info( + "testSecret: Successfully signed into Redshift DB with AWSPENDING secret in %s." + % arn + ) return else: - logger.error("testSecret: Unable to log into database with pending secret of secret ARN %s" % arn) - raise ValueError("Unable to log into database with pending secret of secret ARN %s" % arn) + logger.error( + "testSecret: Unable to log into database with pending secret of secret ARN %s" + % arn + ) + raise ValueError( + "Unable to log into database with pending secret of secret ARN %s" % arn + ) def finish_secret(service_client, arn, token): @@ -264,14 +338,25 @@ def finish_secret(service_client, arn, token): if "AWSCURRENT" in metadata["VersionIdsToStages"][version]: if version == token: # The correct version is already marked as current, return - logger.info("finishSecret: Version %s already marked as AWSCURRENT for %s" % (version, arn)) + logger.info( + "finishSecret: Version %s already marked as AWSCURRENT for %s" + % (version, arn) + ) return current_version = version break # Finalize by staging the secret version current - service_client.update_secret_version_stage(SecretId=arn, VersionStage="AWSCURRENT", MoveToVersionId=token, RemoveFromVersionId=current_version) - logger.info("finishSecret: Successfully set AWSCURRENT stage to version %s for secret %s." % (token, arn)) + service_client.update_secret_version_stage( + SecretId=arn, + VersionStage="AWSCURRENT", + MoveToVersionId=token, + RemoveFromVersionId=current_version, + ) + logger.info( + "finishSecret: Successfully set AWSCURRENT stage to version %s for secret %s." + % (token, arn) + ) def get_connection(secret_dict): @@ -284,22 +369,32 @@ def get_connection(secret_dict): secret_dict (dict): The Secret Dictionary Returns: - Connection: The pgdb.Connection object if successful. None otherwise + Connection: The Redshift Connection object if successful. Raise error otherwise. Raises: KeyError: If the secret json does not contain the expected keys """ # Parse and validate the secret JSON string - port = int(secret_dict['port']) if 'port' in secret_dict else 5439 - dbname = secret_dict['dbname'] if 'dbname' in secret_dict else "dev" + port = int(secret_dict["port"]) if "port" in secret_dict else 5439 + dbname = secret_dict["dbname"] if "dbname" in secret_dict else "dev" # Try to obtain a connection to the db try: - conn = pgdb.connect(host=secret_dict['host'], user=secret_dict['username'], password=secret_dict['password'], database=dbname, port=port, connect_timeout=5) - logger.info("Successfully established connection as user '%s' with host: '%s'" % (secret_dict['username'], secret_dict['host'])) + conn = redshift_connector.connect( + host=secret_dict["host"], + database=dbname, + port=port, + user=secret_dict["username"], + password=secret_dict["password"], + ) + logger.info( + "Successfully established connection as user '%s' with host: '%s'" + % (secret_dict["username"], secret_dict["host"]) + ) return conn - except pg.InternalError: + except Exception as e: + logger.error(e) return None @@ -326,19 +421,18 @@ def get_secret_dict(service_client, arn, stage, token=None): ValueError: If the secret is not valid JSON """ - required_fields = ['host', 'username', 'password'] + required_fields = ["host", "username", "password"] # Only do VersionId validation against the stage if a token is passed in if token: - secret = service_client.get_secret_value(SecretId=arn, VersionId=token, VersionStage=stage) + secret = service_client.get_secret_value( + SecretId=arn, VersionId=token, VersionStage=stage + ) else: secret = service_client.get_secret_value(SecretId=arn, VersionStage=stage) - plaintext = secret['SecretString'] + plaintext = secret["SecretString"] secret_dict = json.loads(plaintext) - # Run validations against the secret - if 'engine' not in secret_dict or secret_dict['engine'] != 'redshift': - raise KeyError("Database engine must be set to 'redshift' in order to use this rotation lambda") for field in required_fields: if field not in secret_dict: raise KeyError("%s key is missing from secret JSON" % field) From d0499caea53ca8f03b9652afceda2b5eea22a907 Mon Sep 17 00:00:00 2001 From: Preshen Goobiah Date: Fri, 5 May 2023 11:29:04 +0200 Subject: [PATCH 2/3] remove formatting --- .../lambda_function.py | 177 +++++------------- 1 file changed, 49 insertions(+), 128 deletions(-) diff --git a/SecretsManagerRedshiftRotationSingleUser/lambda_function.py b/SecretsManagerRedshiftRotationSingleUser/lambda_function.py index 28553f25..93844d0f 100644 --- a/SecretsManagerRedshiftRotationSingleUser/lambda_function.py +++ b/SecretsManagerRedshiftRotationSingleUser/lambda_function.py @@ -44,41 +44,28 @@ def lambda_handler(event, context): KeyError: If the secret json does not contain the expected keys """ - logger.info(event) - arn = event["SecretId"] - token = event["ClientRequestToken"] - step = event["Step"] + arn = event['SecretId'] + token = event['ClientRequestToken'] + step = event['Step'] # Setup the client - service_client = boto3.client("secretsmanager") + service_client = boto3.client('secretsmanager', endpoint_url=os.environ['SECRETS_MANAGER_ENDPOINT']) # Make sure the version is staged correctly metadata = service_client.describe_secret(SecretId=arn) - if "RotationEnabled" in metadata and not metadata["RotationEnabled"]: + if "RotationEnabled" in metadata and not metadata['RotationEnabled']: logger.error("Secret %s is not enabled for rotation" % arn) raise ValueError("Secret %s is not enabled for rotation" % arn) - versions = metadata["VersionIdsToStages"] + versions = metadata['VersionIdsToStages'] if token not in versions: - logger.error( - "Secret version %s has no stage for rotation of secret %s." % (token, arn) - ) - raise ValueError( - "Secret version %s has no stage for rotation of secret %s." % (token, arn) - ) + logger.error("Secret version %s has no stage for rotation of secret %s." % (token, arn)) + raise ValueError("Secret version %s has no stage for rotation of secret %s." % (token, arn)) if "AWSCURRENT" in versions[token]: - logger.info( - "Secret version %s already set as AWSCURRENT for secret %s." % (token, arn) - ) + logger.info("Secret version %s already set as AWSCURRENT for secret %s." % (token, arn)) return elif "AWSPENDING" not in versions[token]: - logger.error( - "Secret version %s not set as AWSPENDING for rotation of secret %s." - % (token, arn) - ) - raise ValueError( - "Secret version %s not set as AWSPENDING for rotation of secret %s." - % (token, arn) - ) + logger.error("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token, arn)) + raise ValueError("Secret version %s not set as AWSPENDING for rotation of secret %s." % (token, arn)) # Call the appropriate step if step == "createSecret": @@ -94,9 +81,7 @@ def lambda_handler(event, context): finish_secret(service_client, arn, token) else: - logger.error( - "lambda_handler: Invalid step parameter %s for secret %s" % (step, arn) - ) + logger.error("lambda_handler: Invalid step parameter %s for secret %s" % (step, arn)) raise ValueError("Invalid step parameter %s for secret %s" % (step, arn)) @@ -128,28 +113,14 @@ def create_secret(service_client, arn, token): logger.info("createSecret: Successfully retrieved secret for %s." % arn) except service_client.exceptions.ResourceNotFoundException: # Get exclude characters from environment variable - exclude_characters = ( - os.environ["EXCLUDE_CHARACTERS"] - if "EXCLUDE_CHARACTERS" in os.environ - else "/@\"'\\:" - ) + exclude_characters = os.environ['EXCLUDE_CHARACTERS'] if 'EXCLUDE_CHARACTERS' in os.environ else '/@"\'\\:' # Generate a random password - passwd = service_client.get_random_password( - ExcludeCharacters=exclude_characters - ) - current_dict["password"] = passwd["RandomPassword"] + passwd = service_client.get_random_password(ExcludeCharacters=exclude_characters) + current_dict['password'] = passwd['RandomPassword'] # Put the secret - service_client.put_secret_value( - SecretId=arn, - ClientRequestToken=token, - SecretString=json.dumps(current_dict), - VersionStages=["AWSPENDING"], - ) - logger.info( - "createSecret: Successfully put secret for ARN %s and version %s." - % (arn, token) - ) + service_client.put_secret_value(SecretId=arn, ClientRequestToken=token, SecretString=json.dumps(current_dict), VersionStages=['AWSPENDING']) + logger.info("createSecret: Successfully put secret for ARN %s and version %s." % (arn, token)) def set_secret(service_client, arn, token): @@ -185,33 +156,18 @@ def set_secret(service_client, arn, token): conn = get_connection(pending_dict) if conn: conn.close() - logger.info( - "setSecret: AWSPENDING secret is already set as password in Redshift DB for secret arn %s." - % arn - ) + logger.info("setSecret: AWSPENDING secret is already set as password in Redshift DB for secret arn %s." % arn) return # Make sure the user from current and pending match - if current_dict["username"] != pending_dict["username"]: - logger.error( - "setSecret: Attempting to modify user %s other than current user %s" - % (pending_dict["username"], current_dict["username"]) - ) - raise ValueError( - "Attempting to modify user %s other than current user %s" - % (pending_dict["username"], current_dict["username"]) - ) + if current_dict['username'] != pending_dict['username']: + logger.error("setSecret: Attempting to modify user %s other than current user %s" % (pending_dict['username'], current_dict['username'])) + raise ValueError("Attempting to modify user %s other than current user %s" % (pending_dict['username'], current_dict['username'])) # Make sure the host from current and pending match - if current_dict["host"] != pending_dict["host"]: - logger.error( - "setSecret: Attempting to modify user for host %s other than current host %s" - % (pending_dict["host"], current_dict["host"]) - ) - raise ValueError( - "Attempting to modify user for host %s other than current host %s" - % (pending_dict["host"], current_dict["host"]) - ) + if current_dict['host'] != pending_dict['host']: + logger.error("setSecret: Attempting to modify user for host %s other than current host %s" % (pending_dict['host'], current_dict['host'])) + raise ValueError("Attempting to modify user for host %s other than current host %s" % (pending_dict['host'], current_dict['host'])) # Now try the current password conn = get_connection(current_dict) @@ -220,40 +176,21 @@ def set_secret(service_client, arn, token): conn = get_connection(previous_dict) # Make sure the user/host from current and pending match - if previous_dict["username"] != pending_dict["username"]: - logger.error( - "setSecret: Attempting to modify user %s other than previous valid user %s" - % (pending_dict["username"], previous_dict["username"]) - ) - raise ValueError( - "Attempting to modify user %s other than previous valid user %s" - % (pending_dict["username"], previous_dict["username"]) - ) - if previous_dict["host"] != pending_dict["host"]: - logger.error( - "setSecret: Attempting to modify user for host %s other than previous host %s" - % (pending_dict["host"], previous_dict["host"]) - ) - raise ValueError( - "Attempting to modify user for host %s other than previous host %s" - % (pending_dict["host"], previous_dict["host"]) - ) + if previous_dict['username'] != pending_dict['username']: + logger.error("setSecret: Attempting to modify user %s other than previous valid user %s" % (pending_dict['username'], previous_dict['username'])) + raise ValueError("Attempting to modify user %s other than previous valid user %s" % (pending_dict['username'], previous_dict['username'])) + if previous_dict['host'] != pending_dict['host']: + logger.error("setSecret: Attempting to modify user for host %s other than previous host %s" % (pending_dict['host'], previous_dict['host'])) + raise ValueError("Attempting to modify user for host %s other than previous host %s" % (pending_dict['host'], previous_dict['host'])) # If we still don't have a connection, raise a ValueError if not conn: - logger.error( - "setSecret: Unable to log into database with previous, current, or pending secret of secret arn %s" - % arn - ) - raise ValueError( - "Unable to log into database with previous, current, or pending secret of secret arn %s" - % arn - ) + logger.error("setSecret: Unable to log into database with previous, current, or pending secret of secret arn %s" % arn) + raise ValueError("Unable to log into database with previous, current, or pending secret of secret arn %s" % arn) # Now set the password to the pending password try: with conn.cursor() as cur: - # Get escaped username via quote_ident cur.execute("SELECT quote_ident(%s)", (pending_dict["username"],)) escaped_username = cur.fetchone()[0] @@ -270,6 +207,7 @@ def set_secret(service_client, arn, token): conn.close() + def test_secret(service_client, arn, token): """Test the pending secret against the database @@ -303,19 +241,11 @@ def test_secret(service_client, arn, token): finally: conn.close() - logger.info( - "testSecret: Successfully signed into Redshift DB with AWSPENDING secret in %s." - % arn - ) + logger.info("testSecret: Successfully signed into Redshift DB with AWSPENDING secret in %s." % arn) return else: - logger.error( - "testSecret: Unable to log into database with pending secret of secret ARN %s" - % arn - ) - raise ValueError( - "Unable to log into database with pending secret of secret ARN %s" % arn - ) + logger.error("testSecret: Unable to log into database with pending secret of secret ARN %s" % arn) + raise ValueError("Unable to log into database with pending secret of secret ARN %s" % arn) def finish_secret(service_client, arn, token): @@ -338,25 +268,14 @@ def finish_secret(service_client, arn, token): if "AWSCURRENT" in metadata["VersionIdsToStages"][version]: if version == token: # The correct version is already marked as current, return - logger.info( - "finishSecret: Version %s already marked as AWSCURRENT for %s" - % (version, arn) - ) + logger.info("finishSecret: Version %s already marked as AWSCURRENT for %s" % (version, arn)) return current_version = version break # Finalize by staging the secret version current - service_client.update_secret_version_stage( - SecretId=arn, - VersionStage="AWSCURRENT", - MoveToVersionId=token, - RemoveFromVersionId=current_version, - ) - logger.info( - "finishSecret: Successfully set AWSCURRENT stage to version %s for secret %s." - % (token, arn) - ) + service_client.update_secret_version_stage(SecretId=arn, VersionStage="AWSCURRENT", MoveToVersionId=token, RemoveFromVersionId=current_version) + logger.info("finishSecret: Successfully set AWSCURRENT stage to version %s for secret %s." % (token, arn)) def get_connection(secret_dict): @@ -369,15 +288,15 @@ def get_connection(secret_dict): secret_dict (dict): The Secret Dictionary Returns: - Connection: The Redshift Connection object if successful. Raise error otherwise. + Connection: The redshift_connector.Connection object if successful. None otherwise Raises: KeyError: If the secret json does not contain the expected keys """ # Parse and validate the secret JSON string - port = int(secret_dict["port"]) if "port" in secret_dict else 5439 - dbname = secret_dict["dbname"] if "dbname" in secret_dict else "dev" + port = int(secret_dict['port']) if 'port' in secret_dict else 5439 + dbname = secret_dict['dbname'] if 'dbname' in secret_dict else "dev" # Try to obtain a connection to the db try: @@ -398,6 +317,7 @@ def get_connection(secret_dict): return None + def get_secret_dict(service_client, arn, stage, token=None): """Gets the secret dictionary corresponding for the secret arn, stage, and token @@ -421,18 +341,19 @@ def get_secret_dict(service_client, arn, stage, token=None): ValueError: If the secret is not valid JSON """ - required_fields = ["host", "username", "password"] + required_fields = ['host', 'username', 'password'] # Only do VersionId validation against the stage if a token is passed in if token: - secret = service_client.get_secret_value( - SecretId=arn, VersionId=token, VersionStage=stage - ) + secret = service_client.get_secret_value(SecretId=arn, VersionId=token, VersionStage=stage) else: secret = service_client.get_secret_value(SecretId=arn, VersionStage=stage) - plaintext = secret["SecretString"] + plaintext = secret['SecretString'] secret_dict = json.loads(plaintext) + # Run validations against the secret + if 'engine' not in secret_dict or secret_dict['engine'] != 'redshift': + raise KeyError("Database engine must be set to 'redshift' in order to use this rotation lambda") for field in required_fields: if field not in secret_dict: raise KeyError("%s key is missing from secret JSON" % field) From d715fe8de3e2a55c01762be2abe53623c203cccf Mon Sep 17 00:00:00 2001 From: Preshen Goobiah Date: Fri, 5 May 2023 12:19:36 +0200 Subject: [PATCH 3/3] use specific exception --- SecretsManagerRedshiftRotationSingleUser/lambda_function.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/SecretsManagerRedshiftRotationSingleUser/lambda_function.py b/SecretsManagerRedshiftRotationSingleUser/lambda_function.py index 93844d0f..377731d6 100644 --- a/SecretsManagerRedshiftRotationSingleUser/lambda_function.py +++ b/SecretsManagerRedshiftRotationSingleUser/lambda_function.py @@ -312,8 +312,8 @@ def get_connection(secret_dict): % (secret_dict["username"], secret_dict["host"]) ) return conn - except Exception as e: - logger.error(e) + except redshift_connector.Error as e: + logger.info("Redshift Connection attempt failed with:\n %s", e) return None