refactor(backend/kernel): replace kernel_call context manager with explicit try/except

Using a ``with`` block for error translation is bad form: ``with``
conventionally signals resource lifecycle (locks, files), so
``with kernel_call("X"):`` hides the fact that the block raises a
mapped exception. Replace with explicit ``try/except Exception as
exc: raise _wrap_kernel_exception("X", exc) from exc`` at every
PyO3 call site.

What changed:

- ``_errors.py``: drop the ``kernel_call`` context manager; export
  ``wrap_kernel_exception(what, exc)`` — a pure function that maps
  a raw exception to a PEP 249 one (KernelError → mapped class via
  ``reraise_kernel_error``; existing Error → passthrough; anything
  else → OperationalError).
- ``client.py``: replace 12 ``with _kernel_call(...):`` blocks with
  inline try/except calling the helper.
- ``result_set.py``: same for the 3 sites (arrow_schema on
  construct, fetch_next_batch in _pull_one_batch, fetch_next_batch
  in _drain).

Behaviour is unchanged — same KernelError → PEP 249 mapping, same
non-KernelError → OperationalError wrapping. Just spelled in a way
that makes control flow visible at the call site and keeps
tracebacks one frame shorter (no ``__exit__`` frame).

87/87 kernel unit tests still pass.

Co-authored-by: Isaac
Signed-off-by: Vikrant Puppala <[email protected]>

Author: vikrantpuppala

src/databricks/sql/backend/kernel/_errors.py

@@ -20,15 +20,24 @@
 These primitives live in their own module so both ``client.py``
 (which orchestrates PyO3 calls) and ``result_set.py`` (which calls
 ``fetch_next_batch`` on the same kernel handles) can share them
-without ``result_set.py`` importing from ``client.py`` — that
-direction would be a layering violation.
+without ``result_set.py`` importing from ``client.py``.
+
+Usage at every PyO3 call site is a plain try/except:
+
+    try:
+        stmt.execute()
+    except Exception as exc:
+        raise wrap_kernel_exception("execute_command", exc) from exc
+
+The helper returns the mapped exception; callers raise it. Plain
+``try/except`` is preferred over a context manager: the control
+flow is visible at the call site, the helper is a pure function
+(trivial to test), and tracebacks don't carry an extra
+``__exit__`` frame.
 """
 
 from __future__ import annotations
 
-import contextlib
-from typing import Iterator
-
 from databricks.sql.exc import (
     DatabaseError,
     Error,
@@ -93,32 +102,24 @@ def reraise_kernel_error(exc: "_kernel.KernelError") -> "Error":
     return new
 
 
-@contextlib.contextmanager
-def kernel_call(what: str) -> Iterator[None]:
-    """Context manager that wraps a span of PyO3 calls so any error
-    crossing the Python/Rust boundary surfaces as a PEP 249
-    exception.
+def wrap_kernel_exception(what: str, exc: BaseException) -> "Error":
+    """Map any exception from a PyO3 call site to a PEP 249 exception.
 
-    ``KernelError`` flows through ``reraise_kernel_error`` (the
-    structured-attribute path). Anything else is wrapped in
-    ``OperationalError`` so DB-API callers see a uniform exception
-    surface and never have to catch native Python exceptions to
-    handle a connector-level failure.
+    - ``KernelError`` → mapped class with structured attrs forwarded.
+    - Already-PEP-249 ``Error`` (e.g. raised by an inner caller that
+      already mapped) → passed through unchanged.
+    - Anything else (``TypeError`` / ``ValueError`` / etc. from PyO3
+      argument conversion, extension-internal errors) → wrapped in
+      ``OperationalError``.
 
-    ``what`` is a short tag used only in the ``OperationalError``
-    message for the non-``KernelError`` path; keep it caller-named
-    (e.g. ``"execute_command"``).
+    Returned, not raised — the caller decides whether to ``raise``
+    or ``raise ... from exc``. ``what`` is a short tag (the calling
+    method name) used only in the ``OperationalError`` message.
     """
-    try:
-        yield
-    except _kernel.KernelError as exc:
-        raise reraise_kernel_error(exc) from exc
-    except Error:
-        # Already a PEP 249 error (e.g. a nested ``kernel_call`` or
-        # the cursor-side guard re-raising one); let it propagate
-        # unchanged.
-        raise
-    except Exception as exc:
-        raise OperationalError(
-            f"Unexpected error from databricks_sql_kernel during {what}: {exc!r}"
-        ) from exc
+    if isinstance(exc, _kernel.KernelError):
+        return reraise_kernel_error(exc)
+    if isinstance(exc, Error):
+        return exc
+    return OperationalError(
+        f"Unexpected error from databricks_sql_kernel during {what}: {exc!r}"
+    )

src/databricks/sql/backend/kernel/client.py

@@ -41,8 +41,8 @@
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.kernel._errors import (
     _kernel,
-    kernel_call as _kernel_call,
     reraise_kernel_error as _reraise_kernel_error,
+    wrap_kernel_exception as _wrap_kernel_exception,
 )
 from databricks.sql.backend.kernel.auth_bridge import kernel_auth_kwargs
 from databricks.sql.backend.kernel.result_set import KernelResultSet
@@ -137,7 +137,7 @@ def open_session(
         if session_configuration:
             session_conf = {k: str(v) for k, v in session_configuration.items()}
         try:
-            with _kernel_call("open_session"):
+            try:
                 self._kernel_session = _kernel.Session(
                     host=self._server_hostname,
                     http_path=self._http_path,
@@ -146,6 +146,8 @@ def open_session(
                     session_conf=session_conf,
                     **self._auth_kwargs,
                 )
+            except Exception as exc:
+                raise _wrap_kernel_exception("open_session", exc) from exc
         finally:
             # Drop the raw access token from the instance once the
             # kernel session is constructed (or failed). The kernel
@@ -222,10 +224,12 @@ def execute_command(
                 "Statement-level query_tags are not yet supported on the kernel backend."
             )
 
-        with _kernel_call("execute_command"):
+        try:
             stmt = self._kernel_session.statement()
+        except Exception as exc:
+            raise _wrap_kernel_exception("execute_command", exc) from exc
         try:
-            with _kernel_call("execute_command"):
+            try:
                 stmt.set_sql(operation)
                 if async_op:
                     async_exec = stmt.submit()
@@ -235,6 +239,8 @@ def execute_command(
                         self._async_handles[command_id.guid] = async_exec
                     return None
                 executed = stmt.execute()
+            except Exception as exc:
+                raise _wrap_kernel_exception("execute_command", exc) from exc
         finally:
             # ``Statement`` is a lifecycle owner separate from the
             # executed handle it produces. Drop it here so the
@@ -253,8 +259,10 @@ def execute_command(
         # can itself raise ``KernelError`` (or, in principle, a PyO3
         # native exception) — wrap the construction so callers see a
         # mapped PEP 249 exception.
-        with _kernel_call("execute_command"):
+        try:
             return self._make_result_set(executed, cursor, command_id)
+        except Exception as exc:
+            raise _wrap_kernel_exception("execute_command", exc) from exc
 
     def cancel_command(self, command_id: CommandId) -> None:
         with self._async_handles_lock:
@@ -266,8 +274,10 @@ def cancel_command(self, command_id: CommandId) -> None:
             # Match the Thrift backend's tolerant behaviour.
             logger.debug("cancel_command: no in-flight async handle for %s", command_id)
             return
-        with _kernel_call("cancel_command"):
+        try:
             handle.cancel()
+        except Exception as exc:
+            raise _wrap_kernel_exception("cancel_command", exc) from exc
 
     def close_command(self, command_id: CommandId) -> None:
         with self._async_handles_lock:
@@ -279,8 +289,10 @@ def close_command(self, command_id: CommandId) -> None:
         if handle is None:
             logger.debug("close_command: no tracked handle for %s", command_id)
             return
-        with _kernel_call("close_command"):
+        try:
             handle.close()
+        except Exception as exc:
+            raise _wrap_kernel_exception("close_command", exc) from exc
 
     def get_query_state(self, command_id: CommandId) -> CommandState:
         with self._async_handles_lock:
@@ -296,8 +308,10 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
             # ran sync and the result was materialised before
             # returning. Terminal by construction.
             return CommandState.SUCCEEDED
-        with _kernel_call("get_query_state"):
+        try:
             state, failure = handle.status()
+        except Exception as exc:
+            raise _wrap_kernel_exception("get_query_state", exc) from exc
         if state == "Failed" and failure is not None:
             # Surface server-reported failure as a database error so
             # the cursor's polling loop terminates with the right
@@ -318,8 +332,10 @@ def get_execution_result(
                 "get_execution_result called for an unknown command_id; "
                 "the kernel backend only tracks async-submitted statements."
             )
-        with _kernel_call("get_execution_result"):
+        try:
             stream = async_exec.await_result()
+        except Exception as exc:
+            raise _wrap_kernel_exception("get_execution_result", exc) from exc
         # The async-exec handle's role ends once it has produced the
         # ``ResultStream`` — keeping it around (and tracked in
         # ``_async_handles``) would leak the server-side
@@ -338,9 +354,11 @@ def get_execution_result(
                 exc,
             )
         # ``KernelResultSet.__init__`` calls ``arrow_schema()`` which
-        # can raise — keep that in the mapped-exception scope.
-        with _kernel_call("get_execution_result"):
+        # can raise — map that to PEP 249 too.
+        try:
             return self._make_result_set(stream, cursor, command_id)
+        except Exception as exc:
+            raise _wrap_kernel_exception("get_execution_result", exc) from exc
 
     # ── Metadata ───────────────────────────────────────────────────
 
@@ -382,9 +400,11 @@ def get_catalogs(
     ) -> "ResultSet":
         if self._kernel_session is None:
             raise InterfaceError("get_catalogs requires an open session.")
-        with _kernel_call("get_catalogs"):
+        try:
             stream = self._kernel_session.metadata().list_catalogs()
             return self._make_result_set(stream, cursor, self._synthetic_command_id())
+        except Exception as exc:
+            raise _wrap_kernel_exception("get_catalogs", exc) from exc
 
     def get_schemas(
         self,
@@ -397,12 +417,14 @@ def get_schemas(
     ) -> "ResultSet":
         if self._kernel_session is None:
             raise InterfaceError("get_schemas requires an open session.")
-        with _kernel_call("get_schemas"):
+        try:
             stream = self._kernel_session.metadata().list_schemas(
                 catalog=catalog_name,
                 schema_pattern=schema_name,
             )
             return self._make_result_set(stream, cursor, self._synthetic_command_id())
+        except Exception as exc:
+            raise _wrap_kernel_exception("get_schemas", exc) from exc
 
     def get_tables(
         self,
@@ -417,7 +439,7 @@ def get_tables(
     ) -> "ResultSet":
         if self._kernel_session is None:
             raise InterfaceError("get_tables requires an open session.")
-        with _kernel_call("get_tables"):
+        try:
             stream = self._kernel_session.metadata().list_tables(
                 catalog=catalog_name,
                 schema_pattern=schema_name,
@@ -448,6 +470,8 @@ def get_tables(
                 cursor,
                 self._synthetic_command_id(),
             )
+        except Exception as exc:
+            raise _wrap_kernel_exception("get_tables", exc) from exc
 
     def get_columns(
         self,
@@ -469,14 +493,16 @@ def get_columns(
             raise ProgrammingError(
                 "get_columns requires catalog_name on the kernel backend."
             )
-        with _kernel_call("get_columns"):
+        try:
             stream = self._kernel_session.metadata().list_columns(
                 catalog=catalog_name,
                 schema_pattern=schema_name,
                 table_pattern=table_name,
                 column_pattern=column_name,
             )
             return self._make_result_set(stream, cursor, self._synthetic_command_id())
+        except Exception as exc:
+            raise _wrap_kernel_exception("get_columns", exc) from exc
 
     # ── Misc ───────────────────────────────────────────────────────
 

src/databricks/sql/backend/kernel/result_set.py

@@ -36,7 +36,7 @@
 
 import pyarrow
 
-from databricks.sql.backend.kernel._errors import kernel_call
+from databricks.sql.backend.kernel._errors import wrap_kernel_exception
 from databricks.sql.backend.kernel.type_mapping import description_from_arrow_schema
 from databricks.sql.backend.types import CommandId, CommandState
 from databricks.sql.result_set import ResultSet
@@ -68,8 +68,10 @@ def __init__(
         arraysize: int,
         buffer_size_bytes: int,
     ):
-        with kernel_call("KernelResultSet.arrow_schema"):
+        try:
             schema = kernel_handle.arrow_schema()
+        except Exception as exc:
+            raise wrap_kernel_exception("KernelResultSet.arrow_schema", exc) from exc
         super().__init__(
             connection=connection,
             backend=backend,
@@ -107,8 +109,10 @@ def _pull_one_batch(self) -> bool:
         is exhausted."""
         if self._exhausted:
             return False
-        with kernel_call("fetch_next_batch"):
+        try:
             batch = self._kernel_handle.fetch_next_batch()
+        except Exception as exc:
+            raise wrap_kernel_exception("fetch_next_batch", exc) from exc
         if batch is None:
             self._exhausted = True
             self.has_more_rows = False
@@ -165,8 +169,10 @@ def _drain(self) -> pyarrow.Table:
             chunks.append(self._buffer.popleft())
         if not self._exhausted:
             while True:
-                with kernel_call("fetch_next_batch"):
+                try:
                     batch = self._kernel_handle.fetch_next_batch()
+                except Exception as exc:
+                    raise wrap_kernel_exception("fetch_next_batch", exc) from exc
                 if batch is None:
                     self._exhausted = True
                     self.has_more_rows = False