stronger typing of Cursor and ExecuteResponse

varun-edachali-dbx · varun-edachali-dbx · commit 980063653f08 · 2025-05-29T03:31:07.000Z
Signed-off-by: varun-edachali-dbx &lt;varun.edachali@databricks.com&gt;
diff --git a/src/databricks/sql/backend/databricks_client.py b/src/databricks/sql/backend/databricks_client.py
@@ -13,6 +13,7 @@
 from abc import ABC, abstractmethod
 from typing import Dict, Tuple, List, Optional, Any, Union
 
+from databricks.sql.client import Cursor
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.backend.types import SessionId, CommandId
 from databricks.sql.utils import ExecuteResponse
@@ -75,7 +76,7 @@ def execute_command(
         max_rows: int,
         max_bytes: int,
         lz4_compression: bool,
-        cursor: Any,
+        cursor: Cursor,
         use_cloud_fetch: bool,
         parameters: List[ttypes.TSparkParameter],
         async_op: bool,
@@ -173,7 +174,7 @@ def get_query_state(self, command_id: CommandId) -> ttypes.TOperationState:
     def get_execution_result(
         self,
         command_id: CommandId,
-        cursor: Any,
+        cursor: Cursor,
     ) -> ExecuteResponse:
         """
         Retrieves the results of a previously executed command.
@@ -201,8 +202,8 @@ def get_catalogs(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
-    ) -> Any:
+        cursor: Cursor,
+    ) -> ExecuteResponse:
         """
         Retrieves a list of available catalogs.
         
@@ -230,10 +231,10 @@ def get_schemas(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: Cursor,
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
-    ) -> Any:
+    ) -> ExecuteResponse:
         """
         Retrieves a list of available schemas.
         
@@ -263,12 +264,12 @@ def get_tables(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: Cursor,
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         table_types: Optional[List[str]] = None,
-    ) -> Any:
+    ) -> ExecuteResponse:
         """
         Retrieves a list of available tables.
         
@@ -300,12 +301,12 @@ def get_columns(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: Cursor,
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         column_name: Optional[str] = None,
-    ) -> Any:
+    ) -> ExecuteResponse:
         """
         Retrieves column metadata for tables.
         
diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py
@@ -7,6 +7,7 @@
 import threading
 from typing import List, Union, Any
 
+from databricks.sql.client import Cursor
 from databricks.sql.thrift_api.TCLIService.ttypes import TOperationState
 from databricks.sql.backend.types import (
     SessionId,
@@ -928,7 +929,7 @@ def execute_command(
         max_rows: int,
         max_bytes: int,
         lz4_compression: bool,
-        cursor: Any,
+        cursor: Cursor,
         use_cloud_fetch=True,
         parameters=[],
         async_op=False,
@@ -986,8 +987,8 @@ def get_catalogs(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
-    ):
+        cursor: Cursor,
+    ) -> ExecuteResponse:
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1006,10 +1007,10 @@ def get_schemas(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: Cursor,
         catalog_name=None,
         schema_name=None,
-    ):
+    ) -> ExecuteResponse:
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1030,12 +1031,12 @@ def get_tables(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: Cursor,
         catalog_name=None,
         schema_name=None,
         table_name=None,
         table_types=None,
-    ):
+    ) -> ExecuteResponse:
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1058,12 +1059,12 @@ def get_columns(
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
-        cursor: Any,
+        cursor: Cursor,
         catalog_name=None,
         schema_name=None,
         table_name=None,
         column_name=None,
-    ):
+    ) -> ExecuteResponse:
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
