apache-airflow-providers-openlineage 1.3.1rc1__py3-none-any.whl

airflow/providers/openlineage/plugins/macros.py

@@ -0,0 +1,60 @@
+# 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.
+from __future__ import annotations
+
+import os
+import typing
+
+from airflow.configuration import conf
+from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter
+
+if typing.TYPE_CHECKING:
+    from airflow.models import TaskInstance
+
+_JOB_NAMESPACE = conf.get("openlineage", "namespace", fallback=os.getenv("OPENLINEAGE_NAMESPACE", "default"))
+
+
+def lineage_run_id(task_instance: TaskInstance):
+    """
+    Macro function which returns the generated run id for a given task.
+
+    This can be used to forward the run id from a task to a child run so the job hierarchy is preserved.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/macros:openlineage`
+    """
+    return OpenLineageAdapter.build_task_instance_run_id(
+        task_instance.task.task_id, task_instance.execution_date, task_instance.try_number
+    )
+
+
+def lineage_parent_id(run_id: str, task_instance: TaskInstance):
+    """
+    Macro function which returns the generated job and run id for a given task.
+
+    This can be used to forward the ids from a task to a child run so the job
+    hierarchy is preserved. Child run can create ParentRunFacet from those ids.
+
+    .. seealso::
+        For more information on how to use this macro, take a look at the guide:
+        :ref:`howto/macros:openlineage`
+    """
+    job_name = OpenLineageAdapter.build_task_instance_run_id(
+        task_instance.task.task_id, task_instance.execution_date, task_instance.try_number
+    )
+    return f"{_JOB_NAMESPACE}/{job_name}/{run_id}"

airflow/providers/openlineage/plugins/openlineage.py

@@ -0,0 +1,51 @@
+# 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.
+from __future__ import annotations
+
+import os
+
+from airflow.configuration import conf
+from airflow.plugins_manager import AirflowPlugin
+from airflow.providers.openlineage.plugins.listener import get_openlineage_listener
+from airflow.providers.openlineage.plugins.macros import lineage_parent_id, lineage_run_id
+
+
+def _is_disabled() -> bool:
+    return (
+        conf.getboolean("openlineage", "disabled", fallback=False)
+        or os.getenv("OPENLINEAGE_DISABLED", "false").lower() == "true"
+        or (
+            conf.get("openlineage", "transport", fallback="") == ""
+            and conf.get("openlineage", "config_path", fallback="") == ""
+            and os.getenv("OPENLINEAGE_URL", "") == ""
+            and os.getenv("OPENLINEAGE_CONFIG", "") == ""
+        )
+    )
+
+
+class OpenLineageProviderPlugin(AirflowPlugin):
+    """
+    Listener that emits numerous Events.
+
+    OpenLineage Plugin provides listener that emits OL events on DAG start,
+    complete and failure and TaskInstances start, complete and failure.
+    """
+
+    name = "OpenLineageProviderPlugin"
+    if not _is_disabled():
+        macros = [lineage_run_id, lineage_parent_id]
+        listeners = [get_openlineage_listener()]

airflow/providers/openlineage/sqlparser.py

@@ -0,0 +1,347 @@
+# 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable
+
+import sqlparse
+from attrs import define
+from openlineage.client.facet import (
+    BaseFacet,
+    ColumnLineageDatasetFacet,
+    ColumnLineageDatasetFacetFieldsAdditional,
+    ColumnLineageDatasetFacetFieldsAdditionalInputFields,
+    ExtractionError,
+    ExtractionErrorRunFacet,
+    SqlJobFacet,
+)
+from openlineage.common.sql import DbTableMeta, SqlMeta, parse
+
+from airflow.providers.openlineage.extractors.base import OperatorLineage
+from airflow.providers.openlineage.utils.sql import (
+    TablesHierarchy,
+    create_information_schema_query,
+    get_table_schemas,
+)
+from airflow.typing_compat import TypedDict
+
+if TYPE_CHECKING:
+    from openlineage.client.run import Dataset
+    from sqlalchemy.engine import Engine
+
+    from airflow.hooks.base import BaseHook
+
+DEFAULT_NAMESPACE = "default"
+DEFAULT_INFORMATION_SCHEMA_COLUMNS = [
+    "table_schema",
+    "table_name",
+    "column_name",
+    "ordinal_position",
+    "udt_name",
+]
+DEFAULT_INFORMATION_SCHEMA_TABLE_NAME = "information_schema.columns"
+
+
+def default_normalize_name_method(name: str) -> str:
+    return name.lower()
+
+
+class GetTableSchemasParams(TypedDict):
+    """get_table_schemas params."""
+
+    normalize_name: Callable[[str], str]
+    is_cross_db: bool
+    information_schema_columns: list[str]
+    information_schema_table: str
+    is_uppercase_names: bool
+    database: str | None
+
+
+@define
+class DatabaseInfo:
+    """
+    Contains database specific information needed to process SQL statement parse result.
+
+    :param scheme: Scheme part of URI in OpenLineage namespace.
+    :param authority: Authority part of URI in OpenLineage namespace.
+        For most cases it should return `{host}:{port}` part of Airflow connection.
+        See: https://github.com/OpenLineage/OpenLineage/blob/main/spec/Naming.md
+    :param database: Takes precedence over parsed database name.
+    :param information_schema_columns: List of columns names from information schema table.
+    :param information_schema_table_name: Information schema table name.
+    :param is_information_schema_cross_db: Specifies if information schema contains
+        cross-database data.
+    :param is_uppercase_names: Specifies if database accepts only uppercase names (e.g. Snowflake).
+    :param normalize_name_method: Method to normalize database, schema and table names.
+        Defaults to `name.lower()`.
+    """
+
+    scheme: str
+    authority: str | None = None
+    database: str | None = None
+    information_schema_columns: list[str] = DEFAULT_INFORMATION_SCHEMA_COLUMNS
+    information_schema_table_name: str = DEFAULT_INFORMATION_SCHEMA_TABLE_NAME
+    is_information_schema_cross_db: bool = False
+    is_uppercase_names: bool = False
+    normalize_name_method: Callable[[str], str] = default_normalize_name_method
+
+
+class SQLParser:
+    """Interface for openlineage-sql.
+
+    :param dialect: dialect specific to the database
+    :param default_schema: schema applied to each table with no schema parsed
+    """
+
+    def __init__(self, dialect: str | None = None, default_schema: str | None = None) -> None:
+        self.dialect = dialect
+        self.default_schema = default_schema
+
+    def parse(self, sql: list[str] | str) -> SqlMeta | None:
+        """Parse a single or a list of SQL statements."""
+        return parse(sql=sql, dialect=self.dialect)
+
+    def parse_table_schemas(
+        self,
+        hook: BaseHook,
+        inputs: list[DbTableMeta],
+        outputs: list[DbTableMeta],
+        database_info: DatabaseInfo,
+        namespace: str = DEFAULT_NAMESPACE,
+        database: str | None = None,
+        sqlalchemy_engine: Engine | None = None,
+    ) -> tuple[list[Dataset], ...]:
+        """Parse schemas for input and output tables."""
+        database_kwargs: GetTableSchemasParams = {
+            "normalize_name": database_info.normalize_name_method,
+            "is_cross_db": database_info.is_information_schema_cross_db,
+            "information_schema_columns": database_info.information_schema_columns,
+            "information_schema_table": database_info.information_schema_table_name,
+            "is_uppercase_names": database_info.is_uppercase_names,
+            "database": database or database_info.database,
+        }
+        return get_table_schemas(
+            hook,
+            namespace,
+            self.default_schema,
+            database or database_info.database,
+            self.create_information_schema_query(
+                tables=inputs, sqlalchemy_engine=sqlalchemy_engine, **database_kwargs
+            )
+            if inputs
+            else None,
+            self.create_information_schema_query(
+                tables=outputs, sqlalchemy_engine=sqlalchemy_engine, **database_kwargs
+            )
+            if outputs
+            else None,
+        )
+
+    def attach_column_lineage(
+        self, datasets: list[Dataset], database: str | None, parse_result: SqlMeta
+    ) -> None:
+        """
+        Attaches column lineage facet to the list of datasets.
+
+        Note that currently each dataset has the same column lineage information set.
+        This would be a matter of change after OpenLineage SQL Parser improvements.
+        """
+        if not len(parse_result.column_lineage):
+            return
+        for dataset in datasets:
+            dataset.facets["columnLineage"] = ColumnLineageDatasetFacet(
+                fields={
+                    column_lineage.descendant.name: ColumnLineageDatasetFacetFieldsAdditional(
+                        inputFields=[
+                            ColumnLineageDatasetFacetFieldsAdditionalInputFields(
+                                namespace=dataset.namespace,
+                                name=".".join(
+                                    filter(
+                                        None,
+                                        (
+                                            column_meta.origin.database or database,
+                                            column_meta.origin.schema or self.default_schema,
+                                            column_meta.origin.name,
+                                        ),
+                                    )
+                                )
+                                if column_meta.origin
+                                else "",
+                                field=column_meta.name,
+                            )
+                            for column_meta in column_lineage.lineage
+                        ],
+                        transformationType="",
+                        transformationDescription="",
+                    )
+                    for column_lineage in parse_result.column_lineage
+                }
+            )
+
+    def generate_openlineage_metadata_from_sql(
+        self,
+        sql: list[str] | str,
+        hook: BaseHook,
+        database_info: DatabaseInfo,
+        database: str | None = None,
+        sqlalchemy_engine: Engine | None = None,
+    ) -> OperatorLineage:
+        """Parses SQL statement(s) and generates OpenLineage metadata.
+
+        Generated OpenLineage metadata contains:
+
+        * input tables with schemas parsed
+        * output tables with schemas parsed
+        * run facets
+        * job facets.
+
+        :param sql: a SQL statement or list of SQL statement to be parsed
+        :param hook: Airflow Hook used to connect to the database
+        :param database_info: database specific information
+        :param database: when passed it takes precedence over parsed database name
+        :param sqlalchemy_engine: when passed, engine's dialect is used to compile SQL queries
+        """
+        job_facets: dict[str, BaseFacet] = {"sql": SqlJobFacet(query=self.normalize_sql(sql))}
+        parse_result = self.parse(self.split_sql_string(sql))
+        if not parse_result:
+            return OperatorLineage(job_facets=job_facets)
+
+        run_facets: dict[str, BaseFacet] = {}
+        if parse_result.errors:
+            run_facets["extractionError"] = ExtractionErrorRunFacet(
+                totalTasks=len(sql) if isinstance(sql, list) else 1,
+                failedTasks=len(parse_result.errors),
+                errors=[
+                    ExtractionError(
+                        errorMessage=error.message,
+                        stackTrace=None,
+                        task=error.origin_statement,
+                        taskNumber=error.index,
+                    )
+                    for error in parse_result.errors
+                ],
+            )
+
+        namespace = self.create_namespace(database_info=database_info)
+        inputs, outputs = self.parse_table_schemas(
+            hook=hook,
+            inputs=parse_result.in_tables,
+            outputs=parse_result.out_tables,
+            namespace=namespace,
+            database=database,
+            database_info=database_info,
+            sqlalchemy_engine=sqlalchemy_engine,
+        )
+
+        self.attach_column_lineage(outputs, database or database_info.database, parse_result)
+
+        return OperatorLineage(
+            inputs=inputs,
+            outputs=outputs,
+            run_facets=run_facets,
+            job_facets=job_facets,
+        )
+
+    @staticmethod
+    def create_namespace(database_info: DatabaseInfo) -> str:
+        return (
+            f"{database_info.scheme}://{database_info.authority}"
+            if database_info.authority
+            else database_info.scheme
+        )
+
+    @classmethod
+    def normalize_sql(cls, sql: list[str] | str) -> str:
+        """Makes sure to return a semicolon-separated SQL statements."""
+        return ";\n".join(stmt.rstrip(" ;\r\n") for stmt in cls.split_sql_string(sql))
+
+    @classmethod
+    def split_sql_string(cls, sql: list[str] | str) -> list[str]:
+        """
+        Split SQL string into list of statements.
+
+        Tries to use `DbApiHook.split_sql_string` if available.
+        Otherwise, uses the same logic.
+        """
+        try:
+            from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+            split_statement = DbApiHook.split_sql_string
+        except (ImportError, AttributeError):
+            # No common.sql Airflow provider available or version is too old.
+            def split_statement(sql: str) -> list[str]:
+                splits = sqlparse.split(sqlparse.format(sql, strip_comments=True))
+                return [s for s in splits if s]
+
+        if isinstance(sql, str):
+            return split_statement(sql)
+        return [obj for stmt in sql for obj in cls.split_sql_string(stmt) if obj != ""]
+
+    @classmethod
+    def create_information_schema_query(
+        cls,
+        tables: list[DbTableMeta],
+        normalize_name: Callable[[str], str],
+        is_cross_db: bool,
+        information_schema_columns,
+        information_schema_table,
+        is_uppercase_names,
+        database: str | None = None,
+        sqlalchemy_engine: Engine | None = None,
+    ) -> str:
+        """Creates SELECT statement to query information schema table."""
+        tables_hierarchy = cls._get_tables_hierarchy(
+            tables,
+            normalize_name=normalize_name,
+            database=database,
+            is_cross_db=is_cross_db,
+        )
+        return create_information_schema_query(
+            columns=information_schema_columns,
+            information_schema_table_name=information_schema_table,
+            tables_hierarchy=tables_hierarchy,
+            uppercase_names=is_uppercase_names,
+            sqlalchemy_engine=sqlalchemy_engine,
+        )
+
+    @staticmethod
+    def _get_tables_hierarchy(
+        tables: list[DbTableMeta],
+        normalize_name: Callable[[str], str],
+        database: str | None = None,
+        is_cross_db: bool = False,
+    ) -> TablesHierarchy:
+        """
+        Creates a hierarchy of database -> schema -> table name.
+
+        This helps to create simpler information schema query grouped by
+        database and schema.
+        :param tables: List of tables.
+        :param normalize_name: A method to normalize all names.
+        :param is_cross_db: If false, set top (database) level to None
+            when creating hierarchy.
+        """
+        hierarchy: TablesHierarchy = {}
+        for table in tables:
+            if is_cross_db:
+                db = table.database or database
+            else:
+                db = None
+            schemas = hierarchy.setdefault(normalize_name(db) if db else db, {})
+            tables = schemas.setdefault(normalize_name(table.schema) if table.schema else None, [])
+            tables.append(table.name)
+        return hierarchy

airflow/providers/openlineage/utils/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/openlineage/utils/sql.py

@@ -0,0 +1,199 @@
+# 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.
+from __future__ import annotations
+
+from collections import defaultdict
+from contextlib import closing
+from enum import IntEnum
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+from attrs import define
+from openlineage.client.facet import SchemaDatasetFacet, SchemaField
+from openlineage.client.run import Dataset
+from sqlalchemy import Column, MetaData, Table, and_, union_all
+
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Engine
+    from sqlalchemy.sql import ClauseElement
+
+    from airflow.hooks.base import BaseHook
+
+
+class ColumnIndex(IntEnum):
+    """Enumerates the indices of columns in information schema view."""
+
+    SCHEMA = 0
+    TABLE_NAME = 1
+    COLUMN_NAME = 2
+    ORDINAL_POSITION = 3
+    # Use 'udt_name' which is the underlying type of column
+    UDT_NAME = 4
+    # Database is optional as 5th column
+    DATABASE = 5
+
+
+TablesHierarchy = Dict[Optional[str], Dict[Optional[str], List[str]]]
+
+
+@define
+class TableSchema:
+    """Temporary object used to construct OpenLineage Dataset."""
+
+    table: str
+    schema: str | None
+    database: str | None
+    fields: list[SchemaField]
+
+    def to_dataset(self, namespace: str, database: str | None = None, schema: str | None = None) -> Dataset:
+        # Prefix the table name with database and schema name using
+        # the format: {database_name}.{table_schema}.{table_name}.
+        name = ".".join(
+            part
+            for part in [self.database or database, self.schema or schema, self.table]
+            if part is not None
+        )
+        return Dataset(
+            namespace=namespace,
+            name=name,
+            facets={"schema": SchemaDatasetFacet(fields=self.fields)} if self.fields else {},
+        )
+
+
+def get_table_schemas(
+    hook: BaseHook,
+    namespace: str,
+    schema: str | None,
+    database: str | None,
+    in_query: str | None,
+    out_query: str | None,
+) -> tuple[list[Dataset], list[Dataset]]:
+    """Query database for table schemas.
+
+    Uses provided hook. Responsibility to provide queries for this function is on particular extractors.
+    If query for input or output table isn't provided, the query is skipped.
+    """
+    # Do not query if we did not get both queries
+    if not in_query and not out_query:
+        return [], []
+
+    with closing(hook.get_conn()) as conn, closing(conn.cursor()) as cursor:
+        if in_query:
+            cursor.execute(in_query)
+            in_datasets = [x.to_dataset(namespace, database, schema) for x in parse_query_result(cursor)]
+        else:
+            in_datasets = []
+        if out_query:
+            cursor.execute(out_query)
+            out_datasets = [x.to_dataset(namespace, database, schema) for x in parse_query_result(cursor)]
+        else:
+            out_datasets = []
+    return in_datasets, out_datasets
+
+
+def parse_query_result(cursor) -> list[TableSchema]:
+    """Fetch results from DB-API 2.0 cursor and creates list of table schemas.
+
+    For each row it creates :class:`TableSchema`.
+    """
+    schemas: dict = {}
+    columns: dict = defaultdict(list)
+    for row in cursor.fetchall():
+        table_schema_name: str = row[ColumnIndex.SCHEMA]
+        table_name: str = row[ColumnIndex.TABLE_NAME]
+        table_column: SchemaField = SchemaField(
+            name=row[ColumnIndex.COLUMN_NAME],
+            type=row[ColumnIndex.UDT_NAME],
+            description=None,
+        )
+        ordinal_position = row[ColumnIndex.ORDINAL_POSITION]
+        try:
+            table_database = row[ColumnIndex.DATABASE]
+        except IndexError:
+            table_database = None
+
+        # Attempt to get table schema
+        table_key = ".".join(filter(None, [table_database, table_schema_name, table_name]))
+
+        schemas[table_key] = TableSchema(
+            table=table_name, schema=table_schema_name, database=table_database, fields=[]
+        )
+        columns[table_key].append((ordinal_position, table_column))
+
+    for schema in schemas.values():
+        table_key = ".".join(filter(None, [schema.database, schema.schema, schema.table]))
+        schema.fields = [x for _, x in sorted(columns[table_key])]
+
+    return list(schemas.values())
+
+
+def create_information_schema_query(
+    columns: list[str],
+    information_schema_table_name: str,
+    tables_hierarchy: TablesHierarchy,
+    uppercase_names: bool = False,
+    sqlalchemy_engine: Engine | None = None,
+) -> str:
+    """Creates query for getting table schemas from information schema."""
+    metadata = MetaData(sqlalchemy_engine)
+    select_statements = []
+    for db, schema_mapping in tables_hierarchy.items():
+        # Information schema table name is expected to be "< information_schema schema >.<view/table name>"
+        # usually "information_schema.columns". In order to use table identifier correct for various table
+        # we need to pass first part of dot-separated identifier as `schema` argument to `sqlalchemy.Table`.
+        if db:
+            # Use database as first part of table identifier.
+            schema = db
+            table_name = information_schema_table_name
+        else:
+            # When no database passed, use schema as first part of table identifier.
+            schema, table_name = information_schema_table_name.split(".")
+        information_schema_table = Table(
+            table_name,
+            metadata,
+            *[Column(column) for column in columns],
+            schema=schema,
+            quote=False,
+        )
+        filter_clauses = create_filter_clauses(schema_mapping, information_schema_table, uppercase_names)
+        select_statements.append(information_schema_table.select().filter(*filter_clauses))
+    return str(
+        union_all(*select_statements).compile(sqlalchemy_engine, compile_kwargs={"literal_binds": True})
+    )
+
+
+def create_filter_clauses(
+    schema_mapping: dict, information_schema_table: Table, uppercase_names: bool = False
+) -> ClauseElement:
+    """
+    Creates comprehensive filter clauses for all tables in one database.
+
+    :param schema_mapping: a dictionary of schema names and list of tables in each
+    :param information_schema_table: `sqlalchemy.Table` instance used to construct clauses
+        For most SQL dbs it contains `table_name` and `table_schema` columns,
+        therefore it is expected the table has them defined.
+    :param uppercase_names: if True use schema and table names uppercase
+    """
+    filter_clauses = []
+    for schema, tables in schema_mapping.items():
+        filter_clause = information_schema_table.c.table_name.in_(
+            name.upper() if uppercase_names else name for name in tables
+        )
+        if schema:
+            schema = schema.upper() if uppercase_names else schema
+            filter_clause = and_(information_schema_table.c.table_schema == schema, filter_clause)
+        filter_clauses.append(filter_clause)
+    return filter_clauses