@@ -104,15 +104,10 @@ def get_table_lineage(self, table_name):
104104 :param table_name: name of the table (without schema)
105105 :return: dict mapping attribute_name -> lineage (or None)
106106 """
107- result = (self & dict (table_name = table_name )).fetch (
108- "attribute_name" , "lineage"
109- )
107+ result = (self & dict (table_name = table_name )).fetch ("attribute_name" , "lineage" )
110108 if len (result [0 ]) == 0 :
111109 return {}
112- return {
113- attr : (lin if lin else None )
114- for attr , lin in zip (result [0 ], result [1 ])
115- }
110+ return {attr : (lin if lin else None ) for attr , lin in zip (result [0 ], result [1 ])}
116111
117112 def delete_table_lineage (self , table_name ):
118113 """
@@ -172,10 +167,14 @@ def compute_lineage_from_dependencies(connection, schema, table_name, attribute_
172167 # Handle alias nodes (numeric string nodes in the graph)
173168 if parent .isdigit ():
174169 # Find the actual parent by traversing through alias
175- for grandparent , gprops in connection .dependencies .parents (parent ).items ():
170+ for grandparent , gprops in connection .dependencies .parents (
171+ parent
172+ ).items ():
176173 if not grandparent .isdigit ():
177174 parent = grandparent
178- parent_attr = gprops .get ("attr_map" , {}).get (attribute_name , parent_attr )
175+ parent_attr = gprops .get ("attr_map" , {}).get (
176+ attribute_name , parent_attr
177+ )
179178 break
180179 parent_schema , parent_table = parse_full_table_name (parent )
181180 return compute_lineage_from_dependencies (
@@ -304,13 +303,16 @@ def get_lineage_for_heading(connection, schema, table_name, heading):
304303 :return: dict mapping attribute_name -> lineage (or None)
305304 """
306305 # Check if ~lineage table exists
307- lineage_table_exists = connection .query (
308- """
306+ lineage_table_exists = (
307+ connection .query (
308+ """
309309 SELECT COUNT(*) FROM information_schema.tables
310310 WHERE table_schema = %s AND table_name = '~lineage'
311311 """ ,
312- args = (schema ,),
313- ).fetchone ()[0 ] > 0
312+ args = (schema ,),
313+ ).fetchone ()[0 ]
314+ > 0
315+ )
314316
315317 if lineage_table_exists :
316318 # Load from ~lineage table
0 commit comments