fix key conflict issue

Author: cbrady

src/schematic/client.py

@@ -247,7 +247,7 @@ def _check_flag_via_api(
                 return cached_value
 
             resp = self.features.check_flag(flag_key, company=company, user=user)
-            if resp is None or resp.data.value is None:
+            if resp is None or resp.data is None or resp.data.value is None:
                 return self._default_response(flag_key, options, REASON_FLAG_NOT_FOUND)
 
             self._safe_cache_set(cache_key, resp.data)
@@ -627,7 +627,7 @@ def _ds_result_to_response(
         """Convert a RulesengineCheckFlagResult (from DataStream) into the
         public CheckFlagResponseData shape."""
         entitlement = (
-            FeatureEntitlement.model_validate(resp.entitlement.model_dump(mode="json"))
+            FeatureEntitlement.model_validate(resp.entitlement.model_dump())
             if resp.entitlement is not None else None
         )
         return CheckFlagResponseData(
@@ -658,7 +658,7 @@ async def _check_flag_via_api(
                 return cached_value
 
             resp = await self.features.check_flag(flag_key, company=company, user=user)
-            if resp is None or resp.data.value is None:
+            if resp is None or resp.data is None or resp.data.value is None:
                 return self._default_response(flag_key, options, REASON_FLAG_NOT_FOUND)
 
             self._safe_cache_set(cache_key, resp.data)

src/schematic/datastream/__init__.py

@@ -2,7 +2,7 @@
 from .datastream_client import DataStreamClient, DataStreamClientOptions
 from .merge import deep_copy_company, deep_copy_user, partial_company, partial_user
 from .rules_engine import RulesEngineClient
-from .types import DataStreamBaseReq, DataStreamError, DataStreamReq, DataStreamResp, EntityType, MessageType
+from .types import DataStreamBaseReq, DataStreamError, DataStreamReq, DataStreamResp, EntityType, KeyConflictError, MessageType
 from .websocket_client import ClientOptions, DatastreamWSClient, convert_api_url_to_websocket_url
 
 __all__ = [
@@ -25,6 +25,7 @@
     "DataStreamReq",
     "DataStreamResp",
     "EntityType",
+    "KeyConflictError",
     "MessageType",
     # WebSocket client
     "ClientOptions",

src/schematic/datastream/datastream_client.py

@@ -16,7 +16,7 @@
 from ..cache import AsyncCacheProvider, AsyncLocalCache
 from .merge import partial_company, partial_user
 from .rules_engine import RulesEngineClient
-from .types import DataStreamBaseReq, DataStreamReq, DataStreamResp, EntityType, MessageType
+from .types import DataStreamBaseReq, DataStreamReq, DataStreamResp, EntityType, KeyConflictError, MessageType
 from .websocket_client import ClientOptions as WSClientOptions, DatastreamWSClient
 
 
@@ -363,7 +363,10 @@ async def get_all_flags(self) -> None:
         try:
             await asyncio.wait_for(asyncio.shield(self._pending_flags), timeout=RESOURCE_TIMEOUT_MS / 1000)
         except asyncio.TimeoutError:
+            fut = self._pending_flags
             self._pending_flags = None
+            if fut is not None and not fut.done():
+                fut.set_result(False)
             raise TimeoutError("Timeout while waiting for flags data")
 
     async def check_flag(
@@ -384,10 +387,20 @@ async def check_flag(
         cached_company: Optional[Any] = None
         cached_user: Optional[Any] = None
 
-        if needs_company:
-            cached_company = await self._get_company_from_cache(company_keys)  # type: ignore[arg-type]
-        if needs_user:
-            cached_user = await self._get_user_from_cache(user_keys)  # type: ignore[arg-type]
+        try:
+            if needs_company:
+                cached_company = await self._get_company_from_cache(company_keys)  # type: ignore[arg-type]
+            if needs_user:
+                cached_user = await self._get_user_from_cache(user_keys)  # type: ignore[arg-type]
+        except KeyConflictError as exc:
+            self._logger.warning("Key conflict for flag %s: %s", flag_key, exc)
+            return RulesengineCheckFlagResult(
+                value=flag.default_value,
+                reason="key conflict",
+                flag_key=flag.key,
+                flag_id=flag.id,
+                err=str(exc),
+            )
 
         # Replicator mode — evaluate with whatever is cached
         if self._replicator_mode:
@@ -670,35 +683,65 @@ async def _send_request(self, request: DataStreamReq) -> None:
     # ------------------------------------------------------------------
 
     async def _get_company_from_cache(self, keys: Dict[str, str]) -> Optional[RulesengineCompany]:
+        matched_id: Optional[str] = None
         for key, value in keys.items():
             ck = self._resource_key_to_cache_key(_PREFIX_COMPANY, key, value)
             try:
                 company_id = await self._company_key_cache.get(ck)
-                self._logger.debug("Company lookup key %s -> %s", ck, company_id)
-                if company_id:
-                    rk = self._resource_id_cache_key(_PREFIX_COMPANY, company_id)
-                    raw = await self._company_cache.get(rk)
-                    self._logger.debug("Company ID key %s -> %s", rk, "hit" if raw is not None else "miss")
-                    if raw is not None:
-                        company = _validate(RulesengineCompany, raw)
-                        return company.model_copy(deep=True)
             except Exception as exc:
                 self._logger.warning("Failed to retrieve company from cache: %s", exc)
+                continue
+            self._logger.debug("Company lookup key %s -> %s", ck, company_id)
+            if not company_id:
+                continue
+            if matched_id is None:
+                matched_id = company_id
+            elif matched_id != company_id:
+                raise KeyConflictError(
+                    f"Company keys match multiple entities: {matched_id} and {company_id}"
+                )
+
+        if matched_id is None:
+            return None
+
+        try:
+            rk = self._resource_id_cache_key(_PREFIX_COMPANY, matched_id)
+            raw = await self._company_cache.get(rk)
+            self._logger.debug("Company ID key %s -> %s", rk, "hit" if raw is not None else "miss")
+            if raw is not None:
+                return _validate(RulesengineCompany, raw).model_copy(deep=True)
+        except Exception as exc:
+            self._logger.warning("Failed to retrieve company from cache: %s", exc)
         return None
 
     async def _get_user_from_cache(self, keys: Dict[str, str]) -> Optional[RulesengineUser]:
+        matched_id: Optional[str] = None
         for key, value in keys.items():
             ck = self._resource_key_to_cache_key(_PREFIX_USER, key, value)
             try:
                 user_id = await self._user_key_cache.get(ck)
-                if user_id:
-                    rk = self._resource_id_cache_key(_PREFIX_USER, user_id)
-                    raw = await self._user_cache.get(rk)
-                    if raw is not None:
-                        user = _validate(RulesengineUser, raw)
-                        return user.model_copy(deep=True)
             except Exception as exc:
                 self._logger.warning("Failed to retrieve user from cache: %s", exc)
+                continue
+            if not user_id:
+                continue
+            if matched_id is None:
+                matched_id = user_id
+            elif matched_id != user_id:
+                raise KeyConflictError(
+                    f"User keys match multiple entities: {matched_id} and {user_id}"
+                )
+
+        if matched_id is None:
+            return None
+
+        try:
+            rk = self._resource_id_cache_key(_PREFIX_USER, matched_id)
+            raw = await self._user_cache.get(rk)
+            if raw is not None:
+                return _validate(RulesengineUser, raw).model_copy(deep=True)
+        except Exception as exc:
+            self._logger.warning("Failed to retrieve user from cache: %s", exc)
         return None
 
     async def _cache_company(self, company: RulesengineCompany) -> None:
@@ -956,7 +999,8 @@ def _make_default_result(
     def _start_replicator_health_check(self) -> None:
         if not self._replicator_health_url:
             return
-        self._health_check_client = httpx.AsyncClient(timeout=REPLICATOR_HEALTH_TIMEOUT_S)
+        if self._health_check_client is None:
+            self._health_check_client = httpx.AsyncClient(timeout=REPLICATOR_HEALTH_TIMEOUT_S)
         self._logger.info(
             "Starting replicator health check: url=%s, interval=%dms",
             self._replicator_health_url, self._replicator_health_check_ms,

src/schematic/datastream/types.py

@@ -70,3 +70,7 @@ class DataStreamError:
     error: str
     keys: Optional[Dict[str, str]] = None
     entity_type: Optional[EntityType] = None
+
+
+class KeyConflictError(Exception):
+    """Raised when lookup keys resolve to multiple distinct entities."""