o-nikolas commented on code in PR #72335:
URL: https://github.com/apache/airflow/pull/72335#discussion_r3962208773
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio_notebook.py:
##########
@@ -173,7 +173,11 @@ def execute(self, context: Context):
region_name=self.hook.conn_region_name,
aws_partition=self.hook.conn_partition,
)
- workflow_name = context["dag"].dag_id # Workflow name is the same as
the dag_id
+ # MWAA Serverless overwrites dag_id with the execution ID for tenant
isolation and
+ # injects the real workflow name as a DAG param, so prefer that when
set. For standard
+ # Workflows and open-source Airflow, dag_id is the workflow name.
+ params = context.get("params") or {}
Review Comment:
Instead of introducing Params to this, which I think complicates things. Can
you just make use of the existing `dag_display_name` attr of dags?
##########
providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_unified_studio_notebook.py:
##########
@@ -173,6 +173,57 @@ def test_execute_passes_dag_id_as_workflow_name(self,
mock_hook_prop):
call_kwargs = mock_hook.start_notebook_run.call_args[1]
assert call_kwargs["workflow_name"] == "my_custom_dag"
+ @patch(HOOK_PATH, new_callable=PropertyMock)
+ def test_execute_prefers_mwaa_serverless_workflow_id_param(self,
mock_hook_prop):
+ """Under MWAA Serverless, dag_id is the per-execution UUID; the real
workflow name is
+ injected as the `mwaa_serverless_workflow_id` DAG param and must be
preferred."""
+ mock_hook = MagicMock()
+ mock_hook_prop.return_value = mock_hook
+ mock_hook.start_notebook_run.return_value = {"id": NOTEBOOK_RUN_ID}
+ mock_hook.wait_for_notebook_run.return_value = {}
+ mock_hook.get_notebook_outputs.return_value = {}
+
+ op = SageMakerUnifiedStudioNotebookOperator(
+ task_id=TASK_ID,
+ notebook_identifier=NOTEBOOK_ID,
+ domain_identifier=DOMAIN_ID,
+ owning_project_identifier=PROJECT_ID,
+ )
+ op.execute(
+ _make_context(
+ dag_id="execution-uuid-12345",
+ params={"mwaa_serverless_workflow_id":
"my-real-workflow-a1b2c3d4"},
+ )
+ )
+
+ call_kwargs = mock_hook.start_notebook_run.call_args[1]
+ assert call_kwargs["workflow_name"] == "my-real-workflow-a1b2c3d4"
+
+ @patch(HOOK_PATH, new_callable=PropertyMock)
+ def test_execute_falls_back_to_dag_id_when_param_empty(self,
mock_hook_prop):
Review Comment:
I think you could likely parametrize this these two tests into one.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio_notebook.py:
##########
@@ -173,7 +173,11 @@ def execute(self, context: Context):
region_name=self.hook.conn_region_name,
aws_partition=self.hook.conn_partition,
)
- workflow_name = context["dag"].dag_id # Workflow name is the same as
the dag_id
+ # MWAA Serverless overwrites dag_id with the execution ID for tenant
isolation and
+ # injects the real workflow name as a DAG param, so prefer that when
set. For standard
+ # Workflows and open-source Airflow, dag_id is the workflow name.
+ params = context.get("params") or {}
+ workflow_name = params.get("mwaa_serverless_workflow_id") or
context["dag"].dag_id
Review Comment:
Can you move this magic string into a constant? This makes it referenceable
in tests too
--
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]
