minor telemetry changes

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

Author: saishreeeee

src/databricks/sql/backend/thrift_backend.py

@@ -895,7 +895,6 @@ def get_execution_result(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
-            session_id_hex=self._session_id_hex,
             buffer_size_bytes=cursor.buffer_size_bytes,
             arraysize=cursor.arraysize,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
@@ -1040,7 +1039,6 @@ def execute_command(
                 max_download_threads=self.max_download_threads,
                 ssl_options=self._ssl_options,
                 has_more_rows=has_more_rows,
-                session_id_hex=self._session_id_hex,
             )
 
     def get_catalogs(
@@ -1079,7 +1077,6 @@ def get_catalogs(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             has_more_rows=has_more_rows,
-            session_id_hex=self._session_id_hex,
         )
 
     def get_schemas(
@@ -1124,7 +1121,6 @@ def get_schemas(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             has_more_rows=has_more_rows,
-            session_id_hex=self._session_id_hex,
         )
 
     def get_tables(
@@ -1173,7 +1169,6 @@ def get_tables(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             has_more_rows=has_more_rows,
-            session_id_hex=self._session_id_hex,
         )
 
     def get_columns(
@@ -1222,7 +1217,6 @@ def get_columns(
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
             has_more_rows=has_more_rows,
-            session_id_hex=self._session_id_hex,
         )
 
     def _handle_execute_response(self, resp, cursor):

src/databricks/sql/result_set.py

@@ -197,7 +197,6 @@ def __init__(
         connection: Connection,
         execute_response: ExecuteResponse,
         thrift_client: ThriftDatabricksClient,
-        session_id_hex: Optional[str],
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
         use_cloud_fetch: bool = True,
@@ -221,7 +220,7 @@ def __init__(
             :param ssl_options: SSL options for cloud fetch
             :param has_more_rows: Whether there are more rows to fetch
         """
-        self.num_downloaded_chunks = 0
+        self.num_chunks = 0
 
         # Initialize ThriftResultSet-specific attributes
         self._use_cloud_fetch = use_cloud_fetch
@@ -241,12 +240,12 @@ def __init__(
                 lz4_compressed=execute_response.lz4_compressed,
                 description=execute_response.description,
                 ssl_options=ssl_options,
-                session_id_hex=session_id_hex,
+                session_id_hex=connection.get_session_id_hex(),
                 statement_id=execute_response.command_id.to_hex_guid(),
-                chunk_id=self.num_downloaded_chunks,
+                chunk_id=self.num_chunks,
             )
             if t_row_set.resultLinks:
-                self.num_downloaded_chunks += len(t_row_set.resultLinks)
+                self.num_chunks += len(t_row_set.resultLinks)
 
         # Call parent constructor with common attributes
         super().__init__(
@@ -279,11 +278,11 @@ def _fill_results_buffer(self):
             arrow_schema_bytes=self._arrow_schema_bytes,
             description=self.description,
             use_cloud_fetch=self._use_cloud_fetch,
-            chunk_id=self.num_downloaded_chunks,
+            chunk_id=self.num_chunks,
         )
         self.results = results
         self.has_more_rows = has_more_rows
-        self.num_downloaded_chunks += result_links_count
+        self.num_chunks += result_links_count
 
     def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]

tests/unit/test_client.py

@@ -127,7 +127,6 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                     connection=connection,
                     execute_response=mock_execute_response,
                     thrift_client=mock_backend,
-                    session_id_hex=Mock(),
                 )
 
                 # Mock execute_command to return our real result set
@@ -187,7 +186,6 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
             connection=mock_connection,
             execute_response=Mock(),
             thrift_client=mock_backend,
-            session_id_hex=Mock(),
         )
         result_set.results = mock_results
 
@@ -217,7 +215,6 @@ def test_closing_result_set_hard_closes_commands(self):
             mock_connection,
             mock_results_response,
             mock_thrift_backend,
-            session_id_hex=Mock(),
         )
         result_set.results = mock_results
 
@@ -266,7 +263,7 @@ def test_negative_fetch_throws_exception(self):
         mock_backend.fetch_results.return_value = (Mock(), False, 0)
 
         result_set = ThriftResultSet(
-            Mock(), Mock(), mock_backend, session_id_hex=Mock()
+            Mock(), Mock(), mock_backend
         )
 
         with self.assertRaises(ValueError) as e:

tests/unit/test_fetches.py

@@ -63,7 +63,6 @@ def make_dummy_result_set_from_initial_results(initial_results):
             ),
             thrift_client=mock_thrift_backend,
             t_row_set=None,
-            session_id_hex=Mock(),
         )
         return rs
 
@@ -108,7 +107,6 @@ def fetch_results(
                 is_staging_operation=False,
             ),
             thrift_client=mock_thrift_backend,
-            session_id_hex=Mock(),
         )
         return rs