Finding source table for columns

[slp-oozzyy]

Hi,

First of all thank you for the great library, it has been a pleasure using it!

I have a question. I am working with a lot of convoluted queries that can have multiple derived tables and multiple nested column name changes. I need to update columns based on its source table. Is there a way to find that information for any given column? Excluding columns that come from calculated columns and given that I have the schema structure for all tables and columns.

I tried to find it via going thru the scope but couldnt managed to do it.

For example:
Lets say I want to find the source table for the 'COL1' at the final select. The answer is from the column INNER_COL1 from the table S1.TABLE1.

SELECT
I.OUTER_COL1 AS COL1
FROM (
SELECT
D1.MID_COL1 AS OUTER_COL1
D1.MID_COL2 AS OUTER_COL2
X.COMBINED_COL AS COMBINED_COL
FROM (
SELECT
T1.INNER_COL1 AS MID_COL1,
T2.INNER_COL2 AS MID_COL2
FROM S1.TABLE1 AS T1
JOIN S2.TABLE2 AS T2
) D1
JOIN (
SELECT
AUX_COL1||AUX_COL2 AS COMBINED_COL,
T2.MID_COL1
FROM S3.AUX_TABLE
JOIN S1.TABLE2 AS T2
) X

) I
JOIN (
SELECT
T4.COL1 as AUX_COL
T4.COL2 as OUTER_COL1
FROM S1.TABLE4 AS T4
) J

Is this possible? Any suggestion is welcomed

Kind Regards
oozzyy

[barakalon]

Scope is indeed the way to do this:

import sqlglot
from sqlglot import expressions as exp
from sqlglot.optimizer.scope import traverse_scope


ast = sqlglot.parse_one("SELECT ...")
physical_columns = [
    c
    for scope in traverse_scope(ast)
    for c in scope.columns
    if isinstance(scope.sources.get(c.table), exp.Table)
]

Now, this only works if the columns are qualified. Qualified means the column is prefixed with the table name. To illustrate - sqlglot can't disambiguate columns in this query without knowing the schema:

unqualified = """
SELECT 
  a,
  b,
FROM physical_table
JOIN (
  SELECT *
  FROM physical_table2
) AS derived_table
"""

If you're dealing with columns that might not be qualified, but you know the columns of the physical tables, sqlglot can qualify the columns for you:

from sqlglot.optimizer.qualify_columns import qualify_columns

schema = {
    "physical_table": {
        "a": "type - doesn't actually matter for qualification...",
    },
    "physical_table2": {
        "b": "-",
    }
}
qualify_columns(unqualified, schema=schema)
# SELECT
#   physical_table.a AS a,
#   derived_table.b AS b
# FROM physical_table
# JOIN (
#   SELECT
#     physical_table2.b AS b
#   FROM physical_table2
# ) AS derived_table

# Now you can get all the physical columns

If you're dealing with columns that might not be qualified, and you don't know the columns of the physical tables, but you can connect to the database, you can fetch the schema on the fly. Happy to chat more about this.

[slp-oozzyy]

Thanks for the advice, I have implemented in today, it works like a charm.

Also, while using scope, is there way to get other expressions? for example you used scope.columns or scope.sources. Lets say I wanted to find all DPipe expressions in a similar way? for ex.

for scope in traverse_scope(ast):
    for dpipe in scope.dpipes:
        if all([ isinstance(scope.sources.get(c.table),exp.Table) for c in dpipe.columns ]):
            **some setting and updating for the dpipe expression**

for this i tried to use find_ancestor to go to the select and use find_all to get the dpipe expression, it kinda works, but wanted to know if there is a better way of doing this

[barakalon]

Would this work? #527

[yuchaofan13]

Would be a good idea to eventually collate a list of all these functionalities in readthedocs.

[barakalon]

Yep, agreed. Plus a module with a collection of common ones.