This is an automated email from the ASF dual-hosted git repository.
mobuchowski pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ae50f958ea2 Surface errors from OpenLineage database info lookups
without losing lineage (#74805)
ae50f958ea2 is described below
commit ae50f958ea22243852d3b58c511aafde615b91f1
Author: Kacper Muda <[email protected]>
AuthorDate: Tue Oct 27 16:35:22 2026 +0100
Surface errors from OpenLineage database info lookups without losing
lineage (#74805)
* Don't silence AttributeError from OpenLineage database info lookup
SQL OpenLineage extraction called hook.get_openlineage_database_info()
inside `except AttributeError` blocks, meant to detect hooks that don't
implement the method. The same handler also swallowed AttributeErrors
raised from inside the method (e.g. a broken optional import in
SnowflakeHook.get_uri()), logging only at debug level. Lineage silently
came out without inputs/outputs and nothing in the task logs said why.
Look the method up with getattr() instead and call it outside any
try/except, so real errors propagate (the extractor already turns them
into a warning and still emits events). Also log failures in SQL hook
lineage at warning level, with details at debug.
Co-Authored-By: Claude <[email protected]>
* Keep partial OpenLineage SQL lineage when database info lookup fails
Reviewers pointed out that letting AttributeError escape lost all lineage
for the task, including GCS outputs and query ids that don't depend on the
failing method, and that hooks without the method were warned about.
A missing method is now detected with getattr and stays silent. An error
raised inside the method is logged as a warning and lineage falls back to
what can still be built. The same applies to the dialect, default schema
and database specific lineage lookups.
Co-Authored-By: Claude <[email protected]>
* Don't silence AttributeError in the common.compat OpenLineage SQL fallback
The fallback used when the OpenLineage provider is too old to ship
get_openlineage_facets_with_sql had the same except AttributeError blocks,
so an error raised inside the hook's database info, dialect or default
schema lookup was reported as a hook without those methods.
Co-Authored-By: Claude <[email protected]>
* Keep OpenLineage lineage when the SQLAlchemy engine or hook helpers fail
A failing get_sqlalchemy_engine() in SQLExecuteQueryOperator still dropped
all lineage of the task, and the new warnings named the hook object and the
wrong method, which made failures hard to attribute to a connection.
Co-Authored-By: Claude <[email protected]>
* Name the connection and both lookups in OpenLineage SQL debug messages
Co-Authored-By: Claude <[email protected]>
* Report a failing OpenLineage dialect lookup once per hook in SQL hook
lineage
A hook whose dialect or default schema lookup fails would log the same
warning for every SQL extra of a task. The SQLParser is now built once per
hook and reused, and the SQL parsing warning names the connection.
Co-Authored-By: Claude <[email protected]>
* Accept a missing connection id in get_sql_parser_for_hook
Co-Authored-By: Claude <[email protected]>
---------
Co-authored-by: Claude <[email protected]>
---
.../common/compat/openlineage/utils/sql.py | 29 ++++---
.../common/compat/openlineage/utils/test_sql.py | 60 +++++++++++++++
.../airflow/providers/common/sql/operators/sql.py | 55 ++++++++-----
.../unit/common/sql/operators/test_sql_execute.py | 89 ++++++++++++++++++++++
.../src/airflow/providers/openlineage/sqlparser.py | 45 +++++++----
.../openlineage/utils/sql_hook_lineage.py | 21 ++++-
.../tests/unit/openlineage/test_sqlparser.py | 53 +++++++++++--
.../openlineage/utils/test_sql_hook_lineage.py | 84 ++++++++++++++++++++
8 files changed, 385 insertions(+), 51 deletions(-)
diff --git
a/providers/common/compat/src/airflow/providers/common/compat/openlineage/utils/sql.py
b/providers/common/compat/src/airflow/providers/common/compat/openlineage/utils/sql.py
index 72e9e1f94f0..40a4e4a0604 100644
---
a/providers/common/compat/src/airflow/providers/common/compat/openlineage/utils/sql.py
+++
b/providers/common/compat/src/airflow/providers/common/compat/openlineage/utils/sql.py
@@ -51,22 +51,31 @@ else:
use_external_connection = True
connection = hook.get_connection(conn_id)
- try:
- database_info = hook.get_openlineage_database_info(connection)
- except AttributeError:
- log.debug("%s has no database info provided", hook)
- database_info = None
+ database_info = None
+ get_database_info = getattr(hook, "get_openlineage_database_info",
None)
+ if get_database_info is not None:
+ try:
+ database_info = get_database_info(connection)
+ except Exception as e:
+ log.warning("Failed to get OpenLineage database info for
%s: %s", conn_id, e)
+ log.debug("OpenLineage failure details:", exc_info=True)
if database_info is None:
+ log.debug("%s has no database info provided", conn_id)
return None
+ get_dialect = getattr(hook, "get_openlineage_database_dialect",
None)
+ get_default_schema = getattr(hook,
"get_openlineage_default_schema", None)
+ if get_dialect is None or get_default_schema is None:
+ log.debug("%s has no database dialect or default schema
provided", conn_id)
+ return None
try:
- sql_parser = SQLParser(
- dialect=hook.get_openlineage_database_dialect(connection),
- default_schema=hook.get_openlineage_default_schema(),
+ sql_parser = SQLParser(dialect=get_dialect(connection),
default_schema=get_default_schema())
+ except Exception as e:
+ log.warning(
+ "Failed to get OpenLineage database dialect or default
schema for %s: %s", conn_id, e
)
- except AttributeError:
- log.debug("%s failed to get database dialect", hook)
+ log.debug("OpenLineage failure details:", exc_info=True)
return None
try:
diff --git
a/providers/common/compat/tests/unit/common/compat/openlineage/utils/test_sql.py
b/providers/common/compat/tests/unit/common/compat/openlineage/utils/test_sql.py
index dfa3eace880..44fd8f5033e 100644
---
a/providers/common/compat/tests/unit/common/compat/openlineage/utils/test_sql.py
+++
b/providers/common/compat/tests/unit/common/compat/openlineage/utils/test_sql.py
@@ -16,8 +16,68 @@
# under the License.
from __future__ import annotations
+import importlib
+from unittest import mock
+
+import pytest
+
+_MODULE = "airflow.providers.common.compat.openlineage.utils.sql"
+
def test_import():
from airflow.providers.common.compat.openlineage.utils.sql import
get_openlineage_facets_with_sql
assert get_openlineage_facets_with_sql is not None
+
+
[email protected]
+def fallback_get_openlineage_facets_with_sql(monkeypatch):
+ """Load the module as if the OpenLineage provider was too old to ship
``get_openlineage_facets_with_sql``."""
+ from airflow.providers.openlineage import sqlparser
+
+ monkeypatch.delattr(sqlparser, "get_openlineage_facets_with_sql")
+ module = importlib.reload(importlib.import_module(_MODULE))
+ yield module.get_openlineage_facets_with_sql
+ monkeypatch.undo()
+ importlib.reload(module)
+
+
[email protected](f"{_MODULE}.log", autospec=True)
+def test_fallback_hook_without_database_info_method_returns_none(
+ mock_log, fallback_get_openlineage_facets_with_sql
+):
+ hook = mock.MagicMock(spec=["get_connection"])
+
+ assert fallback_get_openlineage_facets_with_sql(hook, "SELECT 1", "conn",
None) is None
+ mock_log.warning.assert_not_called()
+
+
[email protected](f"{_MODULE}.log", autospec=True)
+def test_fallback_error_from_database_info_is_logged_as_warning(
+ mock_log, fallback_get_openlineage_facets_with_sql
+):
+ hook = mock.MagicMock(spec=["get_connection",
"get_openlineage_database_info"])
+ hook.get_openlineage_database_info.side_effect = AttributeError("broken
import")
+
+ assert fallback_get_openlineage_facets_with_sql(hook, "SELECT 1", "conn",
None) is None
+ mock_log.warning.assert_called_once()
+ assert mock_log.warning.call_args.args[1] == "conn"
+ assert str(mock_log.warning.call_args.args[2]) == "broken import"
+
+
[email protected](f"{_MODULE}.log", autospec=True)
+def test_fallback_error_from_dialect_is_logged_as_warning(mock_log,
fallback_get_openlineage_facets_with_sql):
+ hook = mock.MagicMock(
+ spec=[
+ "get_connection",
+ "get_openlineage_database_info",
+ "get_openlineage_database_dialect",
+ "get_openlineage_default_schema",
+ ]
+ )
+ hook.get_openlineage_database_dialect.side_effect = AttributeError("broken
import")
+
+ assert fallback_get_openlineage_facets_with_sql(hook, "SELECT 1", "conn",
None) is None
+ mock_log.warning.assert_called_once()
+ assert mock_log.warning.call_args.args[1] == "conn"
+ assert str(mock_log.warning.call_args.args[2]) == "broken import"
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py
b/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py
index 01148ae4507..160485e5ed9 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py
@@ -296,32 +296,47 @@ class BaseSQLOperator(BaseOperator):
# OpenLineage provider release < 1.8.0 - we always use connection
use_external_connection = True
- connection = hook.get_connection(getattr(hook, hook.conn_name_attr))
- try:
- database_info = hook.get_openlineage_database_info(connection)
- except AttributeError:
- self.log.debug("%s has no database info provided", hook)
- database_info = None
+ conn_id = getattr(hook, hook.conn_name_attr)
+ connection = hook.get_connection(conn_id)
+ get_database_info = getattr(hook, "get_openlineage_database_info",
None)
+ database_info = None
+ if get_database_info is not None:
+ try:
+ database_info = get_database_info(connection)
+ except Exception as e:
+ self.log.warning("Failed to get OpenLineage database info for
%s: %s", conn_id, e)
+ self.log.debug("OpenLineage failure details:", exc_info=True)
if database_info is None:
self.log.debug("OpenLineage could not retrieve database
information. Skipping.")
return OperatorLineage()
+ get_dialect = getattr(hook, "get_openlineage_database_dialect", None)
+ get_default_schema = getattr(hook, "get_openlineage_default_schema",
None)
+ if get_dialect is None or get_default_schema is None:
+ self.log.debug("%s has no database dialect or default schema
provided", conn_id)
+ return None
try:
- sql_parser = SQLParser(
- dialect=hook.get_openlineage_database_dialect(connection),
- default_schema=hook.get_openlineage_default_schema(),
+ sql_parser = SQLParser(dialect=get_dialect(connection),
default_schema=get_default_schema())
+ except Exception as e:
+ self.log.warning(
+ "Failed to get OpenLineage database dialect or default schema
for %s: %s", conn_id, e
)
- except AttributeError:
- self.log.debug("%s failed to get database dialect", hook)
+ self.log.debug("OpenLineage failure details:", exc_info=True)
return None
+ sqlalchemy_engine = None
+ try:
+ sqlalchemy_engine = hook.get_sqlalchemy_engine()
+ except Exception as e:
+ self.log.debug("Failed to get sql alchemy engine: %s", e)
+
operator_lineage = sql_parser.generate_openlineage_metadata_from_sql(
sql=sql,
hook=hook,
database_info=database_info,
database=self.database,
- sqlalchemy_engine=hook.get_sqlalchemy_engine(),
+ sqlalchemy_engine=sqlalchemy_engine,
use_connection=use_external_connection,
)
@@ -337,11 +352,17 @@ class BaseSQLOperator(BaseOperator):
operator_lineage = self.get_openlineage_facets_on_start() or
OperatorLineage()
hook = self.get_db_hook()
- try:
- database_specific_lineage =
hook.get_openlineage_database_specific_lineage(task_instance)
- except AttributeError:
- self.log.debug("%s has no database specific lineage provided",
hook)
- database_specific_lineage = None
+ conn_id = getattr(hook, hook.conn_name_attr)
+ get_specific_lineage = getattr(hook,
"get_openlineage_database_specific_lineage", None)
+ database_specific_lineage = None
+ if get_specific_lineage is None:
+ self.log.debug("%s has no database specific lineage provided",
conn_id)
+ else:
+ try:
+ database_specific_lineage = get_specific_lineage(task_instance)
+ except Exception as e:
+ self.log.warning("Failed to get OpenLineage database specific
lineage for %s: %s", conn_id, e)
+ self.log.debug("OpenLineage failure details:", exc_info=True)
if database_specific_lineage is None:
if not self.check_results:
diff --git
a/providers/common/sql/tests/unit/common/sql/operators/test_sql_execute.py
b/providers/common/sql/tests/unit/common/sql/operators/test_sql_execute.py
index 564cfabd1a6..8721034474d 100644
--- a/providers/common/sql/tests/unit/common/sql/operators/test_sql_execute.py
+++ b/providers/common/sql/tests/unit/common/sql/operators/test_sql_execute.py
@@ -368,6 +368,95 @@ FORGOT TO COMMENT"""
)
+class _SQLExecuteQueryOperatorForTest(SQLExecuteQueryOperator):
+ hook: Any
+
+ def get_db_hook(self):
+ return self.hook
+
+
+def _make_openlineage_operator(hook) -> _SQLExecuteQueryOperatorForTest:
+ hook.conn_name_attr = "conn_id"
+ hook.conn_id = "sql_default"
+ op = _SQLExecuteQueryOperatorForTest(task_id=TASK_ID, sql="SELECT 1;")
+ op.hook = hook
+ return op
+
+
[email protected](SQLExecuteQueryOperator, "log", autospec=True)
+def
test_openlineage_hook_without_database_info_method_returns_empty_lineage(mock_log):
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ op = _make_openlineage_operator(mock.MagicMock(spec=["get_connection",
"conn_name_attr", "conn_id"]))
+
+ assert op.get_openlineage_facets_on_start() == OperatorLineage()
+ mock_log.warning.assert_not_called()
+
+
[email protected](SQLExecuteQueryOperator, "log", autospec=True)
+def
test_openlineage_error_from_database_info_is_logged_and_does_not_lose_other_lineage(mock_log):
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ hook = mock.MagicMock(spec=DbApiHook)
+ hook.get_openlineage_database_info.side_effect = AttributeError("broken
import")
+ hook.get_openlineage_database_specific_lineage.return_value =
OperatorLineage(
+ run_facets={"externalQuery": mock.sentinel.facet}
+ )
+ op = _make_openlineage_operator(hook)
+
+ assert op.get_openlineage_facets_on_start() == OperatorLineage()
+ lineage = op.get_openlineage_facets_on_complete(None)
+
+ assert lineage.run_facets == {"externalQuery": mock.sentinel.facet}
+ assert mock_log.warning.call_count == 2
+ assert mock_log.warning.call_args.args[1] == "sql_default"
+ assert str(mock_log.warning.call_args.args[2]) == "broken import"
+
+
[email protected](SQLExecuteQueryOperator, "log", autospec=True)
+def test_openlineage_error_from_dialect_is_logged_with_conn_id(mock_log):
+ hook = mock.MagicMock(spec=DbApiHook)
+ hook.get_openlineage_database_info.return_value =
mock.sentinel.database_info
+ hook.get_openlineage_database_dialect.side_effect = AttributeError("broken
import")
+ op = _make_openlineage_operator(hook)
+
+ assert op.get_openlineage_facets_on_start() is None
+ mock_log.warning.assert_called_once()
+ assert mock_log.warning.call_args.args[1] == "sql_default"
+ assert str(mock_log.warning.call_args.args[2]) == "broken import"
+
+
[email protected](SQLExecuteQueryOperator, "log", autospec=True)
+def
test_openlineage_error_from_database_specific_lineage_is_logged_and_falls_back(mock_log):
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ hook = mock.MagicMock(spec=DbApiHook)
+ hook.get_openlineage_database_info.return_value = None
+ hook.get_openlineage_database_specific_lineage.side_effect =
AttributeError("broken import")
+ op = _make_openlineage_operator(hook)
+
+ assert op.get_openlineage_facets_on_complete(None) == OperatorLineage()
+ mock_log.warning.assert_called_once()
+ assert mock_log.warning.call_args.args[1] == "sql_default"
+ assert str(mock_log.warning.call_args.args[2]) == "broken import"
+
+
[email protected]("airflow.providers.openlineage.sqlparser.SQLParser", autospec=True)
+def
test_openlineage_sqlalchemy_engine_failure_does_not_lose_lineage(mock_sql_parser):
+ from airflow.providers.openlineage.extractors import OperatorLineage
+
+ hook = mock.MagicMock(spec=DbApiHook)
+ hook.get_sqlalchemy_engine.side_effect = RuntimeError("no dialect")
+ hook.get_openlineage_database_info.return_value =
mock.sentinel.database_info
+ op = _make_openlineage_operator(hook)
+ expected = OperatorLineage()
+
mock_sql_parser.return_value.generate_openlineage_metadata_from_sql.return_value
= expected
+
+ assert op.get_openlineage_facets_on_start() is expected
+ generate =
mock_sql_parser.return_value.generate_openlineage_metadata_from_sql
+ assert generate.call_args.kwargs["sqlalchemy_engine"] is None
+
+
def test_with_no_openlineage_provider():
import importlib
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/sqlparser.py
b/providers/openlineage/src/airflow/providers/openlineage/sqlparser.py
index 8ce25f342f4..ef1e6d61b50 100644
--- a/providers/openlineage/src/airflow/providers/openlineage/sqlparser.py
+++ b/providers/openlineage/src/airflow/providers/openlineage/sqlparser.py
@@ -472,6 +472,20 @@ class SQLParser(LoggingMixin):
return hierarchy
+def get_sql_parser_for_hook(hook: DbApiHook, connection: Any, conn_id: str |
None) -> SQLParser | None:
+ get_dialect = getattr(hook, "get_openlineage_database_dialect", None)
+ get_default_schema = getattr(hook, "get_openlineage_default_schema", None)
+ if get_dialect is None or get_default_schema is None:
+ log.debug("%s has no database dialect or default schema provided",
conn_id)
+ return None
+ try:
+ return SQLParser(dialect=get_dialect(connection),
default_schema=get_default_schema())
+ except Exception as e:
+ log.warning("Failed to get OpenLineage database dialect or default
schema for %s: %s", conn_id, e)
+ log.debug("OpenLineage failure details:", exc_info=True)
+ return None
+
+
def get_openlineage_facets_with_sql(
hook: DbApiHook,
sql: str | list[str],
@@ -481,31 +495,32 @@ def get_openlineage_facets_with_sql(
*,
connection: Any | None = None,
database_info: Any | None = None,
+ sql_parser: SQLParser | None = None,
) -> OperatorLineage | None:
- # Accept pre-fetched connection and database_info to avoid redundant
hook.get_connection()
- # calls when processing multiple SQL extras from the same hook. Each
get_connection() call
+ # Accept pre-fetched connection, database_info and sql_parser to avoid
redundant lookups
+ # when processing multiple SQL extras from the same hook. Each
get_connection() call
# hits SecretsManager (miss) then the Airflow API server — passing these
in avoids that cost.
+ # A failing dialect/default schema lookup is then also reported once per
hook, not once per query.
if connection is None:
connection = hook.get_connection(conn_id)
if database_info is None:
- try:
- database_info = hook.get_openlineage_database_info(connection)
- except AttributeError:
- database_info = None
+ get_database_info = getattr(hook, "get_openlineage_database_info",
None)
+ if get_database_info is not None:
+ try:
+ database_info = get_database_info(connection)
+ except Exception as e:
+ log.warning("Failed to get OpenLineage database info for %s:
%s", conn_id, e)
+ log.debug("OpenLineage failure details:", exc_info=True)
if database_info is None:
- log.debug("%s has no database info provided", hook)
+ log.debug("%s has no database info provided", conn_id)
return None
- try:
- sql_parser = SQLParser(
- dialect=hook.get_openlineage_database_dialect(connection),
- default_schema=hook.get_openlineage_default_schema(),
- )
- except AttributeError:
- log.debug("%s failed to get database dialect", hook)
- return None
+ if sql_parser is None:
+ sql_parser = get_sql_parser_for_hook(hook, connection, conn_id)
+ if sql_parser is None:
+ return None
sqlalchemy_engine = None
if use_connection:
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
b/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
index 97ac22ad82a..5b395d270b6 100644
---
a/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
+++
b/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
@@ -47,6 +47,7 @@ from airflow.providers.openlineage.sqlparser import (
SQLParser,
from_table_meta,
get_openlineage_facets_with_sql,
+ get_sql_parser_for_hook,
)
from airflow.providers.openlineage.utils.emission_policy import DatasetFilter,
get_hook_class_name
from airflow.providers.openlineage.utils.utils import _get_parent_run_facet
@@ -141,12 +142,21 @@ def emit_lineage_from_sql_extras(
@cache
def _get_database_info(hook_id: int):
hook, conn_id = _hook_info[hook_id]
+ get_database_info = getattr(hook, "get_openlineage_database_info",
None)
+ if get_database_info is None:
+ return None
try:
- return hook.get_openlineage_database_info(_get_connection(hook_id))
+ return get_database_info(_get_connection(hook_id))
except Exception as e:
- log.debug("Failed to get OpenLineage database info for %s: %s",
conn_id, e)
+ log.warning("Failed to get OpenLineage database info for %s: %s",
conn_id, e)
+ log.debug("OpenLineage failure details:", exc_info=True)
return None
+ @cache
+ def _get_sql_parser(hook_id: int) -> SQLParser | None:
+ hook, conn_id = _hook_info[hook_id]
+ return get_sql_parser_for_hook(hook, _get_connection(hook_id), conn_id)
+
@cache
def _get_namespace(hook_id: int) -> str | None:
db_info = _get_database_info(hook_id)
@@ -170,7 +180,8 @@ def emit_lineage_from_sql_extras(
# Parse SQL to obtain lineage (inputs, outputs, facets)
query_lineage: OperatorLineage | None = None
database_info = _get_database_info(hook_id) if conn_id else None
- if sql_text and conn_id and database_info is not None:
+ sql_parser = _get_sql_parser(hook_id) if conn_id and database_info is
not None else None
+ if sql_text and conn_id and database_info is not None and sql_parser
is not None:
try:
query_lineage = get_openlineage_facets_with_sql(
hook=hook,
@@ -180,9 +191,11 @@ def emit_lineage_from_sql_extras(
use_connection=False, # Temporary solution before we
figure out timeouts for queries
connection=_get_connection(hook_id),
database_info=database_info,
+ sql_parser=sql_parser,
)
except Exception as e:
- log.debug("Failed to parse SQL for query %s: %s", query_count,
e)
+ log.warning("Failed to parse SQL for query %s on %s: %s",
query_count, conn_id, e)
+ log.debug("OpenLineage failure details:", exc_info=True)
# If parsing SQL failed, just attach SQL text as a facet
if query_lineage is None:
diff --git a/providers/openlineage/tests/unit/openlineage/test_sqlparser.py
b/providers/openlineage/tests/unit/openlineage/test_sqlparser.py
index a10cf1446ef..7716d791763 100644
--- a/providers/openlineage/tests/unit/openlineage/test_sqlparser.py
+++ b/providers/openlineage/tests/unit/openlineage/test_sqlparser.py
@@ -24,6 +24,7 @@ from openlineage.client.event_v2 import Dataset
from openlineage.client.facet_v2 import column_lineage_dataset, schema_dataset
from openlineage.common.sql import DbTableMeta
+from airflow.providers.common.sql.hooks.sql import DbApiHook
from airflow.providers.openlineage.sqlparser import (
DatabaseInfo,
GetTableSchemasParams,
@@ -414,17 +415,59 @@ class TestSQLParser:
class TestGetOpenlineageFacetsWithSql:
- def test_returns_none_when_no_database_info(self):
- hook = MagicMock()
- hook.get_openlineage_database_info.side_effect = AttributeError
+ def test_returns_none_when_hook_has_no_database_info_method(self):
+ hook = MagicMock(spec=["get_connection"])
+
+ result = get_openlineage_facets_with_sql(hook=hook, sql="SELECT 1",
conn_id="conn", database=None)
+ assert result is None
+
+ @mock.patch("airflow.providers.openlineage.sqlparser.log", autospec=True)
+ def test_error_from_database_info_is_logged_as_warning(self, mock_log):
+ hook = MagicMock(spec=DbApiHook)
+ hook.get_openlineage_database_info.side_effect =
AttributeError("broken import")
+
+ result = get_openlineage_facets_with_sql(hook=hook, sql="SELECT 1",
conn_id="conn", database=None)
+
+ assert result is None
+ mock_log.warning.assert_called_once()
+ assert mock_log.warning.call_args.args[1] == "conn"
+ assert str(mock_log.warning.call_args.args[2]) == "broken import"
+
+ @mock.patch("airflow.providers.openlineage.sqlparser.log", autospec=True)
+ def test_error_from_dialect_is_logged_as_warning(self, mock_log):
+ hook = MagicMock(spec=DbApiHook)
+ hook.get_openlineage_database_info.return_value =
DatabaseInfo(scheme="myscheme")
+ hook.get_openlineage_database_dialect.side_effect =
AttributeError("broken import")
result = get_openlineage_facets_with_sql(hook=hook, sql="SELECT 1",
conn_id="conn", database=None)
+
assert result is None
+ mock_log.warning.assert_called_once()
+ assert mock_log.warning.call_args.args[1] == "conn"
+ assert str(mock_log.warning.call_args.args[2]) == "broken import"
+
+ def test_prefetched_sql_parser_skips_dialect_lookup(self):
+ hook = MagicMock(spec=DbApiHook)
+ sql_parser = MagicMock(spec=SQLParser)
+
+ result = get_openlineage_facets_with_sql(
+ hook=hook,
+ sql="SELECT 1",
+ conn_id="conn",
+ database=None,
+ use_connection=False,
+ connection=MagicMock(),
+ database_info=DatabaseInfo(scheme="myscheme"),
+ sql_parser=sql_parser,
+ )
+
+ hook.get_openlineage_database_dialect.assert_not_called()
+ sql_parser.generate_openlineage_metadata_from_sql.assert_called_once()
+ assert result is
sql_parser.generate_openlineage_metadata_from_sql.return_value
def test_returns_none_when_no_dialect(self):
- hook = MagicMock()
+ hook = MagicMock(spec=["get_connection",
"get_openlineage_database_info"])
hook.get_openlineage_database_info.return_value =
DatabaseInfo(scheme="myscheme")
- hook.get_openlineage_database_dialect.side_effect = AttributeError
result = get_openlineage_facets_with_sql(hook=hook, sql="SELECT 1",
conn_id="conn", database=None)
assert result is None
diff --git
a/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
b/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
index 31f8a7da56f..b436b8b3f3b 100644
---
a/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
+++
b/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
@@ -25,6 +25,7 @@ from openlineage.client.event_v2 import Dataset as
OpenLineageDataset
from openlineage.client.facet_v2 import external_query_run, sql_job
from airflow.providers.common.sql.hooks.lineage import SqlJobHookLineageExtra
+from airflow.providers.common.sql.hooks.sql import DbApiHook
from airflow.providers.openlineage.extractors.base import OperatorLineage
from airflow.providers.openlineage.sqlparser import SQLParser
from airflow.providers.openlineage.utils.emission_policy import DatasetFilter,
Rule
@@ -76,12 +77,14 @@ class TestEmitLineageFromSqlExtras:
mock.patch(f"{_MODULE}._get_hook_conn_id", return_value="my_conn")
as mock_conn_id,
mock.patch(f"{_MODULE}.SQLParser.create_namespace") as mock_ns,
mock.patch(f"{_MODULE}.get_openlineage_facets_with_sql") as
mock_facets_fn,
+ mock.patch(f"{_MODULE}.get_sql_parser_for_hook") as
mock_get_parser,
mock.patch(f"{_MODULE}._create_ol_event_pair") as mock_build,
mock.patch(f"{_MODULE}.get_openlineage_listener") as mock_listener,
):
self.mock_conn_id = mock_conn_id
self.mock_ns = mock_ns
self.mock_facets_fn = mock_facets_fn
+ self.mock_get_parser = mock_get_parser
self.mock_build = mock_build
self.mock_listener = mock_listener
mock_build.return_value = (mock.sentinel.start_event,
mock.sentinel.end_event)
@@ -221,6 +224,87 @@ class TestEmitLineageFromSqlExtras:
assert call.kwargs["job_facets"]["sql"].query ==
SQLParser.normalize_sql("SELECT broken(")
assert call.kwargs["run_facets"]["externalQuery"].externalQueryId ==
"qid-1"
+ @mock.patch(f"{_MODULE}.log", autospec=True)
+ def test_sql_parsing_failure_is_logged_as_warning(self, mock_log):
+ self.mock_ns.return_value = "ns"
+ self.mock_facets_fn.side_effect = Exception("parse error")
+ mock_ti = mock.MagicMock(dag_id="dag_id", task_id="task_id")
+
+ emit_lineage_from_sql_extras(
+ task_instance=mock_ti,
+ sql_extras=[_make_extra(sql="SELECT broken(", job_id="qid-1")],
+ )
+
+ mock_log.warning.assert_called_once()
+ assert mock_log.warning.call_args.args[1:] == (1, "my_conn",
self.mock_facets_fn.side_effect)
+
+ @mock.patch(f"{_MODULE}.log", autospec=True)
+ def
test_database_info_failure_is_logged_as_warning_and_event_still_emitted(self,
mock_log):
+ hook = mock.MagicMock(spec=DbApiHook)
+ hook.get_openlineage_database_info.side_effect =
AttributeError("broken import")
+ self.mock_ns.return_value = "ns"
+ mock_ti = mock.MagicMock(dag_id="dag_id", task_id="task_id")
+
+ emit_lineage_from_sql_extras(
+ task_instance=mock_ti,
+ sql_extras=[_make_extra(sql="SELECT 1", job_id="qid-1",
hook=hook)],
+ )
+
+ mock_log.warning.assert_called_once()
+ assert mock_log.warning.call_args.args[1] == "my_conn"
+ assert str(mock_log.warning.call_args.args[2]) == "broken import"
+ # Without database info SQL is not parsed, but the query event is
still built and emitted
+ # with just the SQL facet (and no externalQuery facet, as namespace is
unknown).
+ self.mock_facets_fn.assert_not_called()
+ self.mock_build.assert_called_once()
+ call = self.mock_build.call_args
+ assert call.kwargs["inputs"] == []
+ assert call.kwargs["outputs"] == []
+ assert call.kwargs["job_facets"]["sql"].query ==
SQLParser.normalize_sql("SELECT 1")
+ assert "externalQuery" not in call.kwargs["run_facets"]
+
+ def test_sql_parser_is_built_once_per_hook(self):
+ hook = mock.MagicMock(spec=DbApiHook)
+ self.mock_ns.return_value = "ns"
+ mock_ti = mock.MagicMock(dag_id="dag_id", task_id="task_id")
+
+ emit_lineage_from_sql_extras(
+ task_instance=mock_ti,
+ sql_extras=[_make_extra(sql=f"SELECT {i}", hook=hook) for i in
range(3)],
+ )
+
+ self.mock_get_parser.assert_called_once()
+ assert self.mock_facets_fn.call_count == 3
+ for call in self.mock_facets_fn.call_args_list:
+ assert call.kwargs["sql_parser"] is
self.mock_get_parser.return_value
+
+ def test_sql_parsing_is_skipped_when_sql_parser_is_unavailable(self):
+ self.mock_get_parser.return_value = None
+ self.mock_ns.return_value = "ns"
+ mock_ti = mock.MagicMock(dag_id="dag_id", task_id="task_id")
+
+ emit_lineage_from_sql_extras(
+ task_instance=mock_ti,
+ sql_extras=[_make_extra(sql="SELECT 1"), _make_extra(sql="SELECT
2")],
+ )
+
+ self.mock_facets_fn.assert_not_called()
+ assert self.mock_build.call_count == 2
+
+ @mock.patch(f"{_MODULE}.log", autospec=True)
+ def test_hook_without_database_info_method_is_not_a_failure(self,
mock_log):
+ hook = mock.MagicMock(spec=["get_connection"])
+ self.mock_ns.return_value = None
+ mock_ti = mock.MagicMock(dag_id="dag_id", task_id="task_id")
+
+ emit_lineage_from_sql_extras(
+ task_instance=mock_ti,
+ sql_extras=[_make_extra(sql="SELECT 1", job_id="qid-1",
hook=hook)],
+ )
+
+ mock_log.warning.assert_not_called()
+ self.mock_build.assert_called_once()
+
def test_no_namespace_skips_external_query_facet(self):
"""When namespace cannot be resolved, no externalQuery facet is
added."""
self.mock_ns.return_value = None
