Spaces:
Sleeping
Sleeping
| """ | |
| Session manager — lightweight in-memory session metadata. | |
| In production, swap this dict for Redis. Each entry holds only | |
| metadata (not the dataframe itself). The actual data lives on | |
| disk as a Parquet file under DATA_DIR/{session_id}.parquet. | |
| """ | |
| from __future__ import annotations | |
| import time | |
| import threading | |
| from dataclasses import dataclass, field | |
| from typing import Optional | |
| from config import DATA_DIR, SESSION_TTL_MINUTES | |
| class SessionMeta: | |
| session_id: str | |
| file_name: str | |
| file_size_bytes: int | |
| columns: list[dict[str, str]] # [{"name": ..., "dtype": ...}] | |
| row_count: int | |
| status: str = "active" | |
| current_version: int = 0 | |
| created_at: float = field(default_factory=time.time) | |
| last_active: float = field(default_factory=time.time) | |
| def touch(self) -> None: | |
| self.last_active = time.time() | |
| class SessionManager: | |
| """Thread-safe session store.""" | |
| def __init__(self) -> None: | |
| self._sessions: dict[str, SessionMeta] = {} | |
| self._lock = threading.Lock() | |
| def create( | |
| self, | |
| session_id: str, | |
| file_name: str, | |
| file_size_bytes: int, | |
| columns: list[dict[str, str]], | |
| row_count: int, | |
| ) -> SessionMeta: | |
| meta = SessionMeta( | |
| session_id=session_id, | |
| file_name=file_name, | |
| file_size_bytes=file_size_bytes, | |
| columns=columns, | |
| row_count=row_count, | |
| ) | |
| with self._lock: | |
| self._sessions[session_id] = meta | |
| return meta | |
| def get(self, session_id: str) -> Optional[SessionMeta]: | |
| with self._lock: | |
| return self._sessions.get(session_id) | |
| def get_filepath(self, session_id: str) -> str: | |
| """Get the absolute filepath of the current version of the Parquet file.""" | |
| meta = self.get(session_id) | |
| if meta: | |
| version = getattr(meta, "current_version", 0) | |
| if version > 0: | |
| v_path = os.path.join(DATA_DIR, f"{session_id}_v{version}.parquet") | |
| if os.path.exists(v_path): | |
| return v_path | |
| # Version 0 or fallback: check v0 path first | |
| v0_path = os.path.join(DATA_DIR, f"{session_id}_v0.parquet") | |
| if os.path.exists(v0_path): | |
| return v0_path | |
| return os.path.join(DATA_DIR, f"{session_id}.parquet") | |
| def touch(self, session_id: str) -> None: | |
| meta = self.get(session_id) | |
| if meta: | |
| meta.touch() | |
| def remove(self, session_id: str) -> None: | |
| with self._lock: | |
| self._sessions.pop(session_id, None) | |
| def list_active(self) -> list[SessionMeta]: | |
| """Return sessions that haven't expired.""" | |
| now = time.time() | |
| cutoff = now - SESSION_TTL_MINUTES * 60 | |
| with self._lock: | |
| return [ | |
| m for m in self._sessions.values() | |
| if m.last_active > cutoff | |
| ] | |
| def cleanup_expired(self) -> int: | |
| """Remove expired sessions and their Parquet files. Returns count removed.""" | |
| import os | |
| import glob | |
| now = time.time() | |
| cutoff = now - SESSION_TTL_MINUTES * 60 | |
| removed = 0 | |
| with self._lock: | |
| expired = [sid for sid, m in self._sessions.items() if m.last_active <= cutoff] | |
| for sid in expired: | |
| del self._sessions[sid] | |
| pattern = os.path.join(DATA_DIR, f"{sid}*.parquet") | |
| for pq in glob.glob(pattern): | |
| try: | |
| os.remove(pq) | |
| except Exception: | |
| pass | |
| removed += 1 | |
| return removed | |
| # Module-level singleton | |
| session_manager = SessionManager() |