1 import datetime 2 import typing as _t 3 4 from cachelib.base import BaseCache 5 from cachelib.serializers import DynamoDbSerializer 6 7 CREATED_AT_FIELD = "created_at" 8 RESPONSE_FIELD = "response" 9 10 11 class DynamoDbCache(BaseCache): 12 """ 13 Implementation of cachelib.BaseCache that uses an AWS DynamoDb table 14 as the backend. 15 16 Your server process will require dynamodb:GetItem and dynamodb:PutItem 17 IAM permissions on the cache table. 18 19 Limitations: DynamoDB table items are limited to 400 KB in size. Since 20 this class stores cached items in a table, the max size of a cache entry 21 will be slightly less than 400 KB, since the cache key and expiration 22 time fields are also part of the item. 23 24 :param table_name: The name of the DynamoDB table to use 25 :param default_timeout: Set the timeout in seconds after which cache entries 26 expire 27 :param key_field: The name of the hash_key attribute in the DynamoDb 28 table. This must be a string attribute. 29 :param expiration_time_field: The name of the table attribute to store the 30 expiration time in. This will be an int 31 attribute. The timestamp will be stored as 32 seconds past the epoch. If you configure 33 this as the TTL field, then DynamoDB will 34 automatically delete expired entries. 35 :param key_prefix: A prefix that should be added to all keys. 36 37 """ 38 39 serializer = DynamoDbSerializer() 40 41 def __init__( 42 self, 43 table_name: _t.Optional[str] = "python-cache", 44 default_timeout: int = 300, 45 key_field: _t.Optional[str] = "cache_key", 46 expiration_time_field: _t.Optional[str] = "expiration_time", 47 key_prefix: _t.Optional[str] = None, 48 **kwargs: _t.Any 49 ): 50 super().__init__(default_timeout) 51 52 try: 53 import boto3 # type: ignore 54 except ImportError as err: 55 raise RuntimeError("no boto3 module found") from err 56 57 self._table_name = table_name 58 self._key_field = key_field 59 self._expiration_time_field = expiration_time_field 60 self.key_prefix = key_prefix or "" 61 self._dynamo = boto3.resource("dynamodb", **kwargs) 62 self._attr = boto3.dynamodb.conditions.Attr 63 64 try: 65 self._table = self._dynamo.Table(table_name) 66 self._table.load() 67 # catch this exception (triggered if the table doesn't exist) 68 except Exception: 69 table = self._dynamo.create_table( 70 AttributeDefinitions=[ 71 {"AttributeName": key_field, "AttributeType": "S"} 72 ], 73 TableName=table_name, 74 KeySchema=[ 75 {"AttributeName": key_field, "KeyType": "HASH"}, 76 ], 77 BillingMode="PAY_PER_REQUEST", 78 ) 79 table.wait_until_exists() 80 dynamo = boto3.client("dynamodb", **kwargs) 81 dynamo.update_time_to_live( 82 TableName=table_name, 83 TimeToLiveSpecification={ 84 "Enabled": True, 85 "AttributeName": expiration_time_field, 86 }, 87 ) 88 self._table = self._dynamo.Table(table_name) 89 self._table.load() 90 91 def _utcnow(self) -> _t.Any: 92 """Return a tz-aware UTC datetime representing the current time""" 93 return datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc) 94 95 def _get_item(self, key: str, attributes: _t.Optional[list] = None) -> _t.Any: 96 """ 97 Get an item from the cache table, optionally limiting the returned 98 attributes. 99 100 :param key: The cache key of the item to fetch 101 102 :param attributes: An optional list of attributes to fetch. If not 103 given, all attributes are fetched. The 104 expiration_time field will always be added to the 105 list of fetched attributes. 106 :return: The table item for key if it exists and is not expired, else 107 None 108 """ 109 kwargs = {} 110 if attributes: 111 if self._expiration_time_field not in attributes: 112 attributes = list(attributes) + [self._expiration_time_field] 113 kwargs = dict(ProjectionExpression=",".join(attributes)) 114 115 response = self._table.get_item(Key={self._key_field: key}, **kwargs) 116 cache_item = response.get("Item") 117 118 if cache_item: 119 now = int(self._utcnow().timestamp()) 120 if cache_item.get(self._expiration_time_field, now + 100) > now: 121 return cache_item 122 123 return None 124 125 def get(self, key: str) -> _t.Any: 126 """ 127 Get a cache item 128 129 :param key: The cache key of the item to fetch 130 :return: cache value if not expired, else None 131 """ 132 cache_item = self._get_item(self.key_prefix + key) 133 if cache_item: 134 response = cache_item[RESPONSE_FIELD] 135 value = self.serializer.loads(response) 136 return value 137 return None 138 139 def delete(self, key: str) -> bool: 140 """ 141 Deletes an item from the cache. This is a no-op if the item doesn't 142 exist 143 144 :param key: Key of the item to delete. 145 :return: True if the key existed and was deleted 146 """ 147 try: 148 self._table.delete_item( 149 Key={self._key_field: self.key_prefix + key}, 150 ConditionExpression=self._attr(self._key_field).exists(), 151 ) 152 return True 153 except self._dynamo.meta.client.exceptions.ConditionalCheckFailedException: 154 return False 155 156 def _set( 157 self, 158 key: str, 159 value: _t.Any, 160 timeout: _t.Optional[int] = None, 161 overwrite: _t.Optional[bool] = True, 162 ) -> _t.Any: 163 """ 164 Store a cache item, with the option to not overwrite existing items 165 166 :param key: Cache key to use 167 :param value: a serializable object 168 :param timeout: The timeout in seconds for the cached item, to override 169 the default 170 :param overwrite: If true, overwrite any existing cache item with key. 171 If false, the new value will only be stored if no 172 non-expired cache item exists with key. 173 :return: True if the new item was stored. 174 """ 175 timeout = self._normalize_timeout(timeout) 176 now = self._utcnow() 177 178 kwargs = {} 179 if not overwrite: 180 # Cause the put to fail if a non-expired item with this key 181 # already exists 182 183 cond = self._attr(self._key_field).not_exists() | self._attr( 184 self._expiration_time_field 185 ).lte(int(now.timestamp())) 186 kwargs = dict(ConditionExpression=cond) 187 188 try: 189 dump = self.serializer.dumps(value) 190 item = { 191 self._key_field: key, 192 CREATED_AT_FIELD: now.isoformat(), 193 RESPONSE_FIELD: dump, 194 } 195 if timeout > 0: 196 expiration_time = now + datetime.timedelta(seconds=timeout) 197 item[self._expiration_time_field] = int(expiration_time.timestamp()) 198 self._table.put_item(Item=item, **kwargs) 199 return True 200 except Exception: 201 return False 202 203 def set(self, key: str, value: _t.Any, timeout: _t.Optional[int] = None) -> _t.Any: 204 return self._set(self.key_prefix + key, value, timeout=timeout, overwrite=True) 205 206 def add(self, key: str, value: _t.Any, timeout: _t.Optional[int] = None) -> _t.Any: 207 return self._set(self.key_prefix + key, value, timeout=timeout, overwrite=False) 208 209 def has(self, key: str) -> bool: 210 return ( 211 self._get_item(self.key_prefix + key, [self._expiration_time_field]) 212 is not None 213 ) 214 215 def clear(self) -> bool: 216 paginator = self._dynamo.meta.client.get_paginator("scan") 217 pages = paginator.paginate( 218 TableName=self._table_name, ProjectionExpression=self._key_field 219 ) 220 221 with self._table.batch_writer() as batch: 222 for page in pages: 223 for item in page["Items"]: 224 batch.delete_item(Key=item) 225 226 return True 227