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))