feat: Add disable_connection_reuse config method

Disables connection pooling

Author: dustinbyrne

posthog/__init__.py

@@ -23,6 +23,7 @@
     RequiresServerEvaluation as RequiresServerEvaluation,
 )
 from posthog.request import (
+    disable_connection_reuse as disable_connection_reuse,
     enable_keep_alive as enable_keep_alive,
     set_socket_options as set_socket_options,
     SocketOptions as SocketOptions,

posthog/request.py

@@ -85,6 +85,14 @@ def _build_session(socket_options: Optional[SocketOptions] = None) -> requests.S
 
 
 _session = _build_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 set_socket_options(socket_options: Optional[SocketOptions]) -> None:
@@ -95,7 +103,8 @@ def set_socket_options(socket_options: Optional[SocketOptions]) -> None:
         from posthog import set_socket_options
         set_socket_options([(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)])
     """
-    global _session
+    global _session, _socket_options
+    _socket_options = socket_options
     _session = _build_session(socket_options)
 
 
@@ -104,6 +113,12 @@ def enable_keep_alive() -> None:
     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"
 DEFAULT_HOST = US_INGESTION_ENDPOINT
@@ -148,7 +163,7 @@ def post(
             gz.write(data.encode("utf-8"))
         data = buf.getvalue()
 
-    res = _session.post(url, data=data, headers=headers, timeout=timeout)
+    res = _get_session().post(url, data=data, headers=headers, timeout=timeout)
 
     if res.status_code == 200:
         log.debug("data uploaded successfully")
@@ -263,7 +278,7 @@ def get(
     if etag:
         headers["If-None-Match"] = etag
 
-    res = _session.get(full_url, headers=headers, timeout=timeout)
+    res = _get_session().get(full_url, headers=headers, timeout=timeout)
 
     masked_url = _mask_tokens_in_url(full_url)
 

posthog/test/test_request.py

@@ -6,6 +6,7 @@
 import pytest
 import requests
 
+import posthog.request as request_module
 from posthog.request import (
     APIError,
     DatetimeSerializer,
@@ -16,6 +17,7 @@
     batch_post,
     decide,
     determine_server_host,
+    disable_connection_reuse,
     enable_keep_alive,
     get,
     set_socket_options,
@@ -370,3 +372,13 @@ def test_set_socket_options_clears_with_none():
         assert adapter.socket_options is None
     finally:
         set_socket_options(None)
+
+
+def test_disable_connection_reuse_creates_fresh_sessions():
+    try:
+        disable_connection_reuse()
+        session1 = request_module._get_session()
+        session2 = request_module._get_session()
+        assert session1 is not session2
+    finally:
+        request_module._pooling_enabled = True