posthoganalytics 7.0.1__py3-none-any.whl → 7.4.1__py3-none-any.whl

posthoganalytics/flag_definition_cache.py

@@ -0,0 +1,127 @@
+"""
+Flag Definition Cache Provider interface for multi-worker environments.
+
+EXPERIMENTAL: This API may change in future minor version bumps.
+
+This module provides an interface for external caching of feature flag definitions,
+enabling multi-worker environments (Kubernetes, load-balanced servers, serverless
+functions) to share flag definitions and reduce API calls.
+
+Usage:
+
+    from posthoganalytics import Posthog
+    from posthoganalytics.flag_definition_cache import FlagDefinitionCacheProvider
+
+    cache = RedisFlagDefinitionCache(redis_client, "my-team")
+    posthog = Posthog(
+        "<project_api_key>",
+        personal_api_key="<personal_api_key>",
+        flag_definition_cache_provider=cache,
+    )
+"""
+
+from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
+
+from typing_extensions import Required, TypedDict
+
+
+class FlagDefinitionCacheData(TypedDict):
+    """
+    Data structure for cached flag definitions.
+
+    Attributes:
+        flags: List of feature flag definition dictionaries from the API.
+        group_type_mapping: Mapping of group type indices to group names.
+        cohorts: Dictionary of cohort definitions for local evaluation.
+    """
+
+    flags: Required[List[Dict[str, Any]]]
+    group_type_mapping: Required[Dict[str, str]]
+    cohorts: Required[Dict[str, Any]]
+
+
+@runtime_checkable
+class FlagDefinitionCacheProvider(Protocol):
+    """
+    Interface for external caching of feature flag definitions.
+
+    Enables multi-worker environments to share flag definitions, reducing API
+    calls while ensuring all workers have consistent data.
+
+    EXPERIMENTAL: This API may change in future minor version bumps.
+
+    The four methods handle the complete lifecycle of flag definition caching:
+
+    1. `should_fetch_flag_definitions()` - Called before each poll to determine
+       if this worker should fetch new definitions. Use for distributed lock
+       coordination to ensure only one worker fetches at a time.
+
+    2. `get_flag_definitions()` - Called when `should_fetch_flag_definitions()`
+       returns False. Returns cached definitions if available.
+
+    3. `on_flag_definitions_received()` - Called after successfully fetching
+       new definitions from the API. Store the data in your external cache
+       and release any locks.
+
+    4. `shutdown()` - Called when the PostHog client shuts down. Release any
+       distributed locks and clean up resources.
+
+    Error Handling:
+        All methods are wrapped in try/except. Errors will be logged but will
+        never break flag evaluation. On error:
+        - `should_fetch_flag_definitions()` errors default to fetching (fail-safe)
+        - `get_flag_definitions()` errors fall back to API fetch
+        - `on_flag_definitions_received()` errors are logged but flags remain in memory
+        - `shutdown()` errors are logged but shutdown continues
+    """
+
+    def get_flag_definitions(self) -> Optional[FlagDefinitionCacheData]:
+        """
+        Retrieve cached flag definitions.
+
+        Returns:
+            Cached flag definitions if available and valid, None otherwise.
+            Returning None will trigger a fetch from the API if this worker
+            has no flags loaded yet.
+        """
+        ...
+
+    def should_fetch_flag_definitions(self) -> bool:
+        """
+        Determine whether this instance should fetch new flag definitions.
+
+        Use this for distributed lock coordination. Only one worker should
+        return True to avoid thundering herd problems. A typical implementation
+        uses a distributed lock (e.g., Redis SETNX) that expires after the
+        poll interval.
+
+        Returns:
+            True if this instance should fetch from the API, False otherwise.
+            When False, the client will call `get_flag_definitions()` to
+            retrieve cached data instead.
+        """
+        ...
+
+    def on_flag_definitions_received(self, data: FlagDefinitionCacheData) -> None:
+        """
+        Called after successfully receiving new flag definitions from PostHog.
+
+        Use this to store the data in your external cache and release any
+        distributed locks acquired in `should_fetch_flag_definitions()`.
+
+        Args:
+            data: The flag definitions to cache, containing flags,
+                  group_type_mapping, and cohorts.
+        """
+        ...
+
+    def shutdown(self) -> None:
+        """
+        Called when the PostHog client shuts down.
+
+        Use this to release any distributed locks and clean up resources.
+        This method is called even if `should_fetch_flag_definitions()`
+        returned False, so implementations should handle the case where
+        no lock was acquired.
+        """
+        ...

posthoganalytics/request.py

@@ -1,28 +1,163 @@
 import json
 import logging
+import re
+import socket
+from dataclasses import dataclass
 from datetime import date, datetime
 from gzip import GzipFile
 from io import BytesIO
-from typing import Any, Optional, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import requests
 from dateutil.tz import tzutc
+from requests.adapters import HTTPAdapter  # type: ignore[import-untyped]
+from urllib3.connection import HTTPConnection
 from urllib3.util.retry import Retry
 
 from posthoganalytics.utils import remove_trailing_slash
 from posthoganalytics.version import VERSION
 
-# Retry on both connect and read errors
-# by default read errors will only retry idempotent HTTP methods (so not POST)
-adapter = requests.adapters.HTTPAdapter(
-    max_retries=Retry(
-        total=2,
-        connect=2,
-        read=2,
+SocketOptions = List[Tuple[int, int, Union[int, bytes]]]
+
+KEEPALIVE_IDLE_SECONDS = 60
+KEEPALIVE_INTERVAL_SECONDS = 60
+KEEPALIVE_PROBE_COUNT = 3
+
+# TCP keepalive probes idle connections to prevent them from being dropped.
+# SO_KEEPALIVE is cross-platform, but timing options vary:
+# - Linux: TCP_KEEPIDLE, TCP_KEEPINTVL, TCP_KEEPCNT
+# - macOS: only SO_KEEPALIVE (uses system defaults)
+# - Windows: TCP_KEEPIDLE, TCP_KEEPINTVL (since Windows 10 1709)
+KEEP_ALIVE_SOCKET_OPTIONS: SocketOptions = list(
+    HTTPConnection.default_socket_options
+) + [
+    (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+]
+for attr, value in [
+    ("TCP_KEEPIDLE", KEEPALIVE_IDLE_SECONDS),
+    ("TCP_KEEPINTVL", KEEPALIVE_INTERVAL_SECONDS),
+    ("TCP_KEEPCNT", KEEPALIVE_PROBE_COUNT),
+]:
+    if hasattr(socket, attr):
+        KEEP_ALIVE_SOCKET_OPTIONS.append((socket.SOL_TCP, getattr(socket, attr), value))
+
+# Status codes that indicate transient server errors worth retrying
+RETRY_STATUS_FORCELIST = [408, 500, 502, 503, 504]
+
+
+def _mask_tokens_in_url(url: str) -> str:
+    """Mask token values in URLs for safe logging, keeping first 10 chars visible."""
+    return re.sub(r"(token=)([^&]{10})[^&]*", r"\1\2...", url)
+
+
+@dataclass
+class GetResponse:
+    """Response from a GET request with ETag support."""
+
+    data: Any
+    etag: Optional[str] = None
+    not_modified: bool = False
+
+
+class HTTPAdapterWithSocketOptions(HTTPAdapter):
+    """HTTPAdapter with configurable socket options."""
+
+    def __init__(self, *args, socket_options: Optional[SocketOptions] = None, **kwargs):
+        self.socket_options = socket_options
+        super().__init__(*args, **kwargs)
+
+    def init_poolmanager(self, *args, **kwargs):
+        if self.socket_options is not None:
+            kwargs["socket_options"] = self.socket_options
+        super().init_poolmanager(*args, **kwargs)
+
+
+def _build_session(socket_options: Optional[SocketOptions] = None) -> requests.Session:
+    """Build a session for general requests (batch, decide, etc.)."""
+    adapter = HTTPAdapterWithSocketOptions(
+        max_retries=Retry(
+            total=2,
+            connect=2,
+            read=2,
+        ),
+        socket_options=socket_options,
+    )
+    session = requests.Session()
+    session.mount("https://", adapter)
+    return session
+
+
+def _build_flags_session(
+    socket_options: Optional[SocketOptions] = None,
+) -> requests.Session:
+    """
+    Build a session for feature flag requests with POST retries.
+
+    Feature flag requests are idempotent (read-only), so retrying POST
+    requests is safe. This session retries on transient server errors
+    (408, 5xx) and network failures with exponential backoff
+    (0.5s, 1s delays between retries).
+    """
+    adapter = HTTPAdapterWithSocketOptions(
+        max_retries=Retry(
+            total=2,
+            connect=2,
+            read=2,
+            backoff_factor=0.5,
+            status_forcelist=RETRY_STATUS_FORCELIST,
+            allowed_methods=["POST"],
+        ),
+        socket_options=socket_options,
     )
-)
-_session = requests.sessions.Session()
-_session.mount("https://", adapter)
+    session = requests.Session()
+    session.mount("https://", adapter)
+    return session
+
+
+_session = _build_session()
+_flags_session = _build_flags_session()
+_socket_options: Optional[SocketOptions] = None
+_pooling_enabled = True
+
+
+def _get_session() -> requests.Session:
+    if _pooling_enabled:
+        return _session
+    return _build_session(_socket_options)
+
+
+def _get_flags_session() -> requests.Session:
+    if _pooling_enabled:
+        return _flags_session
+    return _build_flags_session(_socket_options)
+
+
+def set_socket_options(socket_options: Optional[SocketOptions]) -> None:
+    """
+    Configure socket options for all HTTP connections.
+
+    Example:
+        from posthoganalytics import set_socket_options
+        set_socket_options([(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)])
+    """
+    global _session, _flags_session, _socket_options
+    if socket_options == _socket_options:
+        return
+    _socket_options = socket_options
+    _session = _build_session(socket_options)
+    _flags_session = _build_flags_session(socket_options)
+
+
+def enable_keep_alive() -> None:
+    """Enable TCP keepalive to prevent idle connections from being dropped."""
+    set_socket_options(KEEP_ALIVE_SOCKET_OPTIONS)
+
+
+def disable_connection_reuse() -> None:
+    """Disable connection reuse, creating a fresh connection for each request."""
+    global _pooling_enabled
+    _pooling_enabled = False
+
 
 US_INGESTION_ENDPOINT = "https://us.i.posthog.com"
 EU_INGESTION_ENDPOINT = "https://eu.i.posthog.com"
@@ -48,6 +183,7 @@ def post(
     path=None,
     gzip: bool = False,
     timeout: int = 15,
+    session: Optional[requests.Session] = None,
     **kwargs,
 ) -> requests.Response:
     """Post the `kwargs` to the API"""
@@ -68,7 +204,9 @@ def post(
             gz.write(data.encode("utf-8"))
         data = buf.getvalue()
 
-    res = _session.post(url, data=data, headers=headers, timeout=timeout)
+    res = (session or _get_session()).post(
+        url, data=data, headers=headers, timeout=timeout
+    )
 
     if res.status_code == 200:
         log.debug("data uploaded successfully")
@@ -124,8 +262,16 @@ def flags(
     timeout: int = 15,
     **kwargs,
 ) -> Any:
-    """Post the `kwargs to the flags API endpoint"""
-    res = post(api_key, host, "/flags/?v=2", gzip, timeout, **kwargs)
+    """Post the kwargs to the flags API endpoint with automatic retries."""
+    res = post(
+        api_key,
+        host,
+        "/flags/?v=2",
+        gzip,
+        timeout,
+        session=_get_flags_session(),
+        **kwargs,
+    )
     return _process_response(
         res, success_message="Feature flags evaluated successfully"
     )
@@ -139,12 +285,13 @@ def remote_config(
     timeout: int = 15,
 ) -> Any:
     """Get remote config flag value from remote_config API endpoint"""
-    return get(
+    response = get(
         personal_api_key,
         f"/api/projects/@current/feature_flags/{key}/remote_config?token={project_api_key}",
         host,
         timeout,
     )
+    return response.data
 
 
 def batch_post(
@@ -162,15 +309,42 @@ def batch_post(
 
 
 def get(
-    api_key: str, url: str, host: Optional[str] = None, timeout: Optional[int] = None
-) -> requests.Response:
-    url = remove_trailing_slash(host or DEFAULT_HOST) + url
-    res = requests.get(
-        url,
-        headers={"Authorization": "Bearer %s" % api_key, "User-Agent": USER_AGENT},
-        timeout=timeout,
+    api_key: str,
+    url: str,
+    host: Optional[str] = None,
+    timeout: Optional[int] = None,
+    etag: Optional[str] = None,
+) -> GetResponse:
+    """
+    Make a GET request with optional ETag support.
+
+    If an etag is provided, sends If-None-Match header. Returns GetResponse with:
+    - not_modified=True and data=None if server returns 304
+    - not_modified=False and data=response if server returns 200
+    """
+    log = logging.getLogger("posthog")
+    full_url = remove_trailing_slash(host or DEFAULT_HOST) + url
+    headers = {"Authorization": "Bearer %s" % api_key, "User-Agent": USER_AGENT}
+
+    if etag:
+        headers["If-None-Match"] = etag
+
+    res = _get_session().get(full_url, headers=headers, timeout=timeout)
+
+    masked_url = _mask_tokens_in_url(full_url)
+
+    # Handle 304 Not Modified
+    if res.status_code == 304:
+        log.debug(f"GET {masked_url} returned 304 Not Modified")
+        response_etag = res.headers.get("ETag")
+        return GetResponse(data=None, etag=response_etag or etag, not_modified=True)
+
+    # Handle normal response
+    data = _process_response(
+        res, success_message=f"GET {masked_url} completed successfully"
     )
-    return _process_response(res, success_message=f"GET {url} completed successfully")
+    response_etag = res.headers.get("ETag")
+    return GetResponse(data=data, etag=response_etag, not_modified=False)
 
 
 class APIError(Exception):
@@ -187,6 +361,12 @@ class QuotaLimitError(APIError):
     pass
 
 
+# Re-export requests exceptions for use in client.py
+# This keeps all requests library imports centralized in this module
+RequestsTimeout = requests.exceptions.Timeout
+RequestsConnectionError = requests.exceptions.ConnectionError
+
+
 class DatetimeSerializer(json.JSONEncoder):
     def default(self, obj: Any):
         if isinstance(obj, (date, datetime)):

posthoganalytics/test/test_client.py

@@ -9,7 +9,7 @@ from parameterized import parameterized
 
 from posthoganalytics.client import Client
 from posthoganalytics.contexts import get_context_session_id, new_context, set_context_session
-from posthoganalytics.request import APIError
+from posthoganalytics.request import APIError, GetResponse
 from posthoganalytics.test.test_utils import FAKE_TEST_API_KEY
 from posthoganalytics.types import FeatureFlag, LegacyFlagMetadata
 from posthoganalytics.version import VERSION
@@ -198,12 +198,6 @@ class TestClient(unittest.TestCase):
             print(capture_call)
             self.assertEqual(capture_call[1]["distinct_id"], "distinct_id")
             self.assertEqual(capture_call[0][0], "$exception")
-            self.assertEqual(
-                capture_call[1]["properties"]["$exception_type"], "Exception"
-            )
-            self.assertEqual(
-                capture_call[1]["properties"]["$exception_message"], "test exception"
-            )
             self.assertEqual(
                 capture_call[1]["properties"]["$exception_list"][0]["mechanism"][
                     "type"
@@ -415,7 +409,9 @@ class TestClient(unittest.TestCase):
             )
             client.feature_flags = [multivariate_flag, basic_flag, false_flag]
 
-            msg_uuid = client.capture("python test event", distinct_id="distinct_id")
+            msg_uuid = client.capture(
+                "python test event", distinct_id="distinct_id", send_feature_flags=True
+            )
             self.assertIsNotNone(msg_uuid)
             self.assertFalse(self.failed)
 
@@ -571,6 +567,7 @@ class TestClient(unittest.TestCase):
                 "python test event",
                 distinct_id="distinct_id",
                 properties={"$feature/beta-feature-local": "my-custom-variant"},
+                send_feature_flags=True,
             )
             self.assertIsNotNone(msg_uuid)
             self.assertFalse(self.failed)
@@ -752,6 +749,178 @@ class TestClient(unittest.TestCase):
 
             self.assertEqual(patch_flags.call_count, 0)
 
+    @mock.patch("posthog.client.flags")
+    def test_capture_with_send_feature_flags_false_and_local_evaluation_doesnt_send_flags(
+        self, patch_flags
+    ):
+        """Test that send_feature_flags=False with local evaluation enabled does NOT send flags"""
+        patch_flags.return_value = {"featureFlags": {"beta-feature": "remote-variant"}}
+
+        multivariate_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "beta-feature-local",
+            "active": True,
+            "rollout_percentage": 100,
+            "filters": {
+                "groups": [
+                    {
+                        "rollout_percentage": 100,
+                    },
+                ],
+                "multivariate": {
+                    "variants": [
+                        {
+                            "key": "first-variant",
+                            "name": "First Variant",
+                            "rollout_percentage": 50,
+                        },
+                        {
+                            "key": "second-variant",
+                            "name": "Second Variant",
+                            "rollout_percentage": 50,
+                        },
+                    ]
+                },
+            },
+        }
+        simple_flag = {
+            "id": 2,
+            "name": "Simple Flag",
+            "key": "simple-flag",
+            "active": True,
+            "filters": {
+                "groups": [
+                    {
+                        "rollout_percentage": 100,
+                    }
+                ],
+            },
+        }
+
+        with mock.patch("posthog.client.batch_post") as mock_post:
+            client = Client(
+                FAKE_TEST_API_KEY,
+                on_error=self.set_fail,
+                personal_api_key=FAKE_TEST_API_KEY,
+                sync_mode=True,
+            )
+            client.feature_flags = [multivariate_flag, simple_flag]
+
+            msg_uuid = client.capture(
+                "python test event",
+                distinct_id="distinct_id",
+                send_feature_flags=False,
+            )
+            self.assertIsNotNone(msg_uuid)
+            self.assertFalse(self.failed)
+
+            # Get the enqueued message from the mock
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            msg = batch_data[0]
+
+            self.assertEqual(msg["event"], "python test event")
+            self.assertEqual(msg["distinct_id"], "distinct_id")
+
+            # CRITICAL: Verify local flags are NOT included in the event
+            self.assertNotIn("$feature/beta-feature-local", msg["properties"])
+            self.assertNotIn("$feature/simple-flag", msg["properties"])
+            self.assertNotIn("$active_feature_flags", msg["properties"])
+
+            # CRITICAL: Verify the /flags API was NOT called
+            self.assertEqual(patch_flags.call_count, 0)
+
+    @mock.patch("posthog.client.flags")
+    def test_capture_with_send_feature_flags_true_and_local_evaluation_uses_local_flags(
+        self, patch_flags
+    ):
+        """Test that send_feature_flags=True with local evaluation enabled uses local flags without API call"""
+        patch_flags.return_value = {"featureFlags": {"remote-flag": "remote-variant"}}
+
+        multivariate_flag = {
+            "id": 1,
+            "name": "Beta Feature",
+            "key": "beta-feature-local",
+            "active": True,
+            "rollout_percentage": 100,
+            "filters": {
+                "groups": [
+                    {
+                        "rollout_percentage": 100,
+                    },
+                ],
+                "multivariate": {
+                    "variants": [
+                        {
+                            "key": "first-variant",
+                            "name": "First Variant",
+                            "rollout_percentage": 50,
+                        },
+                        {
+                            "key": "second-variant",
+                            "name": "Second Variant",
+                            "rollout_percentage": 50,
+                        },
+                    ]
+                },
+            },
+        }
+        simple_flag = {
+            "id": 2,
+            "name": "Simple Flag",
+            "key": "simple-flag",
+            "active": True,
+            "filters": {
+                "groups": [
+                    {
+                        "rollout_percentage": 100,
+                    }
+                ],
+            },
+        }
+
+        with mock.patch("posthog.client.batch_post") as mock_post:
+            client = Client(
+                FAKE_TEST_API_KEY,
+                on_error=self.set_fail,
+                personal_api_key=FAKE_TEST_API_KEY,
+                sync_mode=True,
+            )
+            client.feature_flags = [multivariate_flag, simple_flag]
+
+            msg_uuid = client.capture(
+                "python test event",
+                distinct_id="distinct_id",
+                send_feature_flags=True,
+            )
+            self.assertIsNotNone(msg_uuid)
+            self.assertFalse(self.failed)
+
+            # Get the enqueued message from the mock
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            msg = batch_data[0]
+
+            self.assertEqual(msg["event"], "python test event")
+            self.assertEqual(msg["distinct_id"], "distinct_id")
+
+            # Verify local flags are included in the event
+            self.assertIn("$feature/beta-feature-local", msg["properties"])
+            self.assertIn("$feature/simple-flag", msg["properties"])
+            self.assertEqual(msg["properties"]["$feature/simple-flag"], True)
+
+            # Verify active feature flags are set correctly
+            active_flags = msg["properties"]["$active_feature_flags"]
+            self.assertIn("beta-feature-local", active_flags)
+            self.assertIn("simple-flag", active_flags)
+
+            # The remote flag should NOT be included since we used local evaluation
+            self.assertNotIn("$feature/remote-flag", msg["properties"])
+
+            # CRITICAL: Verify the /flags API was NOT called
+            self.assertEqual(patch_flags.call_count, 0)
+
     @mock.patch("posthog.client.flags")
     def test_capture_with_send_feature_flags_options_only_evaluate_locally_true(
         self, patch_flags
@@ -2095,13 +2264,21 @@ class TestClient(unittest.TestCase):
         self, patch_get, patch_poller
     ):
         """Test that when enable_local_evaluation=False, the poller is not started"""
-        patch_get.return_value = {
-            "flags": [
-                {"id": 1, "name": "Beta Feature", "key": "beta-feature", "active": True}
-            ],
-            "group_type_mapping": {},
-            "cohorts": {},
-        }
+        patch_get.return_value = GetResponse(
+            data={
+                "flags": [
+                    {
+                        "id": 1,
+                        "name": "Beta Feature",
+                        "key": "beta-feature",
+                        "active": True,
+                    }
+                ],
+                "group_type_mapping": {},
+                "cohorts": {},
+            },
+            etag='"test-etag"',
+        )
 
         client = Client(
             FAKE_TEST_API_KEY,
@@ -2123,13 +2300,21 @@ class TestClient(unittest.TestCase):
     @mock.patch("posthog.client.get")
     def test_enable_local_evaluation_true_starts_poller(self, patch_get, patch_poller):
         """Test that when enable_local_evaluation=True (default), the poller is started"""
-        patch_get.return_value = {
-            "flags": [
-                {"id": 1, "name": "Beta Feature", "key": "beta-feature", "active": True}
-            ],
-            "group_type_mapping": {},
-            "cohorts": {},
-        }
+        patch_get.return_value = GetResponse(
+            data={
+                "flags": [
+                    {
+                        "id": 1,
+                        "name": "Beta Feature",
+                        "key": "beta-feature",
+                        "active": True,
+                    }
+                ],
+                "group_type_mapping": {},
+                "cohorts": {},
+            },
+            etag='"test-etag"',
+        )
 
         client = Client(
             FAKE_TEST_API_KEY,