Repository: astronomer/astro-sdk

Issue:
Openlineage support - Add Extractor for `ExportFileOperator`
**Please describe the feature you'd like to see**
We should be able to extract open lineage info from the `ExportFileOperator`.


**Describe the solution you'd like**
- Add a method `get_openlineage_facets` on the Operator (based on [this doc](https://docs.google.com/document/d/1vPsvHejQ24xTbzpz_LYSf0_ixk9oUuBiEUHVEaF9J2U/edit?usp=sharing))
- The "PythonSDKExtractor" built-in https://github.com/astronomer/astro-sdk/issues/898 should be able to work with `ExportFileOperator.get_openlineage_facets`. Test it to make sure it works


**Acceptance Criteria**

- [ ] Post the screenshot of how it looks in the Openlineage/Marquez UI
- [ ] All checks and tests in the CI should pass
- [ ] Unit tests (90% code coverage or more, [once available](https://github.com/astronomer/astro-sdk/issues/191))
- [ ] Integration tests (if the feature relates to a new database or external service)
- [ ] Docstrings in [reStructuredText](https://peps.python.org/pep-0287/) for each of  methods, classes, functions and module-level attributes (including Example DAG on how it should be used)
- [ ] Exception handling in case of errors
- [ ] Logging (are we exposing useful information to the user? e.g. source and destination)
- [ ] Improve the documentation (README, Sphinx, and any other relevant)



Discussion:

Resolving patch:
diff --git a/python-sdk/src/astro/lineage/extractor.py b/python-sdk/src/astro/lineage/extractor.py
--- a/python-sdk/src/astro/lineage/extractor.py
+++ b/python-sdk/src/astro/lineage/extractor.py
@@ -34,6 +34,7 @@ def get_operator_classnames(cls) -> list[str]:
         return [
             "AppendOperator",
             "BaseSQLDecoratedOperator",
+            "ExportFileOperator",
             "LoadFileOperator",
             "MergeOperator",
             "TransformOperator",
diff --git a/python-sdk/src/astro/lineage/facets.py b/python-sdk/src/astro/lineage/facets.py
--- a/python-sdk/src/astro/lineage/facets.py
+++ b/python-sdk/src/astro/lineage/facets.py
@@ -3,7 +3,7 @@
 import attr
 from openlineage.client.facet import BaseFacet
 
-from astro.constants import FileType, MergeConflictStrategy
+from astro.constants import ExportExistsStrategy, FileType, MergeConflictStrategy
 from astro.table import Column, Metadata
 
 
@@ -23,6 +23,23 @@ class InputFileFacet(BaseFacet):
     description: str | None = None
 
 
+@attr.define
+class ExportFileFacet(BaseFacet):
+    """
+    Facet that represents output file for export file
+
+    :param filepath: path of each file.
+    :param file_size: size of the file.
+    :param file_type: type of the file.
+    :param if_exists: Overwrite file if exists.
+    """
+
+    filepath: str
+    file_size: int | None
+    file_type: FileType
+    if_exists: ExportExistsStrategy
+
+
 @attr.define
 class InputFileDatasetFacet(BaseFacet):
     """
diff --git a/python-sdk/src/astro/sql/operators/export_file.py b/python-sdk/src/astro/sql/operators/export_file.py
--- a/python-sdk/src/astro/sql/operators/export_file.py
+++ b/python-sdk/src/astro/sql/operators/export_file.py
@@ -5,11 +5,22 @@
 import pandas as pd
 from airflow.decorators.base import get_unique_task_id
 from airflow.models.xcom_arg import XComArg
+from openlineage.client.facet import (
+    BaseFacet,
+    DataQualityMetricsInputDatasetFacet,
+    DataSourceDatasetFacet,
+    OutputStatisticsOutputDatasetFacet,
+    SchemaDatasetFacet,
+    SchemaField,
+)
+from openlineage.client.run import Dataset as OpenlineageDataset
 
 from astro.airflow.datasets import kwargs_with_datasets
 from astro.constants import ExportExistsStrategy
 from astro.databases import create_database
 from astro.files import File
+from astro.lineage.extractor import OpenLineageFacets
+from astro.lineage.facets import ExportFileFacet
 from astro.sql.operators.base_operator import AstroSQLBaseOperator
 from astro.table import BaseTable, Table
 from astro.utils.typing_compat import Context
@@ -62,6 +73,71 @@ def execute(self, context: Context) -> File:
         else:
             raise FileExistsError(f"{self.output_file.path} file already exists.")
 
+    def get_openlineage_facets(self, task_instance) -> OpenLineageFacets:  # skipcq: PYL-W0613
+        """
+        Collect the input, output, job and run facets for export file operator
+        """
+        input_dataset: list[OpenlineageDataset] = []
+        if isinstance(self.input_data, BaseTable) and self.input_data.openlineage_emit_temp_table_event():
+            input_uri = (
+                f"{self.input_data.openlineage_dataset_namespace()}"
+                f"://{self.input_data.openlineage_dataset_name()}"
+            )
+            input_dataset = [
+                OpenlineageDataset(
+                    namespace=self.input_data.openlineage_dataset_namespace(),
+                    name=self.input_data.openlineage_dataset_name(),
+                    facets={
+                        "dataSource": DataSourceDatasetFacet(
+                            name=self.input_data.openlineage_dataset_name(), uri=input_uri
+                        ),
+                        "schema": SchemaDatasetFacet(
+                            fields=[
+                                SchemaField(
+                                    name=self.input_data.metadata.schema,
+                                    type=self.input_data.metadata.database,
+                                )
+                            ]
+                        ),
+                        "dataQualityMetrics": DataQualityMetricsInputDatasetFacet(
+                            rowCount=self.input_data.row_count, columnMetrics={}
+                        ),
+                    },
+                )
+            ]
+        output_uri = (
+            f"{self.output_file.openlineage_dataset_namespace}"
+            f"://{self.output_file.openlineage_dataset_name}"
+        )
+        output_dataset = [
+            OpenlineageDataset(
+                namespace=self.output_file.openlineage_dataset_namespace,
+                name=self.output_file.openlineage_dataset_name,
+                facets={
+                    "output_file_facet": ExportFileFacet(
+                        filepath=self.output_file.path,
+                        file_size=self.output_file.size,
+                        file_type=self.output_file.type.name,
+                        if_exists=self.if_exists,
+                    ),
+                    "dataSource": DataSourceDatasetFacet(name=self.output_file.path, uri=output_uri),
+                    "outputStatistics": OutputStatisticsOutputDatasetFacet(
+                        rowCount=self.input_data.row_count
+                        if isinstance(self.input_data, BaseTable)
+                        else len(self.input_data),
+                        size=self.output_file.size,
+                    ),
+                },
+            )
+        ]
+
+        run_facets: dict[str, BaseFacet] = {}
+        job_facets: dict[str, BaseFacet] = {}
+
+        return OpenLineageFacets(
+            inputs=input_dataset, outputs=output_dataset, run_facets=run_facets, job_facets=job_facets
+        )
+
 
 def export_file(
     input_data: BaseTable | pd.DataFrame,


Test patch:
diff --git a/python-sdk/tests/extractors/test_extractor.py b/python-sdk/tests/extractors/test_extractor.py
--- a/python-sdk/tests/extractors/test_extractor.py
+++ b/python-sdk/tests/extractors/test_extractor.py
@@ -3,6 +3,7 @@
 from airflow.models.taskinstance import TaskInstance
 from airflow.utils import timezone
 from openlineage.airflow.extractors import TaskMetadata
+from openlineage.client.facet import DataQualityMetricsInputDatasetFacet, OutputStatisticsOutputDatasetFacet
 from openlineage.client.run import Dataset as OpenlineageDataset
 
 from astro import sql as aql
@@ -12,6 +13,7 @@
 from astro.lineage.facets import InputFileDatasetFacet, InputFileFacet, OutputDatabaseDatasetFacet
 from astro.settings import LOAD_FILE_ENABLE_NATIVE_FALLBACK
 from astro.sql import AppendOperator, MergeOperator
+from astro.sql.operators.export_file import ExportFileOperator
 from astro.sql.operators.load_file import LoadFileOperator
 from astro.table import Metadata, Table
 from tests.utils.airflow import create_context
@@ -22,6 +24,25 @@
 TEST_INPUT_DATASET_NAME = "/workspace/sample_pattern"
 TEST_OUTPUT_DATASET_NAMESPACE = "bigquery"
 TEST_OUTPUT_DATASET_NAME = "astronomer-dag-authoring.astro.test-extractor"
+INPUT_STATS_FOR_EXPORT_FILE = [
+    OpenlineageDataset(
+        namespace=TEST_INPUT_DATASET_NAMESPACE,
+        name=TEST_INPUT_DATASET_NAME,
+        facets={
+            "dataQualityMetrics": DataQualityMetricsInputDatasetFacet(rowCount=117, columnMetrics={}),
+        },
+    )
+]
+
+OUTPUT_STATS_FOR_EXPORT_FILE = [
+    OpenlineageDataset(
+        namespace=TEST_OUTPUT_DATASET_NAMESPACE,
+        name=TEST_OUTPUT_DATASET_NAME,
+        facets={
+            "outputStatistics": OutputStatisticsOutputDatasetFacet(rowCount=117, size=65),
+        },
+    )
+]
 INPUT_STATS = [
     OpenlineageDataset(
         namespace=TEST_INPUT_DATASET_NAMESPACE,
@@ -100,6 +121,53 @@ def test_python_sdk_load_file_extract_on_complete():
     assert task_meta.run_facets == {}
 
 
+@pytest.mark.integration
+def test_python_sdk_export_file_extract_on_complete():
+    """
+    Tests that  the custom PythonSDKExtractor is able to process the
+    operator's metadata that needs to be extracted as per OpenLineage
+    for ExportFileOperator.
+    """
+    LoadFileOperator(
+        task_id="load_file",
+        input_file=File(
+            path="https://raw.githubusercontent.com/astronomer/astro-sdk/main/tests/data/imdb_v2.csv"
+        ),
+        output_table=Table(conn_id="sqlite_conn", name="test_extractor"),
+    ).execute({})
+
+    task_id = "export_file"
+    export_file_operator = ExportFileOperator(
+        task_id=task_id,
+        input_data=Table(conn_id="sqlite_conn", name="test_extractor"),
+        output_file=File(
+            path="gs://astro-sdk/workspace/openlineage_export_file.csv",
+            conn_id="bigquery",
+            filetype=FileType.CSV,
+        ),
+        if_exists="replace",
+    )
+
+    task_instance = TaskInstance(task=export_file_operator)
+
+    python_sdk_extractor = PythonSDKExtractor(export_file_operator)
+    task_meta_extract = python_sdk_extractor.extract()
+    assert isinstance(task_meta_extract, TaskMetadata)
+
+    task_meta = python_sdk_extractor.extract_on_complete(task_instance)
+    assert task_meta.name == f"adhoc_airflow.{task_id}"
+    assert (
+        task_meta.inputs[0].facets["dataQualityMetrics"]
+        == INPUT_STATS_FOR_EXPORT_FILE[0].facets["dataQualityMetrics"]
+    )
+    assert (
+        task_meta.outputs[0].facets["outputStatistics"]
+        == OUTPUT_STATS_FOR_EXPORT_FILE[0].facets["outputStatistics"]
+    )
+    assert task_meta.job_facets == {}
+    assert task_meta.run_facets == {}
+
+
 @pytest.mark.integration
 def test_append_op_extract_on_complete():
     """