# -*- coding: utf-8 -*-
"""获取 DWD 所有表的结构"""
import psycopg2
import json

DSN = 'postgresql://local-Python:Neo-local-1991125@100.64.0.4:5432/LLZQ-test'

def get_all_tables(conn, schema='billiards_dwd'):
    """获取所有表及其列"""
    cur = conn.cursor()
    cur.execute("""
        SELECT table_name, column_name, data_type, is_nullable,
               COALESCE(character_maximum_length, numeric_precision) as max_length
        FROM information_schema.columns
        WHERE table_schema = %s
        ORDER BY table_name, ordinal_position
    """, (schema,))
    
    tables = {}
    for row in cur.fetchall():
        table_name, col_name, data_type, nullable, max_len = row
        if table_name not in tables:
            tables[table_name] = []
        tables[table_name].append({
            'column': col_name,
            'type': data_type,
            'nullable': nullable,
            'max_length': max_len
        })
    
    cur.close()
    return tables

def main():
    conn = psycopg2.connect(DSN)
    tables = get_all_tables(conn)
    conn.close()
    
    # 保存到 JSON 文件
    with open('tmp/dwd_schema.json', 'w', encoding='utf-8') as f:
        json.dump(tables, f, ensure_ascii=False, indent=2)
    
    print(f"Found {len(tables)} tables")
    for table_name, columns in sorted(tables.items()):
        print(f"  {table_name}: {len(columns)} columns")

if __name__ == '__main__':
    main()