MINOR: Lineage handle copy queries being skipped (#14855)

This commit is contained in:
Mayur Singal 2024-01-25 10:15:32 +05:30 committed by GitHub
parent 70316b9680
commit 17fb2cabca
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 52 additions and 11 deletions

View File

@ -98,24 +98,30 @@ class LineageParser:
""" """
Get a list of intermediate tables Get a list of intermediate tables
""" """
# These are @lazy_property, not properly being picked up by IDEs. Ignore the warning if self.parser:
return self.retrieve_tables(self.parser.intermediate_tables) # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
return self.retrieve_tables(self.parser.intermediate_tables)
return []
@cached_property @cached_property
def source_tables(self) -> List[Table]: def source_tables(self) -> List[Table]:
""" """
Get a list of source tables Get a list of source tables
""" """
# These are @lazy_property, not properly being picked up by IDEs. Ignore the warning if self.parser:
return self.retrieve_tables(self.parser.source_tables) # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
return self.retrieve_tables(self.parser.source_tables)
return []
@cached_property @cached_property
def target_tables(self) -> List[Table]: def target_tables(self) -> List[Table]:
""" """
Get a list of target tables Get a list of target tables
""" """
# These are @lazy_property, not properly being picked up by IDEs. Ignore the warning if self.parser:
return self.retrieve_tables(self.parser.target_tables) # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
return self.retrieve_tables(self.parser.target_tables)
return []
# pylint: disable=protected-access # pylint: disable=protected-access
@cached_property @cached_property
@ -124,6 +130,8 @@ class LineageParser:
Get a list of tuples of column lineage Get a list of tuples of column lineage
""" """
column_lineage = [] column_lineage = []
if self.parser is None:
return []
try: try:
if self.parser._dialect == SQLPARSE_DIALECT: if self.parser._dialect == SQLPARSE_DIALECT:
return self.parser.get_column_lineage() return self.parser.get_column_lineage()
@ -331,6 +339,8 @@ class LineageParser:
:return: for each table name, list all joins against other tables :return: for each table name, list all joins against other tables
""" """
join_data = defaultdict(list) join_data = defaultdict(list)
if self.parser is None:
return join_data
# These are @lazy_property, not properly being picked up by IDEs. Ignore the warning # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
for statement in self.parser.statements(): for statement in self.parser.statements():
self.stateful_add_joins_from_statement(join_data, sql_statement=statement) self.stateful_add_joins_from_statement(join_data, sql_statement=statement)
@ -378,7 +388,11 @@ class LineageParser:
@staticmethod @staticmethod
def _evaluate_best_parser( def _evaluate_best_parser(
query: str, dialect: Dialect, timeout_seconds: int query: str, dialect: Dialect, timeout_seconds: int
) -> LineageRunner: ) -> Optional[LineageRunner]:
if query is None:
return None
@timeout(seconds=timeout_seconds) @timeout(seconds=timeout_seconds)
def get_sqlfluff_lineage_runner(qry: str, dlct: str) -> LineageRunner: def get_sqlfluff_lineage_runner(qry: str, dlct: str) -> LineageRunner:
lr_dialect = LineageRunner(qry, dialect=dlct) lr_dialect = LineageRunner(qry, dialect=dlct)

View File

@ -206,7 +206,7 @@ class QueryParserTests(TestCase):
""" """
Validate we obtain information from Comon Table Expressions Validate we obtain information from Comon Table Expressions
""" """
query = """CREATE TABLE TESTDB.PUBLIC.TARGET AS query = """CREATE TABLE TESTDB.PUBLIC.TARGET AS
WITH cte_table AS ( WITH cte_table AS (
SELECT SELECT
USERS.ID, USERS.ID,
@ -226,7 +226,6 @@ class QueryParserTests(TestCase):
; ;
""" """
expected_tables = {"testdb.public.users"}
expected_lineage = [ expected_lineage = [
( (
Column("testdb.public.users.id"), Column("testdb.public.users.id"),
@ -258,7 +257,7 @@ class QueryParserTests(TestCase):
""" """
Validate we obtain information from Comon Table Expressions Validate we obtain information from Comon Table Expressions
""" """
query = """CREATE TABLE TESTDB.PUBLIC.TARGET AS query = """CREATE TABLE TESTDB.PUBLIC.TARGET AS
SELECT SELECT
ID, ID,
-- A comment here -- A comment here
@ -295,7 +294,7 @@ class QueryParserTests(TestCase):
""" """
Validate we obtain information from Comon Table Expressions Validate we obtain information from Comon Table Expressions
""" """
query = """CREATE TABLE TESTDB.PUBLIC.TARGET AS query = """CREATE TABLE TESTDB.PUBLIC.TARGET AS
SELECT SELECT
ID AS new_identifier, ID AS new_identifier,
NAME new_name NAME new_name
@ -329,3 +328,31 @@ class QueryParserTests(TestCase):
parser.column_lineage, parser.column_lineage,
expected_lineage, expected_lineage,
) )
def test_copy_query(self):
"""
Validate Copy query is skipped appropriately without any errors
"""
query = """COPY MY_TABLE col1,col2,col3
FROM 's3://bucket/schema/table.csv'
WITH CREDENTIALS ''
REGION 'US-east-2'
"""
expected_lineage = []
expected_tables = set()
parser = LineageParser(query)
tables = {str(table) for table in parser.involved_tables}
self.assertEqual(tables, expected_tables)
self.assertEqual(
parser.column_lineage,
expected_lineage,
)
parser = LineageParser(query, Dialect.MYSQL)
tables = {str(table) for table in parser.involved_tables}
self.assertEqual(tables, expected_tables)
self.assertEqual(
parser.column_lineage,
expected_lineage,
)