mobuchowski commented on code in PR #74805:
URL: https://github.com/apache/airflow/pull/74805#discussion_r4380439896
##########
providers/common/sql/tests/unit/common/sql/operators/test_sql_execute.py:
##########
@@ -368,6 +368,70 @@ def get_db_hook(self):
)
+def test_openlineage_hook_without_database_info_method_returns_empty_lineage():
Review Comment:
This test only checks `on_start() == OperatorLineage()`. It would still pass
if the `getattr` check were removed, because the new `except Exception` would
catch the missing-method `AttributeError` and return the same empty lineage.
The only difference would be a spurious WARNING. Asserting
`mock_log.warning.assert_not_called()` would make it actually guard the
"missing method is not a failure" behaviour, like the matching tests in
`test_sql_hook_lineage.py` and the compat tests.
Small one in the same file: the two other new tests patch the logger with
`with mock.patch.object(SQLExecuteQueryOperatorForTest, "log", autospec=True)`
(L404, L428). The repo guidelines prefer the decorator form, e.g.
`@mock.patch.object(SQLExecuteQueryOperator, "log", autospec=True)`, which the
other new tests in this PR already use. Optional.
##########
providers/openlineage/src/airflow/providers/openlineage/sqlparser.py:
##########
@@ -489,23 +489,29 @@ def get_openlineage_facets_with_sql(
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", hook, e)
+ log.debug("OpenLineage failure details:", exc_info=True)
if database_info is None:
log.debug("%s has no database info provided", hook)
return None
- try:
- sql_parser = SQLParser(
- dialect=hook.get_openlineage_database_dialect(connection),
- default_schema=hook.get_openlineage_default_schema(),
- )
- except AttributeError:
+ 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 failed to get database dialect", hook)
return None
+ try:
+ sql_parser = SQLParser(dialect=get_dialect(connection),
default_schema=get_default_schema())
+ except Exception as e:
+ log.warning("Failed to get OpenLineage database dialect for %s: %s",
hook, e)
Review Comment:
Two small things about the new warnings, which are now the main way users
will diagnose these failures:
1. The `try` covers `get_dialect(connection)`, `get_default_schema()` and
the `SQLParser(...)` constructor, but the message always says "database
dialect". For example, `SnowflakeHook.get_openlineage_default_schema` calls
`_get_conn_params()`. If that fails, the user is pointed at the dialect method,
which isn't the one that broke. Something like "Failed to get OpenLineage
database dialect or default schema" would be accurate.
2. The warnings print the hook object, which renders as
`<airflow.providers.snowflake.hooks.snowflake.SnowflakeHook object at 0x7f…>`.
With several connections you can't tell which one failed.
`sql_hook_lineage._get_database_info` already logs the `conn_id`, and using it
here would make the messages consistent.
The same applies to the copies in `operators/sql.py` (L306, L321, L354) and
in `common/compat/.../openlineage/utils/sql.py`.
##########
providers/common/compat/tests/unit/common/compat/openlineage/utils/test_sql.py:
##########
@@ -16,8 +16,67 @@
# under the License.
from __future__ import annotations
+import importlib
+import sys
+import types
+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():
+ """Load the module as if the OpenLineage provider was too old to ship
``get_openlineage_facets_with_sql``."""
+ old_sqlparser = types.ModuleType("airflow.providers.openlineage.sqlparser")
+ old_sqlparser.SQLParser = mock.MagicMock(name="SQLParser")
+ with mock.patch.dict(sys.modules,
{"airflow.providers.openlineage.sqlparser": old_sqlparser}):
Review Comment:
When `patch.dict(sys.modules, …)` exits, it restores `sys.modules` to its
exact previous contents, so it also **removes** any module first imported
inside the block. The fallback function imports
`airflow.providers.openlineage.utils.utils` (for
`should_use_external_connection`), and that module pulls in
`openlineage.plugins.facets`, `selective_enable` and others. If none of them
were imported before this test, they're imported during it and then dropped
from `sys.modules`. Later tests then get a second copy, so class identity and
`mock.patch` string targets no longer match what earlier imports hold. That
causes failures that depend on test order. I couldn't confirm whether these
modules are already loaded when this test runs. Can you check, e.g. by running
this file alone followed by another openlineage-dependent compat test?
A smaller-blast-radius way to force the fallback is to hide only the
attribute, so `from … import get_openlineage_facets_with_sql` raises
`ImportError`:
```python
@pytest.fixture
def fallback_get_openlineage_facets_with_sql(monkeypatch):
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)
```
This also lets you use the real `SQLParser`, or patch it with `autospec`,
instead of a `MagicMock` with no spec.
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -337,11 +343,16 @@ def get_openlineage_facets_on_complete(self,
task_instance) -> OperatorLineage |
operator_lineage = self.get_openlineage_facets_on_start() or
OperatorLineage()
Review Comment:
One path can still throw away the lineage this PR now preserves. In
`get_openlineage_facets_on_start`, `hook.get_sqlalchemy_engine()` (L330) is
passed straight into `generate_openlineage_metadata_from_sql` with no guard. If
it raises, the exception leaves `on_start`, and `on_complete` (here) never
reaches `get_openlineage_database_specific_lineage` or `_attach_check_facets`.
`DefaultExtractor` then returns `None` for the whole task.
`get_sqlalchemy_engine` raises in at least two situations:
- `AirflowOptionalProviderFeatureException` when SQLAlchemy isn't installed
(it's an optional extra of common-sql now).
- An error from `create_engine` when no SQLAlchemy dialect is installed for
the connection's URI scheme.
`sqlparser.get_openlineage_facets_with_sql` already wraps the same call
(L517–521). It only builds the engine when `use_connection` is set and falls
back to `None` on failure. Suggest the same here:
```python
sqlalchemy_engine = None
if use_external_connection:
try:
sqlalchemy_engine = hook.get_sqlalchemy_engine()
except Exception as e:
self.log.debug("Failed to get sql alchemy engine: %s", e)
```
Alternatively, `on_start` could call `get_openlineage_facets_with_sql` (from
`common.compat`, which common-sql already depends on) instead of keeping a
third copy of this lookup logic. That would fix this gap and stop the copies
drifting apart.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
