From d0cdadbb3f4b89e771e2a05593ef528e622f94e5 Mon Sep 17 00:00:00 2001 From: Aseem Bansal Date: Tue, 10 May 2022 15:48:54 +0530 Subject: [PATCH] fix(sql-parsing): improve error handling (#4862) --- .../src/datahub/utilities/sql_lineage_parser_impl.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/metadata-ingestion/src/datahub/utilities/sql_lineage_parser_impl.py b/metadata-ingestion/src/datahub/utilities/sql_lineage_parser_impl.py index 412c8fcf46..80ea7cc314 100644 --- a/metadata-ingestion/src/datahub/utilities/sql_lineage_parser_impl.py +++ b/metadata-ingestion/src/datahub/utilities/sql_lineage_parser_impl.py @@ -2,7 +2,7 @@ import logging import re import unittest import unittest.mock -from typing import Dict, List, Set +from typing import Dict, List, Optional, Set from sqllineage.core.holders import Column, SQLLineageHolder from sqllineage.exceptions import SQLLineageException @@ -64,7 +64,8 @@ class SqlLineageSQLParserImpl: logger.debug(f"Rewrote original query {original_sql_query} as {sql_query}") self._sql = sql_query - + self._stmt_holders: Optional[List[LineageAnalyzer]] = None + self._sql_holder: Optional[SQLLineageHolder] = None try: self._stmt = [ s @@ -97,6 +98,9 @@ class SqlLineageSQLParserImpl: def get_tables(self) -> List[str]: result: List[str] = list() + if self._sql_holder is None: + logger.error("sql holder not present so cannot get tables") + return result for table in self._sql_holder.source_tables: table_normalized = re.sub(r"^.", "", str(table)) result.append(str(table_normalized)) @@ -115,6 +119,9 @@ class SqlLineageSQLParserImpl: return result def get_columns(self) -> List[str]: + if self._sql_holder is None: + logger.error("sql holder not present so cannot get columns") + return [] graph: DiGraph = self._sql_holder.graph # For mypy attribute checking column_nodes = [n for n in graph.nodes if isinstance(n, Column)] column_graph = graph.subgraph(column_nodes)