fix fetch types

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/types.py

@@ -102,7 +102,6 @@ def from_sea_state(cls, state: str) -> Optional["CommandState"]:
         return state_mapping.get(state, None)
 
 
-
 class BackendType(Enum):
     """
     Enum representing the type of backend
@@ -417,6 +416,7 @@ def to_hex_guid(self) -> str:
         else:
             return str(self.guid)
 
+
 @dataclass
 class ExecuteResponse:
     """Response from executing a SQL command."""
@@ -430,4 +430,4 @@ class ExecuteResponse:
     results_queue: Optional[Any] = None
     has_been_closed_server_side: bool = False
     lz4_compressed: bool = True
-    is_staging_operation: bool = False
+    is_staging_operation: bool = False

src/databricks/sql/result_set.py

@@ -115,12 +115,12 @@ def fetchall(self) -> List[Row]:
         pass
 
     @abstractmethod
-    def fetchmany_arrow(self, size: int) -> Any:
+    def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         """Fetch the next set of rows as an Arrow table."""
         pass
 
     @abstractmethod
-    def fetchall_arrow(self) -> Any:
+    def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all remaining rows as an Arrow table."""
         pass