correct typing

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":
+    ) -> Union["ResultSet", None]:
         pass
 
     @abstractmethod
@@ -62,7 +62,7 @@ def get_execution_result(
         self,
         command_id: CommandId,
         cursor: Any,
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         pass
 
     # == Metadata Operations ==

src/databricks/sql/backend/thrift_backend.py

@@ -944,7 +944,7 @@ def execute_command(
         parameters=[],
         async_op=False,
         enforce_embedded_schema_correctness=False,
-    ) -> "ResultSet":
+    ) -> Union["ResultSet", None]:
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")

src/databricks/sql/client.py

@@ -783,7 +783,7 @@ def execute(
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
         )
 
-        if self.active_result_set.is_staging_operation:
+        if self.active_result_set and self.active_result_set.is_staging_operation:
             self._handle_staging_operation(
                 staging_allowed_local_path=self.backend.staging_allowed_local_path
             )
@@ -879,7 +879,7 @@ def get_async_execution_result(self):
                 self.active_op_handle, self
             )
 
-            if self.active_result_set.is_staging_operation:
+            if self.active_result_set and self.active_result_set.is_staging_operation:
                 self._handle_staging_operation(
                     staging_allowed_local_path=self.backend.staging_allowed_local_path
                 )

src/databricks/sql/result_set.py

@@ -46,6 +46,12 @@ def __iter__(self):
     def rownumber(self):
         return self._next_row_index
 
+    @property
+    @abstractmethod
+    def is_staging_operation(self) -> bool:
+        """Whether this result set represents a staging operation."""
+        pass
+
     # Define abstract methods that concrete implementations must implement
     @abstractmethod
     def _fill_results_buffer(self):
@@ -117,7 +123,7 @@ def __init__(
         self.description = execute_response.description
         self._arrow_schema_bytes = execute_response.arrow_schema_bytes
         self._use_cloud_fetch = use_cloud_fetch
-        self.is_staging_operation = execute_response.is_staging_operation
+        self._is_staging_operation = execute_response.is_staging_operation
 
         # Initialize results queue
         if execute_response.arrow_queue:
@@ -350,3 +356,8 @@ def close(self) -> None:
         finally:
             self.has_been_closed_server_side = True
             self.op_state = ttypes.TOperationState.CLOSED_STATE
+
+    @property
+    def is_staging_operation(self) -> bool:
+        """Whether this result set represents a staging operation."""
+        return self._is_staging_operation