wenshu/python-iceberg/app.py

import io
import os
import logging
from contextlib import asynccontextmanager

import openpyxl
import pyarrow as pa
import pyarrow.csv as pa_csv
import requests
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pyiceberg.catalog.rest import RestCatalog
from pyiceberg.exceptions import NoSuchTableError

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
logger = logging.getLogger(__name__)

# ---------- Patch pyiceberg fsspec S3 factory for Aliyun OSS compatibility ----------
# Aliyun OSS rejects aws-chunked encoding with UNSIGNED-PAYLOAD.
# botocore must use payload_signing_enabled=True so that chunked uploads
# send STREAMING-AWS4-HMAC-SHA256-PAYLOAD instead.
import pyiceberg.io.fsspec as _fsspec_mod

_orig_s3_factory = _fsspec_mod._s3


def _oss_s3_factory(properties):
    """Wrap the original _s3 factory to inject payload_signing_enabled into botocore Config."""
    from s3fs import S3FileSystem
    from pyiceberg.io import (
        S3_ENDPOINT, S3_ACCESS_KEY_ID, S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN,
        S3_REGION, S3_FORCE_VIRTUAL_ADDRESSING, S3_PROXY_URI, S3_CONNECT_TIMEOUT,
        S3_REQUEST_TIMEOUT, S3_ANONYMOUS, S3_PROFILE_NAME,
        AWS_ACCESS_KEY_ID, AWS_REGION, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN, AWS_PROFILE_NAME,
    )
    from pyiceberg.utils.properties import get_first_property_value, strtobool

    client_kwargs = {
        "endpoint_url": properties.get(S3_ENDPOINT),
        "aws_access_key_id": get_first_property_value(properties, S3_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID),
        "aws_secret_access_key": get_first_property_value(properties, S3_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY),
        "aws_session_token": get_first_property_value(properties, S3_SESSION_TOKEN, AWS_SESSION_TOKEN),
        "region_name": get_first_property_value(properties, S3_REGION, AWS_REGION),
    }
    config_kwargs = {}

    if proxy_uri := properties.get(S3_PROXY_URI):
        config_kwargs["proxies"] = {"http": proxy_uri, "https": proxy_uri}
    if connect_timeout := properties.get(S3_CONNECT_TIMEOUT):
        config_kwargs["connect_timeout"] = float(connect_timeout)
    if request_timeout := properties.get(S3_REQUEST_TIMEOUT):
        config_kwargs["read_timeout"] = float(request_timeout)

    # OSS fix: virtual addressing + payload signing (avoids aws-chunked + UNSIGNED-PAYLOAD)
    config_kwargs["s3"] = {
        "addressing_style": "virtual",
        "payload_signing_enabled": True,
    }
    # botocore >= 1.35 adds flexible checksums that trigger aws-chunked encoding;
    # Aliyun OSS does not support this — disable automatic checksum calculation.
    config_kwargs["request_checksum_calculation"] = "when_required"
    config_kwargs["response_checksum_validation"] = "when_required"

    anon = strtobool(properties.get(S3_ANONYMOUS, "false")) if properties.get(S3_ANONYMOUS) else False
    s3_fs_kwargs = {"anon": anon, "client_kwargs": client_kwargs, "config_kwargs": config_kwargs}

    if profile_name := get_first_property_value(properties, S3_PROFILE_NAME, AWS_PROFILE_NAME):
        s3_fs_kwargs["profile"] = profile_name

    return S3FileSystem(**s3_fs_kwargs)


_fsspec_mod.SCHEME_TO_FS["s3"] = _oss_s3_factory
_fsspec_mod.SCHEME_TO_FS["s3a"] = _oss_s3_factory
_fsspec_mod.SCHEME_TO_FS["s3n"] = _oss_s3_factory

# ---------- Config (fetched from Java backend at startup) ----------
BACKEND_URL = os.getenv("BACKEND_URL", "http://localhost:8080")
BACKEND_INTERNAL_HEADER = os.getenv("BACKEND_INTERNAL_HEADER", "X-Internal-Token")
BACKEND_INTERNAL_TOKEN = os.getenv("BACKEND_INTERNAL_TOKEN", "").strip()

CFG: dict[str, str] = {}


def _fetch_config() -> None:
    """Pull config from Spring Boot backend (application.yml is the single source of truth)."""
    global CFG
    headers = {}
    if BACKEND_INTERNAL_TOKEN:
        headers[BACKEND_INTERNAL_HEADER] = BACKEND_INTERNAL_TOKEN

    resp = requests.get(
        f"{BACKEND_URL}/api/data/internal/iceberg-config",
        headers=headers or None,
        timeout=10,
    )
    resp.raise_for_status()
    CFG = resp.json()
    logger.info("Config loaded from backend: polaris_host=%s, catalog=%s, bucket=%s",
                CFG.get("polaris_host"), CFG.get("polaris_catalog"), CFG.get("oss_bucket"))


def get_catalog() -> RestCatalog:
    if not CFG:
        raise RuntimeError("Config not loaded — is the Java backend running?")
    props = {
        "uri": f"{CFG['polaris_host']}/api/catalog",
        "credential": f"{CFG['polaris_client_id']}:{CFG['polaris_client_secret']}",
        "warehouse": CFG["polaris_catalog"],
        "header.Polaris-Realm": CFG["polaris_realm"],
        "scope": "PRINCIPAL_ROLE:ALL",
        "s3.region": CFG["oss_region"],
        # Use fsspec/s3fs (botocore) instead of PyArrow's native S3 (AWS C++ SDK).
        # PyArrow sends aws-chunked encoding with UNSIGNED-PAYLOAD which Aliyun OSS rejects.
        "py-io-impl": "pyiceberg.io.fsspec.FsspecFileIO",
    }
    endpoint = CFG.get("oss_endpoint", "")
    if endpoint:
        props["s3.endpoint"] = endpoint
        props["s3.force-virtual-addressing"] = "true"
    if CFG.get("oss_access_key_id"):
        props["s3.access-key-id"] = CFG["oss_access_key_id"]
    if CFG.get("oss_access_key_secret"):
        props["s3.secret-access-key"] = CFG["oss_access_key_secret"]
    catalog = RestCatalog(name="polaris", **props)
    catalog._session.headers.pop("X-Iceberg-Access-Delegation", None)
    return catalog


def _ensure_namespace(catalog: RestCatalog, database: str) -> None:
    bucket = CFG.get("oss_bucket", "")
    catalog_name = CFG.get("polaris_catalog", "")
    expected_location = f"s3://{bucket}/{catalog_name}/{database}" if bucket else ""

    try:
        ns_props = {"location": expected_location} if expected_location else {}
        catalog.create_namespace(database, ns_props)
    except Exception:
        # Namespace already exists — make sure its location points to the correct bucket
        if expected_location:
            try:
                catalog.update_namespace_properties(database, updates={"location": expected_location})
            except Exception:
                pass


def _write(catalog: RestCatalog, database: str, table: str, arrow_table: pa.Table, mode: str) -> int:
    identifier = f"{database}.{table}"
    _ensure_namespace(catalog, database)
    try:
        tbl = catalog.load_table(identifier)
        if mode == "overwrite":
            tbl.overwrite(arrow_table)
        else:
            tbl.append(arrow_table)
    except NoSuchTableError:
        tbl = catalog.create_table(identifier=identifier, schema=arrow_table.schema)
        tbl.append(arrow_table)
    return len(arrow_table)


# ---------- Startup: fetch config & ensure Polaris grants ----------

def _get_management_token() -> str:
    resp = requests.post(
        f"{CFG['polaris_host']}/api/catalog/v1/oauth/tokens",
        headers={"Polaris-Realm": CFG["polaris_realm"]},
        data={
            "grant_type": "client_credentials",
            "client_id": CFG["polaris_client_id"],
            "client_secret": CFG["polaris_client_secret"],
            "scope": "PRINCIPAL_ROLE:ALL",
        },
        timeout=10,
    )
    resp.raise_for_status()
    return resp.json()["access_token"]


def _ensure_catalog_grants(headers: dict) -> None:
    catalog_name = CFG["polaris_catalog"]
    url = f"{CFG['polaris_host']}/api/management/v1/catalogs/{catalog_name}/catalog-roles/catalog_admin/grants"
    resp = requests.put(url, headers=headers, json={"type": "catalog", "privilege": "CATALOG_MANAGE_CONTENT"}, timeout=10)
    already_exists = resp.status_code == 409 or (resp.status_code == 500 and "duplicate key" in resp.text)
    if resp.status_code == 201:
        logger.info("Granted CATALOG_MANAGE_CONTENT to catalog_admin on %s", catalog_name)
    elif already_exists:
        logger.info("CATALOG_MANAGE_CONTENT already granted on %s", catalog_name)
    else:
        logger.warning("Grant response %s: %s", resp.status_code, resp.text)


def _ensure_catalog_storage_credentials(headers: dict) -> None:
    ak = CFG.get("oss_access_key_id", "")
    sk = CFG.get("oss_access_key_secret", "")
    if not ak or not sk:
        logger.warning("OSS credentials not configured — skipping credential storage")
        return

    catalog_name = CFG["polaris_catalog"]
    bucket = CFG.get("oss_bucket", "")
    base_location = f"s3://{bucket}/{catalog_name}" if bucket else ""

    get_resp = requests.get(
        f"{CFG['polaris_host']}/api/management/v1/catalogs/{catalog_name}",
        headers=headers, timeout=10,
    )
    get_resp.raise_for_status()
    catalog_info = get_resp.json()
    entity_version = catalog_info["entityVersion"]
    storage = catalog_info.get("storageConfigInfo", {})
    current_props = catalog_info.get("properties", {})

    if (storage.get("accessKeyId") == ak
            and (not base_location or current_props.get("default-base-location") == base_location)):
        logger.info("OSS credentials already stored in Polaris catalog %s", catalog_name)
        return

    updated_props = {**current_props}
    if base_location:
        updated_props["default-base-location"] = base_location
    # Ensure DROP with purge can clean up underlying OSS objects for this catalog.
    updated_props["polaris.config.drop-with-purge.enabled"] = "true"

    update_body = {
        "currentEntityVersion": entity_version,
        "properties": updated_props,
        "storageConfigInfo": {
            "storageType": storage.get("storageType", "S3"),
            "allowedLocations": [f"s3://{bucket}/"] if bucket else storage.get("allowedLocations", []),
            "region": storage.get("region", CFG.get("oss_region", "")),
            "endpoint": storage.get("endpoint", CFG.get("oss_endpoint", "")),
            "pathStyleAccess": storage.get("pathStyleAccess", False),
            "stsUnavailable": True,
            "accessKeyId": ak,
            "secretAccessKey": sk,
        },
    }
    put_resp = requests.put(
        f"{CFG['polaris_host']}/api/management/v1/catalogs/{catalog_name}",
        headers=headers, json=update_body, timeout=10,
    )
    if put_resp.status_code == 200:
        logger.info("Stored OSS credentials in Polaris catalog %s", catalog_name)
    else:
        logger.warning("Catalog credential update %s: %s", put_resp.status_code, put_resp.text)


@asynccontextmanager
async def lifespan(app: FastAPI):
    import asyncio
    # Config fetch must succeed — retry until Java backend is reachable
    for attempt in range(30):
        try:
            _fetch_config()
            break
        except Exception as e:
            logger.warning("Config fetch attempt %d failed: %s — retrying in 2s", attempt + 1, e)
            await asyncio.sleep(2)
    else:
        logger.error("Could not fetch config from backend after 30 attempts")

    if CFG:
        try:
            token = _get_management_token()
            headers = {
                "Authorization": f"Bearer {token}",
                "Polaris-Realm": CFG["polaris_realm"],
                "Content-Type": "application/json",
            }
            _ensure_catalog_grants(headers)
            _ensure_catalog_storage_credentials(headers)
        except Exception as e:
            logger.warning("Polaris setup failed (will retry on next restart): %s", e)
    yield


# ---------- App ----------
app = FastAPI(title="Iceberg Data Service", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health")
def health():
    return {"status": "ok"}


# ---------- Excel import endpoint ----------

@app.post("/import/excel")
async def import_excel(
    database: str,
    file: UploadFile = File(...),
    mode: str = Form(default="append"),
):
    try:
        content = await file.read()
        wb = openpyxl.load_workbook(io.BytesIO(content), read_only=True, data_only=True)
        catalog = get_catalog()
        results = []
        seen_names: set[str] = set()
        for sheet_name in wb.sheetnames:
            # 检测同一 Excel 内重复 Sheet 名（Excel 本身禁止，但做防御性检查）
            name_key = sheet_name.strip().lower()
            if name_key in seen_names:
                results.append({"table": sheet_name, "error": f"Excel 中存在重复 Sheet 名「{sheet_name}」"})
                continue
            seen_names.add(name_key)

            ws = wb[sheet_name]
            rows_iter = ws.iter_rows(values_only=True)
            try:
                header_row = next(rows_iter)
            except StopIteration:
                results.append({"table": sheet_name, "skipped": True, "reason": "空 sheet"})
                continue
            headers = [str(h) if h is not None else f"col_{i}" for i, h in enumerate(header_row)]
            if not any(h for h in headers):
                results.append({"table": sheet_name, "skipped": True, "reason": "空表头"})
                continue

            # 检测目标库中是否已存在同名表（overwrite 模式允许覆盖）
            if mode != "overwrite":
                try:
                    catalog.load_table(f"{database}.{sheet_name}")
                    results.append({"table": sheet_name, "error": f"表「{sheet_name}」已存在，请先删除后再导入"})
                    continue
                except NoSuchTableError:
                    pass

            data_rows = list(rows_iter)
            if not data_rows:
                arrays = [pa.array([], type=pa.string()) for _ in headers]
            else:
                cols = [[row[i] if i < len(row) else None for row in data_rows] for i in range(len(headers))]
                arrays = [pa.array(col) for col in cols]
            arrow_table = pa.table(dict(zip(headers, arrays)))
            try:
                row_count = _write(catalog, database, sheet_name, arrow_table, mode)
                results.append({"table": sheet_name, "rows": row_count, "columns": len(headers)})
            except Exception as e:
                results.append({"table": sheet_name, "columns": len(headers), "error": str(e)})
        wb.close()
        return {"success": True, "tables": results}
    except Exception as e:
        logger.error("Excel import error", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/import/csv")
async def import_csv(
    database: str,
    table: str,
    file: UploadFile = File(...),
    mode: str = Form(default="append"),
):
    try:
        content = await file.read()
        arrow_table = pa_csv.read_csv(io.BytesIO(content))
        rows = _write(get_catalog(), database, table, arrow_table, mode)
        return {"success": True, "rows": rows, "columns": len(arrow_table.schema)}
    except Exception as e:
        logger.error("CSV import error", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))



@app.get("/export/csv")
async def export_csv(database: str, table: str):
    try:
        catalog = get_catalog()
        tbl = catalog.load_table(f"{database}.{table}")
        arrow_table = tbl.scan().to_arrow()
        buf = io.BytesIO()
        # Add UTF-8 BOM to make Excel happy with Chinese characters
        buf.write(b"\xef\xbb\xbf")
        pa_csv.write_csv(arrow_table, buf)
        buf.seek(0)
        filename = f"{database}_{table}.csv"
        return StreamingResponse(
            buf,
            media_type="text/csv; charset=utf-8",
            headers={"Content-Disposition": f'attachment; filename="{filename}"'},
        )
    except Exception as e:
        logger.error("CSV export error", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


# ---------- LLM Function Calling 查询接口 ----------

import pyarrow.compute as pc
from typing import Optional


def _safe_json(val):
    """将 PyArrow .as_py() 的值转换为 JSON 可序列化的 Python 类型。"""
    if val is None:
        return None
    if isinstance(val, (int, float, bool, str)):
        return val
    # date / datetime / Decimal 等 → str
    return str(val)


def _to_column_flat(arrow_table: pa.Table, column: str) -> pa.Array:
    """取出指定列并合并 chunk，返回 pa.Array。"""
    col = arrow_table.column(column)
    return col.combine_chunks() if isinstance(col, pa.ChunkedArray) else col


@app.get("/preview")
async def preview(database: str, table: str, limit: int = 5):
    """
    query_sample_data：返回表的前 N 行样本数据。
    响应：{"table": str, "columns": [...], "rows": [[...], ...], "total_returned": int}
    """
    limit = min(max(limit, 1), 500)
    try:
        tbl = get_catalog().load_table(f"{database}.{table}")
        arrow_table = tbl.scan(limit=limit).to_arrow()
        columns = arrow_table.schema.names
        rows = [
            [_safe_json(arrow_table.column(c)[i].as_py()) for c in columns]
            for i in range(len(arrow_table))
        ]
        return {"table": table, "columns": columns, "rows": rows, "total_returned": len(arrow_table)}
    except NoSuchTableError:
        raise HTTPException(status_code=404, detail=f"表 {database}.{table} 不存在")
    except Exception as e:
        logger.error("preview error", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/distinct")
async def distinct_values(database: str, table: str, column: str, limit: int = 20):
    """
    query_distinct_values：返回指定列的不同值及出现次数（按频次降序）。
    响应：{"column": str, "total_distinct": int, "values": [{"val": ..., "count": int}, ...]}
    """
    limit = min(max(limit, 1), 500)
    try:
        tbl = get_catalog().load_table(f"{database}.{table}")
        arrow_table = tbl.scan(selected_fields=(column,)).to_arrow()
        col_flat = _to_column_flat(arrow_table, column)

        total_distinct = pc.count_distinct(col_flat, mode="only_valid").as_py()

        # value_counts() → StructArray {values, counts}；建表后排序
        vc = col_flat.value_counts()
        vc_table = pa.table({
            "val": vc.field("values"),
            "cnt": vc.field("counts"),
        })
        sorted_vc = vc_table.sort_by([("cnt", "descending")]).slice(0, limit)

        values = [
            {"val": _safe_json(sorted_vc.column("val")[i].as_py()),
             "count": sorted_vc.column("cnt")[i].as_py()}
            for i in range(len(sorted_vc))
        ]
        return {"column": column, "total_distinct": total_distinct, "values": values}
    except NoSuchTableError:
        raise HTTPException(status_code=404, detail=f"表 {database}.{table} 不存在")
    except Exception as e:
        logger.error("distinct error", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/statistics")
async def statistics(database: str, table: str, column: Optional[str] = None):
    """
    query_statistics：返回表的行数；若指定列则额外返回 min/max/distinct_count/null_count。
    响应：{"table": str, "total_rows": int, "column"?: str, "min"?: ..., "max"?: ...,
           "distinct_count"?: int, "null_count"?: int}
    """
    try:
        tbl = get_catalog().load_table(f"{database}.{table}")
        result: dict = {"table": table}

        if column:
            arrow_table = tbl.scan(selected_fields=(column,)).to_arrow()
            result["total_rows"] = len(arrow_table)
            result["column"] = column
            col_flat = _to_column_flat(arrow_table, column)
            try:
                result["min"] = _safe_json(pc.min(col_flat).as_py())
                result["max"] = _safe_json(pc.max(col_flat).as_py())
                result["distinct_count"] = pc.count_distinct(col_flat, mode="only_valid").as_py()
                result["null_count"] = col_flat.null_count
            except Exception as stat_err:
                result["stats_error"] = str(stat_err)
        else:
            # 仅统计行数：用 metadata 统计，避免全量扫描
            result["total_rows"] = tbl.scan().count()

        return result
    except NoSuchTableError:
        raise HTTPException(status_code=404, detail=f"表 {database}.{table} 不存在")
    except Exception as e:
        logger.error("statistics error", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))