Annotating Task Lineage

from openlineage.client.event_v2 import Dataset
# Database tablessource_table = Dataset(    namespace="postgres://mydb:5432",    name="public.orders",)target_table = Dataset(    namespace="snowflake://account.snowflakecomputing.com",    name="staging.orders_clean",)
# Filesinput_file = Dataset(    namespace="s3://my-bucket",    name="raw/events/2024-01-01.json",)

from airflow.sdk import Asset
# Using Airflow's native Asset typeorders_asset = Asset(uri="s3://my-bucket/data/orders")

from airflow.datasets import Dataset
# Using Airflow's Dataset type (Airflow 2.4-2.x)orders_dataset = Dataset(uri="s3://my-bucket/data/orders")

from airflow import DAGfrom airflow.operators.bash import BashOperatorfrom openlineage.client.event_v2 import Datasetimport pendulum
# Define your lineage datasetssource_table = Dataset(    namespace="snowflake://account.snowflakecomputing.com",    name="raw.orders",)target_table = Dataset(    namespace="snowflake://account.snowflakecomputing.com",    name="staging.orders_clean",)output_file = Dataset(    namespace="s3://my-bucket",    name="exports/orders.parquet",)
with DAG(    dag_id="etl_with_lineage",    start_date=pendulum.datetime(2024, 1, 1, tz="UTC"),    schedule="@daily",) as dag:
    transform = BashOperator(        task_id="transform_orders",        bash_command="echo 'transforming...'",        inlets=[source_table],           # What this task reads        outlets=[target_table],          # What this task writes    )
    export = BashOperator(        task_id="export_to_s3",        bash_command="echo 'exporting...'",        inlets=[target_table],           # Reads from previous output        outlets=[output_file],           # Writes to S3    )
    transform >> export

from openlineage.client.event_v2 import Dataset
# Multiple source tablescustomers = Dataset(namespace="postgres://crm:5432", name="public.customers")orders = Dataset(namespace="postgres://sales:5432", name="public.orders")products = Dataset(namespace="postgres://inventory:5432", name="public.products")
# Multiple output tablesdaily_summary = Dataset(namespace="snowflake://account", name="analytics.daily_summary")customer_metrics = Dataset(namespace="snowflake://account", name="analytics.customer_metrics")
aggregate_task = PythonOperator(    task_id="build_daily_aggregates",    python_callable=build_aggregates,    inlets=[customers, orders, products],      # All inputs    outlets=[daily_summary, customer_metrics], # All outputs)

from airflow.models import BaseOperator
class MyCustomOperator(BaseOperator):    def __init__(self, source_table: str, target_table: str, **kwargs):        super().__init__(**kwargs)        self.source_table = source_table        self.target_table = target_table
    def execute(self, context):        # ... perform the actual work ...        self.log.info(f"Processing {self.source_table} -> {self.target_table}")
    def get_openlineage_facets_on_complete(self, task_instance):        """Return lineage after successful execution."""        from openlineage.client.event_v2 import Dataset        from airflow.providers.openlineage.extractors import OperatorLineage
        return OperatorLineage(            inputs=[Dataset(namespace="warehouse://db", name=self.source_table)],            outputs=[Dataset(namespace="warehouse://db", name=self.target_table)],        )

from airflow.models import BaseOperatorfrom openlineage.client.event_v2 import Dataset
class MyCustomOperator(BaseOperator):    def __init__(self, source_table: str, target_table: str, **kwargs):        super().__init__(**kwargs)        self.source_table = source_table        self.target_table = target_table
    def execute(self, context):        # Set lineage dynamically based on operator parameters        self.inlets = [            Dataset(namespace="warehouse://db", name=self.source_table)        ]        self.outlets = [            Dataset(namespace="warehouse://db", name=self.target_table)        ]
        # ... perform the actual work ...        self.log.info(f"Processing {self.source_table} -> {self.target_table}")

from openlineage.client.event_v2 import Dataset
# Snowflakefrom openlineage.client.naming.snowflake import SnowflakeDatasetNaming
naming = SnowflakeDatasetNaming(    account_identifier="myorg-myaccount",    database="mydb",    schema="myschema",    table="mytable",)dataset = Dataset(namespace=naming.get_namespace(), name=naming.get_name())# -> namespace: "snowflake://myorg-myaccount", name: "mydb.myschema.mytable"
# BigQueryfrom openlineage.client.naming.bigquery import BigQueryDatasetNaming
naming = BigQueryDatasetNaming(    project="my-project",    dataset="my_dataset",    table="my_table",)dataset = Dataset(namespace=naming.get_namespace(), name=naming.get_name())# -> namespace: "bigquery", name: "my-project.my_dataset.my_table"
# S3from openlineage.client.naming.s3 import S3DatasetNaming
naming = S3DatasetNaming(bucket="my-bucket", key="path/to/file.parquet")dataset = Dataset(namespace=naming.get_namespace(), name=naming.get_name())# -> namespace: "s3://my-bucket", name: "path/to/file.parquet"
# PostgreSQLfrom openlineage.client.naming.postgres import PostgresDatasetNaming
naming = PostgresDatasetNaming(    host="localhost",    port=5432,    database="mydb",    schema="public",    table="users",)dataset = Dataset(namespace=naming.get_namespace(), name=naming.get_name())# -> namespace: "postgres://localhost:5432", name: "mydb.public.users"

from openlineage.client.event_v2 import Datasetfrom openlineage.client.naming.snowflake import SnowflakeDatasetNaming
def snowflake_dataset(schema: str, table: str) -> Dataset:    """Create a Snowflake Dataset using the naming helper."""    naming = SnowflakeDatasetNaming(        account_identifier="mycompany",        database="analytics",        schema=schema,        table=table,    )    return Dataset(namespace=naming.get_namespace(), name=naming.get_name())
# Usagesource = snowflake_dataset("raw", "orders")target = snowflake_dataset("staging", "orders_clean")

transform = SqlOperator(    task_id="transform_orders",    sql="...",    # Lineage: Reads raw orders, joins with customers, writes to staging    inlets=[        snowflake_dataset("raw", "orders"),        snowflake_dataset("raw", "customers"),    ],    outlets=[        snowflake_dataset("staging", "order_details"),    ],)