o-nikolas commented on code in PR #73929:
URL: https://github.com/apache/airflow/pull/73929#discussion_r4149132804
##########
providers/amazon/tests/system/amazon/aws/example_athena_spark.py:
##########
@@ -100,15 +127,26 @@ def stop_athena_spark_session(session_id: str) -> None:
stop_session = stop_athena_spark_session(session_id)
+ delete_bucket = S3DeleteBucketOperator(
+ task_id="delete_bucket",
+ bucket_name=bucket_name,
+ force_delete=True,
+ trigger_rule=TriggerRule.ALL_DONE,
+ )
+
chain(
# TEST SETUP
test_context,
+ create_bucket,
+ setup_work_group,
session_id,
idle_session_id,
# TEST BODY
run_spark_calculation,
# TEST TEARDOWN
stop_session,
+ delete_work_group(work_group),
Review Comment:
We very frequently have issues during cleanup with trying to do back-to-back
deletes without enough time in between. If the session is stopped, do we need
to wait explicitly before deleting a work group?
##########
providers/amazon/tests/system/amazon/aws/example_athena_spark.py:
##########
@@ -40,15 +41,34 @@
# Compatibility for Airflow < 3.1
from airflow.utils.trigger_rule import TriggerRule # type:
ignore[no-redef,attr-defined]
-from system.amazon.aws.utils import SystemTestContextBuilder
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
DAG_ID = "example_athena_spark"
-# The Spark workgroup is preconfigured test infrastructure; this DAG creates
only the session.
-# Test runners can override the default by exporting ATHENA_SPARK_WORK_GROUP.
-ATHENA_SPARK_WORK_GROUP_KEY = "ATHENA_SPARK_WORK_GROUP"
+# Athena rejects a PySpark work group without an execution role, so the role
is preconfigured
+# test infrastructure. The results bucket and the work group are created here.
+EXECUTION_ROLE_ARN_KEY = "EXECUTION_ROLE_ARN"
Review Comment:
I think we usually just call these `ROLE_ARN_KEY = "ROLE_ARN"`, could change
to that for consistency, but very minor nit.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
