Add OpenLineage system tests for new emission_policy feature

providers/openlineage/tests/system/openlineage/example_openlineage_policy_conf_locked_dag.py

@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+DAG exercising the OpenLineage emission policy authoring API.
+
+A locked conf rule cannot be overridden by authoring.
+
+Required global conf: `{"scope": {"dag_id": "openlineage_policy_conf_locked_dag"}, "locked": true, "controls": {"include_source_code": false}}`
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow import DAG
+from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy
+from airflow.providers.standard.operators.bash import BashOperator
+
+from system.openlineage.expected_events import get_expected_event_file_path
+from system.openlineage.operator import OpenLineageTestOperator
+
+DAG_ID = "openlineage_policy_conf_locked_dag"
+
+with DAG(
+    dag_id=DAG_ID,
+    start_date=datetime(2021, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args={"retries": 0},
+) as dag:
+    t_locked = BashOperator(task_id="t_locked_blocked", bash_command="echo locked_test")
+    check_events = OpenLineageTestOperator(
+        task_id="check_events", file_path=get_expected_event_file_path(DAG_ID)
+    )
+    t_locked >> check_events
+
+extend_global_openlineage_emission_policy(t_locked, include_source_code=True)
+
+
+from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)

providers/openlineage/tests/system/openlineage/example_openlineage_policy_conf_source_code_dag.py

@@ -0,0 +1,59 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+DAG exercising the OpenLineage emission policy authoring API.
+
+An unlocked conf rule suppresses source code; authoring re-enables it per-task.
+
+Required global conf: `{"scope": {"dag_id": "openlineage_policy_conf_source_code_dag"}, "controls": {"include_source_code": false}}`
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow import DAG
+from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy
+from airflow.providers.standard.operators.bash import BashOperator
+
+from system.openlineage.expected_events import get_expected_event_file_path
+from system.openlineage.operator import OpenLineageTestOperator
+
+DAG_ID = "openlineage_policy_conf_source_code_dag"
+
+with DAG(
+    dag_id=DAG_ID,
+    start_date=datetime(2021, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args={"retries": 0},
+) as dag:
+    t_conf_suppressed = BashOperator(task_id="t_conf_suppressed", bash_command="echo conf_test")
+    t_auth_override = extend_global_openlineage_emission_policy(
+        BashOperator(task_id="t_authoring_override", bash_command="echo auth_override"),
+        include_source_code=True,
+    )
+    check_events = OpenLineageTestOperator(
+        task_id="check_events", file_path=get_expected_event_file_path(DAG_ID)
+    )
+    t_conf_suppressed >> t_auth_override >> check_events
+
+
+from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)

providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_emit_false_dag.py

@@ -0,0 +1,61 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+DAG exercising the OpenLineage emission policy authoring API.
+
+Dag-level ``emit=False`` silences everything (task + Dag events).
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow import DAG
+from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy
+from airflow.providers.standard.operators.python import PythonOperator
+
+from system.openlineage.expected_events import get_expected_event_file_path
+from system.openlineage.operator import OpenLineageTestOperator
+
+
+def _say_hello():
+    print("hello")
+
+
+DAG_ID = "openlineage_policy_dag_emit_false_dag"
+
+with DAG(
+    dag_id=DAG_ID,
+    start_date=datetime(2021, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args={"retries": 0},
+) as dag:
+    task_a = PythonOperator(task_id="task_a", python_callable=_say_hello)
+    task_b = PythonOperator(task_id="task_b", python_callable=_say_hello)
+    check_events = OpenLineageTestOperator(
+        task_id="check_events", file_path=get_expected_event_file_path(DAG_ID)
+    )
+    task_a >> task_b >> check_events
+
+extend_global_openlineage_emission_policy(dag, emit=False)
+
+
+from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)

providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_events_false_dag.py

@@ -0,0 +1,68 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+DAG exercising the OpenLineage emission policy authoring API.
+
+Dag-level ``emit_dag_events=False`` suppresses DAG-run events only; task events
+must remain fully intact (source code, inputs/outputs, ...). The task registers
+hook lineage so the emitted task event carries real inputs/outputs to assert on.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow import DAG
+from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy
+from airflow.providers.standard.operators.python import PythonOperator
+
+from system.openlineage.expected_events import get_expected_event_file_path
+from system.openlineage.operator import OpenLineageTestOperator
+
+
+def _register_hook_lineage():
+    from airflow.lineage.hook import get_hook_lineage_collector
+
+    collector = get_hook_lineage_collector()
+    collector.add_input_asset(context=None, uri="file://host1/in1.txt")
+    collector.add_input_asset(context=None, uri="file://host1/in2.txt")
+    collector.add_output_asset(context=None, uri="file://host2/out1.txt")
+    collector.add_output_asset(context=None, uri="file://host2/out2.txt")
+
+
+DAG_ID = "openlineage_policy_dag_events_false_dag"
+
+with DAG(
+    dag_id=DAG_ID,
+    start_date=datetime(2021, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args={"retries": 0},
+) as dag:
+    run = PythonOperator(task_id="run", python_callable=_register_hook_lineage)
+    check_events = OpenLineageTestOperator(
+        task_id="check_events", file_path=get_expected_event_file_path(DAG_ID)
+    )
+    run >> check_events
+
+extend_global_openlineage_emission_policy(dag, emit_dag_events=False)
+
+
+from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)

providers/openlineage/tests/system/openlineage/example_openlineage_policy_dag_override_task_dag.py

@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+DAG exercising the OpenLineage emission policy authoring API.
+
+A dag-level flag (``include_source_code=False``) is overridden per-task.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow import DAG
+from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy
+from airflow.providers.standard.operators.bash import BashOperator
+
+from system.openlineage.expected_events import get_expected_event_file_path
+from system.openlineage.operator import OpenLineageTestOperator
+
+DAG_ID = "openlineage_policy_dag_override_task_dag"
+
+with DAG(
+    dag_id=DAG_ID,
+    start_date=datetime(2021, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args={"retries": 0},
+) as dag:
+    t_default = BashOperator(task_id="t_default", bash_command="echo dag_default")
+    t_override = BashOperator(task_id="t_override", bash_command="echo task_override")
+    check_events = OpenLineageTestOperator(
+        task_id="check_events", file_path=get_expected_event_file_path(DAG_ID)
+    )
+    t_default >> t_override >> check_events
+
+extend_global_openlineage_emission_policy(dag, include_source_code=False)
+extend_global_openlineage_emission_policy(t_override, include_source_code=True)
+
+
+from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)

providers/openlineage/tests/system/openlineage/example_openlineage_policy_extract_metadata_dag.py

@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+DAG exercising the OpenLineage emission policy authoring API.
+
+Per-task ``extract_operator_metadata=False`` gives empty inputs/outputs.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from openlineage.client.event_v2 import Dataset
+
+from airflow import DAG
+from airflow.providers.common.compat.sdk import BaseOperator
+from airflow.providers.openlineage.api.emission_policy import extend_global_openlineage_emission_policy
+from airflow.providers.openlineage.extractors.base import OperatorLineage
+
+from system.openlineage.expected_events import get_expected_event_file_path
+from system.openlineage.operator import OpenLineageTestOperator
+
+
+class _PolicyCustomOp(BaseOperator):
+    """Custom operator returning fixed inputs/outputs via the OL extractor API."""
+
+    def execute(self, context):
+        print("custom op ran")
+
+    def get_openlineage_facets_on_complete(self, task_instance):
+        return OperatorLineage(
+            inputs=[Dataset("s3://ol-test-in", "in.csv")],
+            outputs=[Dataset("s3://ol-test-out", "out.csv")],
+        )
+
+
+DAG_ID = "openlineage_policy_extract_metadata_dag"
+
+with DAG(
+    dag_id=DAG_ID,
+    start_date=datetime(2021, 1, 1),
+    schedule=None,
+    catchup=False,
+    default_args={"retries": 0},
+) as dag:
+    t_with_extract = _PolicyCustomOp(task_id="t_with_extract")
+    t_no_extract = extend_global_openlineage_emission_policy(
+        _PolicyCustomOp(task_id="t_no_extract"),
+        extract_operator_metadata=False,
+    )
+    check_events = OpenLineageTestOperator(
+        task_id="check_events", file_path=get_expected_event_file_path(DAG_ID)
+    )
+    t_with_extract >> t_no_extract >> check_events
+
+
+from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)