fix(backend/kernel): retain parent Statement for async-submitted commands

The kernel's `Statement.close()` invalidates any executed handle the
Statement produced (see `databricks-sql-kernel/src/statement/validity.rs`
and `mutable::close`). Today `KernelDatabricksClient.execute_command`
closes `stmt` in `finally` regardless of `async_op` — which is correct
for the sync path (`stmt.execute()` returns a fully-materialised result
before the close) but wrong for the async path: as soon as the user
calls `cursor.execute_async(...)` and tries to poll or fetch, the
async handle is already invalidated and the next operation raises
"executed-statement handle has been invalidated by a re-execute on its
parent Statement".

This makes the entire async surface
(`execute_async` + `is_query_pending`/`get_query_state` +
`get_async_execution_result`) effectively non-functional on the kernel
backend today.

Fix: when `async_op=True`, retain the parent Statement in a new
`_async_statements` dict alongside `_async_handles`, and close it in
`close_command`, `close_session`, and `get_execution_result` (which
already closes the async handle once the result stream is produced).
The sync path is unchanged.

Discovered via the python-comparator audit harness running the
`STATEMENT_ASYNC` suite — the kernel side produced 0 HTTP requests
because every async cursor operation short-circuited locally on the
invalidated handle.

Test:
  conn = sql.connect(..., use_kernel=True)
  cur = conn.cursor()
  cur.execute_async("SELECT 1 AS x")
  while cur.is_query_pending(): time.sleep(0.2)
  cur.get_async_execution_result()
  assert cur.fetchall() == [(1,)]
Previously raised KernelError("InvalidStatementHandle: …") on
`is_query_pending`; now succeeds.

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

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

@@ -129,6 +129,15 @@ def __init__(
         # Guarded by ``_async_handles_lock`` so concurrent cursors on the
         # same connection don't race on submit / close / close-session.
         self._async_handles: Dict[str, Any] = {}
+        # Parent ``Statement`` objects kept alive alongside async handles.
+        # On the kernel, ``Statement.close()`` flips the validity flag on
+        # the produced executed handle (see kernel
+        # ``statement::mutable::close``), so we cannot close the
+        # Statement immediately after ``submit()`` as we do for sync
+        # ``execute()``. Instead retain it here and close it in
+        # ``close_command`` / ``close_session`` after the async handle
+        # has finished its work.
+        self._async_statements: Dict[str, Any] = {}
         # CommandId.guids of async commands that have already been
         # closed (via ``close_command`` or ``close_session``). Lets
         # ``get_query_state`` report ``CLOSED`` for them rather than
@@ -197,7 +206,9 @@ def close_session(self, session_id: SessionId) -> None:
         # server-side CloseStatement before the session goes away.
         with self._async_handles_lock:
             tracked = list(self._async_handles.items())
+            tracked_stmts = list(self._async_statements.items())
             self._async_handles.clear()
+            self._async_statements.clear()
             for guid, _ in tracked:
                 self._closed_commands.add(guid)
         for _, handle in tracked:
@@ -211,6 +222,16 @@ def close_session(self, session_id: SessionId) -> None:
                 logger.warning(
                     "Error closing async handle during session close: %s", exc
                 )
+        # Now drop the parent Statements that were keeping those handles
+        # alive. Same non-fatal close semantics — close errors are not
+        # actionable at session-close time.
+        for _, stmt in tracked_stmts:
+            try:
+                stmt.close()
+            except Exception as exc:
+                logger.warning(
+                    "Error closing async statement during session close: %s", exc
+                )
         try:
             self._kernel_session.close()
         except Exception as exc:
@@ -249,6 +270,11 @@ def execute_command(
             stmt = self._kernel_session.statement()
         except Exception as exc:
             raise _wrap_kernel_exception("execute_command", exc) from exc
+        # ``async_op`` keeps ``stmt`` alive (tracked in
+        # ``_async_statements`` and closed by ``close_command``); the sync
+        # path drops it in finally. ``close_stmt`` is the post-success
+        # decision flag — it stays True on sync, flips to False on async.
+        close_stmt = True
         try:
             try:
                 stmt.set_sql(operation)
@@ -262,21 +288,26 @@ def execute_command(
                     cursor.active_command_id = command_id
                     with self._async_handles_lock:
                         self._async_handles[command_id.guid] = async_exec
+                        # Closing the kernel ``Statement`` invalidates the
+                        # async handle (see kernel validity flag). Retain
+                        # the Statement here and close it on
+                        # ``close_command`` / ``close_session``.
+                        self._async_statements[command_id.guid] = stmt
+                    close_stmt = False
                     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
-            # parent doesn't keep the handle alive longer than the
-            # caller expects. Swallow all close errors (including
-            # PyO3 native exceptions) — a failed stmt.close() is
-            # not actionable for the caller.
-            try:
-                stmt.close()
-            except Exception:
-                pass
+            if close_stmt:
+                # Sync path: ``Statement`` is a lifecycle owner separate
+                # from the executed handle. Drop it here so the parent
+                # doesn't outlive its caller. Swallow close errors —
+                # they're not actionable.
+                try:
+                    stmt.close()
+                except Exception:
+                    pass
 
         command_id = CommandId.from_sea_statement_id(executed.statement_id)
         cursor.active_command_id = command_id
@@ -307,17 +338,34 @@ def cancel_command(self, command_id: CommandId) -> None:
     def close_command(self, command_id: CommandId) -> None:
         with self._async_handles_lock:
             handle = self._async_handles.pop(command_id.guid, None)
+            stmt = self._async_statements.pop(command_id.guid, None)
             if handle is not None:
                 # Record the close so ``get_query_state`` can report
                 # ``CLOSED`` (not ``SUCCEEDED``) for this command.
                 self._closed_commands.add(command_id.guid)
         if handle is None:
             logger.debug("close_command: no tracked handle for %s", command_id)
+            # Still drop the parent Statement if somehow tracked without
+            # the handle — keeps the invariant clean even on bookkeeping
+            # races.
+            if stmt is not None:
+                try:
+                    stmt.close()
+                except Exception:
+                    pass
             return
         try:
             handle.close()
         except Exception as exc:
             raise _wrap_kernel_exception("close_command", exc) from exc
+        finally:
+            # Now safe to close the parent Statement — the executed
+            # handle has finished its lifecycle.
+            if stmt is not None:
+                try:
+                    stmt.close()
+                except Exception:
+                    pass
 
     def get_query_state(self, command_id: CommandId) -> CommandState:
         with self._async_handles_lock:
@@ -378,6 +426,7 @@ def get_execution_result(
         # it wraps. Drop tracking and fire-and-forget the close.
         with self._async_handles_lock:
             self._async_handles.pop(command_id.guid, None)
+            stmt = self._async_statements.pop(command_id.guid, None)
             self._closed_commands.add(command_id.guid)
         try:
             async_exec.close()
@@ -387,6 +436,18 @@ def get_execution_result(
                 command_id,
                 exc,
             )
+        # The parent Statement is no longer needed once the async handle
+        # has produced its ResultStream. Close to release server-side
+        # tracking; matches the sync path's eager Statement close.
+        if stmt is not None:
+            try:
+                stmt.close()
+            except Exception as exc:
+                logger.warning(
+                    "Error closing async statement after await_result for %s: %s",
+                    command_id,
+                    exc,
+                )
         # ``KernelResultSet.__init__`` calls ``arrow_schema()`` which
         # can raise — map that to PEP 249 too.
         try: