diff --git a/ldclient/impl/integrations/consul/async_consul_feature_store.py b/ldclient/impl/integrations/consul/async_consul_feature_store.py new file mode 100644 index 00000000..1858d937 --- /dev/null +++ b/ldclient/impl/integrations/consul/async_consul_feature_store.py @@ -0,0 +1,162 @@ +import json +from typing import Any, Dict, Mapping, Optional + +from ldclient import log +from ldclient.interfaces import AsyncFeatureStoreCore, DiagnosticDescription +from ldclient.versioned_data_kind import VersionedDataKind + +have_async_consul = False +try: + import consul.aio + + have_async_consul = True +except ImportError: + pass + +# +# Internal implementation of the async Consul feature store. +# +# It uses the same Consul KV layout as the synchronous Consul feature store, so an async and a +# synchronous SDK can share one Consul instance. +# +# Implementation notes: +# +# * Feature flags, segments, and any other kind of entity the LaunchDarkly client may wish +# to store, are stored as individual items with the key "{prefix}/features/{flag-key}", +# "{prefix}/segments/{segment-key}", etc. +# +# * The special key "{prefix}/$inited" indicates that the store contains a complete data set. +# +# * Since Consul has limited support for transactions (they can't contain more than 64 +# operations), the init method-- which replaces the entire data store-- is not guaranteed to +# be atomic, so there can be a race condition if another process is adding new data via +# Upsert. To minimize this, we don't delete all the data at the start; instead, we update +# the items we've received, and then delete all other items. That could potentially result in +# deleting new data from another process, but that would be the case anyway if the Init +# happened to execute later than the Upsert; we are relying on the fact that normally the +# process that did the Init will also receive the new data shortly and do its own Upsert. +# + + +class _AsyncConsulFeatureStoreCore(DiagnosticDescription, AsyncFeatureStoreCore): + """Async Consul implementation of :class:`ldclient.interfaces.AsyncFeatureStoreCore`. + + It stores data in the same Consul KV layout as the synchronous Consul feature store, so an async + and a synchronous SDK can share one Consul instance. + """ + + def __init__(self, host: Optional[str], port: Optional[int], prefix: Optional[str], consul_opts: Optional[dict]): + if not have_async_consul: + raise NotImplementedError("Cannot use async Consul feature store because the py-consul package is not installed") + opts = dict(consul_opts or {}) + if host is not None: + opts['host'] = host + if port is not None: + opts['port'] = port + self._opts = opts + self._prefix = ("launchdarkly" if prefix is None else prefix) + "/" + self._client: Optional[Any] = None + + def _get_client(self): + # py-consul's asyncio client builds its aiohttp session when it is constructed, which needs a + # running event loop. We create the client lazily on the first store operation so the factory + # method can be called synchronously while building configuration. + if self._client is None: + self._client = consul.aio.Consul(**self._opts) + return self._client + + async def is_available(self) -> bool: + try: + await self._get_client().kv.get(self._inited_key()) + return True + except BaseException: + return False + + async def init_internal(self, all_data: Mapping[VersionedDataKind, Mapping[str, dict]]) -> None: + client = self._get_client() + + # Start by reading the existing keys; we will later delete any of these that weren't in all_data. + index, keys = await client.kv.get(self._prefix, recurse=True, keys=True) + unused_old_keys = set(keys or []) + + num_items = 0 + inited_key = self._inited_key() + unused_old_keys.discard(inited_key) + + # Insert or update every provided item. Note that this Consul client doesn't support batch + # operations (the "txn" method), so we'll write them one at a time. + for kind, items in all_data.items(): + for key, item in items.items(): + encoded_item = json.dumps(item) + db_key = self._item_key(kind, item['key']) + await client.kv.put(db_key, encoded_item) + unused_old_keys.discard(db_key) + num_items = num_items + 1 + + # Now delete any previously existing items whose keys were not in the current data + for key in unused_old_keys: + await client.kv.delete(key) + + # Now set the special key that we check in initialized_internal() + await client.kv.put(inited_key, "") + + log.info('Initialized async Consul store with %d items', num_items) + + async def get_internal(self, kind: VersionedDataKind, key: str) -> Optional[dict]: + index, resp = await self._get_client().kv.get(self._item_key(kind, key)) + return None if resp is None else json.loads(resp['Value'].decode('utf-8')) + + async def get_all_internal(self, kind: VersionedDataKind) -> Mapping[str, dict]: + items_out: Dict[str, dict] = {} + index, results = await self._get_client().kv.get(self._kind_key(kind), recurse=True) + for result in results or []: + item = json.loads(result['Value'].decode('utf-8')) + items_out[item['key']] = item + return items_out + + async def upsert_internal(self, kind: VersionedDataKind, new_item: dict) -> dict: + client = self._get_client() + key = self._item_key(kind, new_item['key']) + encoded_item = json.dumps(new_item) + + # We will potentially keep retrying indefinitely until someone's write succeeds + while True: + index, old_value = await client.kv.get(key) + if old_value is None: + mod_index = 0 + else: + old_item = json.loads(old_value['Value'].decode('utf-8')) + # Check whether the item is stale. If so, don't do the update (and return the existing item to + # AsyncCachingStoreWrapper so it can be cached) + if old_item['version'] >= new_item['version']: + return old_item + mod_index = old_value['ModifyIndex'] + + # Otherwise, try to write. We will do a compare-and-set operation, so the write will only succeed if + # the key's ModifyIndex is still equal to the previous value. If the previous ModifyIndex was zero, + # it means the key did not previously exist and the write will only succeed if it still doesn't exist. + success = await client.kv.put(key, encoded_item, cas=mod_index) + if success: + return new_item + + log.debug('Concurrent modification detected, retrying') + + async def initialized_internal(self) -> bool: + index, resp = await self._get_client().kv.get(self._inited_key()) + return resp is not None + + async def close(self) -> None: + if self._client is not None: + await self._client.close() + + def describe_configuration(self, config) -> str: + return 'Consul' + + def _kind_key(self, kind: VersionedDataKind) -> str: + return self._prefix + kind.namespace + + def _item_key(self, kind: VersionedDataKind, key: str) -> str: + return self._kind_key(kind) + '/' + key + + def _inited_key(self) -> str: + return self._prefix + '$inited' diff --git a/ldclient/integrations/__init__.py b/ldclient/integrations/__init__.py index 3f729589..a69b18bd 100644 --- a/ldclient/integrations/__init__.py +++ b/ldclient/integrations/__init__.py @@ -71,6 +71,54 @@ def new_feature_store( core = _ConsulFeatureStoreCore(host, port, prefix, consul_opts) return CachingStoreWrapper(core, caching) + @staticmethod + def async_feature_store( + host: Optional[str] = None, port: Optional[int] = None, prefix: Optional[str] = None, consul_opts: Optional[dict] = None, caching: CacheConfig = CacheConfig.default() + ): + """Creates an async Consul-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 ``py-consul`` package with its ``asyncio`` + extra (``py-consul[asyncio]``). Then, put the object returned by this method into the + ``feature_store`` property of your client configuration when constructing an ``AsyncLDClient``. + :: + + from ldclient.integrations import Consul + store = Consul.async_feature_store() + config = Config(feature_store=store) + + The data layout matches :func:`new_feature_store`, so an async and a synchronous SDK can share + one Consul instance. + + :param host: hostname of the Consul server (uses ``localhost`` if omitted) + :param port: port of the Consul server (uses 8500 if omitted) + :param prefix: a namespace prefix to be prepended to all Consul keys + :param consul_opts: optional parameters for configuring the Consul client, if you need + to set any of them besides host and port, as defined in the + `py-consul API `_ + :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.consul.async_consul_feature_store import ( + _AsyncConsulFeatureStoreCore + ) + core = _AsyncConsulFeatureStoreCore(host, port, prefix, consul_opts) + wrapper = AsyncCachingStoreWrapper(core, caching) + wrapper._core = core # exposed for testing + return wrapper + class DynamoDB: """Provides factory methods for integrations between the LaunchDarkly SDK and DynamoDB.""" diff --git a/ldclient/testing/integrations/test_async_consul_feature_store.py b/ldclient/testing/integrations/test_async_consul_feature_store.py new file mode 100644 index 00000000..35030afa --- /dev/null +++ b/ldclient/testing/integrations/test_async_consul_feature_store.py @@ -0,0 +1,119 @@ +""" +Integration tests for the async Consul feature store (_AsyncConsulFeatureStoreCore wrapped by +AsyncCachingStoreWrapper). + +These tests require a real Consul instance running on localhost:8500. They are skipped when the +py-consul 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 Consul in +ldclient.testing.test_async_feature_store_helpers. +""" + +import pytest + +from ldclient.async_feature_store_helpers import AsyncCachingStoreWrapper +from ldclient.feature_store import CacheConfig +from ldclient.integrations import Consul +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 +from ldclient.versioned_data_kind import FEATURES + +have_async_consul = False +try: + import consul.aio + + have_async_consul = True +except ImportError: + pass + +pytestmark = pytest.mark.skipif( + not have_async_consul, + reason="skipping async Consul tests because py-consul package is not installed" +) + +DEFAULT_PREFIX = 'launchdarkly' + + +def clear_data(prefix): + # A synchronous client is enough for test setup and teardown. + client = consul.Consul() + index, keys = client.kv.get((prefix or DEFAULT_PREFIX) + "/", recurse=True, keys=True) + for key in keys or []: + client.kv.delete(key) + + +class AsyncConsulFeatureStoreTester(AsyncFeatureStoreTester): + def __init__(self, prefix=None, caching=None): + self.prefix = prefix + self.caching = caching if caching is not None else CacheConfig.disabled() + + async def create_feature_store(self) -> AsyncFeatureStore: + return Consul.async_feature_store(prefix=self.prefix, caching=self.caching) + + +@pytest.mark.skipif(skip_database_tests, reason="skipping database tests") +class TestAsyncConsulFeatureStore(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 AsyncConsulFeatureStoreTester(prefix, caching) + + @pytest.fixture(autouse=True) + def clear_data_before_each(self, tester): + clear_data(tester.prefix) + + +@pytest.mark.asyncio +@pytest.mark.skipif(skip_database_tests, reason="skipping database tests") +async def test_stores_with_different_prefixes_are_independent(): + clear_data("a") + clear_data("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 = Consul.async_feature_store(prefix="a") + store_b = Consul.async_feature_store(prefix="b") + 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() + + +@pytest.mark.asyncio +@pytest.mark.skipif(skip_database_tests, reason="skipping database tests") +async def test_available_and_monitoring(): + store = Consul.async_feature_store() + try: + assert store.is_monitoring_enabled() is True + assert await store.is_available() is True + finally: + await store.close() + + +def test_async_feature_store_is_caching_wrapper(): + store = Consul.async_feature_store() + assert isinstance(store, AsyncCachingStoreWrapper) + + +# Consul does not support Big Segments. diff --git a/pyproject.toml b/pyproject.toml index ba144385..de715e2b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ Documentation = "https://launchdarkly-python-sdk.readthedocs.io/en/latest/" async = ["aiohttp>=3.9,<4"] redis = ["redis>=2.10.5"] consul = ["python-consul>=1.0.1"] +async-consul = ["py-consul[asyncio]>=1.7.1"] dynamodb = ["boto3>=1.9.71"] test-filesource = ["pyyaml>=5.3.1", "watchdog>=3.0.0"] @@ -61,6 +62,7 @@ dev = [ "types-redis>=4.0", "types-setuptools>=68.0", "aiohttp>=3.8.0", + "py-consul[asyncio]>=1.7.1", ] contract-tests = [ "Flask<4",