Merge remote-tracking branch 'origin/sea-migration' into backend-interface

Author: varun-edachali-dbx

.github/CODEOWNERS

@@ -2,4 +2,4 @@
 # the repo. Unless a later match takes precedence, these
 # users will be requested for review when someone opens a
 # pull request.
-*       @deeksha-db @samikshya-db @jprakash-db @yunbodeng-db @jackyhu-db @benc-db
+*       @deeksha-db @samikshya-db @jprakash-db @jackyhu-db @madhav-db @gopalldb @jayantsing-db @vikrantpuppala @shivam2680

.github/workflows/code-quality-checks.yml

@@ -3,9 +3,13 @@ on:
   push:
     branches:
       - main
+      - sea-migration
+      - telemetry
   pull_request:
     branches:
       - main
+      - sea-migration
+      - telemetry
 jobs:
   run-unit-tests:
     runs-on: ubuntu-latest

.github/workflows/integration.yml

@@ -4,6 +4,11 @@ on:
     paths-ignore:
       - "**.MD"
       - "**.md"
+  pull_request:
+    branches:
+      - main
+      - sea-migration
+      - telemetry
 
 jobs:
   run-e2e-tests:

src/databricks/sql/backend/thrift_backend.py

@@ -132,6 +132,13 @@ def __init__(
         # max_download_threads
         #  Number of threads for handling cloud fetch downloads. Defaults to 10
 
+        logger.debug(
+            "ThriftBackend.__init__(server_hostname=%s, port=%s, http_path=%s)",
+            server_hostname,
+            port,
+            http_path,
+        )
+
         port = port or 443
         if kwargs.get("_connection_uri"):
             uri = kwargs.get("_connection_uri")
@@ -403,6 +410,8 @@ def attempt_request(attempt):
 
                 # TODO: don't use exception handling for GOS polling...
 
+                logger.error("ThriftBackend.attempt_request: HTTPError: %s", err)
+
                 gos_name = TCLIServiceClient.GetOperationStatus.__name__
                 if method.__name__ == gos_name:
                     delay_default = (
@@ -447,6 +456,7 @@ def attempt_request(attempt):
                     else:
                         logger.warning(log_string)
             except Exception as err:
+                logger.error("ThriftBackend.attempt_request: Exception: %s", err)
                 error = err
                 retry_delay = extract_retry_delay(attempt)
                 error_message = (
@@ -904,6 +914,12 @@ def execute_command(
     ):
         assert session_handle is not None
 
+        logger.debug(
+            "ThriftBackend.execute_command(operation=%s, session_handle=%s)",
+            operation,
+            session_handle,
+        )
+
         spark_arrow_types = ttypes.TSparkArrowTypes(
             timestampAsArrow=self._use_arrow_native_timestamps,
             decimalAsArrow=self._use_arrow_native_decimals,
@@ -1090,6 +1106,7 @@ def fetch_results(
         return queue, resp.hasMoreRows
 
     def close_command(self, op_handle):
+        logger.debug("ThriftBackend.close_command(op_handle=%s)", op_handle)
         req = ttypes.TCloseOperationReq(operationHandle=op_handle)
         resp = self.make_request(self._client.CloseOperation, req)
         return resp.status

src/databricks/sql/client.py

@@ -216,6 +216,12 @@ def read(self) -> Optional[OAuthToken]:
         # use_cloud_fetch
         # Enable use of cloud fetch to extract large query results in parallel via cloud storage
 
+        logger.debug(
+            "Connection.__init__(server_hostname=%s, http_path=%s)",
+            server_hostname,
+            http_path,
+        )
+
         if access_token:
             access_token_kv = {"access_token": access_token}
             kwargs = {**kwargs, **access_token_kv}
@@ -278,7 +284,13 @@ def __enter__(self) -> "Connection":
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self.close()
+        try:
+            self.close()
+        except BaseException as e:
+            logger.warning(f"Exception during connection close in __exit__: {e}")
+            if exc_type is None:
+                raise
+        return False
 
     def __del__(self):
         if self.open:
@@ -401,7 +413,14 @@ def __enter__(self) -> "Cursor":
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self.close()
+        try:
+            logger.debug("Cursor context manager exiting, calling close()")
+            self.close()
+        except BaseException as e:
+            logger.warning(f"Exception during cursor close in __exit__: {e}")
+            if exc_type is None:
+                raise
+        return False
 
     def __iter__(self):
         if self.active_result_set:
@@ -732,6 +751,9 @@ def execute(
 
         :returns self
         """
+        logger.debug(
+            "Cursor.execute(operation=%s, parameters=%s)", operation, parameters
+        )
 
         param_approach = self._determine_parameter_approach(parameters)
         if param_approach == ParameterApproach.NONE:
@@ -1112,7 +1134,21 @@ def cancel(self) -> None:
     def close(self) -> None:
         """Close cursor"""
         self.open = False
-        self.active_op_handle = None
+
+        # Close active operation handle if it exists
+        if self.active_op_handle:
+            try:
+                self.thrift_backend.close_command(self.active_op_handle)
+            except RequestError as e:
+                if isinstance(e.args[1], CursorAlreadyClosedError):
+                    logger.info("Operation was canceled by a prior request")
+                else:
+                    logging.warning(f"Error closing operation handle: {e}")
+            except Exception as e:
+                logging.warning(f"Error closing operation handle: {e}")
+            finally:
+                self.active_op_handle = None
+
         if self.active_result_set:
             self._close_and_clear_active_result_set()
 

tests/e2e/test_driver.py

@@ -50,7 +50,7 @@
 
 from tests.e2e.common.uc_volume_tests import PySQLUCVolumeTestSuiteMixin
 
-from databricks.sql.exc import SessionAlreadyClosedError
+from databricks.sql.exc import SessionAlreadyClosedError, CursorAlreadyClosedError
 
 log = logging.getLogger(__name__)
 
@@ -820,7 +820,6 @@ def test_close_connection_closes_cursors(self):
             ars = cursor.active_result_set
 
             # We must manually run this check because thrift_backend always forces `has_been_closed_server_side` to True
-
             # Cursor op state should be open before connection is closed
             status_request = ttypes.TGetOperationStatusReq(
                 operationHandle=ars.command_id, getProgressUpdate=False
@@ -847,9 +846,104 @@ def test_closing_a_closed_connection_doesnt_fail(self, caplog):
         with self.connection() as conn:
             # First .close() call is explicit here
             conn.close()
-
         assert "Session appears to have been closed already" in caplog.text
 
+        conn = None
+        try:
+            with pytest.raises(KeyboardInterrupt):
+                with self.connection() as c:
+                    conn = c
+                    raise KeyboardInterrupt("Simulated interrupt")
+        finally:
+            if conn is not None:
+                assert not conn.open, "Connection should be closed after KeyboardInterrupt"
+
+    def test_cursor_close_properly_closes_operation(self):
+        """Test that Cursor.close() properly closes the active operation handle on the server."""
+        with self.connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute("SELECT 1 AS test")
+                assert cursor.active_op_handle is not None
+                cursor.close()
+                assert cursor.active_op_handle is None
+                assert not cursor.open
+            finally:
+                if cursor.open:
+                    cursor.close()
+
+        conn = None
+        cursor = None
+        try:
+            with self.connection() as c:
+                conn = c
+                with pytest.raises(KeyboardInterrupt):
+                    with conn.cursor() as cur:
+                        cursor = cur
+                        raise KeyboardInterrupt("Simulated interrupt")
+        finally:
+            if cursor is not None:
+                assert not cursor.open, "Cursor should be closed after KeyboardInterrupt"
+
+    def test_nested_cursor_context_managers(self):
+        """Test that nested cursor context managers properly close operations on the server."""
+        with self.connection() as conn:
+            with conn.cursor() as cursor1:
+                cursor1.execute("SELECT 1 AS test1")
+                assert cursor1.active_op_handle is not None
+
+                with conn.cursor() as cursor2:
+                    cursor2.execute("SELECT 2 AS test2")
+                    assert cursor2.active_op_handle is not None
+
+                # After inner context manager exit, cursor2 should be not open
+                assert not cursor2.open
+                assert cursor2.active_op_handle is None
+
+            # After outer context manager exit, cursor1 should be not open
+            assert not cursor1.open
+            assert cursor1.active_op_handle is None
+
+    def test_cursor_error_handling(self):
+        """Test that cursor close handles errors properly to prevent orphaned operations."""
+        with self.connection() as conn:
+            cursor = conn.cursor()
+
+            cursor.execute("SELECT 1 AS test")
+
+            op_handle = cursor.active_op_handle
+
+            assert op_handle is not None
+
+            # Manually close the operation to simulate server-side closure
+            conn.thrift_backend.close_command(op_handle)
+
+            cursor.close()
+
+            assert not cursor.open
+
+    def test_result_set_close(self):
+        """Test that ResultSet.close() properly closes operations on the server and handles state correctly."""
+        with self.connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute("SELECT * FROM RANGE(10)")
+
+                result_set = cursor.active_result_set
+                assert result_set is not None
+
+                initial_op_state = result_set.op_state
+
+                result_set.close()
+
+                assert result_set.op_state == result_set.thrift_backend.CLOSED_OP_STATE
+                assert result_set.op_state != initial_op_state
+
+                # Closing the result set again should be a no-op and not raise exceptions
+                result_set.close()
+            finally:
+                cursor.close()
+
 
 # use a RetrySuite to encapsulate these tests which we'll typically want to run together; however keep
 # the 429/503 subsuites separate since they execute under different circumstances.