set statement type to query for chunk download

saishreeeee · saishreeeee · commit a35d3e9904d6 · 2025-07-16T15:41:09.000+05:30
Signed-off-by: Sai Shree Pradhan &lt;saishree.pradhan@databricks.com&gt;
diff --git a/src/databricks/sql/backend/sea/queue.py b/src/databricks/sql/backend/sea/queue.py
@@ -135,14 +135,12 @@ def __init__(
         super().__init__(
             max_download_threads=max_download_threads,
             ssl_options=ssl_options,
-            # TODO: fix these arguments when telemetry is implemented in SEA
-            session_id_hex=None,
             statement_id=statement_id,
-            statement_type=StatementType.NONE,
             chunk_id=0,
             schema_bytes=None,
             lz4_compressed=lz4_compressed,
             description=description,
+            session_id_hex=None,  # TODO: fix this argument when telemetry is implemented in SEA
         )
 
         self._sea_client = sea_client
diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py
@@ -889,7 +889,6 @@ def get_execution_result(
             arrow_schema_bytes=schema_bytes,
             result_format=t_result_set_metadata_resp.resultFormat,
         )
-        execute_response.command_id.set_statement_type(StatementType.QUERY)
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1029,8 +1028,6 @@ def execute_command(
             if resp.directResults and resp.directResults.resultSet:
                 t_row_set = resp.directResults.resultSet.results
 
-            execute_response.command_id.set_statement_type(StatementType.QUERY)
-
             return ThriftResultSet(
                 connection=cursor.connection,
                 execute_response=execute_response,
@@ -1072,8 +1069,6 @@ def get_catalogs(
         if resp.directResults and resp.directResults.resultSet:
             t_row_set = resp.directResults.resultSet.results
 
-        execute_response.command_id.set_statement_type(StatementType.METADATA)
-
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
@@ -1121,8 +1116,6 @@ def get_schemas(
         if resp.directResults and resp.directResults.resultSet:
             t_row_set = resp.directResults.resultSet.results
 
-        execute_response.command_id.set_statement_type(StatementType.METADATA)
-
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
@@ -1174,8 +1167,6 @@ def get_tables(
         if resp.directResults and resp.directResults.resultSet:
             t_row_set = resp.directResults.resultSet.results
 
-        execute_response.command_id.set_statement_type(StatementType.METADATA)
-
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
@@ -1227,8 +1218,6 @@ def get_columns(
         if resp.directResults and resp.directResults.resultSet:
             t_row_set = resp.directResults.resultSet.results
 
-        execute_response.command_id.set_statement_type(StatementType.METADATA)
-
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
@@ -1315,7 +1304,6 @@ def fetch_results(
             ssl_options=self._ssl_options,
             session_id_hex=self._session_id_hex,
             statement_id=command_id.to_hex_guid(),
-            statement_type=command_id.statement_type,
             chunk_id=chunk_id,
         )
 
diff --git a/src/databricks/sql/backend/types.py b/src/databricks/sql/backend/types.py
@@ -301,7 +301,6 @@ def __init__(
         self.operation_type = operation_type
         self.has_result_set = has_result_set
         self.modified_row_count = modified_row_count
-        self._statement_type = StatementType.NONE
 
     def __str__(self) -> str:
         """
@@ -413,19 +412,6 @@ def to_hex_guid(self) -> str:
         else:
             return str(self.guid)
 
-    def set_statement_type(self, statement_type: StatementType):
-        """
-        Set the statement type for this command.
-        """
-        self._statement_type = statement_type
-
-    @property
-    def statement_type(self) -> StatementType:
-        """
-        Get the statement type for this command.
-        """
-        return self._statement_type
-
 
 @dataclass
 class ExecuteResponse:
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
@@ -708,7 +708,7 @@ def _handle_staging_operation(
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
-    @log_latency()
+    @log_latency(StatementType.SQL)
     def _handle_staging_put(
         self, presigned_url: str, local_file: str, headers: Optional[dict] = None
     ):
@@ -717,7 +717,6 @@ def _handle_staging_put(
         Raise an exception if request fails. Returns no data.
         """
 
-        self.statement_type = StatementType.SQL
         if local_file is None:
             raise ProgrammingError(
                 "Cannot perform PUT without specifying a local_file",
@@ -749,7 +748,7 @@ def _handle_staging_put(
                 + "but not yet applied on the server. It's possible this command may fail later."
             )
 
-    @log_latency()
+    @log_latency(StatementType.SQL)
     def _handle_staging_get(
         self, local_file: str, presigned_url: str, headers: Optional[dict] = None
     ):
@@ -758,7 +757,6 @@ def _handle_staging_get(
         Raise an exception if request fails. Returns no data.
         """
 
-        self.statement_type = StatementType.SQL
         if local_file is None:
             raise ProgrammingError(
                 "Cannot perform GET without specifying a local_file",
@@ -778,13 +776,12 @@ def _handle_staging_get(
         with open(local_file, "wb") as fp:
             fp.write(r.content)
 
-    @log_latency()
+    @log_latency(StatementType.SQL)
     def _handle_staging_remove(
         self, presigned_url: str, headers: Optional[dict] = None
     ):
         """Make an HTTP DELETE request to the presigned_url"""
 
-        self.statement_type = StatementType.SQL
         r = requests.delete(url=presigned_url, headers=headers)
 
         if not r.ok:
@@ -793,7 +790,7 @@ def _handle_staging_remove(
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
-    @log_latency()
+    @log_latency(StatementType.QUERY)
     def execute(
         self,
         operation: str,
@@ -832,7 +829,6 @@ def execute(
         :returns self
         """
 
-        self.statement_type = StatementType.QUERY
         logger.debug(
             "Cursor.execute(operation=%s, parameters=%s)", operation, parameters
         )
@@ -879,7 +875,7 @@ def execute(
 
         return self
 
-    @log_latency()
+    @log_latency(StatementType.QUERY)
     def execute_async(
         self,
         operation: str,
@@ -895,7 +891,6 @@ def execute_async(
         :return:
         """
 
-        self.statement_type = StatementType.QUERY
         param_approach = self._determine_parameter_approach(parameters)
         if param_approach == ParameterApproach.NONE:
             prepared_params = NO_NATIVE_PARAMS
@@ -999,14 +994,13 @@ def executemany(self, operation, seq_of_parameters):
             self.execute(operation, parameters)
         return self
 
-    @log_latency()
+    @log_latency(StatementType.METADATA)
     def catalogs(self) -> "Cursor":
         """
         Get all available catalogs.
 
         :returns self
         """
-        self.statement_type = StatementType.METADATA
         self._check_not_closed()
         self._close_and_clear_active_result_set()
         self.active_result_set = self.backend.get_catalogs(
@@ -1017,7 +1011,7 @@ def catalogs(self) -> "Cursor":
         )
         return self
 
-    @log_latency()
+    @log_latency(StatementType.METADATA)
     def schemas(
         self, catalog_name: Optional[str] = None, schema_name: Optional[str] = None
     ) -> "Cursor":
@@ -1027,7 +1021,6 @@ def schemas(
         Names can contain % wildcards.
         :returns self
         """
-        self.statement_type = StatementType.METADATA
         self._check_not_closed()
         self._close_and_clear_active_result_set()
         self.active_result_set = self.backend.get_schemas(
@@ -1040,7 +1033,7 @@ def schemas(
         )
         return self
 
-    @log_latency()
+    @log_latency(StatementType.METADATA)
     def tables(
         self,
         catalog_name: Optional[str] = None,
@@ -1054,7 +1047,6 @@ def tables(
         Names can contain % wildcards.
         :returns self
         """
-        self.statement_type = StatementType.METADATA
         self._check_not_closed()
         self._close_and_clear_active_result_set()
 
@@ -1070,7 +1062,7 @@ def tables(
         )
         return self
 
-    @log_latency()
+    @log_latency(StatementType.METADATA)
     def columns(
         self,
         catalog_name: Optional[str] = None,
@@ -1084,7 +1076,6 @@ def columns(
         Names can contain % wildcards.
         :returns self
         """
-        self.statement_type = StatementType.METADATA
         self._check_not_closed()
         self._close_and_clear_active_result_set()
 
diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py
@@ -24,7 +24,6 @@ def __init__(
         ssl_options: SSLOptions,
         session_id_hex: Optional[str],
         statement_id: str,
-        statement_type: StatementType,
         chunk_id: int,
     ):
         self._pending_links: List[Tuple[int, TSparkArrowResultLink]] = []
@@ -48,7 +47,6 @@ def __init__(
         self._ssl_options = ssl_options
         self.session_id_hex = session_id_hex
         self.statement_id = statement_id
-        self.statement_type = statement_type
 
     def get_next_downloaded_file(
         self, next_row_offset: int
@@ -111,7 +109,6 @@ def _schedule_downloads(self):
                 chunk_id=chunk_id,
                 session_id_hex=self.session_id_hex,
                 statement_id=self.statement_id,
-                statement_type=self.statement_type,
             )
             task = self._thread_pool.submit(handler.run)
             self._download_tasks.append(task)
diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py
@@ -72,17 +72,15 @@ def __init__(
         chunk_id: int,
         session_id_hex: Optional[str],
         statement_id: str,
-        statement_type: StatementType,
     ):
         self.settings = settings
         self.link = link
         self._ssl_options = ssl_options
         self.chunk_id = chunk_id
         self.session_id_hex = session_id_hex
         self.statement_id = statement_id
-        self.statement_type = statement_type
 
-    @log_latency()
+    @log_latency(StatementType.QUERY)
     def run(self) -> DownloadedFile:
         """
         Download the file described in the cloud fetch link.
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
@@ -217,7 +217,6 @@ def __init__(
             :param ssl_options: SSL options for cloud fetch
             :param is_direct_results: Whether there are more rows to fetch
         """
-        self.statement_type = execute_response.command_id.statement_type
         self.num_downloaded_chunks = 0
 
         # Initialize ThriftResultSet-specific attributes
@@ -240,7 +239,6 @@ def __init__(
                 ssl_options=ssl_options,
                 session_id_hex=session_id_hex,
                 statement_id=execute_response.command_id.to_hex_guid(),
-                statement_type=self.statement_type,
                 chunk_id=self.num_downloaded_chunks,
             )
             if t_row_set and t_row_set.resultLinks:
diff --git a/src/databricks/sql/telemetry/latency_logger.py b/src/databricks/sql/telemetry/latency_logger.py
@@ -43,9 +43,6 @@ def get_retry_count(self):
     def get_chunk_id(self):
         pass
 
-    def get_statement_type(self):
-        pass
-
 
 class CursorExtractor(TelemetryExtractor):
     """
@@ -86,9 +83,6 @@ def get_retry_count(self) -> int:
     def get_chunk_id(self):
         return None
 
-    def get_statement_type(self):
-        return self.statement_type
-
 
 class ResultSetDownloadHandlerExtractor(TelemetryExtractor):
     """
@@ -114,9 +108,6 @@ def get_retry_count(self) -> Optional[int]:
     def get_chunk_id(self) -> Optional[int]:
         return self._obj.chunk_id
 
-    def get_statement_type(self):
-        return self.statement_type
-
 
 def get_extractor(obj):
     """
@@ -144,7 +135,7 @@ def get_extractor(obj):
         return None
 
 
-def log_latency():
+def log_latency(statement_type: StatementType = StatementType.NONE):
     """
     Decorator for logging execution latency and telemetry information.
 
@@ -159,7 +150,7 @@ def log_latency():
     - Sends the telemetry data asynchronously via TelemetryClient
 
     Usage:
-        @log_latency()
+        @log_latency(StatementType.QUERY)
         def execute(self, query):
             # Method implementation
             pass
@@ -199,7 +190,7 @@ def _safe_call(func_to_call):
                     statement_id = _safe_call(extractor.get_statement_id)
 
                     sql_exec_event = SqlExecutionEvent(
-                        statement_type=_safe_call(extractor.get_statement_type),
+                        statement_type=statement_type,
                         is_compressed=_safe_call(extractor.get_is_compressed),
                         execution_result=_safe_call(
                             extractor.get_execution_result_format
diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py
diff --git a/tests/unit/test_cloud_fetch_queue.py b/tests/unit/test_cloud_fetch_queue.py
diff --git a/tests/unit/test_download_manager.py b/tests/unit/test_download_manager.py
diff --git a/tests/unit/test_downloader.py b/tests/unit/test_downloader.py
