diff --git a/pyathena/model.py b/pyathena/model.py index b751ccbc..44ba17c6 100644 --- a/pyathena/model.py +++ b/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") diff --git a/tests/pyathena/spark/test_async_spark_cursor.py b/tests/pyathena/spark/test_async_spark_cursor.py index 3d727b5b..e2031dcc 100644 --- a/tests/pyathena/spark/test_async_spark_cursor.py +++ b/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 diff --git a/tests/pyathena/spark/test_spark_cursor.py b/tests/pyathena/spark/test_spark_cursor.py index 689f61ca..85ac2b5d 100644 --- a/tests/pyathena/spark/test_spark_cursor.py +++ b/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)