acryl-datahub-airflow-plugin 1.3.1.4__py3-none-any.whl → 1.3.1.5__py3-none-any.whl

datahub_airflow_plugin/example_dags/__init__.py

@@ -0,0 +1,32 @@
+"""
+DataHub Airflow Plugin Example DAGs
+
+This directory contains example DAGs demonstrating various features of the DataHub Airflow plugin.
+
+## Directory Structure
+
+- **airflow2/**: Example DAGs for Airflow 2.x with compatibility layers
+- **airflow3/**: Example DAGs for Airflow 3.0+ using native syntax (no compatibility layers)
+- Root directory: Legacy example DAGs with compatibility layers (deprecated, use airflow2/ or airflow3/)
+
+## Choosing the Right Examples
+
+- If you're using **Airflow 3.0+**, refer to examples in `airflow3/`
+- If you're using **Airflow 2.x**, refer to examples in `airflow2/`
+- For production DAGs that need to work across multiple Airflow versions, see `airflow2/` for compatibility patterns
+
+## Available Examples
+
+### Lineage Collection
+- `lineage_backend_demo.py`: Basic lineage collection using inlets/outlets
+- `lineage_backend_taskflow_demo.py`: Lineage collection with TaskFlow API
+
+### Data Ingestion
+- `snowflake_sample_dag.py`: Ingest Snowflake metadata into DataHub
+- `mysql_sample_dag.py`: Ingest MySQL metadata into DataHub
+- `generic_recipe_sample_dag.py`: Run any DataHub recipe from Airflow
+
+### Advanced Features
+- `lineage_emission_dag.py`: Custom lineage emission with DatahubEmitterOperator
+- `graph_usage_sample_dag.py`: Complex DAG graph with multiple dependencies
+"""

datahub_airflow_plugin/example_dags/airflow2/__init__.py

@@ -0,0 +1,8 @@
+"""
+Airflow 2.x Example DAGs
+
+This directory contains example DAGs for Airflow 2.x.
+These DAGs use compatibility layers to work across multiple Airflow 2.x versions.
+
+For Airflow 3.0+ examples, see the ../airflow3/ directory.
+"""

datahub_airflow_plugin/example_dags/airflow2/generic_recipe_sample_dag.py

@@ -0,0 +1,54 @@
+"""Generic DataHub Ingest via Recipe
+
+This example demonstrates how to load any configuration file and run a
+DataHub ingestion pipeline within an Airflow DAG.
+"""
+
+from datetime import timedelta
+
+from airflow import DAG
+from airflow.operators.python import PythonOperator
+
+from datahub.configuration.config_loader import load_config_file
+from datahub.ingestion.run.pipeline import Pipeline
+from datahub_airflow_plugin._airflow_version_specific import (
+    days_ago,
+    get_airflow_compatible_dag_kwargs,
+)
+
+default_args = {
+    "owner": "airflow",
+    "depends_on_past": False,
+    "email": ["jdoe@example.com"],
+    "email_on_failure": False,
+    "email_on_retry": False,
+    "retries": 1,
+    "retry_delay": timedelta(minutes=5),
+    "execution_timeout": timedelta(minutes=120),
+}
+
+
+def datahub_recipe():
+    # Note that this will also resolve environment variables in the recipe.
+    config = load_config_file("path/to/recipe.yml")
+
+    pipeline = Pipeline.create(config)
+    pipeline.run()
+    pipeline.raise_from_status()
+
+
+with DAG(
+    "datahub_ingest_using_recipe",
+    **get_airflow_compatible_dag_kwargs(
+        default_args=default_args,
+        description="An example DAG which runs a DataHub ingestion recipe",
+        start_date=days_ago(2),
+        schedule_interval=timedelta(days=1),
+        catchup=False,
+        default_view="tree",
+    ),
+) as dag:
+    ingest_task = PythonOperator(
+        task_id="ingest_using_recipe",
+        python_callable=datahub_recipe,
+    )

datahub_airflow_plugin/example_dags/airflow2/graph_usage_sample_dag.py

@@ -0,0 +1,43 @@
+"""This example DAG demonstrates how to create and use a DataHubGraph client."""
+
+from datetime import timedelta
+
+import pendulum
+from airflow.decorators import (  # type: ignore[attr-defined] # Decorators not available in all Airflow versions
+    dag,
+    task,
+)
+
+from datahub.ingestion.graph.client import DataHubGraph, RemovedStatusFilter
+from datahub_airflow_plugin._airflow_version_specific import (
+    get_airflow_compatible_dag_kwargs,
+)
+from datahub_airflow_plugin.hooks.datahub import DatahubRestHook
+
+dag_decorator_kwargs = get_airflow_compatible_dag_kwargs(
+    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    schedule_interval=timedelta(days=1),
+    catchup=False,
+)
+
+
+@dag(**dag_decorator_kwargs)
+def datahub_graph_usage_sample_dag():
+    @task()
+    def use_the_graph():
+        graph: DataHubGraph = DatahubRestHook("my_datahub_rest_conn_id").make_graph()
+        graph.test_connection()
+
+        # Example usage: Find all soft-deleted BigQuery DEV entities
+        # in DataHub, and hard delete them.
+        for urn in graph.get_urns_by_filter(
+            platform="bigquery",
+            env="DEV",
+            status=RemovedStatusFilter.ONLY_SOFT_DELETED,
+        ):
+            graph.hard_delete_entity(urn)
+
+    use_the_graph()
+
+
+datahub_graph_usage_sample_dag()

datahub_airflow_plugin/example_dags/airflow2/lineage_backend_demo.py

@@ -0,0 +1,69 @@
+"""Lineage Backend
+
+An example DAG demonstrating the usage of DataHub's Airflow lineage backend.
+"""
+
+from datetime import timedelta
+
+from airflow import DAG
+from airflow.operators.bash import BashOperator
+
+from datahub_airflow_plugin._airflow_version_specific import days_ago
+from datahub_airflow_plugin.entities import Dataset, Urn
+
+default_args = {
+    "owner": "airflow",
+    "depends_on_past": False,
+    "email": ["jdoe@example.com"],
+    "email_on_failure": False,
+    "execution_timeout": timedelta(minutes=5),
+}
+
+
+# Create DAG arguments conditionally for Airflow version compatibility
+import airflow  # noqa: E402
+
+dag_kwargs = {
+    "dag_id": "datahub_lineage_backend_demo",
+    "default_args": default_args,
+    "description": "An example DAG demonstrating the usage of DataHub's Airflow lineage backend.",
+    "start_date": days_ago(2),
+    "tags": ["example_tag"],
+    "catchup": False,
+}
+
+# Handle schedule parameter change in Airflow 3.0
+if hasattr(airflow, "__version__") and airflow.__version__.startswith(
+    ("3.", "2.10", "2.9", "2.8", "2.7")
+):
+    # Use schedule for newer Airflow versions (2.7+)
+    dag_kwargs["schedule"] = timedelta(days=1)
+else:
+    # Use schedule_interval for older versions
+    dag_kwargs["schedule_interval"] = timedelta(days=1)
+
+# Add default_view only for older Airflow versions that support it
+if hasattr(airflow, "__version__") and not airflow.__version__.startswith("3."):
+    dag_kwargs["default_view"] = "tree"
+
+with DAG(**dag_kwargs) as dag:
+    task1 = BashOperator(
+        task_id="run_data_task",
+        dag=dag,
+        bash_command="echo 'This is where you might run your data tooling.'",
+        inlets=[
+            Dataset(platform="snowflake", name="mydb.schema.tableA"),
+            Dataset(platform="snowflake", name="mydb.schema.tableB", env="DEV"),
+            Dataset(
+                platform="snowflake",
+                name="mydb.schema.tableC",
+                platform_instance="cloud",
+            ),
+            # You can also put dataset URNs in the inlets/outlets lists.
+            Urn(
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)"
+            ),
+            Urn("urn:li:dataJob:(urn:li:dataFlow:(airflow,dag1,prod),task1)"),
+        ],
+        outlets=[Dataset("snowflake", "mydb.schema.tableD")],
+    )

datahub_airflow_plugin/example_dags/airflow2/lineage_backend_taskflow_demo.py

@@ -0,0 +1,69 @@
+"""Lineage Backend
+
+An example DAG demonstrating the usage of DataHub's Airflow lineage backend using the TaskFlow API.
+"""
+
+from datetime import timedelta
+
+from airflow.decorators import dag, task  # type: ignore[attr-defined]
+
+from datahub_airflow_plugin._airflow_version_specific import days_ago
+from datahub_airflow_plugin.entities import Dataset, Urn
+
+default_args = {
+    "owner": "airflow",
+    "depends_on_past": False,
+    "email": ["jdoe@example.com"],
+    "email_on_failure": False,
+    "execution_timeout": timedelta(minutes=5),
+}
+
+
+# Create DAG decorator arguments conditionally for Airflow version compatibility
+import airflow  # noqa: E402
+
+dag_decorator_kwargs = {
+    "default_args": default_args,
+    "description": "An example DAG demonstrating the usage of DataHub's Airflow lineage backend using the TaskFlow API.",
+    "start_date": days_ago(2),
+    "tags": ["example_tag"],
+    "catchup": False,
+}
+
+# Handle schedule parameter change in Airflow 3.0
+if hasattr(airflow, "__version__") and airflow.__version__.startswith(
+    ("3.", "2.10", "2.9", "2.8", "2.7")
+):
+    # Use schedule for newer Airflow versions (2.7+)
+    dag_decorator_kwargs["schedule"] = timedelta(days=1)
+else:
+    # Use schedule_interval for older versions
+    dag_decorator_kwargs["schedule_interval"] = timedelta(days=1)
+
+# Add default_view only for older Airflow versions that support it
+if hasattr(airflow, "__version__") and not airflow.__version__.startswith("3."):
+    dag_decorator_kwargs["default_view"] = "tree"
+
+
+@dag(**dag_decorator_kwargs)
+def datahub_lineage_backend_taskflow_demo():
+    @task(
+        inlets=[
+            Dataset("snowflake", "mydb.schema.tableA"),
+            Dataset("snowflake", "mydb.schema.tableB", "DEV"),
+            # You can also put dataset URNs in the inlets/outlets lists.
+            Urn(
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)"
+            ),
+            Urn("urn:li:dataJob:(urn:li:dataFlow:(airflow,dag1,prod),task1)"),
+        ],
+        outlets=[Dataset("snowflake", "mydb.schema.tableD")],
+    )
+    def run_data_task():
+        # This is where you might run your data tooling.
+        pass
+
+    run_data_task()
+
+
+datahub_lineage_backend_taskflow_dag = datahub_lineage_backend_taskflow_demo()

datahub_airflow_plugin/example_dags/airflow2/lineage_emission_dag.py

@@ -0,0 +1,81 @@
+"""Lineage Emission
+
+This example demonstrates how to emit lineage to DataHub within an Airflow DAG.
+"""
+
+from datetime import timedelta
+
+from airflow import DAG
+from airflow.operators.bash import BashOperator
+
+import datahub.emitter.mce_builder as builder
+from datahub_airflow_plugin._airflow_version_specific import days_ago
+from datahub_airflow_plugin.operators.datahub import DatahubEmitterOperator
+
+default_args = {
+    "owner": "airflow",
+    "depends_on_past": False,
+    "email": ["jdoe@example.com"],
+    "email_on_failure": False,
+    "email_on_retry": False,
+    "retries": 1,
+    "retry_delay": timedelta(minutes=5),
+    "execution_timeout": timedelta(minutes=120),
+}
+
+
+# Create DAG arguments conditionally for Airflow version compatibility
+import airflow  # noqa: E402
+
+dag_kwargs = {
+    "dag_id": "datahub_lineage_emission_example",
+    "default_args": default_args,
+    "description": "An example DAG demonstrating lineage emission within an Airflow DAG.",
+    "start_date": days_ago(2),
+    "catchup": False,
+}
+
+# Handle schedule parameter change in Airflow 3.0
+if hasattr(airflow, "__version__") and airflow.__version__.startswith(
+    ("3.", "2.10", "2.9", "2.8", "2.7")
+):
+    # Use schedule for newer Airflow versions (2.7+)
+    dag_kwargs["schedule"] = timedelta(days=1)
+else:
+    # Use schedule_interval for older versions
+    dag_kwargs["schedule_interval"] = timedelta(days=1)
+
+# Add default_view only for older Airflow versions that support it
+if hasattr(airflow, "__version__") and not airflow.__version__.startswith("3."):
+    dag_kwargs["default_view"] = "tree"
+
+with DAG(**dag_kwargs) as dag:
+    transformation_task = BashOperator(
+        task_id="transformation_task",
+        dag=dag,
+        bash_command="echo 'This is where you might run your data tooling.'",
+    )
+
+    emit_lineage_task = DatahubEmitterOperator(
+        task_id="emit_lineage",
+        datahub_conn_id="datahub_rest_default",
+        mces=[
+            builder.make_lineage_mce(
+                upstream_urns=[
+                    builder.make_dataset_urn(
+                        platform="snowflake", name="mydb.schema.tableA"
+                    ),
+                    builder.make_dataset_urn_with_platform_instance(
+                        platform="snowflake",
+                        name="mydb.schema.tableB",
+                        platform_instance="cloud",
+                    ),
+                ],
+                downstream_urn=builder.make_dataset_urn(
+                    platform="snowflake", name="mydb.schema.tableC", env="DEV"
+                ),
+            )
+        ],
+    )
+
+    transformation_task >> emit_lineage_task

datahub_airflow_plugin/example_dags/airflow2/mysql_sample_dag.py

@@ -0,0 +1,68 @@
+"""MySQL DataHub Ingest DAG
+
+This example demonstrates how to ingest metadata from MySQL into DataHub
+from within an Airflow DAG. Note that the DB connection configuration is
+embedded within the code.
+"""
+
+from datetime import datetime, timedelta
+
+from airflow import DAG
+from airflow.operators.python import PythonVirtualenvOperator
+
+from datahub_airflow_plugin._airflow_version_specific import (
+    get_airflow_compatible_dag_kwargs,
+)
+
+
+def ingest_from_mysql():
+    from datahub.ingestion.run.pipeline import Pipeline
+
+    pipeline = Pipeline.create(
+        # This configuration is analogous to a recipe configuration.
+        {
+            "source": {
+                "type": "mysql",
+                "config": {
+                    # If you want to use Airflow connections, take a look at the snowflake_sample_dag.py example.
+                    "username": "user",
+                    "password": "pass",
+                    "database": "db_name",
+                    "host_port": "localhost:3306",
+                },
+            },
+            "sink": {
+                "type": "datahub-rest",
+                "config": {"server": "http://localhost:8080"},
+            },
+        }
+    )
+    pipeline.run()
+    pipeline.pretty_print_summary()
+    pipeline.raise_from_status()
+
+
+with DAG(
+    "datahub_mysql_ingest",
+    **get_airflow_compatible_dag_kwargs(
+        default_args={
+            "owner": "airflow",
+        },
+        description="An example DAG which ingests metadata from MySQL to DataHub",
+        start_date=datetime(2022, 1, 1),
+        schedule_interval=timedelta(days=1),
+        catchup=False,
+        default_view="tree",
+    ),
+) as dag:
+    # While it is also possible to use the PythonOperator, we recommend using
+    # the PythonVirtualenvOperator to ensure that there are no dependency
+    # conflicts between DataHub and the rest of your Airflow environment.
+    ingest_task = PythonVirtualenvOperator(
+        task_id="ingest_from_mysql",
+        requirements=[
+            "acryl-datahub[mysql]",
+        ],
+        system_site_packages=False,
+        python_callable=ingest_from_mysql,
+    )

datahub_airflow_plugin/example_dags/airflow2/snowflake_sample_dag.py

@@ -0,0 +1,99 @@
+"""Snowflake DataHub Ingest DAG
+
+This example demonstrates how to ingest metadata from Snowflake into DataHub
+from within an Airflow DAG. In contrast to the MySQL example, this DAG
+pulls the DB connection configuration from Airflow's connection store.
+"""
+
+from datetime import datetime, timedelta
+
+from airflow import DAG
+
+from datahub_airflow_plugin._airflow_version_specific import (
+    IS_AIRFLOW_3_OR_HIGHER,
+    get_airflow_compatible_dag_kwargs,
+)
+
+if IS_AIRFLOW_3_OR_HIGHER:
+    from airflow.hooks.base_hook import BaseHook  # type: ignore[attr-defined]
+else:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined]
+from airflow.operators.python import PythonVirtualenvOperator
+
+
+def ingest_from_snowflake(snowflake_credentials, datahub_gms_server):
+    from datahub.ingestion.run.pipeline import Pipeline
+
+    pipeline = Pipeline.create(
+        # This configuration is analogous to a recipe configuration.
+        {
+            "source": {
+                "type": "snowflake",
+                "config": {
+                    **snowflake_credentials,
+                    # Other Snowflake config can be added here.
+                    "profiling": {"enabled": False},
+                },
+            },
+            # Other ingestion features, like transformers, are also supported.
+            # "transformers": [
+            #     {
+            #         "type": "simple_add_dataset_ownership",
+            #         "config": {
+            #             "owner_urns": [
+            #                 "urn:li:corpuser:example",
+            #             ]
+            #         },
+            #     }
+            # ],
+            "sink": {
+                "type": "datahub-rest",
+                "config": {"server": datahub_gms_server},
+            },
+        }
+    )
+    pipeline.run()
+    pipeline.pretty_print_summary()
+    pipeline.raise_from_status()
+
+
+with DAG(
+    "datahub_snowflake_ingest",
+    **get_airflow_compatible_dag_kwargs(
+        default_args={
+            "owner": "airflow",
+        },
+        description="An example DAG which ingests metadata from Snowflake to DataHub",
+        start_date=datetime(2022, 1, 1),
+        schedule_interval=timedelta(days=1),
+        catchup=False,
+        default_view="tree",
+    ),
+) as dag:
+    # This example pulls credentials from Airflow's connection store.
+    # For this to work, you must have previously configured these connections in Airflow.
+    # See the Airflow docs for details: https://airflow.apache.org/docs/apache-airflow/stable/howto/connection.html
+    snowflake_conn = BaseHook.get_connection("snowflake_admin_default")
+    datahub_conn = BaseHook.get_connection("datahub_rest_default")
+
+    # While it is also possible to use the PythonOperator, we recommend using
+    # the PythonVirtualenvOperator to ensure that there are no dependency
+    # conflicts between DataHub and the rest of your Airflow environment.
+    ingest_task = PythonVirtualenvOperator(
+        task_id="ingest_from_snowflake",
+        requirements=[
+            "acryl-datahub[snowflake]",
+        ],
+        system_site_packages=False,
+        python_callable=ingest_from_snowflake,
+        op_kwargs={
+            "snowflake_credentials": {
+                "username": snowflake_conn.login,
+                "password": snowflake_conn.password,
+                "account_id": snowflake_conn.extra_dejson["account"],
+                "warehouse": snowflake_conn.extra_dejson.get("warehouse"),
+                "role": snowflake_conn.extra_dejson.get("role"),
+            },
+            "datahub_gms_server": datahub_conn.host,
+        },
+    )

datahub_airflow_plugin/example_dags/airflow3/__init__.py

@@ -0,0 +1,8 @@
+"""
+Airflow 3.0+ Example DAGs
+
+This directory contains example DAGs specifically for Airflow 3.0+.
+These DAGs use native Airflow 3.0 syntax without compatibility layers.
+
+For Airflow 2.x examples, see the ../airflow2/ directory.
+"""

datahub_airflow_plugin/example_dags/airflow3/lineage_backend_demo.py

@@ -0,0 +1,51 @@
+"""Lineage Backend
+
+An example DAG demonstrating the usage of DataHub's Airflow lineage backend.
+
+This is the Airflow 3.0+ version.
+"""
+
+from datetime import datetime, timedelta
+
+from airflow import DAG
+from airflow.operators.bash import BashOperator
+
+from datahub_airflow_plugin.entities import Dataset, Urn
+
+default_args = {
+    "owner": "airflow",
+    "depends_on_past": False,
+    "email": ["jdoe@example.com"],
+    "email_on_failure": False,
+    "execution_timeout": timedelta(minutes=5),
+}
+
+with DAG(
+    dag_id="datahub_lineage_backend_demo",
+    default_args=default_args,
+    description="An example DAG demonstrating the usage of DataHub's Airflow lineage backend.",
+    start_date=datetime(2023, 1, 1),
+    schedule=timedelta(days=1),
+    tags=["example_tag"],
+    catchup=False,
+) as dag:
+    task1 = BashOperator(
+        task_id="run_data_task",
+        dag=dag,
+        bash_command="echo 'This is where you might run your data tooling.'",
+        inlets=[
+            Dataset(platform="snowflake", name="mydb.schema.tableA"),
+            Dataset(platform="snowflake", name="mydb.schema.tableB", env="DEV"),
+            Dataset(
+                platform="snowflake",
+                name="mydb.schema.tableC",
+                platform_instance="cloud",
+            ),
+            # You can also put dataset URNs in the inlets/outlets lists.
+            Urn(
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)"
+            ),
+            Urn("urn:li:dataJob:(urn:li:dataFlow:(airflow,dag1,prod),task1)"),
+        ],
+        outlets=[Dataset("snowflake", "mydb.schema.tableD")],
+    )

datahub_airflow_plugin/example_dags/airflow3/lineage_backend_taskflow_demo.py

@@ -0,0 +1,51 @@
+"""Lineage Backend
+
+An example DAG demonstrating the usage of DataHub's Airflow lineage backend using the TaskFlow API.
+
+This is the Airflow 3.0+ version.
+"""
+
+from datetime import datetime, timedelta
+
+from airflow.decorators import dag, task  # type: ignore[attr-defined]
+
+from datahub_airflow_plugin.entities import Dataset, Urn
+
+default_args = {
+    "owner": "airflow",
+    "depends_on_past": False,
+    "email": ["jdoe@example.com"],
+    "email_on_failure": False,
+    "execution_timeout": timedelta(minutes=5),
+}
+
+
+@dag(
+    default_args=default_args,
+    description="An example DAG demonstrating the usage of DataHub's Airflow lineage backend using the TaskFlow API.",
+    start_date=datetime(2023, 1, 1),
+    schedule=timedelta(days=1),
+    tags=["example_tag"],
+    catchup=False,
+)
+def datahub_lineage_backend_taskflow_demo():
+    @task(
+        inlets=[
+            Dataset("snowflake", "mydb.schema.tableA"),
+            Dataset("snowflake", "mydb.schema.tableB", "DEV"),
+            # You can also put dataset URNs in the inlets/outlets lists.
+            Urn(
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableC,PROD)"
+            ),
+            Urn("urn:li:dataJob:(urn:li:dataFlow:(airflow,dag1,prod),task1)"),
+        ],
+        outlets=[Dataset("snowflake", "mydb.schema.tableD")],
+    )
+    def run_data_task():
+        # This is where you might run your data tooling.
+        pass
+
+    run_data_task()
+
+
+datahub_lineage_backend_taskflow_dag = datahub_lineage_backend_taskflow_demo()