How to rename tables and columns in a query (safely)

[ibestvina]

What would be the best way to rename the tables and columns in a query, ensuring that the query is still functionally the same?

E.g., given a query:

SELECT
  old_a.old_x AS aliased_x,
  old_a.old_y AS old_y,
  aliased_b.old_q
FROM
  old_a
CROSS JOIN (
  SELECT old_z AS old_q FROM old_b
) AS aliased_b

and renaming dictionaries (which I will ensure are bijections):

table_renames = {'old_a': 'new_a', 'old_b': 'new_b'}
column_renames = {'old_a': {'old_x': 'new_x', 'old_y': 'new_y'}, 'old_b': {'old_z': 'new_z', 'old_q': 'new_q'}}

Additionally I can provide the schema dict if needed:

schema={"old_a": {"old_x": "int", "old_y": "int"}, "old_b": {"old_z": "int", "old_q": "int"}}

I would like the resulting query to be:

SELECT
  new_a.new_x AS aliased_x,
  new_a.new_y AS old_y,
  aliased_b.old_q
FROM
  new_a
CROSS JOIN (
  SELECT new_z AS old_q FROM new_b
) AS aliased_b

I am mainly worried about how to make sure that old_q doesn't get renamed, as it's an alias, not the actual column name, even though it is present in the dictionary. In the above example, optimization removes the subquery, so the renaming is actually not hard afterwards. But I assume optimization will not always be able to do so, and some more complex examples might, even after optimization, still have edge cases like the one above.

For reference, this is what I have done so far, but I am very unhappy with it, firstly because I do not understand if it's actually correct in general, and second because it runs the whole optimization, and Ideally I would not have to run it, or at least I could only run the necessary steps (but I do not know which ones are necessary).

original_code = """
SELECT
  old_a.old_x AS aliased_x,
  old_a.old_y AS old_y,
  aliased_b.old_q
FROM
  old_a
CROSS JOIN (
  SELECT old_z AS old_q FROM old_b
) AS aliased_b
"""
table_renames = {"old_a": "new_a", "old_b": "new_b"}
column_renames = {
    "old_a": {"old_x": "new_x", "old_y": "new_y"},
    "old_b": {"old_z": "new_z", "old_q": "new_q"},
}
schema = {
    "old_a": {"old_x": "int", "old_y": "int"},
    "old_b": {"old_z": "int", "old_q": "int"},
}

optimized = optimize(sqlglot.parse_one(original_code), schema=schema, read="sqlite")
    

def table_renaming(node):
    if isinstance(node, exp.Table):
        node_id = node.find(exp.Identifier)
        if node_id is None:
            return node
        table_name = node_id.this
        if table_name is None:
            return node
        if table_name in table_renames:
            node_id.set("this", table_renames[table_name])
    return node

def column_renaming(node):
    if isinstance(node, exp.Column):
        column_id = node.find(exp.Identifier)
        if column_id is None:
            return node
        column_name = column_id.this
        if column_name is None:
            return node
        table_name = node.table
        if table_name is None:
            return node
        if column_name in column_renames.get(table_name, {}):
            column_id.set("this", column_renames[table_name][column_name])
    return node


renamed_columns = optimized.transform(column_renaming)
renamed_columns_and_tables = renamed_columns.transform(table_renaming)
print(renamed_columns_and_tables.sql(pretty=True))

[tobymao]

you can run the optimizer with just qualify columns and qualify tables

[ibestvina]

Thank you! I thought of doing something like that, but I'm still not sure if my approach above is sound. Is there anything else I should traverse while transforming? Any special handling of aliases?

[tobymao]

you can use traverse_scope to see all references i think