Mulang' Onando
Mulang' Onando

Reputation: 11

Recursively traversing the sqlglot AST and replacing all values with new values from a datastore

I am completely new to using sqlglot library and I find it difficult to get going with the documentation. Seems quite convoluted. Here is my problem:

given a SQL query, e.g.

SELECT T.game_platform_id FROM ( SELECT T2.game_platform_id, MAX(T2.num_sales) FROM region AS T1 INNER JOIN region_sales AS T2 ON T1.id = T2.region_id WHERE T1.region_name = 'Japan' ) t

I want to traverse the AST for this query. I know I can get the AST from the sqlglot functions like parse_one(), and scope etc. But accessing elements in the tree is not so well documented. So for the query above, I want to be a able to access the region_name = 'Japan' part of the query, and compare the value "Japan" to see if it exists in my schema, then replace the value with the closest match in the database. Like wise if there was a function applied to it e.g. Max, Min, AVg, Upper, Lower etc, I want to capture those and restore them in the final query.

My attempt is currently a random attempt and does not help solve the problem; the challenge is being able to access the elements for the expression in the sqlglot AST tree.

def visit_ast(ast, aliases: Dict[str, str], data_source_id: str):  

    # Base cases
    if not ast or not isinstance(ast, (Expression, Condition, Literal)):
        return

    if isinstance(ast,  (EQ, Like, AggFunc)):  # 1. Here : I am not sure is these are the classes to target
        col_expr = ast.args["this"] 
        tbl, col = aliases[col_expr.table], col_expr.output_name
        val = ast.args["expression"].output_name
        
        returned = closest_value(val, tbl, col, data_source_id) # A function to return the closes value for replacement
        
        # 2. Here: How do I alter the values in place within the AST
    
    for child in getattr(ast, 'args', {}).values():  # recurse over the expression's children, in place
        await visit_ast(child, aliases, datasource_id)

Upvotes: 1

Views: 676

Answers (0)

Related Questions