diff --git a/ldclient/impl/integrations/dynamodb/async_dynamodb_feature_store.py b/ldclient/impl/integrations/dynamodb/async_dynamodb_feature_store.py new file mode 100644 index 00000000..fc810eca --- /dev/null +++ b/ldclient/impl/integrations/dynamodb/async_dynamodb_feature_store.py @@ -0,0 +1,220 @@ +import asyncio +import json +from contextlib import AsyncExitStack +from typing import Any, Mapping, Optional, cast + +from ldclient.impl.util import log +from ldclient.interfaces import AsyncFeatureStoreCore, DiagnosticDescription +from ldclient.versioned_data_kind import VersionedDataKind + +have_aioboto3 = False +try: + import aioboto3 + + have_aioboto3 = True +except ImportError: + pass + + +# +# Internal implementation of the async DynamoDB feature store. +# +# Implementation notes: +# +# * This store uses the same table layout and key schema as the synchronous DynamoDB feature store, +# so an async and a synchronous SDK can share one table. Feature flags, segments, and any other kind +# of entity are all put in the same table. The two required attributes are "key" (present in all +# storeable entities) and "namespace" (used to disambiguate between flags and segments). +# +# * Because of DynamoDB's restrictions on attribute values (e.g. empty strings are not allowed), the +# standard DynamoDB marshaling with one attribute per object property is not used. Instead, the +# entire object is serialized to JSON and stored in a single attribute, "item". The "version" +# property is also stored as a separate attribute since it is used for updates. +# +# * Since DynamoDB has no transactions, init() - which replaces the entire data store - is not +# atomic, so there can be a race condition if another process is adding new data via upsert(). To +# minimize this, we do not delete all the data at the start; instead, we update the items we have +# received, and then delete all other items. That could delete new data from another process, but +# that would happen anyway if the init() ran later than the upsert(); we rely on the fact that the +# process that did the init() will normally receive the new data shortly and do its own upsert(). +# +# * DynamoDB has a maximum item size of 400KB. Since each feature flag or user segment is stored as +# a single item, this mechanism will not work for extremely large flags or segments. +# +# * aioboto3 clients are async context managers, so unlike the synchronous boto3 client they cannot +# be created in __init__. The client is created and entered lazily on first use inside the running +# event loop, kept for the lifetime of the store, and released in close(). +# + + +class _AsyncDynamoDBFeatureStoreCore(DiagnosticDescription, AsyncFeatureStoreCore): + """Async DynamoDB implementation of :class:`ldclient.interfaces.AsyncFeatureStoreCore`. + + It stores data in the same DynamoDB table layout and key schema as the synchronous DynamoDB + feature store, so an async and a synchronous SDK can share one table. + """ + + PARTITION_KEY = 'namespace' + SORT_KEY = 'key' + VERSION_ATTRIBUTE = 'version' + ITEM_JSON_ATTRIBUTE = 'item' + + def __init__(self, table_name: str, prefix: Optional[str], dynamodb_opts: Mapping[str, Any]): + if not have_aioboto3: + raise NotImplementedError("Cannot use async DynamoDB feature store because aioboto3 package is not installed") + self._table_name = table_name + self._prefix = (prefix + ":") if prefix else "" + self._dynamodb_opts = dict(dynamodb_opts) + self._session = aioboto3.Session() + self._exit_stack = AsyncExitStack() + self._client: Optional[Any] = None + # Guards lazy client creation so two concurrent coroutines cannot each enter a client. + self._client_lock = asyncio.Lock() + + async def _get_client(self) -> Any: + if self._client is not None: + return self._client + async with self._client_lock: + if self._client is None: + self._client = await self._exit_stack.enter_async_context(self._session.client('dynamodb', **self._dynamodb_opts)) + return self._client + + async def is_available(self) -> bool: + try: + inited_key = self._inited_key() + client = await self._get_client() + await self._get_item_by_keys(client, inited_key, inited_key) + return True + except BaseException: + return False + + async def init_internal(self, all_data: Mapping[VersionedDataKind, Mapping[str, dict]]) -> None: + client = await self._get_client() + # Start by reading the existing keys; we will later delete any of these that were not in all_data. + unused_old_keys = await self._read_existing_keys(client, all_data.keys()) + requests = [] + num_items = 0 + inited_key = self._inited_key() + + # Insert or update every provided item + for kind, items in all_data.items(): + for key, item in items.items(): + encoded_item = self._marshal_item(kind, item) + requests.append({'PutRequest': {'Item': encoded_item}}) + combined_key = (self._namespace_for_kind(kind), key) + unused_old_keys.discard(combined_key) + num_items = num_items + 1 + + # Now delete any previously existing items whose keys were not in the current data + for combined_key in unused_old_keys: + if combined_key[0] != inited_key: + requests.append({'DeleteRequest': {'Key': self._make_keys(combined_key[0], combined_key[1])}}) + + # Now set the special key that we check in initialized_internal() + requests.append({'PutRequest': {'Item': self._make_keys(inited_key, inited_key)}}) + + await _AsyncDynamoDBHelpers.batch_write_requests(client, self._table_name, requests) + log.info('Initialized table %s with %d items', self._table_name, num_items) + + async def get_internal(self, kind: VersionedDataKind, key: str) -> Optional[dict]: + client = await self._get_client() + resp = await self._get_item_by_keys(client, self._namespace_for_kind(kind), key) + return self._unmarshal_item(resp.get('Item')) + + async def get_all_internal(self, kind: VersionedDataKind) -> Mapping[str, dict]: + client = await self._get_client() + items_out = {} + paginator = client.get_paginator('query') + async for resp in paginator.paginate(**self._make_query_for_kind(kind)): + for item in resp['Items']: + # Every stored item carries the JSON attribute, so _unmarshal_item never returns None here. + item_out = cast(dict, self._unmarshal_item(item)) + items_out[item_out['key']] = item_out + return items_out + + async def upsert_internal(self, kind: VersionedDataKind, item: dict) -> dict: + client = await self._get_client() + encoded_item = self._marshal_item(kind, item) + try: + req = { + 'TableName': self._table_name, + 'Item': encoded_item, + 'ConditionExpression': 'attribute_not_exists(#namespace) or attribute_not_exists(#key) or :version > #version', + 'ExpressionAttributeNames': {'#namespace': self.PARTITION_KEY, '#key': self.SORT_KEY, '#version': self.VERSION_ATTRIBUTE}, + 'ExpressionAttributeValues': {':version': {'N': str(item['version'])}}, + } + await client.put_item(**req) + except client.exceptions.ConditionalCheckFailedException: + # The item was not updated because there's a newer item in the database. We must now + # read the item that's in the database and return it, so the wrapper can cache it. + return cast(dict, await self.get_internal(kind, item['key'])) + return item + + async def initialized_internal(self) -> bool: + client = await self._get_client() + resp = await self._get_item_by_keys(client, self._inited_key(), self._inited_key()) + return resp.get('Item') is not None and len(resp['Item']) > 0 + + async def close(self) -> None: + await self._exit_stack.aclose() + self._client = None + + def describe_configuration(self, config) -> str: + return 'DynamoDB' + + def _prefixed_namespace(self, base: str) -> str: + return self._prefix + base + + def _namespace_for_kind(self, kind: VersionedDataKind) -> str: + return self._prefixed_namespace(kind.namespace) + + def _inited_key(self) -> str: + return self._prefixed_namespace('$inited') + + def _make_keys(self, namespace: str, key: str) -> dict: + return {self.PARTITION_KEY: {'S': namespace}, self.SORT_KEY: {'S': key}} + + def _make_query_for_kind(self, kind: VersionedDataKind) -> dict: + return { + 'TableName': self._table_name, + 'ConsistentRead': True, + 'KeyConditions': {self.PARTITION_KEY: {'AttributeValueList': [{'S': self._namespace_for_kind(kind)}], 'ComparisonOperator': 'EQ'}}, + } + + async def _get_item_by_keys(self, client: Any, namespace: str, key: str) -> dict: + return await client.get_item(TableName=self._table_name, Key=self._make_keys(namespace, key)) + + async def _read_existing_keys(self, client: Any, kinds) -> set: + keys: set = set() + for kind in kinds: + req = self._make_query_for_kind(kind) + req['ProjectionExpression'] = '#namespace, #key' + req['ExpressionAttributeNames'] = {'#namespace': self.PARTITION_KEY, '#key': self.SORT_KEY} + paginator = client.get_paginator('query') + async for resp in paginator.paginate(**req): + for item in resp['Items']: + namespace = item[self.PARTITION_KEY]['S'] + key = item[self.SORT_KEY]['S'] + keys.add((namespace, key)) + return keys + + def _marshal_item(self, kind: VersionedDataKind, item: dict) -> dict: + json_str = json.dumps(item) + ret = self._make_keys(self._namespace_for_kind(kind), item['key']) + ret[self.VERSION_ATTRIBUTE] = {'N': str(item['version'])} + ret[self.ITEM_JSON_ATTRIBUTE] = {'S': json_str} + return ret + + def _unmarshal_item(self, item: Optional[dict]) -> Optional[dict]: + if item is None: + return None + json_attr = item.get(self.ITEM_JSON_ATTRIBUTE) + return None if json_attr is None else json.loads(json_attr['S']) + + +class _AsyncDynamoDBHelpers: + @staticmethod + async def batch_write_requests(client: Any, table_name: str, requests: list) -> None: + batch_size = 25 + for batch in (requests[i: i + batch_size] for i in range(0, len(requests), batch_size)): + await client.batch_write_item(RequestItems={table_name: batch}) diff --git a/ldclient/integrations/__init__.py b/ldclient/integrations/__init__.py index 3f729589..5d49e5cb 100644 --- a/ldclient/integrations/__init__.py +++ b/ldclient/integrations/__init__.py @@ -147,6 +147,60 @@ def new_big_segment_store(table_name: str, prefix: Optional[str] = None, dynamod """ return _DynamoDBBigSegmentStore(table_name, prefix, dynamodb_opts) + @staticmethod + def async_feature_store(table_name: str, prefix: Optional[str] = None, dynamodb_opts: Mapping[str, Any] = {}, caching: CacheConfig = CacheConfig.default()): + """Creates an async DynamoDB-backed implementation of :class:`~ldclient.interfaces.AsyncFeatureStore`. + + .. caution:: + This feature is experimental and should NOT be considered ready for production + use. It may change or be removed without notice and is not subject to backwards + compatibility guarantees. Pin to a specific minor version and review the changelog + before upgrading. + + For more details about how and why you can use a persistent feature store, see the + `SDK reference guide `_. + + To use this method, you must first install the ``aioboto3`` package. Then, put the object + returned by this method into the ``feature_store`` property of your client configuration + when constructing an ``AsyncLDClient``. + :: + + from ldclient.config import Config + from ldclient.integrations import DynamoDB + store = DynamoDB.async_feature_store("my-table-name") + config = Config(feature_store=store) + + The data layout matches :func:`new_feature_store`, so an async and a synchronous SDK can share + one DynamoDB table. + + Note that the DynamoDB table must already exist; the LaunchDarkly SDK does not create the table + automatically, because it has no way of knowing what additional properties (such as permissions + and throughput) you would want it to have. The table must have a partition key called + "namespace" and a sort key called "key", both with a string type. + + By default, the DynamoDB client will try to get your AWS credentials and region name from + environment variables and/or local configuration files, as described in the AWS SDK documentation. + You may also pass configuration settings in ``dynamodb_opts``. + + :param table_name: the name of an existing DynamoDB table + :param prefix: an optional namespace prefix to be prepended to all DynamoDB keys + :param dynamodb_opts: optional parameters for configuring the DynamoDB client, forwarded to + ``aioboto3.Session.client`` + :param caching: specifies whether local caching should be enabled and if so, + sets the cache properties; defaults to :func:`ldclient.feature_store.CacheConfig.default()`. + See :class:`ldclient.feature_store.CacheConfig`. + """ + from ldclient.async_feature_store_helpers import ( + AsyncCachingStoreWrapper + ) + from ldclient.impl.integrations.dynamodb.async_dynamodb_feature_store import ( + _AsyncDynamoDBFeatureStoreCore + ) + core = _AsyncDynamoDBFeatureStoreCore(table_name, prefix, dynamodb_opts) + wrapper = AsyncCachingStoreWrapper(core, caching) + wrapper._core = core # exposed for testing + return wrapper + class Redis: """Provides factory methods for integrations between the LaunchDarkly SDK and Redis.""" diff --git a/ldclient/testing/integrations/test_async_dynamodb_feature_store.py b/ldclient/testing/integrations/test_async_dynamodb_feature_store.py new file mode 100644 index 00000000..8a797f1a --- /dev/null +++ b/ldclient/testing/integrations/test_async_dynamodb_feature_store.py @@ -0,0 +1,210 @@ +""" +Integration tests for the async DynamoDB feature store (_AsyncDynamoDBFeatureStoreCore wrapped by +AsyncCachingStoreWrapper). + +These tests require a local DynamoDB instance running on localhost:8000. They are skipped when the +aioboto3 package is not installed or when the LD_SKIP_DATABASE_TESTS environment variable is set to +'1'. The caching-wrapper logic itself is covered without DynamoDB in +ldclient.testing.test_async_feature_store_helpers. Table setup and teardown use the synchronous +boto3 client, so the database-backed tests also need boto3. +""" + +import time + +import pytest + +from ldclient.async_feature_store_helpers import AsyncCachingStoreWrapper +from ldclient.feature_store import CacheConfig +from ldclient.impl.integrations.dynamodb.async_dynamodb_feature_store import ( + _AsyncDynamoDBFeatureStoreCore +) +from ldclient.integrations import DynamoDB +from ldclient.interfaces import AsyncFeatureStore +from ldclient.testing.async_feature_store_test_base import ( + AsyncFeatureStoreTestBase, + AsyncFeatureStoreTester +) +from ldclient.testing.test_util import skip_database_tests + +have_aioboto3 = False +try: + import aioboto3 + + have_aioboto3 = True +except ImportError: + pass + +try: + import boto3 + + have_boto3 = True +except ImportError: + have_boto3 = False + +pytestmark = pytest.mark.skipif( + not have_aioboto3, + reason="skipping async DynamoDB tests because aioboto3 package is not installed" +) + + +class DynamoDBTestHelper: + table_name = 'LD_DYNAMODB_TEST_TABLE' + table_created = False + options = {'aws_access_key_id': 'key', 'aws_secret_access_key': 'secret', 'endpoint_url': 'http://localhost:8000', 'region_name': 'us-east-1'} # not used by local DynamoDB, but still required + + @staticmethod + def make_client(): + """Return a synchronous boto3 client for test setup and teardown.""" + return boto3.client('dynamodb', **DynamoDBTestHelper.options) + + @staticmethod + def clear_data_for_prefix(prefix): + client = DynamoDBTestHelper.make_client() + delete_requests = [] + req = { + 'TableName': DynamoDBTestHelper.table_name, + 'ConsistentRead': True, + 'ProjectionExpression': '#namespace, #key', + 'ExpressionAttributeNames': {'#namespace': _AsyncDynamoDBFeatureStoreCore.PARTITION_KEY, '#key': _AsyncDynamoDBFeatureStoreCore.SORT_KEY}, + } + for resp in client.get_paginator('scan').paginate(**req): + for item in resp['Items']: + delete_requests.append({'DeleteRequest': {'Key': item}}) + _sync_batch_write_requests(client, DynamoDBTestHelper.table_name, delete_requests) + + @staticmethod + def ensure_table_created(): + if DynamoDBTestHelper.table_created: + return + DynamoDBTestHelper.table_created = True + client = DynamoDBTestHelper.make_client() + try: + client.describe_table(TableName=DynamoDBTestHelper.table_name) + return + except client.exceptions.ResourceNotFoundException: + pass + req = { + 'TableName': DynamoDBTestHelper.table_name, + 'KeySchema': [ + {'AttributeName': _AsyncDynamoDBFeatureStoreCore.PARTITION_KEY, 'KeyType': 'HASH'}, + {'AttributeName': _AsyncDynamoDBFeatureStoreCore.SORT_KEY, 'KeyType': 'RANGE'}, + ], + 'AttributeDefinitions': [ + {'AttributeName': _AsyncDynamoDBFeatureStoreCore.PARTITION_KEY, 'AttributeType': 'S'}, + {'AttributeName': _AsyncDynamoDBFeatureStoreCore.SORT_KEY, 'AttributeType': 'S'}, + ], + 'ProvisionedThroughput': {'ReadCapacityUnits': 1, 'WriteCapacityUnits': 1}, + } + client.create_table(**req) + while True: + try: + client.describe_table(TableName=DynamoDBTestHelper.table_name) + return + except client.exceptions.ResourceNotFoundException: + time.sleep(0.5) + + +def _sync_batch_write_requests(client, table_name, requests): + batch_size = 25 + for batch in (requests[i: i + batch_size] for i in range(0, len(requests), batch_size)): + client.batch_write_item(RequestItems={table_name: batch}) + + +class AsyncDynamoDBFeatureStoreTester(AsyncFeatureStoreTester): + def __init__(self, prefix=None, caching=None): + self.prefix = prefix + self.caching = caching if caching is not None else CacheConfig.disabled() + DynamoDBTestHelper.ensure_table_created() + + async def create_feature_store(self) -> AsyncFeatureStore: + return DynamoDB.async_feature_store(DynamoDBTestHelper.table_name, prefix=self.prefix, caching=self.caching, dynamodb_opts=DynamoDBTestHelper.options) + + +@pytest.mark.skipif(skip_database_tests, reason="skipping database tests") +@pytest.mark.skipif(not have_boto3, reason="skipping: boto3 not available for test setup") +class TestAsyncDynamoDBFeatureStore(AsyncFeatureStoreTestBase): + @pytest.fixture(params=[(False, False), (True, False), (False, True), (True, True)]) + def tester(self, request): + specify_prefix, use_caching = request.param + prefix = "testprefix" if specify_prefix else None + caching = CacheConfig.default() if use_caching else CacheConfig.disabled() + return AsyncDynamoDBFeatureStoreTester(prefix, caching) + + @pytest.fixture(autouse=True) + def clear_data_before_each(self, tester): + DynamoDBTestHelper.clear_data_for_prefix(tester.prefix) + + +@pytest.mark.asyncio +@pytest.mark.skipif(skip_database_tests, reason="skipping database tests") +@pytest.mark.skipif(not have_boto3, reason="skipping: boto3 not available for test setup") +async def test_available_and_monitoring(): + DynamoDBTestHelper.ensure_table_created() + store = DynamoDB.async_feature_store(DynamoDBTestHelper.table_name, dynamodb_opts=DynamoDBTestHelper.options) + try: + assert store.is_monitoring_enabled() is True + assert await store.is_available() is True + finally: + await store.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif(skip_database_tests, reason="skipping database tests") +async def test_detects_nonexistent_store(): + options = dict(DynamoDBTestHelper.options) + options['endpoint_url'] = 'http://i-mean-what-are-the-odds' + store = DynamoDB.async_feature_store(DynamoDBTestHelper.table_name, dynamodb_opts=options) + try: + assert store.is_monitoring_enabled() is True + assert await store.is_available() is False + finally: + await store.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif(skip_database_tests, reason="skipping database tests") +@pytest.mark.skipif(not have_boto3, reason="skipping: boto3 not available for test setup") +async def test_stores_with_different_prefixes_are_independent(): + from ldclient.versioned_data_kind import FEATURES + + DynamoDBTestHelper.ensure_table_created() + DynamoDBTestHelper.clear_data_for_prefix("a") + DynamoDBTestHelper.clear_data_for_prefix("b") + + flag_a1 = {'key': 'flagA1', 'version': 1} + flag_a2 = {'key': 'flagA2', 'version': 1} + flag_b1 = {'key': 'flagB1', 'version': 1} + flag_b2 = {'key': 'flagB2', 'version': 1} + + store_a = DynamoDB.async_feature_store(DynamoDBTestHelper.table_name, prefix="a", dynamodb_opts=DynamoDBTestHelper.options) + store_b = DynamoDB.async_feature_store(DynamoDBTestHelper.table_name, prefix="b", dynamodb_opts=DynamoDBTestHelper.options) + try: + await store_a.init({FEATURES: {'flagA1': flag_a1}}) + await store_a.upsert(FEATURES, flag_a2) + + await store_b.init({FEATURES: {'flagB1': flag_b1}}) + await store_b.upsert(FEATURES, flag_b2) + + assert await store_a.get(FEATURES, 'flagA1') == FEATURES.decode(flag_a1) + assert await store_a.get(FEATURES, 'flagB1') is None + assert await store_a.all(FEATURES) == {'flagA1': FEATURES.decode(flag_a1), 'flagA2': FEATURES.decode(flag_a2)} + + assert await store_b.get(FEATURES, 'flagB1') == FEATURES.decode(flag_b1) + assert await store_b.get(FEATURES, 'flagA1') is None + assert await store_b.all(FEATURES) == {'flagB1': FEATURES.decode(flag_b1), 'flagB2': FEATURES.decode(flag_b2)} + finally: + await store_a.close() + await store_b.close() + + +def test_async_feature_store_is_caching_wrapper(): + store = DynamoDB.async_feature_store(DynamoDBTestHelper.table_name) + assert isinstance(store, AsyncCachingStoreWrapper) + + +def test_constructing_without_aioboto3_raises(monkeypatch): + import ldclient.impl.integrations.dynamodb.async_dynamodb_feature_store as mod + + monkeypatch.setattr(mod, "have_aioboto3", False) + with pytest.raises(NotImplementedError): + mod._AsyncDynamoDBFeatureStoreCore(DynamoDBTestHelper.table_name, None, {}) diff --git a/pyproject.toml b/pyproject.toml index ba144385..31b753b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ async = ["aiohttp>=3.9,<4"] redis = ["redis>=2.10.5"] consul = ["python-consul>=1.0.1"] dynamodb = ["boto3>=1.9.71"] +async-dynamodb = ["aioboto3>=11.0"] test-filesource = ["pyyaml>=5.3.1", "watchdog>=3.0.0"] [dependency-groups] @@ -49,6 +50,7 @@ dev = [ "pytest-asyncio>=0.23", "redis>=4.2.0,<9.0.0", "boto3>=1.9.71,<2.0.0", + "aioboto3>=11.0", "coverage>=4.4", "jsonpickle>1.4.1", "pytest-cov>=2.4.0",