add strong return type for execute_command

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

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -83,7 +83,7 @@ def execute_command(
         parameters: List[ttypes.TSparkParameter],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
-    ) -> Any:
+    ) -> Optional[ExecuteResponse]:
         """
         Executes a SQL command or query within the specified session.
 

src/databricks/sql/backend/thrift_backend.py

@@ -5,7 +5,7 @@
 import time
 import uuid
 import threading
-from typing import List, Union, Any, TYPE_CHECKING
+from typing import List, Optional, Union, Any, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
@@ -929,7 +929,7 @@ def execute_command(
         parameters=[],
         async_op=False,
         enforce_embedded_schema_correctness=False,
-    ):
+    ) -> Optional[ExecuteResponse]:
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -974,6 +974,7 @@ def execute_command(
 
         if async_op:
             self._handle_execute_response_async(resp, cursor)
+            return None
         else:
             return self._handle_execute_response(resp, cursor)
 

src/databricks/sql/client.py

@@ -783,6 +783,8 @@ def execute(
             async_op=False,
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
         )
+        assert execute_response is not None  # async_op = False above
+
         self.active_result_set = ResultSet(
             self.connection,
             execute_response,