removed default arg for batch size

Signed-off-by: Sai Shree Pradhan <[email protected]>

Author: saishreeeee

src/databricks/sql/client.py

@@ -251,7 +251,9 @@ def read(self) -> Optional[OAuthToken]:
         self.telemetry_enabled = (
             self.client_telemetry_enabled and self.server_telemetry_enabled
         )
-        self.telemetry_batch_size = kwargs.get("telemetry_batch_size")
+        self.telemetry_batch_size = kwargs.get(
+            "telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE
+        )
 
         user_agent_entry = kwargs.get("user_agent_entry")
         if user_agent_entry is None:

src/databricks/sql/telemetry/telemetry_client.py

@@ -132,22 +132,18 @@ class TelemetryClient(BaseTelemetryClient):
     TELEMETRY_AUTHENTICATED_PATH = "/telemetry-ext"
     TELEMETRY_UNAUTHENTICATED_PATH = "/telemetry-unauth"
 
-    DEFAULT_BATCH_SIZE = 100
-
     def __init__(
         self,
         telemetry_enabled,
         session_id_hex,
         auth_provider,
         host_url,
         executor,
-        batch_size=None,
+        batch_size,
     ):
         logger.debug("Initializing TelemetryClient for connection: %s", session_id_hex)
         self._telemetry_enabled = telemetry_enabled
-        self._batch_size = (
-            batch_size if batch_size is not None else self.DEFAULT_BATCH_SIZE
-        )
+        self._batch_size = batch_size
         self._session_id_hex = session_id_hex
         self._auth_provider = auth_provider
         self._user_agent = None
@@ -332,6 +328,8 @@ class TelemetryClientFactory:
     _flush_event = threading.Event()
     _flush_interval_seconds = 90
 
+    DEFAULT_BATCH_SIZE = 100
+
     @classmethod
     def _initialize(cls):
         """Initialize the factory if not already initialized"""
@@ -403,7 +401,7 @@ def initialize_telemetry_client(
         session_id_hex,
         auth_provider,
         host_url,
-        batch_size=None,
+        batch_size,
     ):
         """Initialize a telemetry client for a specific connection if telemetry is enabled"""
         try:

tests/unit/test_telemetry.py

@@ -31,6 +31,7 @@ def mock_telemetry_client():
         auth_provider=auth_provider,
         host_url="test-host.com",
         executor=executor,
+        batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
     )
 
 
@@ -208,7 +209,8 @@ def test_client_lifecycle_flow(self):
             telemetry_enabled=True,
             session_id_hex=session_id_hex,
             auth_provider=auth_provider,
-            host_url="test-host.com"
+            host_url="test-host.com",
+            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
@@ -232,7 +234,8 @@ def test_disabled_telemetry_flow(self):
             telemetry_enabled=False,
             session_id_hex=session_id_hex,
             auth_provider=None,
-            host_url="test-host.com"
+            host_url="test-host.com",
+            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
@@ -249,7 +252,8 @@ def test_factory_error_handling(self):
                 telemetry_enabled=True,
                 session_id_hex=session_id,
                 auth_provider=AccessTokenAuthProvider("token"),
-                host_url="test-host.com"
+                host_url="test-host.com",
+                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
             )
 
         # Should fall back to NoopTelemetryClient
@@ -267,7 +271,8 @@ def test_factory_shutdown_flow(self):
                 telemetry_enabled=True,
                 session_id_hex=session,
                 auth_provider=AccessTokenAuthProvider("token"),
-                host_url="test-host.com"
+                host_url="test-host.com",
+                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
             )
 
         # Factory should be initialized