Parsing presto queries to get columns from WHERE and JOIN clauses

[vjdbricks]

I am trying to parse set of presto queries and extract the columns involved in WHERE and JOIN clause along with the corresponding table names. I was able to do it for few queries using sqlglot library but, when i try to run on the volume of 200K+ queries running into errors mostly with None type.

My presto queries are from a Production system and have all the possible complexity.

Below is the code I have written for this objective. It would be really helpful if someone could take a look at it provide any feedback.

def extract_where_columns(expression):
    """Extract columns from WHERE clause."""
    where_columns_details = {}
    for whr in expression.find_all(exp.Where):
        left_op = whr.this.left.this.name
        right_op = whr.this.right.this.name
        table_name = f"{whr.parent_select.args['from'].this.db}.{whr.parent_select.args['from'].this.name}"
        where_columns_details.setdefault(table_name, set()).update([left_op, right_op])
    return where_columns_details


def extract_join_columns(node):
    """Extract columns involved in JOIN operations."""
    join_table_columns = {}
    for join in node.args["joins"]:
        main_table = node.args['from'].name
        join_table_columns.setdefault(main_table, set())

        join_table = join.find(exp.Table).name
        join_table_columns.setdefault(join_table, set())

        for eql in join.args["on"].find_all(exp.EQ):
            join_table_columns[eql.left.table].add(eql.left.name)
            join_table_columns[eql.right.table].add(eql.right.name)
    return join_table_columns


def extract_scopes(ast):
    """Extract table scopes."""
    scope_tables_mapping = {}
    uniq_scopes = set(traverse_scope(ast))
    for scope in uniq_scopes:
        for k, v in scope.sources.items():
            if isinstance(v, sqlglot.optimizer.scope.Scope):
                expression = sqlglot.parse_one(v.expression.sql(), read="presto")
                for ttt in set(expression.find_all(exp.Table)):
                    scope_tables_mapping.setdefault(k, set()).add(f"{ttt.db}.{ttt.name}")
    return scope_tables_mapping


    ast = sqlglot.parse_one(query)

    where_columns_details = extract_where_columns(ast)
    join_table_columns = extract_join_columns(ast)
    scope_tables_mapping = extract_scopes(ast)

    actual_columns = {}
    for alias, columns in join_table_columns.items():
        actual_names = scope_tables_mapping.get(alias)
        if actual_names:
            for actual_name in actual_names:
                actual_columns[actual_name] = columns

 

    table_info = {}
    for key in set(where_columns_details.keys()).union(join_table_columns.keys()):
        table_info[key] = {
            'Where': list(where_columns_details.get(key, [])),
            'Join': list(join_table_columns.get(key, []))
        }

[tobymao]

you need to make your code safer and handle nulls

[vjdbricks]

Thanks for reviewing it @tobymao. Could you please suggest if any specific change that i need to make?

[tobymao]

no, this is up to you to debug, and i don't have enough information to help.