vincbeck commented on code in PR #66976:
URL: https://github.com/apache/airflow/pull/66976#discussion_r3274958797
##########
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 will probably fail when it is run in a distributed environment like
ECS. You cannot assume all tasks from a Dag are executed in the same environment
--
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]
