added statement type to command id

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/backend/thrift_backend.py

@@ -9,12 +9,11 @@
 from uuid import UUID
 
 from databricks.sql.result_set import ThriftResultSet
-
+from databricks.sql.telemetry.models.event import StatementType
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
     from databricks.sql.result_set import ResultSet
-    from databricks.sql.telemetry.models.event import StatementType
 
 from databricks.sql.backend.types import (
     CommandState,
@@ -833,7 +832,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
         return execute_response, is_direct_results
 
     def get_execution_result(
-        self, command_id: CommandId, cursor: "Cursor", statement_type: StatementType
+        self, command_id: CommandId, cursor: "Cursor"
     ) -> "ResultSet":
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle:
@@ -889,6 +888,7 @@ 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,
@@ -902,7 +902,6 @@ def get_execution_result(
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
             session_id_hex=self._session_id_hex,
-            statement_type=statement_type,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -968,7 +967,6 @@ def execute_command(
         max_bytes: int,
         lz4_compression: bool,
         cursor: Cursor,
-        statement_type: StatementType,
         use_cloud_fetch=True,
         parameters=[],
         async_op=False,
@@ -1030,6 +1028,8 @@ 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,
@@ -1042,7 +1042,6 @@ def execute_command(
                 ssl_options=self._ssl_options,
                 is_direct_results=is_direct_results,
                 session_id_hex=self._session_id_hex,
-                statement_type=statement_type,
             )
 
     def get_catalogs(
@@ -1051,7 +1050,6 @@ def get_catalogs(
         max_rows: int,
         max_bytes: int,
         cursor: "Cursor",
-        statement_type: StatementType,
     ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
@@ -1073,6 +1071,8 @@ 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,
@@ -1085,7 +1085,6 @@ def get_catalogs(
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
             session_id_hex=self._session_id_hex,
-            statement_id=statement_type,
         )
 
     def get_schemas(
@@ -1094,7 +1093,6 @@ def get_schemas(
         max_rows: int,
         max_bytes: int,
         cursor: Cursor,
-        statement_type: StatementType,
         catalog_name=None,
         schema_name=None,
     ) -> "ResultSet":
@@ -1122,6 +1120,8 @@ 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,
@@ -1134,7 +1134,6 @@ def get_schemas(
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
             session_id_hex=self._session_id_hex,
-            statement_type=statement_type,
         )
 
     def get_tables(
@@ -1143,7 +1142,6 @@ def get_tables(
         max_rows: int,
         max_bytes: int,
         cursor: Cursor,
-        statement_type: StatementType,
         catalog_name=None,
         schema_name=None,
         table_name=None,
@@ -1175,6 +1173,8 @@ 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,
@@ -1187,7 +1187,6 @@ def get_tables(
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
             session_id_hex=self._session_id_hex,
-            statement_type=statement_type,
         )
 
     def get_columns(
@@ -1196,7 +1195,6 @@ def get_columns(
         max_rows: int,
         max_bytes: int,
         cursor: Cursor,
-        statement_type: StatementType,
         catalog_name=None,
         schema_name=None,
         table_name=None,
@@ -1228,6 +1226,8 @@ 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,
@@ -1240,7 +1240,6 @@ def get_columns(
             ssl_options=self._ssl_options,
             is_direct_results=is_direct_results,
             session_id_hex=self._session_id_hex,
-            statement_type=statement_type,
         )
 
     def _handle_execute_response(self, resp, cursor):
@@ -1275,7 +1274,6 @@ def fetch_results(
         lz4_compressed: bool,
         arrow_schema_bytes,
         description,
-        statement_type,
         chunk_id: int,
         use_cloud_fetch=True,
     ):
@@ -1316,7 +1314,7 @@ def fetch_results(
             ssl_options=self._ssl_options,
             session_id_hex=self._session_id_hex,
             statement_id=command_id.to_hex_guid(),
-            statement_type=statement_type,
+            statement_type=command_id.statement_type,
             chunk_id=chunk_id,
         )
 

src/databricks/sql/backend/types.py

@@ -4,6 +4,7 @@
 import logging
 
 from databricks.sql.backend.utils.guid_utils import guid_to_hex_id
+from databricks.sql.telemetry.models.enums import StatementType
 from databricks.sql.thrift_api.TCLIService import ttypes
 
 logger = logging.getLogger(__name__)
@@ -281,6 +282,7 @@ def __init__(
         operation_type: Optional[int] = None,
         has_result_set: bool = False,
         modified_row_count: Optional[int] = None,
+        statement_type: Optional[StatementType] = None,
     ):
         """
         Initialize a CommandId.
@@ -300,6 +302,7 @@ def __init__(
         self.operation_type = operation_type
         self.has_result_set = has_result_set
         self.modified_row_count = modified_row_count
+        self._statement_type = statement_type
 
     def __str__(self) -> str:
         """
@@ -411,6 +414,19 @@ 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) -> Optional[StatementType]:
+        """
+        Get the statement type for this command.
+        """
+        return self._statement_type
+
 
 @dataclass
 class ExecuteResponse:

src/databricks/sql/client.py

@@ -870,7 +870,6 @@ def execute(
             async_op=False,
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
             row_limit=self.row_limit,
-            statement_type=self.statement_type,
         )
 
         if self.active_result_set and self.active_result_set.is_staging_operation:
@@ -930,7 +929,6 @@ def execute_async(
             async_op=True,
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
             row_limit=self.row_limit,
-            statement_type=self.statement_type,
         )
 
         return self
@@ -971,7 +969,7 @@ def get_async_execution_result(self):
         operation_state = self.get_query_state()
         if operation_state == CommandState.SUCCEEDED:
             self.active_result_set = self.backend.get_execution_result(
-                self.active_command_id, cursor=self, statement_type=self.statement_type
+                self.active_command_id, self
             )
 
             if self.active_result_set and self.active_result_set.is_staging_operation:
@@ -1016,7 +1014,6 @@ def catalogs(self) -> "Cursor":
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
             cursor=self,
-            statement_type=self.statement_type,
         )
         return self
 
@@ -1040,7 +1037,6 @@ def schemas(
             cursor=self,
             catalog_name=catalog_name,
             schema_name=schema_name,
-            statement_type=self.statement_type,
         )
         return self
 
@@ -1071,7 +1067,6 @@ def tables(
             schema_name=schema_name,
             table_name=table_name,
             table_types=table_types,
-            statement_type=self.statement_type,
         )
         return self
 
@@ -1102,7 +1097,6 @@ def columns(
             schema_name=schema_name,
             table_name=table_name,
             column_name=column_name,
-            statement_type=self.statement_type,
         )
         return self
 

src/databricks/sql/result_set.py

@@ -194,7 +194,6 @@ def __init__(
         execute_response: "ExecuteResponse",
         thrift_client: "ThriftDatabricksClient",
         session_id_hex: Optional[str],
-        statement_type: StatementType,
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
         use_cloud_fetch: bool = True,
@@ -218,7 +217,7 @@ def __init__(
             :param ssl_options: SSL options for cloud fetch
             :param is_direct_results: Whether there are more rows to fetch
         """
-        self.statement_type = statement_type
+        self.statement_type = execute_response.command_id.statement_type
         self.chunk_id = 0
 
         # Initialize ThriftResultSet-specific attributes
@@ -241,7 +240,7 @@ def __init__(
                 ssl_options=ssl_options,
                 session_id_hex=session_id_hex,
                 statement_id=execute_response.command_id.to_hex_guid(),
-                statement_type=statement_type,
+                statement_type=self.statement_type,
                 chunk_id=self.chunk_id,
             )
             if t_row_set and t_row_set.resultLinks:
@@ -278,7 +277,6 @@ def _fill_results_buffer(self):
             arrow_schema_bytes=self._arrow_schema_bytes,
             description=self.description,
             use_cloud_fetch=self._use_cloud_fetch,
-            statement_type=self.statement_type,
             chunk_id=self.chunk_id,
         )
         self.results = results

src/databricks/sql/telemetry/latency_logger.py

@@ -144,7 +144,7 @@ def get_extractor(obj):
         return None
 
 
-def log_latency(statement_type: StatementType = StatementType.NONE):
+def log_latency():
     """
     Decorator for logging execution latency and telemetry information.
 
@@ -158,11 +158,8 @@ def log_latency(statement_type: StatementType = StatementType.NONE):
     - Creates a SqlExecutionEvent with execution details
     - Sends the telemetry data asynchronously via TelemetryClient
 
-    Args:
-        statement_type (StatementType): The type of SQL statement being executed.
-
     Usage:
-        @log_latency(StatementType.SQL)
+        @log_latency()
         def execute(self, query):
             # Method implementation
             pass

tests/unit/test_client.py

@@ -129,7 +129,6 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                     execute_response=mock_execute_response,
                     thrift_client=mock_backend,
                     session_id_hex=Mock(),
-                    statement_type=Mock(),
                 )
 
                 # Mock execute_command to return our real result set
@@ -197,7 +196,6 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
             execute_response=Mock(),
             thrift_client=mock_backend,
             session_id_hex=Mock(),
-            statement_type=Mock(),
         )
         result_set.results = mock_results
 
@@ -225,7 +223,7 @@ def test_closing_result_set_hard_closes_commands(self):
 
         mock_thrift_backend.fetch_results.return_value = (Mock(), False, 0)
         result_set = ThriftResultSet(
-            mock_connection, mock_results_response, mock_thrift_backend, session_id_hex=Mock(), statement_type=Mock(),)
+            mock_connection, mock_results_response, mock_thrift_backend, session_id_hex=Mock())
 
         result_set.close()
 
@@ -271,7 +269,7 @@ def test_negative_fetch_throws_exception(self):
         mock_backend = Mock()
         mock_backend.fetch_results.return_value = (Mock(), False, 0)
 
-        result_set = ThriftResultSet(Mock(), Mock(), mock_backend, session_id_hex=Mock(), statement_type=Mock())
+        result_set = ThriftResultSet(Mock(), Mock(), mock_backend, session_id_hex=Mock())
 
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)