apps/api/main_hf.py

"""
Hugging Face Spaces 专用启动文件
精简版本，只包含核心查询功能
"""
import os
import subprocess
import sys
from pathlib import Path

os.environ["DATABASE_URL"] = "sqlite+aiosqlite:///data/customs_data.db"

from contextlib import asynccontextmanager
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse

# 只导入核心trade查询路由
from apps.api.routers import trade

# 获取项目根目录
BASE_DIR = Path(__file__).resolve().parent.parent.parent
STATIC_DIR = BASE_DIR / "apps" / "api" / "static"
DB_PATH = BASE_DIR / "data" / "customs_data.db"

@asynccontextmanager
async def lifespan(app: FastAPI):
    print("🚀 Starting Customs Data Query System...")
    
    # 每次启动都重建数据库（确保schema正确）
    if DB_PATH.exists():
        print(f"🗑️  Removing old database: {DB_PATH}")
        DB_PATH.unlink()
    
    print("📦 Initializing database...")
    try:
        subprocess.run([sys.executable, "init_sqlite.py"], check=True, cwd=str(BASE_DIR))
        print("✅ Database initialized successfully")
    except Exception as e:
        print(f"❌ Failed to initialize database: {e}")
        raise
    
    yield
    print("👋 Shutting down...")

app = FastAPI(
    title="海关数据查询系统",
    description="巴西海关贸易数据查询服务 - 包含29万条2023年1月真实数据",
    version="1.0.0",
    lifespan=lifespan
)

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

# 挂载静态文件
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

@app.get("/", tags=["UI"])
async def serve_ui():
    """返回前端页面"""
    return FileResponse(str(STATIC_DIR / "index.html"))

@app.get("/api/v1/health", tags=["System"])
async def health_check():
    """健康检查"""
    import sqlite3
    from pathlib import Path
    db_path = Path("data/customs_data.db")
    db_exists = db_path.exists()
    record_count = 0
    if db_exists:
        conn = sqlite3.connect(str(db_path))
        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM standard_trade_records")
        record_count = cursor.fetchone()[0]
        conn.close()
    
    return {
        "status": "ok", 
        "service": "Customs Data API",
        "database": "SQLite",
        "records": f"290,177 (sample: {record_count})",
        "db_exists": db_exists,
        "db_path": str(db_path.absolute())
    }

# 只包含核心trade查询路由
app.include_router(trade.router, prefix="/api/v1/trade", tags=["Trade Search"])

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)