o-nikolas commented on code in PR #66976:
URL: https://github.com/apache/airflow/pull/66976#discussion_r3270512477
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_notebook.py:
##########
@@ -54,16 +59,31 @@
PROJECT_ID_KEY = "PROJECT_ID"
NOTEBOOK_ID_KEY = "NOTEBOOK_ID"
NOTEBOOK_B_ID_KEY = "NOTEBOOK_B_ID"
+DATAZONE_ROLE_ARN_KEY = "DATAZONE_ROLE_ARN"
sys_test_context_task = (
SystemTestContextBuilder()
.add_variable(DOMAIN_ID_KEY)
.add_variable(PROJECT_ID_KEY)
.add_variable(NOTEBOOK_ID_KEY)
.add_variable(NOTEBOOK_B_ID_KEY)
+ .add_variable(DATAZONE_ROLE_ARN_KEY)
.build()
)
+DATAZONE_CONN_ID = "aws_datazone_notebook"
+
+
+@task
+def setup_datazone_connection(role_arn: str) -> str:
+ """Configure an Airflow connection that assumes the DataZone environment
role."""
+ conn_extra = urllib.parse.quote_plus(
+ json.dumps({"role_arn": role_arn, "assume_role_method": "assume_role"})
+ )
+ os.environ[f"AIRFLOW_CONN_{DATAZONE_CONN_ID.upper()}"] =
f"aws://?extra={conn_extra}"
Review Comment:
This is only going to work for local executor/`dag.test`. Executors like
ECS, Batch and Lambda that run each task in it's own container/vm will not be
able to read this env var that was set in another container. The connection
must be added to the DB for this to work.
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_notebook.py:
##########
@@ -101,20 +121,13 @@
)
# [END howto_operator_sagemaker_unified_studio_notebook]
- # [START howto_operator_sagemaker_unified_studio_notebook_deferrable]
- run_notebook_deferrable = SageMakerUnifiedStudioNotebookOperator(
- task_id="notebook-deferrable-task",
- notebook_identifier=notebook_id,
- domain_identifier=domain_id,
- owning_project_identifier=project_id,
- deferrable=True, # optional
- waiter_delay=10, # optional
- )
- # [END howto_operator_sagemaker_unified_studio_notebook_deferrable]
+ # NOTE: Deferrable mode skipped — the operator does not propagate
aws_conn_id
+ # to the trigger (upstream bug). Tracked separately.
Review Comment:
Is it separate? It looks like the bug fix is in this PR?
--
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]
