Update backend/app/api/datasets.py

backend/app/api/datasets.py

@@ -1,30 +1,75 @@
-import json
+from json import dumps
+from pathlib import Path
+
 import pandas as pd
 from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
 from sqlalchemy.orm import Session
+
 from backend.app.db import get_db
 from backend.app.repositories.dataset_repo import create_dataset, get_dataset
-from backend.app.schemas.dataset import DatasetProfileResponse, DatasetUploadResponse
 from backend.app.services.profiling_service import profile_dataframe
-from backend.app.services.storage_service import dumps, save_upload
+from backend.app.utils.ids import make_dataset_id
+
+router = APIRouter(tags=["datasets"])
 
-router = APIRouter(prefix="/datasets", tags=["datasets"])
+UPLOAD_DIR = Path("/data/uploads")
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
 
-@router.post("/upload", response_model=DatasetUploadResponse)
+
+@router.post("/datasets/upload")
 async def upload_dataset(file: UploadFile = File(...), db: Session = Depends(get_db)):
-    if not file.filename.lower().endswith(".csv"):
-        raise HTTPException(status_code=400, detail="Only CSV files are supported.")
-    content = await file.read()
-    dataset_id, path = save_upload(file.filename, content)
-    df = pd.read_csv(path)
+    if not file.filename:
+        raise HTTPException(status_code=400, detail="Missing file name")
+
+    suffix = Path(file.filename).suffix.lower()
+    if suffix not in {".csv", ".xlsx", ".xls"}:
+        raise HTTPException(status_code=400, detail="Only CSV and Excel files are supported")
+
+    dataset_id = make_dataset_id()
+
+    if suffix == ".csv":
+        content = await file.read()
+        path = UPLOAD_DIR / f"{dataset_id}.csv"
+        path.write_bytes(content)
+        df = pd.read_csv(path)
+    else:
+        content = await file.read()
+        path = UPLOAD_DIR / f"{dataset_id}{suffix}"
+        path.write_bytes(content)
+        df = pd.read_excel(path)
+
     profile = profile_dataframe(df)
-    create_dataset(db, id=dataset_id, name=file.filename, file_path=str(path), row_count=len(df), column_count=len(df.columns), schema_json=dumps({"columns": list(df.columns)}), profile_json=dumps(profile))
-    return DatasetUploadResponse(dataset_id=dataset_id, name=file.filename, row_count=len(df), column_count=len(df.columns))
 
-@router.get("/{dataset_id}/profile", response_model=DatasetProfileResponse)
+    create_dataset(
+        db,
+        id=dataset_id,
+        name=file.filename,
+        file_path=str(path),
+        row_count=len(df),
+        column_count=len(df.columns),
+        schema_json=dumps({"columns": list(df.columns)}),
+        profile_json=dumps(profile),
+    )
+
+    return {
+        "dataset_id": dataset_id,
+        "name": file.filename,
+        "row_count": int(len(df)),
+        "column_count": int(len(df.columns)),
+    }
+
+
+@router.get("/datasets/{dataset_id}/profile")
 def dataset_profile(dataset_id: str, db: Session = Depends(get_db)):
     dataset = get_dataset(db, dataset_id)
     if not dataset:
-        raise HTTPException(status_code=404, detail="Dataset not found.")
-    profile = json.loads(dataset.profile_json)
-    return DatasetProfileResponse(dataset_id=dataset_id, columns=profile["columns"], recommended_algorithms=profile["recommended_algorithms"])
+        raise HTTPException(status_code=404, detail="Dataset not found")
+
+    return {
+        "dataset_id": dataset.id,
+        "name": dataset.name,
+        "row_count": dataset.row_count,
+        "column_count": dataset.column_count,
+        "schema_json": dataset.schema_json,
+        "profile_json": dataset.profile_json,
+    }