add caching for ssm parameters and datadog values

almdrlib/environment.py

@@ -46,16 +46,19 @@
 import boto3
 import os
 import botocore
+import cachetools
+import logging
+
+
+logger = logging.getLogger(__name__)
 
 
 class AlEnvException(Exception):
     pass
 
-
 class AlEnvAwsConfigurationException(AlEnvException):
     pass
 
-
 class AlEnvConfigurationTableUnavailableException(AlEnvException):
     pass
 
@@ -64,32 +67,23 @@ class AlmdrlibSourceNotEnabledError(AlEnvException):
 
 
 class AlEnv:
+
+    dynamodb = None
+    ssm = None
+    table = None
+    table_name = None
+
     def __init__(self, application_name, client=None, source=("dynamodb",)):
         self.application_name = application_name
         self.client = client
         self.source = source if type(source) in [tuple, list] else (source,)
-        self.region = AlEnv._get_region()
-        self.stack_name = AlEnv._get_stack_name()
-        self.table_name = AlEnv._table_name(self.region, self.stack_name)
-        try:
-            self.dynamodb = boto3.resource('dynamodb')
-            self.ssm = boto3.client('ssm')
-        except (botocore.exceptions.NoRegionError, botocore.exceptions.NoCredentialsError) as e:
-            raise AlEnvAwsConfigurationException(f'Please validate your AWS configuration') from e
-        if "dynamodb" in source:
-            try:
-                self.table = self.dynamodb.Table(self.table_name)
-                self._table_date_time = self.table.creation_date_time
-            except botocore.exceptions.ClientError as e:
-                if e.response['Error']['Code'] == 'ResourceNotFoundException':
-                    raise AlEnvConfigurationTableUnavailableException(self.table_name)
-                else:
-                    raise AlEnvException() from e
+
+        self._setup_aws_resources(source)
 
     def get(self, key, default=None, format='decoded', type=None):
         if "dynamodb" not in self.source:
             raise AlmdrlibSourceNotEnabledError("dynamodb is not enabled for this environment")
-        fetched_value = self.table.get_item(Key={"key": self._make_ddb_key(key)}).get('Item', {}).get('value')
+        fetched_value = AlEnv._get_cached_dynamodb_item(self._make_ddb_key(key))
         converted = AlEnv._convert(fetched_value, format, type)
         if converted is not None:
             return converted
@@ -99,16 +93,10 @@ def get(self, key, default=None, format='decoded', type=None):
     def get_parameter(self, key, default=None, decrypt=False):
         if "ssm" not in self.source:
             raise AlmdrlibSourceNotEnabledError("ssm is not enabled for this environment")
-        try:
-            parameter = self.ssm.get_parameter(Name=self._make_ssm_key(key), WithDecryption=decrypt)
-        except self.ssm.exceptions.ParameterNotFound:
-            return default
-        except botocore.exceptions.ClientError as e:
-            raise AlEnvException() from e
-        return parameter["Parameter"]["Value"]
+        return AlEnv._get_cached_ssm_parameter(self._make_ssm_key(key), default, decrypt)
 
     def _make_ssm_key(self, option_key):
-        return f"/deployments/{self.stack_name}/{self._get_region()}/env-settings/{self.application_name}/{self._make_client_option_key(option_key)}"
+        return f"/deployments/{self._get_stack_name()}/{self._get_region()}/env-settings/{self.application_name}/{self._make_client_option_key(option_key)}"
 
     def _make_ddb_key(self, option_key):
         return f"{self.application_name}.{self._make_client_option_key(option_key)}"
@@ -118,7 +106,50 @@ def _make_client_option_key(self, option_key):
             return option_key
         else:
             return f"{self.client}.{option_key}"
+
+    @staticmethod
+    @cachetools.cached(cache=cachetools.TTLCache(maxsize=16, ttl=3600))
+    def _get_cached_ssm_parameter(ssm_key, default=None, decrypt=False):
+        try:
+            logger.debug(f"Query SSM for parameter {ssm_key}")
+            parameter = AlEnv.ssm.get_parameter(Name=ssm_key, WithDecryption=decrypt)
+        except AlEnv.ssm.exceptions.ParameterNotFound:
+            return default
+        except botocore.exceptions.ClientError as e:
+            raise AlEnvException() from e
+        return parameter["Parameter"]["Value"]
+
+    @staticmethod
+    @cachetools.cached(cache=cachetools.TTLCache(maxsize=16, ttl=3600))
+    def _get_cached_dynamodb_item(key):
+        logger.debug(f"Query DynamoDB table {AlEnv.table_name} for item {key}")
+        return AlEnv.table.get_item(Key={"key": key}).get('Item', {}).get('value')
+
+    @staticmethod
+    def _setup_aws_resources(source):
+        try:
+            if "dynamodb" in source:
+                if AlEnv.dynamodb is None:
+                    AlEnv.dynamodb = boto3.resource('dynamodb')
+                if AlEnv.table is None:
+                    AlEnv._setup_dynamodb_table()
+            AlEnv.ssm = boto3.client('ssm')
+        except (botocore.exceptions.NoRegionError, botocore.exceptions.NoCredentialsError) as e:
+            raise AlEnvAwsConfigurationException(f'Please validate your AWS configuration') from e
 
+    @staticmethod
+    def _setup_dynamodb_table():        
+        region = AlEnv._get_region()
+        stack_name = AlEnv._get_stack_name()
+        AlEnv.table_name = AlEnv._table_name(region, stack_name)
+        try:
+            AlEnv.table = AlEnv.dynamodb.Table(AlEnv.table_name)
+        except botocore.exceptions.ClientError as e:
+            if e.response['Error']['Code'] == 'ResourceNotFoundException':
+                raise AlEnvConfigurationTableUnavailableException(AlEnv.table_name)
+            else:
+                raise AlEnvException() from e
+
     @staticmethod
     def _convert(value, format, type):
         if format == 'raw' and value:

requirements.txt

@@ -6,6 +6,7 @@ pyyaml==5.4.1
 jsonschema[format_nongpl]==3.2.0
 m2r2==0.3.2
 boto3>=1.16.57
+cachetools>=5.3.0
 markupsafe==2.0.1
 sphinxcontrib-contentui
 sphinx_paramlinks

tests/test_environment.py

@@ -61,15 +61,14 @@ def get_parameter(**kwargs):
         else:
             raise MockBotoSSM.exceptions.ParameterNotFound()
 
-
 class TestAlEnv(unittest.TestCase):
     def test_something(self):
         boto3.resource = MagicMock(return_value=MockBotoDDB)
         boto3.client = MagicMock(return_value=MockBotoSSM)
         os.environ['ALERTLOGIC_STACK_REGION'] = 'us-west-1'
         os.environ['ALERTLOGIC_STACK_NAME'] = 'production'
         env = AlEnv("someapplication", source="dynamodb")
-        assert env.table_name == 'us-west-1.production.dev.global.settings'
+        assert AlEnv.table_name == 'us-west-1.production.dev.global.settings'
         assert env.get("strkey") == 'strvalue'
         assert env.get("strkey", format='raw') == '"strvalue"'
         assert env.get("intkey") == '1'