SQLGlot库全面解析

pip install sqlglot[all]  # 安装完整功能包（含所有方言支持）

import sqlglot
from sqlglot.dialects import mysql, postgres
# 设置全局默认方言
sqlglot.dialect = "mysql"
# 或针对特定会话设置
with sqlglot.dialect_context("postgres"):
    # 此代码块内使用postgresql方言
    pass

from sqlglot import parse_one
sql = """
with daily_metrics as (
    select 
        date_trunc('day', event_time) as day,
        product_id,
        count(distinct user_id) as dau
    from events
    where event_type = 'click'
    group by 1, 2
)
select 
    a.day,
    a.product_id,
    a.dau,
    b.sales,
    round(a.dau / b.sales, 2) as conversion_rate
from daily_metrics a
join (
    select 
        date_trunc('day', order_time) as day,
        product_id,
        sum(amount) as sales
    from orders
    group by 1, 2
) b on a.day = b.day and a.product_id = b.product_id
where a.day > current_date - interval '7' day
order by conversion_rate desc
"""
ast = parse_one(sql)
print(f"ast节点数: {len(ast.find_all())}")
print(f"cte数量: {len(ast.args['with'].expressions)}")

import sqlglot
mysql_sql = """
select 
    user_id,
    group_concat(distinct product_id order by purchase_date separator ',') as products
from purchases
where status = 'completed'
group by user_id
having count(distinct order_id) > 3
"""
bq_sql = sqlglot.transpile(
    mysql_sql,
    read="mysql",
    write="bigquery",
    pretty=true
)[0]
print(bq_sql)

select
    user_id,
    string_agg(distinct cast(product_id as string), ',' order by purchase_date) as products
from
    purchases
where
    status = 'completed'
group by
    user_id
having
    count(distinct order_id) > 3

from sqlglot import parse_one, optimize
sql = """
select 
    a.department,
    a.avg_salary,
    (select avg(salary) 
     from employees 
     where department = a.department 
     and hire_date > date_add(current_date, interval -5 year)) as junior_avg
from (
    select 
        department, 
        avg(salary) as avg_salary
    from employees
    group by department
) a
"""
optimized = optimize(
    parse_one(sql),
    schema={
        "employees": {
            "columns": ["id", "name", "department", "salary", "hire_date"],
            "indexes": ["department", "hire_date"]
        }
    }
)
print(optimized.sql(pretty=true))

with anon_1 as (
    select
        department,
        avg(salary) as avg_salary
    from
        employees
    group by
        department
)
select
    a.department,
    a.avg_salary,
    (
        select
            avg(salary)
        from
            employees
        where
            department = a.department
            and hire_date > date_add(current_date, interval -5 year)
    ) as junior_avg
from
    anon_1 a

from sqlglot import exp, select, func
def build_funnel_query(events, date_column="event_date"):
    query = select(
        f"date_trunc('day', {date_column}) as day"
    ).with_alias("base_query")
    for i, event in enumerate(events):
        filter_expr = exp.condition(f"event_type = '{event['type']}'")
        if "filters" in event:
            filter_expr = filter_expr.and_(exp.condition(event["filters"]))
        count_expr = (
            exp.count(distinct=true)
            .of(exp.column("user_id"))
            .where(filter_expr)
            .alias(f"step_{i+1}_count")
        )
        query = query.add(count_expr)
    return (
        query.from_("events")
        .group_by("day")
        .order_by("day")
        .sql(dialect="snowflake")
    )
funnel_steps = [
    {"type": "page_view", "name": "页面访问"},
    {"type": "add_to_cart", "name": "加入购物车"},
    {"type": "checkout_start", "name": "开始结账"},
    {"type": "purchase", "name": "完成购买", "filters": "status = 'success' and amount > 0"}
]
print(build_funnel_query(funnel_steps))

from sqlglot import parse, transform, exp
def id_mask_transform(expression):
    if isinstance(expression, exp.column) and expression.name == "id_card":
        return exp.func(
            "concat",
            exp.literal.string("****"),
            exp.func("substr", expression, 11, 4)
        ).alias("id_card")
    return expression
sql = "select name, id_card, phone from users where age > 18"
ast = parse(sql)
transformed_ast = transform(ast, step=id_mask_transform)
print(transformed_ast.sql(dialect="mysql"))

select name, concat('****', substr(id_card, 11, 4)) as id_card, phone from users where age > 18

import sqlglot
from timeit import timeit
def compare_dialects(sql, dialects=["mysql", "postgres", "spark"]):
    results = {}
    for dialect in dialects:
        try:
            parsed = sqlglot.parse_one(sql)
            generated = parsed.sql(dialect=dialect)
            # 模拟执行时间（实际应连接数据库执行）
            exec_time = timeit(lambda: parse_one(generated), number=100)
            results[dialect] = {
                "sql": generated,
                "parse_time": exec_time,
                "length": len(generated)
            }
        except exception as e:
            results[dialect] = {"error": str(e)}
    return results
query = """
select 
    user_id,
    sum(case when event_type = 'click' then 1 else 0 end) as clicks,
    sum(case when event_type = 'view' then 1 else 0 end) as views
from events
group by user_id
"""
print(compare_dialects(query))

from sqlglot import parse_one
from collections import defaultdict
def analyze_sql_pattern(sql):
    ast = parse_one(sql)
    pattern_stats = defaultdict(int)
    # 统计join类型
    for join in ast.find_all(exp.join):
        join_type = join.args.get("join_type", "inner").upper()
        pattern_stats[f"join_{join_type}"] += 1
    # 统计聚合函数
    for func in ast.find_all(exp.func):
        if func.name.upper() in ["sum", "avg", "count", "max", "min"]:
            pattern_stats[f"agg_{func.name.upper()}"] += 1
    return dict(pattern_stats)
complex_query = """
select 
    u.id,
    u.name,
    count(distinct o.order_id) as order_count,
    sum(o.amount) as total_amount,
    avg(o.amount) as avg_order_value
from users u
left join orders o on u.id = o.user_id
where u.status = 'active'
group by u.id, u.name
having count(distinct o.order_id) > 5
"""
print(analyze_sql_pattern(complex_query))

import pandas as pd
from sqlglot import parse_one
def sql_to_dataframe(sql, data):
    ast = parse_one(sql)
    # 实际实现需要解析ast并转换为pandas操作
    # 此处仅为概念演示
    if "select * from" in sql.upper():
        return pd.dataframe(data)
    elif "where" in sql.upper():
        condition = sql.split("where")[1].split("group by")[0].strip()
        # 简化处理，实际需解析条件表达式
        filtered_data = {k: v for k, v in data.items() if eval(condition, {}, v)}
        return pd.dataframe(filtered_data)
    return pd.dataframe()
sample_data = {
    "id": [1, 2, 3],
    "name": ["alice", "bob", "charlie"],
    "age": [25, 30, 35]
}
df = sql_to_dataframe("select name, age from sample where age > 28", sample_data)
print(df)

from functools import lru_cache
from sqlglot import parse_one
@lru_cache(maxsize=1000)
def cached_parse(sql):
    return parse_one(sql)
# 重复解析相同sql时将直接从缓存获取

from sqlglot import exp
# 预定义常用表达式
common_expressions = {
    "recent_7_days": exp.condition(
        "event_time >= date_sub(current_date, interval 7 day)"
    ),
    "active_users": exp.condition(
        "last_active_date >= date_sub(current_date, interval 30 day)"
    )
}
def build_query_with_patterns(base_sql, patterns):
    ast = parse_one(base_sql)
    for alias, expr in patterns.items():
        ast = ast.with_cte(exp.cte(alias, exp.select().from_("dummy").where(expr)))
    return ast.sql()

from concurrent.futures import threadpoolexecutor
import sqlglot
def parallel_parse(sql_list, max_workers=4):
    with threadpoolexecutor(max_workers=max_workers) as executor:
        results = list(executor.map(sqlglot.parse_one, sql_list))
    return results
large_sql_batch = ["select * from table{}".format(i) for i in range(100)]
parsed_asts = parallel_parse(large_sql_batch)