phanikumv commented on code in PR #66976:
URL: https://github.com/apache/airflow/pull/66976#discussion_r3252367756
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_notebook.py:
##########
@@ -77,6 +82,14 @@
project_id = test_context[PROJECT_ID_KEY]
notebook_id = test_context[NOTEBOOK_ID_KEY]
notebook_b_id = test_context[NOTEBOOK_B_ID_KEY]
+ # The DataZone StartNotebookRun API requires the caller to be the
environment's
+ # native user role. Configure a connection via environment variable.
+ # SystemTestContextBuilder exports DATAZONE_ROLE_ARN to os.environ during
variable fetching.
+ DATAZONE_CONN_ID = "aws_datazone_notebook"
+ conn_extra = urllib.parse.quote_plus(
+ json.dumps({"role_arn": os.environ[DATAZONE_ROLE_ARN_KEY],
"assume_role_method": "assume_role"})
+ )
+ os.environ[f"AIRFLOW_CONN_{DATAZONE_CONN_ID.upper()}"] =
f"aws://?extra={conn_extra}"
Review Comment:
Suggest to move this block to a task.
```python
@task
def setup_datazone_connection(role_arn: str) -> str:
conn_id = "aws_datazone_notebook"
extra = urllib.parse.quote_plus(
json.dumps({"role_arn": role_arn, "assume_role_method":
"assume_role"})
)
os.environ[f"AIRFLOW_CONN_{conn_id.upper()}"] = f"aws://?extra={extra}"
return conn_id
datazone_conn_id =
setup_datazone_connection(test_context[DATAZONE_ROLE_ARN_KEY])
# then in each operator/sensor: aws_conn_id=datazone_conn_id
```
This defers the role-ARN read to task runtime (XCom-resolved from
`test_context`), writes `AIRFLOW_CONN_*` inside the same in-process worker that
runs the notebook operators.
--
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]
