feat: Add async Consul persistent feature store

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'

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 <https://docs.launchdarkly.com/sdk/concepts/data-stores>`_.
+
+        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 <https://py-consul.readthedocs.io/en/latest/#consul>`_
+        :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."""

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.

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",