Adllow telemetry URL to be overridden with env var (#6946)

dagster-io · Mar 12, 2022 · 1244427 · 1244427
1 parent 94c8e7f
commit 1244427
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/python_modules/dagster/dagster/core/telemetry.py b/python_modules/dagster/dagster/core/telemetry.py
@@ -6,7 +6,7 @@
 See class TelemetryEntry for logged fields.
 
 For local development:
-  Spin up local telemetry server and set DAGSTER_TELEMETRY_URL = 'http://localhost:3000/actions'
+  Spin up local telemetry server and set the environment variable DAGSTER_TELEMETRY_URL = 'http://localhost:3000/actions'
   To test RotatingFileHandler, can set MAX_BYTES = 500
 """
 

diff --git a/python_modules/dagster/dagster/core/telemetry_upload.py b/python_modules/dagster/dagster/core/telemetry_upload.py
@@ -6,7 +6,9 @@
 
 from .telemetry import MAX_BYTES, get_dir_from_dagster_home
 
-DAGSTER_TELEMETRY_URL = "http://telemetry.dagster.io/actions"
+
+def get_dagster_telemetry_url():
+    return os.getenv("DAGSTER_TELEMETRY_URL", default="http://telemetry.dagster.io/actions")
 
 
 def is_running_in_test():
@@ -116,7 +118,7 @@ def _upload_logs(dagster_log_dir, log_size, dagster_log_queue_dir, raise_errors)
 
                     data = zlib.compress(byte, zlib.Z_BEST_COMPRESSION)
                     headers = {"content-encoding": "gzip"}
-                    r = requests.post(DAGSTER_TELEMETRY_URL, data=data, headers=headers)
+                    r = requests.post(get_dagster_telemetry_url(), data=data, headers=headers)
                     if r.status_code == 200:
                         success = True
                     retry_num += 1

diff --git a/python_modules/dagster/dagster_tests/core_tests/test_telemetry_upload.py b/python_modules/dagster/dagster_tests/core_tests/test_telemetry_upload.py
@@ -8,7 +8,7 @@
 
 from dagster.cli.pipeline import pipeline_execute_command
 from dagster.core.telemetry import cleanup_telemetry_logger
-from dagster.core.telemetry_upload import DAGSTER_TELEMETRY_URL, upload_logs
+from dagster.core.telemetry_upload import get_dagster_telemetry_url, upload_logs
 from dagster.core.test_utils import environ, instance_for_test
 from dagster.utils import pushd, script_relative_path
 
@@ -29,7 +29,7 @@ def test_dagster_telemetry_upload(env):
     for handler in logger.handlers:
         logger.removeHandler(handler)
 
-    responses.add(responses.POST, DAGSTER_TELEMETRY_URL)
+    responses.add(responses.POST, get_dagster_telemetry_url())
 
     with instance_for_test(overrides={"telemetry": {"enabled": True}}):
         with environ(env):
@@ -53,7 +53,7 @@ def side_effect(_):
             cleanup_telemetry_logger()
 
             upload_logs(mock_stop_event, raise_errors=True)
-            assert responses.assert_call_count(DAGSTER_TELEMETRY_URL, 1)
+            assert responses.assert_call_count(get_dagster_telemetry_url(), 1)
 
 
 @pytest.mark.parametrize(
@@ -77,4 +77,4 @@ def test_dagster_telemetry_no_test_env_upload(env):
                 )
 
             upload_logs(mock.MagicMock())
-            assert responses.assert_call_count(DAGSTER_TELEMETRY_URL, 0)
+            assert responses.assert_call_count(get_dagster_telemetry_url(), 0)