From 6f80ca6ae6c1945c52885910bcd3266469cad609 Mon Sep 17 00:00:00 2001 From: Kacper Muda Date: Mon, 3 Feb 2025 11:09:20 +0100 Subject: [PATCH] fix: OL sql parsing add try-except for sqlalchemy engine Signed-off-by: Kacper Muda --- .../providers/common/compat/openlineage/utils/sql.py | 8 +++++++- .../src/airflow/providers/openlineage/sqlparser.py | 8 +++++++- 2 files changed, 14 insertions(+), 2 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 9a9618d7a2eb6..72e9e1f94f0ba 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 @@ -69,12 +69,18 @@ def get_openlineage_facets_with_sql( log.debug("%s failed to get database dialect", hook) return None + try: + sqlalchemy_engine = hook.get_sqlalchemy_engine() + except Exception as e: + log.debug("Failed to get sql alchemy engine: %s", e) + sqlalchemy_engine = None + operator_lineage = sql_parser.generate_openlineage_metadata_from_sql( sql=sql, hook=hook, database_info=database_info, database=database, - sqlalchemy_engine=hook.get_sqlalchemy_engine(), + sqlalchemy_engine=sqlalchemy_engine, use_connection=use_external_connection, ) diff --git a/providers/openlineage/src/airflow/providers/openlineage/sqlparser.py b/providers/openlineage/src/airflow/providers/openlineage/sqlparser.py index e45a05b7c6458..7edb01cabd22a 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/sqlparser.py +++ b/providers/openlineage/src/airflow/providers/openlineage/sqlparser.py @@ -482,12 +482,18 @@ def get_openlineage_facets_with_sql( log.debug("%s failed to get database dialect", hook) return None + try: + sqlalchemy_engine = hook.get_sqlalchemy_engine() + except Exception as e: + log.debug("Failed to get sql alchemy engine: %s", e) + sqlalchemy_engine = None + operator_lineage = sql_parser.generate_openlineage_metadata_from_sql( sql=sql, hook=hook, database_info=database_info, database=database, - sqlalchemy_engine=hook.get_sqlalchemy_engine(), + sqlalchemy_engine=sqlalchemy_engine, use_connection=should_use_external_connection(hook), )