Merge branch 'main' into add-exposure-service

zhukaihan · web-flow · commit f75c590b7d63 · 2025-11-25T10:16:15.000-08:00
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
@@ -10,7 +10,7 @@ jobs:
         uses: actions/checkout@v2
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
 
       - name: Set up pdoc
         run:  pip install pdoc3
diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml
@@ -32,7 +32,7 @@ jobs:
           fetch-depth: 0
 
       - name: Set up Python 3.8
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: 3.8
 
diff --git a/.github/workflows/publish-to-test-pypi.yml b/.github/workflows/publish-to-test-pypi.yml
@@ -22,7 +22,7 @@ jobs:
       - uses: actions/checkout@v3
 
       - name: Set up Python 3.8
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: 3.8
 
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -14,7 +14,7 @@ jobs:
         uses: actions/checkout@v3
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
           cache: 'pip'
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,12 @@
 
 <!--next-version-placeholder-->
 
+## v1.9.0 (2025-11-20)
+
+### Feature
+
+* Allow configs to take in user provided logger instances for client logging behavior ([#65](https://github.com/amplitude/experiment-python-server/issues/65)) ([`688dc70`](https://github.com/amplitude/experiment-python-server/commit/688dc70a32763bcac91bde1fc8abcbb31b2d26de))
+
 ## v1.8.3 (2025-11-10)
 
 ### Fix
diff --git a/src/amplitude_experiment/local/client.py b/src/amplitude_experiment/local/client.py
@@ -1,4 +1,3 @@
-import logging
 from threading import Lock
 from typing import Any, List, Dict, Set
 
@@ -49,17 +48,13 @@ def __init__(self, api_key: str, config: LocalEvaluationConfig = None):
             instance = Amplitude(config.assignment_config.api_key, config.assignment_config)
             self.assignment_service = AssignmentService(instance, AssignmentFilter(
                 config.assignment_config.cache_capacity), config.assignment_config.send_evaluated_props)
-        
         # Exposure service is always instantiated, using deployment key if no api key provided
         self.exposure_service = None
         if config and config.exposure_config:
             exposure_config = config.exposure_config
             exposure_instance = Amplitude(exposure_config.api_key, exposure_config)
             self.exposure_service = ExposureService(exposure_instance, ExposureFilter(exposure_config.cache_capacity))
-        self.logger = logging.getLogger("Amplitude")
-        self.logger.addHandler(logging.StreamHandler())
-        if self.config.debug:
-            self.logger.setLevel(logging.DEBUG)
+        self.logger = self.config.logger
         self.__setup_connection_pool()
         self.lock = Lock()
         self.cohort_storage = InMemoryCohortStorage()
diff --git a/src/amplitude_experiment/local/config.py b/src/amplitude_experiment/local/config.py
@@ -1,3 +1,6 @@
+import logging
+import sys
+
 from ..assignment import AssignmentConfig
 from ..exposure import ExposureConfig
 from ..cohort.cohort_sync_config import CohortSyncConfig, DEFAULT_COHORT_SYNC_URL, EU_COHORT_SYNC_URL
@@ -23,7 +26,8 @@ def __init__(self, debug: bool = False,
                  stream_flag_conn_timeout: int = 1500,
                  assignment_config: AssignmentConfig = None,
                  exposure_config: ExposureConfig = None,
-                 cohort_sync_config: CohortSyncConfig = None):
+                 cohort_sync_config: CohortSyncConfig = None,
+                 logger: logging.Logger = None):
         """
         Initialize a config
            Parameters:
@@ -37,6 +41,8 @@ def __init__(self, debug: bool = False,
                 assignment_config (AssignmentConfig): The assignment configuration. @deprecated use exposure_config instead.
                 exposure_config (ExposureConfig): The exposure configuration.
                 cohort_sync_config (CohortSyncConfig): The cohort sync configuration.
+                logger (logging.Logger): Optional logger instance. If provided, this logger will be used instead of
+                  creating a new one. The debug flag only applies when no logger is provided.
 
            Returns:
                The config object
@@ -61,3 +67,16 @@ def __init__(self, debug: bool = False,
         self.stream_flag_conn_timeout = stream_flag_conn_timeout
         self.assignment_config = assignment_config
         self.exposure_config = exposure_config
+        # Set up logger: use provided logger or create default one
+        if logger is None:
+            self.logger = logging.getLogger("Amplitude")
+            # Only add handler if logger doesn't already have one
+            if not self.logger.handlers:
+                handler = logging.StreamHandler(sys.stderr)
+                self.logger.addHandler(handler)
+            # Set log level: DEBUG if debug=True, otherwise WARNING
+            # Only apply debug flag to default logger, not user-provided loggers
+            log_level = logging.DEBUG if self.debug else logging.WARNING
+            self.logger.setLevel(log_level)
+        else:
+            self.logger = logger
diff --git a/src/amplitude_experiment/remote/client.py b/src/amplitude_experiment/remote/client.py
@@ -1,11 +1,11 @@
 import json
-import logging
 import threading
 import time
 from time import sleep
 from typing import Any, Dict
 
 from .config import RemoteEvaluationConfig
+from .fetch_options import FetchOptions
 from ..connection_pool import HTTPConnectionPool
 from ..exception import FetchException
 from ..user import User
@@ -32,26 +32,24 @@ def __init__(self, api_key, config=None):
             raise ValueError("Experiment API key is empty")
         self.api_key = api_key
         self.config = config or RemoteEvaluationConfig()
-        self.logger = logging.getLogger("Amplitude")
-        self.logger.addHandler(logging.StreamHandler())
-        if self.config.debug:
-            self.logger.setLevel(logging.DEBUG)
+        self.logger = self.config.logger
         self.__setup_connection_pool()
 
-    def fetch_v2(self, user: User):
+    def fetch_v2(self, user: User, fetch_options: FetchOptions = None):
         """
         Fetch all variants for a user synchronously. This method will automatically retry if configured, and throw if
         all retries fail. This function differs from fetch as it will return a default variant object if the flag
         was evaluated but the user was not assigned (i.e. off).
 
             Parameters:
                 user (User): The Experiment User to fetch variants for.
+                fetch_options (FetchOptions): The Fetch Options
 
             Returns:
                 Variants Dictionary.
         """
         try:
-            return self.__fetch_internal(user)
+            return self.__fetch_internal(user, fetch_options)
         except Exception as e:
             self.logger.error(f"[Experiment] Failed to fetch variants: {e}")
             raise e
@@ -67,17 +65,18 @@ def fetch_async_v2(self, user: User, callback=None):
         thread.start()
 
     @deprecated("Use fetch_v2")
-    def fetch(self, user: User):
+    def fetch(self, user: User, fetch_options: FetchOptions = None):
         """
         Fetch all variants for a user synchronous. This method will automatically retry if configured.
             Parameters:
                 user (User): The Experiment User
+                fetch_options (FetchOptions): The Fetch Options
 
             Returns:
                 Variants Dictionary.
         """
         try:
-            variants = self.fetch_v2(user)
+            variants = self.fetch_v2(user, fetch_options)
             return self.__filter_default_variants(variants)
         except Exception:
             return {}
@@ -107,16 +106,16 @@ def __fetch_async_internal(self, user, callback):
                 callback(user, {}, e)
             return {}
 
-    def __fetch_internal(self, user):
+    def __fetch_internal(self, user, fetch_options: FetchOptions = None):
         self.logger.debug(f"[Experiment] Fetching variants for user: {user}")
         try:
-            return self.__do_fetch(user)
+            return self.__do_fetch(user, fetch_options)
         except Exception as e:
             self.logger.error(f"[Experiment] Fetch failed: {e}")
             if self.__should_retry_fetch(e):
-                return self.__retry_fetch(user)
+                return self.__retry_fetch(user, fetch_options)
 
-    def __retry_fetch(self, user):
+    def __retry_fetch(self, user, fetch_options: FetchOptions = None):
         if self.config.fetch_retries == 0:
             return {}
         self.logger.debug("[Experiment] Retrying fetch")
@@ -125,21 +124,26 @@ def __retry_fetch(self, user):
         for i in range(self.config.fetch_retries):
             sleep(delay_millis / 1000.0)
             try:
-                return self.__do_fetch(user)
+                return self.__do_fetch(user, fetch_options)
             except Exception as e:
                 self.logger.error(f"[Experiment] Retry failed: {e}")
                 err = e
             delay_millis = min(delay_millis * self.config.fetch_retry_backoff_scalar,
                                self.config.fetch_retry_backoff_max_millis)
         raise err
 
-    def __do_fetch(self, user):
+    def __do_fetch(self, user, fetch_options: FetchOptions = None):
         start = time.time()
         user_context = self.__add_context(user)
         headers = {
             'Authorization': f"Api-Key {self.api_key}",
             'Content-Type': 'application/json;charset=utf-8'
         }
+        if fetch_options and fetch_options.tracksAssignment is not None:
+            headers['X-Amp-Exp-Track'] = "track" if fetch_options.tracksAssignment else "no-track"
+        if fetch_options and fetch_options.tracksExposure is not None:
+            headers['X-Amp-Exp-Exposure-Track'] = "track" if fetch_options.tracksExposure else "no-track"
+
         conn = self._connection_pool.acquire()
         body = user_context.to_json().encode('utf8')
         if len(body) > 8000:
diff --git a/src/amplitude_experiment/remote/config.py b/src/amplitude_experiment/remote/config.py
@@ -2,6 +2,9 @@
 
 DEFAULT_SERVER_URL = 'https://api.lab.amplitude.com'
 EU_SERVER_URL = 'https://api.lab.eu.amplitude.com'
+import logging
+import sys
+
 
 class RemoteEvaluationConfig:
     """Experiment Remote Client Configuration"""
@@ -14,7 +17,8 @@ def __init__(self, debug=False,
                  fetch_retry_backoff_max_millis=10000,
                  fetch_retry_backoff_scalar=1.5,
                  fetch_retry_timeout_millis=10000,
-                 server_zone: ServerZone = ServerZone.US):
+                 server_zone: ServerZone = ServerZone.US,
+                 logger=None):
         """
         Initialize a config
             Parameters:
@@ -30,6 +34,8 @@ def __init__(self, debug=False,
                 fetch_retry_backoff_scalar (float): Scales the minimum backoff exponentially.
                 fetch_retry_timeout_millis (int): The request timeout for retrying fetch requests.
                 server_zone (str): Select the Amplitude data center to get flags and variants from, US or EU.
+                logger (logging.Logger): Optional logger instance. If provided, this logger will be used instead of
+                  creating a new one. The debug flag only applies when no logger is provided.
 
             Returns:
                 The config object
@@ -46,3 +52,16 @@ def __init__(self, debug=False,
         if server_url == DEFAULT_SERVER_URL and server_zone == ServerZone.EU:
             self.server_url = EU_SERVER_URL
 
+        # Set up logger: use provided logger or create default one
+        if logger is None:
+            self.logger = logging.getLogger("Amplitude")
+            # Only add handler if logger doesn't already have one
+            if not self.logger.handlers:
+                handler = logging.StreamHandler(sys.stderr)
+                self.logger.addHandler(handler)
+            # Set log level: DEBUG if debug=True, otherwise WARNING
+            # Only apply debug flag to default logger, not user-provided loggers
+            log_level = logging.DEBUG if self.debug else logging.WARNING
+            self.logger.setLevel(log_level)
+        else:
+            self.logger = logger
diff --git a/src/amplitude_experiment/remote/fetch_options.py b/src/amplitude_experiment/remote/fetch_options.py
@@ -0,0 +1,14 @@
+from typing import Optional
+class FetchOptions:
+    def __init__(self, tracksAssignment: Optional[bool] = None, tracksExposure: Optional[bool] = None):
+        """
+        Fetch Options
+            Parameters:
+                tracksAssignment (Optional[bool]): Whether to track the assignment. The default None uses the server's default behavior (track the assignment event).
+                tracksExposure (Optional[bool]): Whether to track the exposure. The default None uses the server's default behavior (don't track the exposure event).
+        """
+        self.tracksAssignment = tracksAssignment
+        self.tracksExposure = tracksExposure
+
+    def __str__(self):
+        return f"FetchOptions(tracksAssignment={self.tracksAssignment}, tracksExposure={self.tracksExposure})"
diff --git a/src/amplitude_experiment/version.py b/src/amplitude_experiment/version.py
@@ -1 +1 @@
-__version__ = "1.8.3"
+__version__ = "1.9.0"
diff --git a/tests/local/config_test.py b/tests/local/config_test.py
@@ -0,0 +1,70 @@
+import io
+import logging
+import sys
+import unittest
+
+from src.amplitude_experiment.local.config import LocalEvaluationConfig
+
+
+class LocalEvaluationConfigLoggerTestCase(unittest.TestCase):
+    """Tests for LocalEvaluationConfig logger configuration"""
+
+    def setUp(self):
+        """Clear existing handlers from the Amplitude logger before each test"""
+        logger = logging.getLogger("Amplitude")
+        logger.handlers.clear()
+        logger.setLevel(logging.NOTSET)
+
+    def tearDown(self):
+        """Clean up handlers after each test"""
+        logger = logging.getLogger("Amplitude")
+        logger.handlers.clear()
+
+    def test_default_logger_has_warning_level(self):
+        """Test that default logger has WARNING level and stderr handler when debug=False"""
+        config = LocalEvaluationConfig(debug=False)
+        self.assertEqual(config.logger.level, logging.WARNING)
+        self.assertEqual(len(config.logger.handlers), 1)
+        handler = config.logger.handlers[0]
+        self.assertIsInstance(handler, logging.StreamHandler)
+        self.assertEqual(handler.stream, sys.stderr)
+
+    def test_default_logger_has_debug_level_when_debug_true(self):
+        """Test that default logger has DEBUG level when debug=True"""
+        config = LocalEvaluationConfig(debug=True)
+        self.assertEqual(config.logger.level, logging.DEBUG)
+
+    def test_custom_logger_is_used(self):
+        """Test that provided custom logger is used"""
+        custom_logger = logging.getLogger("CustomLogger")
+        config = LocalEvaluationConfig(logger=custom_logger)
+        self.assertEqual(config.logger, custom_logger)
+        self.assertEqual(config.logger.name, "CustomLogger")
+
+    def test_custom_logger_level_not_modified_by_debug_flag(self):
+        """Test that custom logger level is not modified by debug flag"""
+        # Test with debug=True
+        custom_logger = logging.getLogger("CustomLoggerDebug")
+        custom_logger.setLevel(logging.ERROR)
+        config = LocalEvaluationConfig(debug=True, logger=custom_logger)
+        # Logger level should remain unchanged (ERROR), not modified to DEBUG
+        self.assertEqual(config.logger.level, logging.ERROR)
+        
+        # Test with debug=False
+        custom_logger2 = logging.getLogger("CustomLoggerWarning")
+        custom_logger2.setLevel(logging.INFO)
+        config2 = LocalEvaluationConfig(debug=False, logger=custom_logger2)
+        # Logger level should remain unchanged (INFO), not modified to WARNING
+        self.assertEqual(config2.logger.level, logging.INFO)
+
+    def test_default_logger_only_one_handler_added(self):
+        """Test that only one handler is added to default logger"""
+        config1 = LocalEvaluationConfig()
+        config2 = LocalEvaluationConfig()
+        # Both should use the same logger instance (singleton)
+        logger = logging.getLogger("Amplitude")
+        # Should only have one handler even after creating multiple configs
+        self.assertEqual(len(logger.handlers), 1)
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/remote/client_test.py b/tests/remote/client_test.py
diff --git a/tests/remote/config_test.py b/tests/remote/config_test.py
