Merge branch 'main' into sea-http-client

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -790,7 +790,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
         direct_results = resp.directResults
         has_been_closed_server_side = direct_results and direct_results.closeOperation
 
-        is_direct_results = (
+        has_more_rows = (
             (not direct_results)
             or (not direct_results.resultSet)
             or direct_results.resultSet.hasMoreRows
@@ -831,7 +831,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
             result_format=t_result_set_metadata_resp.resultFormat,
         )
 
-        return execute_response, is_direct_results
+        return execute_response, has_more_rows
 
     def get_execution_result(
         self, command_id: CommandId, cursor: Cursor
@@ -876,7 +876,7 @@ def get_execution_result(
 
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
         is_staging_operation = t_result_set_metadata_resp.isStagingOperation
-        is_direct_results = resp.hasMoreRows
+        has_more_rows = resp.hasMoreRows
 
         status = CommandState.from_thrift_state(resp.status) or CommandState.RUNNING
 
@@ -902,7 +902,7 @@ def get_execution_result(
             t_row_set=resp.results,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            has_more_rows=has_more_rows,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -1021,7 +1021,7 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response, is_direct_results = self._handle_execute_response(
+            execute_response, has_more_rows = self._handle_execute_response(
                 resp, cursor
             )
 
@@ -1039,8 +1039,7 @@ def execute_command(
                 t_row_set=t_row_set,
                 max_download_threads=self.max_download_threads,
                 ssl_options=self._ssl_options,
-                is_direct_results=is_direct_results,
-                session_id_hex=self._session_id_hex,
+                has_more_rows=has_more_rows,
             )
 
     def get_catalogs(
@@ -1062,9 +1061,7 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1080,8 +1077,7 @@ def get_catalogs(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
-            session_id_hex=self._session_id_hex,
+            has_more_rows=has_more_rows,
         )
 
     def get_schemas(
@@ -1109,9 +1105,7 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1127,8 +1121,7 @@ def get_schemas(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
-            session_id_hex=self._session_id_hex,
+            has_more_rows=has_more_rows,
         )
 
     def get_tables(
@@ -1160,9 +1153,7 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1178,8 +1169,7 @@ def get_tables(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
-            session_id_hex=self._session_id_hex,
+            has_more_rows=has_more_rows,
         )
 
     def get_columns(
@@ -1211,9 +1201,7 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
-            resp, cursor
-        )
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1229,8 +1217,7 @@ def get_columns(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
-            session_id_hex=self._session_id_hex,
+            has_more_rows=has_more_rows,
         )
 
     def _handle_execute_response(self, resp, cursor):

src/databricks/sql/common/http.py

@@ -5,8 +5,10 @@
 import threading
 from dataclasses import dataclass
 from contextlib import contextmanager
-from typing import Generator
+from typing import Generator, Optional
 import logging
+from requests.adapters import HTTPAdapter
+from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
 
 logger = logging.getLogger(__name__)
 
@@ -81,3 +83,70 @@ def execute(
 
     def close(self):
         self.session.close()
+
+
+class TelemetryHTTPAdapter(HTTPAdapter):
+    """
+    Custom HTTP adapter to prepare our DatabricksRetryPolicy before each request.
+    This ensures the retry timer is started and the command type is set correctly,
+    allowing the policy to manage its state for the duration of the request retries.
+    """
+
+    def send(self, request, **kwargs):
+        self.max_retries.command_type = CommandType.OTHER
+        self.max_retries.start_retry_timer()
+        return super().send(request, **kwargs)
+
+
+class TelemetryHttpClient:  # TODO: Unify all the http clients in the PySQL Connector
+    """Singleton HTTP client for sending telemetry data."""
+
+    _instance: Optional["TelemetryHttpClient"] = None
+    _lock = threading.Lock()
+
+    TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT = 3
+    TELEMETRY_RETRY_DELAY_MIN = 1.0
+    TELEMETRY_RETRY_DELAY_MAX = 10.0
+    TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION = 30.0
+
+    def __init__(self):
+        """Initializes the session and mounts the custom retry adapter."""
+        retry_policy = DatabricksRetryPolicy(
+            delay_min=self.TELEMETRY_RETRY_DELAY_MIN,
+            delay_max=self.TELEMETRY_RETRY_DELAY_MAX,
+            stop_after_attempts_count=self.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_COUNT,
+            stop_after_attempts_duration=self.TELEMETRY_RETRY_STOP_AFTER_ATTEMPTS_DURATION,
+            delay_default=1.0,
+            force_dangerous_codes=[],
+        )
+        adapter = TelemetryHTTPAdapter(max_retries=retry_policy)
+        self.session = requests.Session()
+        self.session.mount("https://", adapter)
+        self.session.mount("http://", adapter)
+
+    @classmethod
+    def get_instance(cls) -> "TelemetryHttpClient":
+        """Get the singleton instance of the TelemetryHttpClient."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    logger.debug("Initializing singleton TelemetryHttpClient")
+                    cls._instance = TelemetryHttpClient()
+        return cls._instance
+
+    def post(self, url: str, **kwargs) -> requests.Response:
+        """
+        Executes a POST request using the configured session.
+
+        This is a blocking call intended to be run in a background thread.
+        """
+        logger.debug("Executing telemetry POST request to: %s", url)
+        return self.session.post(url, **kwargs)
+
+    def close(self):
+        """Closes the underlying requests.Session."""
+        logger.debug("Closing TelemetryHttpClient session.")
+        self.session.close()
+        # Clear the instance to allow for re-initialization if needed
+        with TelemetryHttpClient._lock:
+            TelemetryHttpClient._instance = None

src/databricks/sql/exc.py

@@ -2,8 +2,6 @@
 import logging
 
 logger = logging.getLogger(__name__)
-from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
-
 
 ### PEP-249 Mandated ###
 # https://peps.python.org/pep-0249/#exceptions
@@ -22,6 +20,8 @@ def __init__(
 
         error_name = self.__class__.__name__
         if session_id_hex:
+            from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
+
             telemetry_client = TelemetryClientFactory.get_telemetry_client(
                 session_id_hex
             )

src/databricks/sql/result_set.py

@@ -43,7 +43,7 @@ def __init__(
         command_id: CommandId,
         status: CommandState,
         has_been_closed_server_side: bool = False,
-        is_direct_results: bool = False,
+        has_more_rows: bool = False,
         results_queue=None,
         description: List[Tuple] = [],
         is_staging_operation: bool = False,
@@ -61,7 +61,7 @@ def __init__(
             :param command_id: The command ID
             :param status: The command status
             :param has_been_closed_server_side: Whether the command has been closed on the server
-            :param is_direct_results: Whether the command has more rows
+            :param has_more_rows: Whether the command has more rows
             :param results_queue: The results queue
             :param description: column description of the results
             :param is_staging_operation: Whether the command is a staging operation
@@ -76,7 +76,7 @@ def __init__(
         self.command_id = command_id
         self.status = status
         self.has_been_closed_server_side = has_been_closed_server_side
-        self.is_direct_results = is_direct_results
+        self.has_more_rows = has_more_rows
         self.results = results_queue
         self._is_staging_operation = is_staging_operation
         self.lz4_compressed = lz4_compressed
@@ -170,7 +170,11 @@ def close(self) -> None:
         been closed on the server for some other reason, issue a request to the server to close it.
         """
         try:
-            self.results.close()
+            if self.results is not None:
+                self.results.close()
+            else:
+                logger.warning("result set close: queue not initialized")
+
             if (
                 self.status != CommandState.CLOSED
                 and not self.has_been_closed_server_side
@@ -193,14 +197,13 @@ 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,
         t_row_set=None,
         max_download_threads: int = 10,
         ssl_options=None,
-        is_direct_results: bool = True,
+        has_more_rows: bool = True,
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
@@ -215,13 +218,13 @@ def __init__(
             :param t_row_set: The TRowSet containing result data (if available)
             :param max_download_threads: Maximum number of download threads for cloud fetch
             :param ssl_options: SSL options for cloud fetch
-            :param is_direct_results: Whether there are more rows to 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
-        self.is_direct_results = is_direct_results
+        self.has_more_rows = has_more_rows
 
         # Build the results queue if t_row_set is provided
         results_queue = None
@@ -237,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__(
@@ -253,7 +256,7 @@ def __init__(
             command_id=execute_response.command_id,
             status=execute_response.status,
             has_been_closed_server_side=execute_response.has_been_closed_server_side,
-            is_direct_results=is_direct_results,
+            has_more_rows=has_more_rows,
             results_queue=results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
@@ -266,7 +269,7 @@ def __init__(
             self._fill_results_buffer()
 
     def _fill_results_buffer(self):
-        results, is_direct_results, result_links_count = self.backend.fetch_results(
+        results, has_more_rows, result_links_count = self.backend.fetch_results(
             command_id=self.command_id,
             max_rows=self.arraysize,
             max_bytes=self.buffer_size_bytes,
@@ -275,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.is_direct_results = is_direct_results
-        self.num_downloaded_chunks += result_links_count
+        self.has_more_rows = has_more_rows
+        self.num_chunks += result_links_count
 
     def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]
@@ -326,7 +329,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         while (
             n_remaining_rows > 0
             and not self.has_been_closed_server_side
-            and self.is_direct_results
+            and self.has_more_rows
         ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
@@ -351,7 +354,7 @@ def fetchmany_columnar(self, size: int):
         while (
             n_remaining_rows > 0
             and not self.has_been_closed_server_side
-            and self.is_direct_results
+            and self.has_more_rows
         ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
@@ -366,7 +369,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
         partial_result_chunks = [results]
-        while not self.has_been_closed_server_side and self.is_direct_results:
+        while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             if isinstance(results, ColumnTable) and isinstance(
@@ -392,7 +395,7 @@ def fetchall_columnar(self):
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
 
-        while not self.has_been_closed_server_side and self.is_direct_results:
+        while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             results = self.merge_columnar(results, partial_results)