""" JSON extractor for structured data files. Supports only CSV, XLS, XLSX file types for JSON extraction. Returns error for unsupported file types. """ import io import warnings from pathlib import Path from typing import Any, Dict, Optional, Union import pandas as pd from logger import get_logger logger = get_logger(__name__) # Supported file extensions for JSON extraction SUPPORTED_EXTENSIONS = {'.csv', '.xls', '.xlsx'} # Resource limits to prevent abuse from app.core.config import settings as _app_settings MAX_FILE_SIZE_BYTES = _app_settings.MAX_FILE_SIZE_BYTES MAX_CSV_ROWS = 100000 MAX_EXCEL_ROWS = 50000 MAX_MEMORY_ROWS = 100000 def _validate_file_size(size: int) -> Optional[str]: if size > MAX_FILE_SIZE_BYTES: return f"File size {size} bytes exceeds limit of {MAX_FILE_SIZE_BYTES} bytes" return None def _check_memory_usage(rows: int, cols: int) -> Optional[str]: approx_mb = (rows * cols * 50) / (1024 * 1024) if rows * cols > MAX_MEMORY_ROWS * 20: return f"Data size too large (approx {approx_mb:.1f} MB). Too many cells: {rows}x{cols}" return None def extract_json_from_file( file_path: Union[str, Path], file_data: Optional[bytes] = None ) -> Dict[str, Any]: """ Extract JSON data from structured files (CSV, XLS, XLSX). Parameters ---------- file_path : Union[str, Path] Path to the file or filename with extension file_data : Optional[bytes] Raw file data (for stream processing) Returns ------- Dict[str, Any] Extracted data or error information """ ext = Path(file_path).suffix.lower() if ext not in SUPPORTED_EXTENSIONS: return { "error": f"Unsupported file type: {ext}. Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}", "file_type": ext } # Validate file size if we have raw data if file_data is not None: size_error = _validate_file_size(len(file_data)) if size_error: return {"error": size_error, "file_type": ext} elif Path(file_path).exists(): size_error = _validate_file_size(Path(file_path).stat().st_size) if size_error: return {"error": size_error, "file_type": ext} try: with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) if ext == '.csv': if file_data: stream = io.BytesIO(file_data) # Peek for encoding detection sample = stream.read(1024) stream.seek(0) df = pd.read_csv(stream, nrows=MAX_CSV_ROWS + 1, low_memory=False) else: df = pd.read_csv(file_path, nrows=MAX_CSV_ROWS + 1, low_memory=False) else: if file_data: df = pd.read_excel(io.BytesIO(file_data), engine='openpyxl' if ext == '.xlsx' else 'xlrd') else: df = pd.read_excel(file_path, engine='openpyxl' if ext == '.xlsx' else 'xlrd') # Enforce row limits to prevent memory exhaustion max_rows = MAX_EXCEL_ROWS if ext != '.csv' else MAX_CSV_ROWS if len(df) > max_rows: return { "error": f"File contains {len(df)} rows, exceeds limit of {max_rows}", "file_type": ext, "row_count": len(df) } # Additional memory guard mem_error = _check_memory_usage(len(df), len(df.columns)) if mem_error: return {"error": mem_error, "file_type": ext} # Efficient conversion to JSON-serializable format result = { "success": True, "file_type": ext, "data": { "columns": list(df.columns), "rows": df.where(pd.notnull(df), None).to_dict(orient='records'), "shape": [len(df), len(df.columns)], "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()} } } logger.info("Extracted JSON from %s: %d rows, %d cols", ext, len(df), len(df.columns)) return result except pd.errors.EmptyDataError: return { "error": f"File is empty or has no data", "file_type": ext } except MemoryError: return { "error": "Out of memory processing file", "file_type": ext } except Exception as exc: logger.exception("JSON extraction failed for %s", ext) return { "error": f"Processing failed: {str(exc)}", "file_type": ext, "exception_type": type(exc).__name__ } def is_supported_file_type(file_path: Union[str, Path]) -> bool: """ Check if file type is supported for JSON extraction. Parameters ---------- file_path : Union[str, Path] Path to the file or filename with extension Returns ------- bool True if supported, False otherwise """ extension = Path(file_path).suffix.lower() return extension in SUPPORTED_EXTENSIONS def get_supported_extensions() -> list: """ Get list of supported file extensions for JSON extraction. Returns ------- list Supported file extensions """ return sorted(SUPPORTED_EXTENSIONS)