# -*- coding: utf-8 -*-
"""
添加剩余的 DWD 列
"""
import psycopg2

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

# 修正后的表名
MISSING_COLUMNS = {
    'billiards_dwd.dwd_assistant_trash_event': [
        ('tenant_id', 'BIGINT', '租户ID'),
    ],
    'billiards_dwd.dwd_store_goods_sale': [
        ('coupon_share_money', 'NUMERIC(18,2)', '优惠券分摊金额'),
    ],
}

def add_column(conn, full_table, col_name, col_type, comment):
    sql = f'ALTER TABLE {full_table} ADD COLUMN IF NOT EXISTS "{col_name}" {col_type}'
    comment_sql = f"COMMENT ON COLUMN {full_table}.\"{col_name}\" IS '{comment}'"
    with conn.cursor() as cur:
        cur.execute(sql)
        cur.execute(comment_sql)
    conn.commit()
    print(f"  [OK] {full_table}.{col_name} ({col_type})")

def main():
    conn = psycopg2.connect(DSN)
    for full_table, columns in MISSING_COLUMNS.items():
        for col_name, col_type, comment in columns:
            add_column(conn, full_table, col_name, col_type, comment)
    conn.close()
    print("Done!")

if __name__ == '__main__':
    main()