Create main.py

main.py

@@ -0,0 +1,385 @@
+"""
+FastAPI application for HF Uploader Dashboard
+"""
+import os
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, Depends, HTTPException, UploadFile, File
+from fastapi.responses import FileResponse, JSONResponse
+from fastapi.staticfiles import StaticFiles
+from sqlalchemy.orm import Session
+from datetime import datetime
+
+from database import init_db, get_db, SessionLocal
+from models import (
+    HFConfig,
+    UploadQueue,
+    UploadStatusEnum,
+    ProcessingState,
+    ProcessingStatusEnum,
+    UploadErrorLog,
+)
+from schemas import (
+    HFConfigUpdate,
+    HFConfigResponse,
+    UploadQueueItem,
+    UploadQueueStats,
+    ProcessingStateResponse,
+    ProcessRequest,
+    UploadRequest,
+    FilePreviewRequest,
+    FilePreviewResponse,
+    RateLimitStatus,
+    UploadStatus,
+    ErrorLogResponse,
+)
+from file_processor import FileProcessor
+from hf_uploader import HFUploader
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global instances
+file_processor = FileProcessor(processed_files_dir="processed_files")
+hf_uploader_instance = None
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Lifespan context manager for startup and shutdown"""
+    # Startup
+    logger.info("Initializing database...")
+    init_db()
+    logger.info("Application started")
+    yield
+    # Shutdown
+    logger.info("Application shutting down...")
+
+
+# Create FastAPI app
+app = FastAPI(
+    title="HF Uploader Dashboard",
+    description="Hugging Face Dataset Uploader with Rate Limiting",
+    version="1.0.0",
+    lifespan=lifespan,
+)
+
+
+# ============================================================================
+# Configuration Endpoints
+# ============================================================================
+
+
+@app.get("/api/config", response_model=HFConfigResponse)
+async def get_config(db: Session = Depends(get_db)):
+    """Get current HF configuration"""
+    config = db.query(HFConfig).first()
+    if not config:
+        # Create default config
+        config = HFConfig()
+        db.add(config)
+        db.commit()
+        db.refresh(config)
+    return config
+
+
+@app.post("/api/config", response_model=HFConfigResponse)
+async def update_config(
+    config_update: HFConfigUpdate, db: Session = Depends(get_db)
+):
+    """Update HF configuration"""
+    config = db.query(HFConfig).first()
+    if not config:
+        config = HFConfig()
+        db.add(config)
+
+    # Update fields
+    update_data = config_update.model_dump(exclude_unset=True)
+    for field, value in update_data.items():
+        setattr(config, field, value)
+
+    config.updated_at = datetime.utcnow()
+    db.commit()
+    db.refresh(config)
+
+    # Update global uploader instance
+    global hf_uploader_instance
+    if config.hf_token:
+        hf_uploader_instance = HFUploader(config.hf_token, config.target_repo)
+
+    return config
+
+
+# ============================================================================
+# Upload Queue Endpoints
+# ============================================================================
+
+
+@app.get("/api/queue", response_model=list[UploadQueueItem])
+async def get_upload_queue(
+    status: str = None, db: Session = Depends(get_db)
+):
+    """Get upload queue items"""
+    query = db.query(UploadQueue)
+    if status:
+        query = query.filter(UploadQueue.status == status)
+    return query.all()
+
+
+@app.get("/api/queue/stats", response_model=UploadQueueStats)
+async def get_queue_stats(db: Session = Depends(get_db)):
+    """Get upload queue statistics"""
+    all_items = db.query(UploadQueue).all()
+    return UploadQueueStats(
+        pending=len([i for i in all_items if i.status == UploadStatusEnum.PENDING]),
+        uploading=len(
+            [i for i in all_items if i.status == UploadStatusEnum.UPLOADING]
+        ),
+        completed=len(
+            [i for i in all_items if i.status == UploadStatusEnum.COMPLETED]
+        ),
+        failed=len([i for i in all_items if i.status == UploadStatusEnum.FAILED]),
+        total=len(all_items),
+    )
+
+
+@app.post("/api/queue/add")
+async def add_to_queue(db: Session = Depends(get_db)):
+    """Add processed files to upload queue"""
+    processed_files = file_processor.get_processed_files()
+    added_count = 0
+
+    for filename in processed_files:
+        # Check if already in queue
+        existing = (
+            db.query(UploadQueue)
+            .filter(UploadQueue.file_name == filename)
+            .first()
+        )
+        if existing:
+            continue
+
+        file_path = str(file_processor.processed_files_dir / filename)
+        file_size = file_processor.get_file_size(filename) or 0
+
+        queue_item = UploadQueue(
+            file_name=filename,
+            file_path=file_path,
+            file_size=file_size,
+            status=UploadStatusEnum.PENDING,
+        )
+        db.add(queue_item)
+        added_count += 1
+
+    db.commit()
+    return {"added": added_count, "total": len(processed_files)}
+
+
+@app.post("/api/queue/clear")
+async def clear_queue(db: Session = Depends(get_db)):
+    """Clear upload queue"""
+    db.query(UploadQueue).delete()
+    db.commit()
+    return {"message": "Queue cleared"}
+
+
+# ============================================================================
+# Processing Endpoints
+# ============================================================================
+
+
+@app.get("/api/processing/state", response_model=ProcessingStateResponse)
+async def get_processing_state(db: Session = Depends(get_db)):
+    """Get current processing state"""
+    state = db.query(ProcessingState).first()
+    if not state:
+        state = ProcessingState(status=ProcessingStatusEnum.IDLE)
+        db.add(state)
+        db.commit()
+        db.refresh(state)
+    return state
+
+
+@app.post("/api/processing/start")
+async def start_processing(
+    request: ProcessRequest, db: Session = Depends(get_db)
+):
+    """Start dataset processing"""
+    config = db.query(HFConfig).first()
+    if not config:
+        raise HTTPException(status_code=400, detail="Configuration not set")
+
+    if not config.hf_token:
+        raise HTTPException(status_code=400, detail="HF token not configured")
+
+    # Run processing in background
+    asyncio.create_task(
+        file_processor.process_datasets(
+            all_repo_id=config.source_all_repo,
+            ato_repo_id=config.source_ato_repo,
+            hf_token=config.hf_token,
+            max_files=request.max_files,
+            db=db,
+        )
+    )
+
+    return {"message": "Processing started"}
+
+
+@app.get("/api/processing/files")
+async def get_processed_files():
+    """Get list of processed files"""
+    files = file_processor.get_processed_files()
+    return {"files": files, "count": len(files)}
+
+
+# ============================================================================
+# Upload Endpoints
+# ============================================================================
+
+
+@app.post("/api/upload/start")
+async def start_upload(request: UploadRequest, db: Session = Depends(get_db)):
+    """Start uploading files"""
+    global hf_uploader_instance
+
+    config = db.query(HFConfig).first()
+    if not config or not config.hf_token:
+        raise HTTPException(status_code=400, detail="HF token not configured")
+
+    if not hf_uploader_instance:
+        hf_uploader_instance = HFUploader(config.hf_token, config.target_repo)
+
+    # Get files to upload
+    files_to_upload = (
+        db.query(UploadQueue)
+        .filter(UploadQueue.id.in_(request.file_ids))
+        .all()
+    )
+
+    if not files_to_upload:
+        raise HTTPException(status_code=400, detail="No files found")
+
+    # Convert to upload format
+    upload_files = [
+        {
+            "id": f.id,
+            "file_name": f.file_name,
+            "file_path": f.file_path,
+        }
+        for f in files_to_upload
+    ]
+
+    # Run upload in background
+    asyncio.create_task(
+        hf_uploader_instance.upload_files_batch(
+            files=upload_files,
+            db=db,
+            batch_size=config.upload_batch_size,
+        )
+    )
+
+    return {"message": "Upload started", "file_count": len(upload_files)}
+
+
+@app.get("/api/upload/status", response_model=UploadStatus)
+async def get_upload_status(db: Session = Depends(get_db)):
+    """Get current upload status"""
+    global hf_uploader_instance
+
+    config = db.query(HFConfig).first()
+    if not config or not config.hf_token:
+        raise HTTPException(status_code=400, detail="HF token not configured")
+
+    if not hf_uploader_instance:
+        hf_uploader_instance = HFUploader(config.hf_token, config.target_repo)
+
+    return await hf_uploader_instance.get_upload_status(db)
+
+
+@app.post("/api/upload/retry/{item_id}")
+async def retry_upload(item_id: int, db: Session = Depends(get_db)):
+    """Retry uploading a failed file"""
+    queue_item = db.query(UploadQueue).filter(UploadQueue.id == item_id).first()
+    if not queue_item:
+        raise HTTPException(status_code=404, detail="Item not found")
+
+    queue_item.status = UploadStatusEnum.PENDING
+    queue_item.retry_count = 0
+    db.commit()
+
+    return {"message": "Retry scheduled"}
+
+
+# ============================================================================
+# File Preview Endpoints
+# ============================================================================
+
+
+@app.post("/api/file/preview", response_model=FilePreviewResponse)
+async def preview_file(request: FilePreviewRequest):
+    """Get file preview"""
+    content = file_processor.get_file_content(request.filename)
+    if not content:
+        raise HTTPException(status_code=404, detail="File not found")
+
+    file_size = file_processor.get_file_size(request.filename) or 0
+
+    return FilePreviewResponse(
+        filename=request.filename, size=file_size, content=content
+    )
+
+
+# ============================================================================
+# Error Log Endpoints
+# ============================================================================
+
+
+@app.get("/api/errors", response_model=list[ErrorLogResponse])
+async def get_error_logs(limit: int = 50, db: Session = Depends(get_db)):
+    """Get error logs"""
+    return (
+        db.query(UploadErrorLog)
+        .order_by(UploadErrorLog.created_at.desc())
+        .limit(limit)
+        .all()
+    )
+
+
+@app.delete("/api/errors")
+async def clear_error_logs(db: Session = Depends(get_db)):
+    """Clear error logs"""
+    db.query(UploadErrorLog).delete()
+    db.commit()
+    return {"message": "Error logs cleared"}
+
+
+# ============================================================================
+# Static Files
+# ============================================================================
+
+
+# Create static directory if it doesn't exist
+os.makedirs("static", exist_ok=True)
+
+# Mount static files
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+
+# ============================================================================
+# Root Endpoint
+# ============================================================================
+
+
+@app.get("/")
+async def root():
+    """Root endpoint - serve dashboard"""
+    return FileResponse("static/index.html")
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)