enforce ResultSet return in exec commands in backend client

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -42,7 +42,7 @@ def execute_command(
         parameters: List[ttypes.TSparkParameter],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
-    ) -> "ResultSet":  # Changed return type to ResultSet
+    ) -> "ResultSet":
         pass
 
     @abstractmethod
@@ -73,7 +73,7 @@ def get_catalogs(
         max_rows: int,
         max_bytes: int,
         cursor: Any,
-    ) -> Any:
+    ) -> "ResultSet":
         pass
 
     @abstractmethod
@@ -85,7 +85,7 @@ def get_schemas(
         cursor: Any,
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
-    ) -> Any:
+    ) -> "ResultSet":
         pass
 
     @abstractmethod
@@ -99,7 +99,7 @@ def get_tables(
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         table_types: Optional[List[str]] = None,
-    ) -> Any:
+    ) -> "ResultSet":
         pass
 
     @abstractmethod
@@ -113,7 +113,7 @@ def get_columns(
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         column_name: Optional[str] = None,
-    ) -> Any:
+    ) -> "ResultSet":
         pass
 
     # == Utility Methods ==

src/databricks/sql/backend/thrift_backend.py

@@ -48,7 +48,7 @@
 )
 from databricks.sql.types import SSLOptions
 from databricks.sql.backend.databricks_client import DatabricksClient
-from databricks.sql.result_set import ThriftResultSet
+from databricks.sql.result_set import ResultSet, ThriftResultSet
 
 logger = logging.getLogger(__name__)
 
@@ -945,7 +945,7 @@ def execute_command(
         parameters=[],
         async_op=False,
         enforce_embedded_schema_correctness=False,
-    ):
+    ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1009,7 +1009,7 @@ def get_catalogs(
         max_rows: int,
         max_bytes: int,
         cursor: Any,
-    ):
+    ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1041,7 +1041,7 @@ def get_schemas(
         cursor: Any,
         catalog_name=None,
         schema_name=None,
-    ):
+    ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1077,7 +1077,7 @@ def get_tables(
         schema_name=None,
         table_name=None,
         table_types=None,
-    ):
+    ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1115,7 +1115,7 @@ def get_columns(
         schema_name=None,
         table_name=None,
         column_name=None,
-    ):
+    ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")