How to determine column datatypes in a CTE with SELECT based on an initial schema?

[ragunandhant]

@tobymao
I'm working on a project where I have a predefined database schema (similar to the example below) and I need to parse SQL queries to understand the resulting datatypes of columns, especially when CTEs (Common Table Expressions) and SELECT statements are involved.

My Goal: Given an initial schema and a SQL query, I want to programmatically determine the datatype of each column in the final result set of that query. This is relatively straightforward for direct table selections, but it becomes more complex with transformations and expressions within SELECT or CTEs.

schema = {
    "users": {
        "id": "INT",
        "name": "VARCHAR(100)",
        "email": "VARCHAR(255)",
        "created_at": "TIMESTAMP",
        "is_active": "BOOLEAN"
    }
}

sql = "WITH active_users AS (SELECT id, name FROM users WHERE is_active = true) SELECT name, id FROM active_users"

In the above sql example, the active_users CTE selects id (INT) and name (VARCHAR(100)). The final SELECT then reorders them. I'd like to be able to output something like:

{
"name": "VARCHAR(100)",
"id": "INT"
}

[georgesittas]

Hey, I think you should look into qualify and annotate_types:

sqlglot/tests/test_optimizer.py

Line 880 in 5a476a5

def test_annotate_types(self):

.

[ragunandhant]

hi @georgesittas but i have tried with

import sqlglot as sg
from sqlglot import exp
from sqlglot.optimizer.annotate_types import annotate_types
schema = {
    "users": {
        "id": "INTEGER",
        "name": "VARCHAR(100)",
        "email": "VARCHAR(255)",
        "created_at": "TIMESTAMP",
        "is_active": "BOOLEAN"
    }
}

sql = "SELECT id, name, email FROM users WHERE is_active = true"
parsed = sg.parse_one(sql)
annotated = annotate_types(parsed, schema=schema ,dialect="tsql")

print(annotated.expressions)

output :

[Column(
  this=Identifier(this=id, quoted=False, _type=DataType(this=Type.UNKNOWN)),
  _type=DataType(this=Type.UNKNOWN)), Column(
  this=Identifier(this=name, quoted=False, _type=DataType(this=Type.UNKNOWN)),
  _type=DataType(this=Type.UNKNOWN)), Column(
  this=Identifier(this=email, quoted=False, _type=DataType(this=Type.UNKNOWN)),
  _type=DataType(this=Type.UNKNOWN))]

could explain why this happening

[georgesittas]

You need to qualify the query, otherwise annotate_types won't work.

[ragunandhant]

thanks