Last active
January 3, 2020 13:11
-
-
Save antdking/3228bdd0313a78f0859274f8d624f6a1 to your computer and use it in GitHub Desktop.
Manage loading in values from parameter store into the environment
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import logging | |
import os | |
from collections import defaultdict | |
from typing import Dict | |
LOG = logging.getLogger(__name__) | |
class ParameterStoreEnvirontment: | |
detection_prefix = "ssm-pse:" | |
def __init__(self, ssm_client): | |
self.client = ssm_client | |
def detect_and_inject( | |
self, | |
with_decryption: bool = True, | |
default: str = None, | |
injection_target: Dict[str, str] = os.environ, | |
) -> None: | |
"""Scans a dictionary for values that look like SSM PSE values, and loads | |
them in from Parameter Store. | |
""" | |
env_to_path_map = self.detect_env_for_loading(injection_target) | |
new_environment = self.fetch( | |
env_to_path_map, with_decryption=with_decryption, default=default | |
) | |
for key, value in new_environment.items(): | |
injection_target[key] = value | |
return | |
def fetch( | |
self, input: Dict[str, str], with_decryption: bool = True, default: str = None, | |
) -> Dict[str, str]: | |
# We need a list, as multiple keys can map to the same values | |
inverted_input = defaultdict(list) | |
for var, path in input.items(): | |
inverted_input[path].append(var) | |
response = self.client.get_parameters( | |
Names=list(inverted_input.keys()), WithDecryption=with_decryption, | |
) | |
self._validate_response(response, allow_missing=default is not None) | |
return { | |
var: param["Value"] | |
for param in response["Parameters"] | |
for var in inverted_input[param["Name"]] | |
} | |
def detect_env_for_loading( | |
self, detection_target: Dict[str, str] = os.environ, | |
) -> Dict[str, str]: | |
""" | |
Search an object for anything contains our expected prefix, and looks like | |
a Parameter Store path. | |
""" | |
output = {} | |
for key, value in detection_target.items(): | |
if value.startswith(self.detection_prefix): | |
ps_path = value[len(self.detection_prefix) :] | |
# just some extra sanity to make sure it's kinda like a path (PS needs this) | |
if ps_path.startswith("/"): | |
output[key] = ps_path | |
else: | |
LOG.warning( | |
"Environment value detected with PSE Prefix: {}={}".format( | |
key, value | |
) | |
) | |
return output | |
@staticmethod | |
def _validate_response(response, allow_missing: bool) -> None: | |
if response["InvalidParameters"]: | |
LOG.warning( | |
"Invalid Parameters detected: %(parameters)r", | |
{"parameters": response["InvalidParameters"]}, | |
) | |
if default is None: | |
raise ValueError("Invalid Parameters detected without a `default` set.") | |
return |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment