From 628d9c1e93f62c6e12734f02db036c28d84558f8 Mon Sep 17 00:00:00 2001 From: Maciej Obuchowski Date: Sat, 8 Mar 2025 16:19:32 +0100 Subject: [PATCH] add OpenLineage configuration injection to LivyOperator Signed-off-by: Maciej Obuchowski --- generated/provider_dependencies.json | 2 + providers/apache/livy/README.rst | 30 +-- providers/apache/livy/pyproject.toml | 2 + .../apache/livy/get_provider_info.py | 1 + .../providers/apache/livy/operators/livy.py | 26 +++ .../unit/apache/livy/operators/test_livy.py | 193 ++++++++++++++++++ 6 files changed, 240 insertions(+), 14 deletions(-) diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 2c5fc991291c8..5448935341688 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -246,6 +246,7 @@ "apache.livy": { "deps": [ "aiohttp>=3.9.2", + "apache-airflow-providers-common-compat>=1.5.0", "apache-airflow-providers-http>=5.1.0", "apache-airflow>=2.9.0", "asgiref>=2.3.0" @@ -253,6 +254,7 @@ "devel-deps": [], "plugins": [], "cross-providers-deps": [ + "common.compat", "http" ], "excluded-python-versions": [], diff --git a/providers/apache/livy/README.rst b/providers/apache/livy/README.rst index 99d821ab6602b..c7d93253fa391 100644 --- a/providers/apache/livy/README.rst +++ b/providers/apache/livy/README.rst @@ -50,14 +50,15 @@ The package supports the following python versions: 3.9,3.10,3.11,3.12 Requirements ------------ -================================= ================== -PIP package Version required -================================= ================== -``apache-airflow`` ``>=2.9.0`` -``apache-airflow-providers-http`` ``>=5.1.0`` -``aiohttp`` ``>=3.9.2`` -``asgiref`` ``>=2.3.0`` -================================= ================== +========================================== ================== +PIP package Version required +========================================== ================== +``apache-airflow`` ``>=2.9.0`` +``apache-airflow-providers-http`` ``>=5.1.0`` +``apache-airflow-providers-common-compat`` ``>=1.5.0`` +``aiohttp`` ``>=3.9.2`` +``asgiref`` ``>=2.3.0`` +========================================== ================== Cross provider package dependencies ----------------------------------- @@ -69,14 +70,15 @@ You can install such cross-provider dependencies when installing from PyPI. For .. code-block:: bash - pip install apache-airflow-providers-apache-livy[http] + pip install apache-airflow-providers-apache-livy[common.compat] -================================================================================================ ======== -Dependent package Extra -================================================================================================ ======== -`apache-airflow-providers-http `_ ``http`` -================================================================================================ ======== +================================================================================================================== ================= +Dependent package Extra +================================================================================================================== ================= +`apache-airflow-providers-common-compat `_ ``common.compat`` +`apache-airflow-providers-http `_ ``http`` +================================================================================================================== ================= The changelog for the provider package can be found in the `changelog `_. diff --git a/providers/apache/livy/pyproject.toml b/providers/apache/livy/pyproject.toml index 8ce6be120be6c..a3a4f3a8aeb3d 100644 --- a/providers/apache/livy/pyproject.toml +++ b/providers/apache/livy/pyproject.toml @@ -59,6 +59,7 @@ requires-python = "~=3.9" dependencies = [ "apache-airflow>=2.9.0", "apache-airflow-providers-http>=5.1.0", + "apache-airflow-providers-common-compat>=1.5.0", "aiohttp>=3.9.2", "asgiref>=2.3.0", ] @@ -68,6 +69,7 @@ dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", "apache-airflow-providers-http", # Additional devel dependencies (do not remove this line and add extra development dependencies) ] diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/get_provider_info.py b/providers/apache/livy/src/airflow/providers/apache/livy/get_provider_info.py index ea926ef9035d1..dca704ae47bda 100644 --- a/providers/apache/livy/src/airflow/providers/apache/livy/get_provider_info.py +++ b/providers/apache/livy/src/airflow/providers/apache/livy/get_provider_info.py @@ -104,6 +104,7 @@ def get_provider_info(): "dependencies": [ "apache-airflow>=2.9.0", "apache-airflow-providers-http>=5.1.0", + "apache-airflow-providers-common-compat>=1.5.0", "aiohttp>=3.9.2", "asgiref>=2.3.0", ], diff --git a/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py b/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py index 746ea55cceeaf..12159955acf33 100644 --- a/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py +++ b/providers/apache/livy/src/airflow/providers/apache/livy/operators/livy.py @@ -28,6 +28,10 @@ from airflow.models import BaseOperator from airflow.providers.apache.livy.hooks.livy import BatchState, LivyHook from airflow.providers.apache.livy.triggers.livy import LivyTrigger +from airflow.providers.common.compat.openlineage.utils.spark import ( + inject_parent_job_information_into_spark_properties, + inject_transport_information_into_spark_properties, +) if TYPE_CHECKING: from airflow.utils.context import Context @@ -94,10 +98,19 @@ def __init__( extra_headers: dict[str, Any] | None = None, retry_args: dict[str, Any] | None = None, deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + openlineage_inject_parent_job_info: bool = conf.getboolean( + "openlineage", "spark_inject_parent_job_info", fallback=False + ), + openlineage_inject_transport_info: bool = conf.getboolean( + "openlineage", "spark_inject_transport_info", fallback=False + ), **kwargs: Any, ) -> None: super().__init__(**kwargs) + if conf is None: + conf = {} + spark_params = { # Prepare spark parameters, it will be templated later. "file": file, @@ -128,6 +141,8 @@ def __init__( self._batch_id: int | str | None = None self.retry_args = retry_args self.deferrable = deferrable + self.openlineage_inject_parent_job_info = openlineage_inject_parent_job_info + self.openlineage_inject_transport_info = openlineage_inject_transport_info @cached_property def hook(self) -> LivyHook: @@ -145,6 +160,17 @@ def hook(self) -> LivyHook: ) def execute(self, context: Context) -> Any: + if self.openlineage_inject_parent_job_info: + self.log.debug("Injecting parent job information into Spark properties") + self.spark_params["conf"] = inject_parent_job_information_into_spark_properties( + self.spark_params["conf"], context + ) + if self.openlineage_inject_transport_info: + self.log.debug("Injecting transport information into Spark properties") + self.spark_params["conf"] = inject_transport_information_into_spark_properties( + self.spark_params["conf"], context + ) + _batch_id: int | str = self.hook.post_batch(**self.spark_params) self._batch_id = _batch_id self.log.info("Generated batch-id is %s", self._batch_id) diff --git a/providers/apache/livy/tests/unit/apache/livy/operators/test_livy.py b/providers/apache/livy/tests/unit/apache/livy/operators/test_livy.py index d83af04180401..be03438fe94bd 100644 --- a/providers/apache/livy/tests/unit/apache/livy/operators/test_livy.py +++ b/providers/apache/livy/tests/unit/apache/livy/operators/test_livy.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import logging from unittest.mock import MagicMock, patch import pytest @@ -362,6 +363,198 @@ def test_execute_complete_timeout(self, mock_delete, mock_post): mock_delete.assert_called_once_with(BATCH_ID) self.mock_context["ti"].xcom_push.assert_not_called() + @patch.object(LivyOperator, "hook", new_callable=MagicMock) + @patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener") + def test_inject_simple_openlineage_config_to_spark(self, mock_get_openlineage_listener, mock_hook): + # Given / When + from openlineage.client.transport.http import ( + ApiKeyTokenProvider, + HttpCompression, + HttpConfig, + HttpTransport, + ) + + mock_get_openlineage_listener.return_value.adapter.get_or_create_openlineage_client.return_value.transport = HttpTransport( + config=HttpConfig( + url="http://localhost:5000", + endpoint="api/v2/lineage", + timeout=5050, + auth=ApiKeyTokenProvider({"api_key": "12345"}), + compression=HttpCompression.GZIP, + custom_headers={"X-OpenLineage-Custom-Header": "airflow"}, + ) + ) + + operator = LivyOperator( + file="sparkapp", + livy_conn_id="livy_default", + polling_interval=1, + dag=self.dag, + task_id="livy_example", + conf={}, + deferrable=False, + openlineage_inject_parent_job_info=False, + openlineage_inject_transport_info=True, + ) + operator.hook.get_batch_state.return_value = BatchState.SUCCESS + operator.hook.TERMINAL_STATES = [BatchState.SUCCESS] + operator.execute(MagicMock()) + + assert operator.spark_params["conf"] == { + "spark.openlineage.transport.type": "http", + "spark.openlineage.transport.url": "http://localhost:5000", + "spark.openlineage.transport.endpoint": "api/v2/lineage", + "spark.openlineage.transport.timeoutInMillis": "5050000", + "spark.openlineage.transport.compression": "gzip", + "spark.openlineage.transport.auth.type": "api_key", + "spark.openlineage.transport.auth.apiKey": "Bearer 12345", + "spark.openlineage.transport.headers.X-OpenLineage-Custom-Header": "airflow", + } + + @patch.object(LivyOperator, "hook", new_callable=MagicMock) + @patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.get_batch_state") + @patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener") + def test_inject_composite_openlineage_config_to_spark( + self, mock_get_openlineage_listener, mock_get_batch_state, mock_hook + ): + # Given / When + from openlineage.client.transport.composite import CompositeConfig, CompositeTransport + + mock_get_openlineage_listener.return_value.adapter.get_or_create_openlineage_client.return_value.transport = CompositeTransport( + CompositeConfig.from_dict( + { + "transports": { + "test1": { + "type": "http", + "url": "http://localhost:5000", + "endpoint": "api/v2/lineage", + "timeout": 5050, + "auth": { + "type": "api_key", + "api_key": "12345", + }, + "compression": "gzip", + "custom_headers": {"X-OpenLineage-Custom-Header": "airflow"}, + }, + "test2": {"type": "http", "url": "https://example.com:1234"}, + "test3": {"type": "console"}, + } + } + ) + ) + + mock_ti = MagicMock() + mock_ti.dag_id = "test_dag_id" + mock_ti.task_id = "spark_submit_job" + mock_ti.try_number = 1 + mock_ti.dag_run.logical_date = DEFAULT_DATE + mock_ti.dag_run.run_after = DEFAULT_DATE + mock_ti.logical_date = DEFAULT_DATE + mock_ti.map_index = -1 + mock_get_batch_state.return_value = BatchState.SUCCESS + + operator = LivyOperator( + file="sparkapp", + livy_conn_id="spark_default", + polling_interval=1, + dag=self.dag, + task_id="livy_example", + deferrable=False, + openlineage_inject_parent_job_info=True, + openlineage_inject_transport_info=True, + ) + operator.hook.get_batch_state.return_value = BatchState.SUCCESS + operator.hook.TERMINAL_STATES = [BatchState.SUCCESS] + + operator.execute({"ti": mock_ti}) + + assert operator.spark_params["conf"] == { + "spark.openlineage.parentJobName": "test_dag_id.spark_submit_job", + "spark.openlineage.parentJobNamespace": "default", + "spark.openlineage.parentRunId": "01595753-6400-710b-8a12-9e978335a56d", + "spark.openlineage.transport.type": "composite", + "spark.openlineage.transport.continueOnFailure": "True", + "spark.openlineage.transport.transports.test1.type": "http", + "spark.openlineage.transport.transports.test1.url": "http://localhost:5000", + "spark.openlineage.transport.transports.test1.endpoint": "api/v2/lineage", + "spark.openlineage.transport.transports.test1.timeoutInMillis": "5050000", + "spark.openlineage.transport.transports.test1.auth.type": "api_key", + "spark.openlineage.transport.transports.test1.auth.apiKey": "Bearer 12345", + "spark.openlineage.transport.transports.test1.compression": "gzip", + "spark.openlineage.transport.transports.test1.headers.X-OpenLineage-Custom-Header": "airflow", + "spark.openlineage.transport.transports.test2.type": "http", + "spark.openlineage.transport.transports.test2.url": "https://example.com:1234", + "spark.openlineage.transport.transports.test2.endpoint": "api/v1/lineage", + "spark.openlineage.transport.transports.test2.timeoutInMillis": "5000", + } + + @patch.object(LivyOperator, "hook", new_callable=MagicMock) + @patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.get_batch_state") + @patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener") + def test_inject_openlineage_composite_config_wrong_transport_to_spark( + self, mock_get_openlineage_listener, mock_get_batch_state, mock_hook, caplog + ): + # Given / When + from openlineage.client.transport.composite import CompositeConfig, CompositeTransport + + mock_get_openlineage_listener.return_value.adapter.get_or_create_openlineage_client.return_value.transport = CompositeTransport( + CompositeConfig.from_dict({"transports": {"test1": {"type": "console"}}}) + ) + mock_get_batch_state.return_value = BatchState.SUCCESS + + with caplog.at_level(logging.INFO): + operator = LivyOperator( + file="sparkapp", + livy_conn_id="livy_default", + polling_interval=1, + dag=self.dag, + task_id="livy_example", + deferrable=False, + openlineage_inject_parent_job_info=False, + openlineage_inject_transport_info=True, + ) + operator.hook.get_batch_state.return_value = BatchState.SUCCESS + operator.hook.TERMINAL_STATES = [BatchState.SUCCESS] + operator.execute(MagicMock()) + + assert ( + "OpenLineage transport type `composite` does not contain http transport. Skipping injection of OpenLineage transport information into Spark properties." + in caplog.text + ) + assert operator.spark_params["conf"] == {} + + @patch.object(LivyOperator, "hook", new_callable=MagicMock) + @patch("airflow.providers.apache.livy.hooks.livy.LivyAsyncHook.get_batch_state") + @patch("airflow.providers.openlineage.utils.spark.get_openlineage_listener") + def test_inject_openlineage_simple_config_wrong_transport_to_spark( + self, mock_get_openlineage_listener, mock_get_batch_state, mock_hook, caplog + ): + # Given / When + from openlineage.client.transport.console import ConsoleConfig, ConsoleTransport + + mock_get_openlineage_listener.return_value.adapter.get_or_create_openlineage_client.return_value.transport = ConsoleTransport( + config=ConsoleConfig() + ) + mock_hook.get_batch_state.return_value = BatchState.SUCCESS + + with caplog.at_level(logging.INFO): + operator = LivyOperator( + file="sparkapp", + livy_conn_id="livy_default", + polling_interval=1, + dag=self.dag, + task_id="livy_example", + deferrable=False, + openlineage_inject_parent_job_info=False, + openlineage_inject_transport_info=True, + ) + operator.hook.get_batch_state.return_value = BatchState.SUCCESS + operator.hook.TERMINAL_STATES = [BatchState.SUCCESS] + operator.execute(MagicMock()) + + assert "OpenLineage transport type `console` does not support automatic injection of OpenLineage transport information into Spark properties." + assert operator.spark_params["conf"] == {} + @pytest.mark.db_test def test_spark_params_templating(create_task_instance_of_operator, session):