(airflow) branch main updated: Fix SageMaker Unified Studio trigger hook config (#72453)

[vincbeck]

This is an automated email from the ASF dual-hosted git repository.

vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 760ee9a2d39 Fix SageMaker Unified Studio trigger hook config (#72453)
760ee9a2d39 is described below

commit 760ee9a2d399374114f06080b5dd6be4c07e7b2d
Author: Alice <78815181+weiyu1...@users.noreply.github.com>
AuthorDate: Fri Sep 4 07:30:24 2026 -0500

    Fix SageMaker Unified Studio trigger hook config (#72453)
---
 .../operators/sagemaker_unified_studio_notebook.py    |  3 +++
 .../test_sagemaker_unified_studio_notebook.py         | 19 +++++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git 
a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio_notebook.py
 
b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio_notebook.py
index e10ed90a412..66ef0e41aae 100644
--- 
a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio_notebook.py
+++ 
b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio_notebook.py
@@ -195,6 +195,9 @@ class 
SageMakerUnifiedStudioNotebookOperator(AwsBaseOperator[SageMakerUnifiedStu
                     owning_project_identifier=self.owning_project_identifier,
                     waiter_delay=self.waiter_delay,
                     timeout_configuration=self.timeout_configuration,
+                    region_name=self.region_name,
+                    verify=self.verify,
+                    botocore_config=self.botocore_config,
                 ),
                 method_name="execute_complete",
             )
diff --git 
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_unified_studio_notebook.py
 
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_unified_studio_notebook.py
index 1233afb490c..7b75bf8cc70 100644
--- 
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_unified_studio_notebook.py
+++ 
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_unified_studio_notebook.py
@@ -264,6 +264,13 @@ class TestSageMakerUnifiedStudioNotebookOperator:
         mock_hook_prop.return_value = mock_hook
         mock_hook.start_notebook_run.return_value = {"id": NOTEBOOK_RUN_ID}
 
+        botocore_config = {
+            "retries": {
+                "max_attempts": 7,
+                "mode": "standard",
+            }
+        }
+
         op = SageMakerUnifiedStudioNotebookOperator(
             task_id=TASK_ID,
             notebook_identifier=NOTEBOOK_ID,
@@ -272,18 +279,30 @@ class TestSageMakerUnifiedStudioNotebookOperator:
             deferrable=True,
             waiter_delay=20,
             timeout_configuration={"runTimeoutInMinutes": 120},
+            region_name="us-west-2",
+            verify="/tmp/custom-ca.pem",
+            botocore_config=botocore_config,
         )
 
         with pytest.raises(TaskDeferred) as exc_info:
             op.execute(_make_context())
 
         trigger = exc_info.value.trigger
+
         assert isinstance(trigger, SageMakerUnifiedStudioNotebookTrigger)
         assert trigger.notebook_run_id == NOTEBOOK_RUN_ID
         assert trigger.domain_identifier == DOMAIN_ID
         assert trigger.owning_project_identifier == PROJECT_ID
         assert trigger.waiter_delay == 20
         assert trigger.timeout_configuration == {"runTimeoutInMinutes": 120}
+
+        _, kwargs = trigger.serialize()
+
+        assert kwargs["region_name"] == "us-west-2"
+        assert kwargs["verify"] == "/tmp/custom-ca.pem"
+        assert kwargs["botocore_config"] == botocore_config
+        assert kwargs["waiter_max_attempts"] == 360
+
         assert exc_info.value.method_name == "execute_complete"
         mock_hook.wait_for_notebook_run.assert_not_called()