Terminate session when canceling calculation execution

pyathena/model.py

@@ -304,9 +304,8 @@ def __init__(self, response: Dict[str, Any]) -> None:
         self._description: Optional[str] = response.get("Description")
         self._working_directory: Optional[str] = response.get("WorkingDirectory")
 
-        result = response.get("Result")
-        if not result:
-            raise DataError("KeyError `Result`")
+        # If cancelled, the result does not exist.
+        result = response.get("Result", {})
         self._std_out_s3_uri: Optional[str] = result.get("StdOutS3Uri")
         self._std_error_s3_uri: Optional[str] = result.get("StdErrorS3Uri")
         self._result_s3_uri: Optional[str] = result.get("ResultS3Uri")

tests/pyathena/spark/test_async_spark_cursor.py

@@ -117,5 +117,9 @@ def test_cancel(self, async_spark_cursor):
         )
         time.sleep(randint(5, 10))
         async_spark_cursor.cancel(query_id).result()
+
+        # TODO: Calculation execution is not canceled unless session is terminated
+        async_spark_cursor.close()
+
         calculation_execution = future.result()
         assert calculation_execution.state == AthenaCalculationExecution.STATE_CANCELED

tests/pyathena/spark/test_spark_cursor.py

@@ -130,6 +130,9 @@ def cancel(c):
             time.sleep(randint(5, 10))
             c.cancel()
 
+            # TODO: Calculation execution is not canceled unless session is terminated
+            c.close()
+
         with ThreadPoolExecutor(max_workers=1) as executor:
             executor.submit(cancel, spark_cursor)