|
| 1 | +import asyncio |
| 2 | +import json |
| 3 | +from contextlib import AsyncExitStack |
| 4 | +from typing import Any, Mapping, Optional, cast |
| 5 | + |
| 6 | +from ldclient.impl.util import log |
| 7 | +from ldclient.interfaces import AsyncFeatureStoreCore, DiagnosticDescription |
| 8 | +from ldclient.versioned_data_kind import VersionedDataKind |
| 9 | + |
| 10 | +have_aioboto3 = False |
| 11 | +try: |
| 12 | + import aioboto3 |
| 13 | + |
| 14 | + have_aioboto3 = True |
| 15 | +except ImportError: |
| 16 | + pass |
| 17 | + |
| 18 | + |
| 19 | +# |
| 20 | +# Internal implementation of the async DynamoDB feature store. |
| 21 | +# |
| 22 | +# Implementation notes: |
| 23 | +# |
| 24 | +# * Feature flags, segments, and any other kind of entity are all put in the same table. The two |
| 25 | +# required attributes are "key" (present in all storeable entities) and "namespace" (used to |
| 26 | +# disambiguate between flags and segments). |
| 27 | +# |
| 28 | +# * Because of DynamoDB's restrictions on attribute values (e.g. empty strings are not allowed), the |
| 29 | +# standard DynamoDB marshaling with one attribute per object property is not used. Instead, the |
| 30 | +# entire object is serialized to JSON and stored in a single attribute, "item". The "version" |
| 31 | +# property is also stored as a separate attribute since it is used for updates. |
| 32 | +# |
| 33 | +# * Since DynamoDB has no transactions, init() - which replaces the entire data store - is not |
| 34 | +# atomic, so there can be a race condition if another process is adding new data via upsert(). To |
| 35 | +# minimize this, we do not delete all the data at the start; instead, we update the items we have |
| 36 | +# received, and then delete all other items. That could delete new data from another process, but |
| 37 | +# that would happen anyway if the init() ran later than the upsert(); we rely on the fact that the |
| 38 | +# process that did the init() will normally receive the new data shortly and do its own upsert(). |
| 39 | +# |
| 40 | +# * DynamoDB has a maximum item size of 400KB. Since each feature flag or user segment is stored as |
| 41 | +# a single item, this mechanism will not work for extremely large flags or segments. |
| 42 | +# |
| 43 | +# * aioboto3 clients are async context managers, so unlike the synchronous boto3 client they cannot |
| 44 | +# be created in __init__. The client is created and entered lazily on first use inside the running |
| 45 | +# event loop, kept for the lifetime of the store, and released in close(). |
| 46 | +# |
| 47 | + |
| 48 | + |
| 49 | +class _AsyncDynamoDBFeatureStoreCore(DiagnosticDescription, AsyncFeatureStoreCore): |
| 50 | + PARTITION_KEY = 'namespace' |
| 51 | + SORT_KEY = 'key' |
| 52 | + VERSION_ATTRIBUTE = 'version' |
| 53 | + ITEM_JSON_ATTRIBUTE = 'item' |
| 54 | + |
| 55 | + def __init__(self, table_name: str, prefix: Optional[str], dynamodb_opts: Mapping[str, Any]): |
| 56 | + if not have_aioboto3: |
| 57 | + raise NotImplementedError("Cannot use async DynamoDB feature store because aioboto3 package is not installed") |
| 58 | + self._table_name = table_name |
| 59 | + self._prefix = (prefix + ":") if prefix else "" |
| 60 | + self._dynamodb_opts = dict(dynamodb_opts) |
| 61 | + self._session = aioboto3.Session() |
| 62 | + self._exit_stack = AsyncExitStack() |
| 63 | + self._client: Optional[Any] = None |
| 64 | + self._closed = False |
| 65 | + # Guards lazy client creation and close so they cannot interleave: a |
| 66 | + # close must not race a client that is still being created. |
| 67 | + self._client_lock = asyncio.Lock() |
| 68 | + |
| 69 | + async def _get_client(self) -> Any: |
| 70 | + if self._client is not None: |
| 71 | + return self._client |
| 72 | + async with self._client_lock: |
| 73 | + if self._closed: |
| 74 | + raise RuntimeError("DynamoDB feature store is closed") |
| 75 | + if self._client is None: |
| 76 | + self._client = await self._exit_stack.enter_async_context(self._session.client('dynamodb', **self._dynamodb_opts)) |
| 77 | + return self._client |
| 78 | + |
| 79 | + async def is_available(self) -> bool: |
| 80 | + try: |
| 81 | + inited_key = self._inited_key() |
| 82 | + client = await self._get_client() |
| 83 | + await self._get_item_by_keys(client, inited_key, inited_key) |
| 84 | + return True |
| 85 | + except BaseException: |
| 86 | + return False |
| 87 | + |
| 88 | + async def init_internal(self, all_data: Mapping[VersionedDataKind, Mapping[str, dict]]) -> None: |
| 89 | + client = await self._get_client() |
| 90 | + # Start by reading the existing keys; we will later delete any of these that were not in all_data. |
| 91 | + unused_old_keys = await self._read_existing_keys(client, all_data.keys()) |
| 92 | + requests = [] |
| 93 | + num_items = 0 |
| 94 | + inited_key = self._inited_key() |
| 95 | + |
| 96 | + # Insert or update every provided item |
| 97 | + for kind, items in all_data.items(): |
| 98 | + for key, item in items.items(): |
| 99 | + encoded_item = self._marshal_item(kind, item) |
| 100 | + requests.append({'PutRequest': {'Item': encoded_item}}) |
| 101 | + combined_key = (self._namespace_for_kind(kind), key) |
| 102 | + unused_old_keys.discard(combined_key) |
| 103 | + num_items = num_items + 1 |
| 104 | + |
| 105 | + # Now delete any previously existing items whose keys were not in the current data |
| 106 | + for combined_key in unused_old_keys: |
| 107 | + if combined_key[0] != inited_key: |
| 108 | + requests.append({'DeleteRequest': {'Key': self._make_keys(combined_key[0], combined_key[1])}}) |
| 109 | + |
| 110 | + # Now set the special key that we check in initialized_internal() |
| 111 | + requests.append({'PutRequest': {'Item': self._make_keys(inited_key, inited_key)}}) |
| 112 | + |
| 113 | + await _AsyncDynamoDBHelpers.batch_write_requests(client, self._table_name, requests) |
| 114 | + log.info('Initialized table %s with %d items', self._table_name, num_items) |
| 115 | + |
| 116 | + async def get_internal(self, kind: VersionedDataKind, key: str) -> Optional[dict]: |
| 117 | + client = await self._get_client() |
| 118 | + resp = await self._get_item_by_keys(client, self._namespace_for_kind(kind), key) |
| 119 | + return self._unmarshal_item(resp.get('Item')) |
| 120 | + |
| 121 | + async def get_all_internal(self, kind: VersionedDataKind) -> Mapping[str, dict]: |
| 122 | + client = await self._get_client() |
| 123 | + items_out = {} |
| 124 | + paginator = client.get_paginator('query') |
| 125 | + async for resp in paginator.paginate(**self._make_query_for_kind(kind)): |
| 126 | + for item in resp['Items']: |
| 127 | + # Every stored item carries the JSON attribute, so _unmarshal_item never returns None here. |
| 128 | + item_out = cast(dict, self._unmarshal_item(item)) |
| 129 | + items_out[item_out['key']] = item_out |
| 130 | + return items_out |
| 131 | + |
| 132 | + async def upsert_internal(self, kind: VersionedDataKind, item: dict) -> dict: |
| 133 | + client = await self._get_client() |
| 134 | + encoded_item = self._marshal_item(kind, item) |
| 135 | + try: |
| 136 | + req = { |
| 137 | + 'TableName': self._table_name, |
| 138 | + 'Item': encoded_item, |
| 139 | + 'ConditionExpression': 'attribute_not_exists(#namespace) or attribute_not_exists(#key) or :version > #version', |
| 140 | + 'ExpressionAttributeNames': {'#namespace': self.PARTITION_KEY, '#key': self.SORT_KEY, '#version': self.VERSION_ATTRIBUTE}, |
| 141 | + 'ExpressionAttributeValues': {':version': {'N': str(item['version'])}}, |
| 142 | + } |
| 143 | + await client.put_item(**req) |
| 144 | + except client.exceptions.ConditionalCheckFailedException: |
| 145 | + # The item was not updated because there's a newer item in the database. We must now |
| 146 | + # read the item that's in the database and return it, so the wrapper can cache it. |
| 147 | + return cast(dict, await self.get_internal(kind, item['key'])) |
| 148 | + return item |
| 149 | + |
| 150 | + async def initialized_internal(self) -> bool: |
| 151 | + client = await self._get_client() |
| 152 | + resp = await self._get_item_by_keys(client, self._inited_key(), self._inited_key()) |
| 153 | + return resp.get('Item') is not None and len(resp['Item']) > 0 |
| 154 | + |
| 155 | + async def close(self) -> None: |
| 156 | + async with self._client_lock: |
| 157 | + self._closed = True |
| 158 | + await self._exit_stack.aclose() |
| 159 | + self._client = None |
| 160 | + |
| 161 | + def describe_configuration(self, config) -> str: |
| 162 | + return 'DynamoDB' |
| 163 | + |
| 164 | + def _prefixed_namespace(self, base: str) -> str: |
| 165 | + return self._prefix + base |
| 166 | + |
| 167 | + def _namespace_for_kind(self, kind: VersionedDataKind) -> str: |
| 168 | + return self._prefixed_namespace(kind.namespace) |
| 169 | + |
| 170 | + def _inited_key(self) -> str: |
| 171 | + return self._prefixed_namespace('$inited') |
| 172 | + |
| 173 | + def _make_keys(self, namespace: str, key: str) -> dict: |
| 174 | + return {self.PARTITION_KEY: {'S': namespace}, self.SORT_KEY: {'S': key}} |
| 175 | + |
| 176 | + def _make_query_for_kind(self, kind: VersionedDataKind) -> dict: |
| 177 | + return { |
| 178 | + 'TableName': self._table_name, |
| 179 | + 'ConsistentRead': True, |
| 180 | + 'KeyConditions': {self.PARTITION_KEY: {'AttributeValueList': [{'S': self._namespace_for_kind(kind)}], 'ComparisonOperator': 'EQ'}}, |
| 181 | + } |
| 182 | + |
| 183 | + async def _get_item_by_keys(self, client: Any, namespace: str, key: str) -> dict: |
| 184 | + return await client.get_item(TableName=self._table_name, Key=self._make_keys(namespace, key)) |
| 185 | + |
| 186 | + async def _read_existing_keys(self, client: Any, kinds) -> set: |
| 187 | + keys: set = set() |
| 188 | + for kind in kinds: |
| 189 | + req = self._make_query_for_kind(kind) |
| 190 | + req['ProjectionExpression'] = '#namespace, #key' |
| 191 | + req['ExpressionAttributeNames'] = {'#namespace': self.PARTITION_KEY, '#key': self.SORT_KEY} |
| 192 | + paginator = client.get_paginator('query') |
| 193 | + async for resp in paginator.paginate(**req): |
| 194 | + for item in resp['Items']: |
| 195 | + namespace = item[self.PARTITION_KEY]['S'] |
| 196 | + key = item[self.SORT_KEY]['S'] |
| 197 | + keys.add((namespace, key)) |
| 198 | + return keys |
| 199 | + |
| 200 | + def _marshal_item(self, kind: VersionedDataKind, item: dict) -> dict: |
| 201 | + json_str = json.dumps(item) |
| 202 | + ret = self._make_keys(self._namespace_for_kind(kind), item['key']) |
| 203 | + ret[self.VERSION_ATTRIBUTE] = {'N': str(item['version'])} |
| 204 | + ret[self.ITEM_JSON_ATTRIBUTE] = {'S': json_str} |
| 205 | + return ret |
| 206 | + |
| 207 | + def _unmarshal_item(self, item: Optional[dict]) -> Optional[dict]: |
| 208 | + if item is None: |
| 209 | + return None |
| 210 | + json_attr = item.get(self.ITEM_JSON_ATTRIBUTE) |
| 211 | + return None if json_attr is None else json.loads(json_attr['S']) |
| 212 | + |
| 213 | + |
| 214 | +class _AsyncDynamoDBHelpers: |
| 215 | + @staticmethod |
| 216 | + async def batch_write_requests(client: Any, table_name: str, requests: list) -> None: |
| 217 | + batch_size = 25 |
| 218 | + for batch in (requests[i: i + batch_size] for i in range(0, len(requests), batch_size)): |
| 219 | + await client.batch_write_item(RequestItems={table_name: batch}) |
0 commit comments