diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/triggers/livy.py b/providers/apache/livy/src/airflow/providers/apache/livy/triggers/livy.py index 20a8e0382d7d9..2d0cb99a64240 100644 --- a/providers/apache/livy/src/airflow/providers/apache/livy/triggers/livy.py +++ b/providers/apache/livy/src/airflow/providers/apache/livy/triggers/livy.py @@ -42,6 +42,7 @@ class LivyTrigger(BaseTrigger): depends on the option that's being modified. :param extra_headers: A dictionary of headers passed to the HTTP request to livy. :param livy_hook_async: LivyAsyncHook object + :param endpoint_prefix: Optional URL prefix for the Livy API endpoint. """ def __init__( @@ -80,6 +81,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "extra_headers": self._extra_headers, "livy_hook_async": self._livy_hook_async, "execution_timeout": self._execution_timeout, + "endpoint_prefix": self._endpoint_prefix, }, ) diff --git a/providers/apache/livy/tests/unit/apache/livy/triggers/test_livy.py b/providers/apache/livy/tests/unit/apache/livy/triggers/test_livy.py index df85a84bac338..1fc92a2c7b1ac 100644 --- a/providers/apache/livy/tests/unit/apache/livy/triggers/test_livy.py +++ b/providers/apache/livy/tests/unit/apache/livy/triggers/test_livy.py @@ -48,8 +48,21 @@ def test_livy_trigger_serialization(self): "extra_headers": None, "livy_hook_async": None, "execution_timeout": None, + "endpoint_prefix": None, } + def test_livy_trigger_serialize_round_trip_endpoint_prefix(self): + trigger = LivyTrigger( + batch_id=1, + spark_params={}, + livy_conn_id=LivyHook.default_conn_name, + polling_interval=0, + endpoint_prefix="/custom", + ) + _, kwargs = trigger.serialize() + restored = LivyTrigger(**kwargs) + assert restored._endpoint_prefix == "/custom" + @pytest.mark.asyncio @mock.patch("airflow.providers.apache.livy.triggers.livy.LivyTrigger.poll_for_termination") async def test_livy_trigger_run_with_no_poll_interval(self, mock_poll_for_termination):