Fix #323: fixes in table_usage.py (#2576)

Co-authored-by: Mayur SIngal <mayursingal@Mayurs-MacBook-Pro.local>
Co-authored-by: Ayush Shah <ayush@getcollate.io>
This commit is contained in:
Mayur Singal 2022-02-07 00:04:56 +05:30 committed by GitHub
parent 6b92aa3554
commit e1d7d644b3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -32,12 +32,9 @@ def get_table_column_join(table, table_aliases, joins):
joined_with = [] joined_with = []
for join in joins: for join in joins:
try: try:
if "." in join: if "." not in join:
if join.count(".") < 3: continue
jtable, column = join.split(".") jtable, column = join.split(".")[-2:]
else:
jtable, column = join.split(".")[2:]
if table == jtable or jtable in table_aliases: if table == jtable or jtable in table_aliases:
table_column = TableColumn( table_column = TableColumn(
table=table_aliases[jtable] if jtable in table_aliases else jtable, table=table_aliases[jtable] if jtable in table_aliases else jtable,
@ -86,13 +83,14 @@ class TableUsageStage(Stage[QueryParserData]):
return cls(ctx, config, metadata_config) return cls(ctx, config, metadata_config)
def stage_record(self, record: QueryParserData) -> None: def stage_record(self, record: QueryParserData) -> None:
if record is not None: if record is None:
for table in record.tables: return None
table_usage_count = None for table in record.tables:
if table in self.table_usage.keys(): try:
table_usage_count = self.table_usage.get(table) table_usage_count = self.table_usage.get(table)
if table_usage_count is not None:
table_usage_count.count = table_usage_count.count + 1 table_usage_count.count = table_usage_count.count + 1
if "join" in record.columns: if record.columns.get("join") is not None:
table_usage_count.joins.append( table_usage_count.joins.append(
get_table_column_join( get_table_column_join(
table, record.tables_aliases, record.columns["join"] table, record.tables_aliases, record.columns["join"]
@ -100,7 +98,7 @@ class TableUsageStage(Stage[QueryParserData]):
) )
else: else:
joins = [] joins = []
if "join" in record.columns: if record.columns.get("join") is not None:
tbl_column_join = get_table_column_join( tbl_column_join = get_table_column_join(
table, record.tables_aliases, record.columns["join"] table, record.tables_aliases, record.columns["join"]
) )
@ -114,7 +112,9 @@ class TableUsageStage(Stage[QueryParserData]):
joins=joins, joins=joins,
service_name=record.service_name, service_name=record.service_name,
) )
self.table_usage[table] = table_usage_count except Exception as exc:
logger.error("Error in staging record {}".format(exc))
self.table_usage[table] = table_usage_count
def get_status(self): def get_status(self):
return self.status return self.status