packages/core/olap.py

import clickhouse_connect
from packages.core.config import settings

def get_clickhouse_client():
    """
    获取 ClickHouse 客户端实例。
    """
    # 解析 CLICKHOUSE_URL
    # 示例: http://default:password@localhost:8123
    import urllib.parse
    parsed = urllib.parse.urlparse(settings.CLICKHOUSE_URL)
    
    host = parsed.hostname or "localhost"
    port = parsed.port or 8123
    username = parsed.username or "default"
    password = parsed.password or ""
    
    client = clickhouse_connect.get_client(
        host=host,
        port=port,
        username=username,
        password=password,
        database="customs_data"
    )
    return client

def init_clickhouse_schema():
    """
    初始化 ClickHouse 数据库与表结构
    设计“国家+月份”分区表
    """
    client = get_clickhouse_client()
    
    # 确保数据库存在 (通过 client 创建时指定的 database="customs_data" 可能需要事先建库)
    client.command("CREATE DATABASE IF NOT EXISTS customs_data")
    
    # 创建贸易明细大宽表
    # 使用 ReplacingMergeTree 支持按 record_id 去重更新
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS customs_data.trade_records
    (
        record_id String,
        source_country String,
        trade_direction String,
        trade_date Date,
        trade_month UInt32 MATERIALIZED toYYYYMM(trade_date),
        importer_name Nullable(String),
        exporter_name Nullable(String),
        hs_code String,
        product_name String,
        amount Float64,
        currency String,
        weight Float64,
        weight_unit String,
        origin_country String,
        destination_country String,
        departure_port Nullable(String),
        arrival_port Nullable(String),
        transport_mode String,
        created_at DateTime DEFAULT now(),
        updated_at DateTime DEFAULT now()
    )
    ENGINE = ReplacingMergeTree(updated_at)
    PARTITION BY (source_country, trade_month)
    ORDER BY (trade_direction, hs_code, trade_date, record_id)
    """
    
    client.command(create_table_sql)
    print("ClickHouse schema initialized successfully.")