PM-pinou
8fa36b774e
Key changes: - New: data_loader.py SQLite persistence with drop_sqlite_table - New: db_utils.py retry_on_lock decorator (3 retries, exponential backoff) - New: tool_registry.py with 27 MCP tools (filter_and_cluster, compute_scores, etc.) - New: tls_ref.py for TLS cipher/reference data - New: import_tlsdb.py management command - New: scripts/start_server.py for portable runtime - New: migrations 0003-0007 for SQLite table, display_id, llm fields - Changed: views.py unified pipeline worker, retry_run, display_id everywhere - Changed: models.py with display_id auto-assignment, run_type, sqlite_table, llm_thinking, tool_calls_json - Changed: urls.py added retry_run route - Changed: session_store.py robust JSON persistence - Changed: AGENTS.md v7 fix summary added - Changed: templates — globe rewrite (inertia/polar flip), auto.html (thinking/tool accordions), base.html (toast/config), all pages use display_id - Changed: run.bat PYTHONUTF8=1 - Deleted: entity_detector.py, entity_aggregator.py (replaced by filter_and_cluster clustering pipeline) - Test: 92/92 unit tests passing
355 lines
13 KiB
Python
355 lines
13 KiB
Python
"""Thread-safe in-memory session store for datasets and results.
|
||
|
||
Singleton pattern with RLock guard. Stores:
|
||
- Datasets: dataset_id -> (LazyFrame, schema, metadata)
|
||
- Cluster results: result_id -> (labels, n_clusters, params, parent_dataset_id)
|
||
- Feature results: result_id -> (features, parent_cluster_result_id)
|
||
"""
|
||
import json
|
||
import os
|
||
import threading
|
||
import gc
|
||
import uuid
|
||
from pathlib import Path
|
||
from typing import Any, Optional
|
||
|
||
import polars as pl
|
||
|
||
|
||
# Approximate per-element byte sizes for memory estimation
|
||
_DTYPE_BYTE_SIZES = {
|
||
# Polars dtype strings
|
||
'Int8': 1, 'Int16': 2, 'Int32': 4, 'Int64': 8,
|
||
'UInt8': 1, 'UInt16': 2, 'UInt32': 4, 'UInt64': 8,
|
||
'Float32': 4, 'Float64': 8,
|
||
'Boolean': 1, 'Bool': 1,
|
||
'Utf8': 50, 'String': 50, 'Categorical': 8,
|
||
'Date': 4, 'Datetime': 8, 'Time': 8, 'Duration': 8,
|
||
'List': 16, 'Object': 16, 'Null': 0,
|
||
# Polars __str__ representations
|
||
'i32': 4, 'i64': 8, 'u32': 4, 'u64': 8,
|
||
'f32': 4, 'f64': 8, 'str': 50, 'bool': 1,
|
||
}
|
||
|
||
|
||
def _estimate_dtype_bytes(dtype: Any) -> int:
|
||
"""Guess per-element byte size from a Polars dtype."""
|
||
key = str(dtype)
|
||
for prefix, size in _DTYPE_BYTE_SIZES.items():
|
||
if key.startswith(prefix) or key == prefix:
|
||
return size
|
||
return 8 # fallback
|
||
|
||
|
||
def _generate_id(prefix: str = 'ds') -> str:
|
||
"""Generate a short unique identifier."""
|
||
return f"{prefix}_{uuid.uuid4().hex[:12]}"
|
||
|
||
|
||
def _get_store_path() -> Path:
|
||
"""Return the persistent JSON file path for dataset metadata.
|
||
|
||
Uses ``%APPDATA%/TianXuan/.session_store.json`` (Windows) or
|
||
``~/.tianxuan_session_store.json`` (other platforms).
|
||
"""
|
||
appdata = os.environ.get('APPDATA', os.path.expanduser('~'))
|
||
store_dir = Path(appdata) / 'TianXuan'
|
||
store_dir.mkdir(parents=True, exist_ok=True)
|
||
return store_dir / '.session_store.json'
|
||
|
||
|
||
class SessionStore:
|
||
"""Singleton in-memory store for datasets and analysis results.
|
||
|
||
Thread-safe via RLock. Datasets store Polars LazyFrames (query plans,
|
||
not materialised data), so clone/shallow-copy is effectively free.
|
||
|
||
Dataset metadata is persisted to ``_get_store_path()`` so that after a
|
||
server restart datasets whose CSV files still exist can be re-loaded.
|
||
"""
|
||
|
||
_instance: Optional['SessionStore'] = None
|
||
_lock: threading.RLock = threading.RLock()
|
||
|
||
def __new__(cls) -> 'SessionStore':
|
||
if cls._instance is None:
|
||
with cls._lock:
|
||
if cls._instance is None:
|
||
cls._instance = super().__new__(cls)
|
||
cls._instance._stores: dict[str, dict[str, Any]] = {}
|
||
cls._instance._datalock = threading.RLock()
|
||
return cls._instance
|
||
|
||
# ── persistence ────────────────────────────────────────────────────
|
||
|
||
def _save_to_disk(self) -> None:
|
||
"""Persist current dataset metadata to JSON file.
|
||
|
||
Only metadata (schema, csv_glob, row_count, …) is saved — never the
|
||
LazyFrame itself, which is an in-memory query plan.
|
||
"""
|
||
with self._datalock:
|
||
serializable = {}
|
||
for ds_id, entry in self._stores.items():
|
||
if entry.get('type') != 'dataset':
|
||
continue
|
||
meta = entry.get('metadata', {})
|
||
serializable[ds_id] = {
|
||
'schema': entry.get('schema'),
|
||
'metadata': meta,
|
||
'csv_glob': meta.get('csv_glob'),
|
||
}
|
||
path = _get_store_path()
|
||
path.write_text(
|
||
json.dumps(serializable, ensure_ascii=False, indent=2),
|
||
encoding='utf-8',
|
||
)
|
||
|
||
def restore_from_disk(self) -> int:
|
||
"""Re-register datasets from persisted metadata.
|
||
|
||
Attempts two strategies in order:
|
||
1. **SQLite table** — when ``metadata.sqlite_table`` exists (fast,
|
||
always available, no CSV dependency).
|
||
2. **CSV glob** — falls back to ``load_csv_directory`` for datasets
|
||
that were stored before the SQLite migration or whose SQLite table
|
||
was dropped.
|
||
|
||
Returns the number of successfully restored datasets.
|
||
"""
|
||
path = _get_store_path()
|
||
if not path.exists():
|
||
return 0
|
||
try:
|
||
data = json.loads(path.read_text(encoding='utf-8'))
|
||
except (json.JSONDecodeError, OSError):
|
||
return 0
|
||
|
||
# Lazy imports to avoid circular dependency at module level
|
||
from analysis.data_loader import load_csv_directory, load_from_db # fmt: skip
|
||
|
||
restored = 0
|
||
for ds_id, info in data.items():
|
||
meta = dict(info.get('metadata', {}))
|
||
schema = dict(info.get('schema', {}))
|
||
restored_ok = False
|
||
|
||
# Strategy 1: SQLite table (fast path)
|
||
sqlite_table = meta.get('sqlite_table')
|
||
if sqlite_table:
|
||
try:
|
||
lf = load_from_db(sqlite_table)
|
||
if lf is not None:
|
||
self.store_dataset(ds_id, lf, schema=schema, metadata=meta)
|
||
restored += 1
|
||
restored_ok = True
|
||
except Exception:
|
||
pass # fall through to CSV restore
|
||
|
||
if restored_ok:
|
||
continue
|
||
|
||
# Strategy 2: CSV glob (legacy / fallback)
|
||
csv_glob = info.get('csv_glob')
|
||
if not csv_glob:
|
||
continue
|
||
try:
|
||
lf, schema_csv, row_count, file_count, memory_mb = load_csv_directory(csv_glob)
|
||
# Merge schema from CSV (may be more accurate than persisted schema)
|
||
self.store_dataset(ds_id, lf, schema=schema_csv, metadata=meta)
|
||
restored += 1
|
||
except Exception:
|
||
# File(s) no longer exist or schema changed — skip gracefully
|
||
continue
|
||
return restored
|
||
|
||
# ── dataset operations ──────────────────────────────────────────────
|
||
|
||
def store_dataset(
|
||
self,
|
||
dataset_id: str,
|
||
lazyframe: pl.LazyFrame,
|
||
schema: Optional[dict] = None,
|
||
metadata: Optional[dict] = None,
|
||
parent_id: Optional[str] = None,
|
||
) -> None:
|
||
"""Store a LazyFrame under *dataset_id*.
|
||
|
||
Args:
|
||
dataset_id: Unique key for the dataset.
|
||
lazyframe: Polars LazyFrame (query plan, not materialised).
|
||
schema: Column schema dict {name: dtype_string}.
|
||
metadata: Optional dict (row_count, file_count, …).
|
||
parent_id: If this is a derived dataset (filter/clone/…), the
|
||
source dataset_id.
|
||
"""
|
||
with self._datalock:
|
||
self._stores[dataset_id] = {
|
||
'type': 'dataset',
|
||
'lazyframe': lazyframe,
|
||
'schema': schema or {},
|
||
'metadata': metadata or {},
|
||
'parent_id': parent_id,
|
||
}
|
||
self._save_to_disk()
|
||
|
||
def get_dataset(self, dataset_id: str) -> Optional[dict]:
|
||
"""Retrieve a stored dataset entry (or *None*)."""
|
||
with self._datalock:
|
||
entry = self._stores.get(dataset_id)
|
||
if entry and entry.get('type') == 'dataset':
|
||
return entry
|
||
return None
|
||
|
||
def drop_dataset(self, dataset_id: str) -> bool:
|
||
"""Remove a dataset from the store and force garbage collection.
|
||
|
||
Returns True if the dataset existed.
|
||
"""
|
||
with self._datalock:
|
||
existed = dataset_id in self._stores and self._stores[dataset_id].get('type') == 'dataset'
|
||
if existed:
|
||
del self._stores[dataset_id]
|
||
if existed:
|
||
gc.collect()
|
||
self._save_to_disk()
|
||
return existed
|
||
|
||
def clone_dataset(self, dataset_id: str) -> Optional[str]:
|
||
"""Shallow clone — LazyFrame is a query plan so no data is copied.
|
||
|
||
Returns the new dataset_id, or *None* if source does not exist.
|
||
"""
|
||
with self._datalock:
|
||
entry = self.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return None
|
||
new_id = _generate_id('clone')
|
||
self._stores[new_id] = {
|
||
'type': 'dataset',
|
||
'lazyframe': entry['lazyframe'], # same plan pointer
|
||
'schema': dict(entry['schema']),
|
||
'metadata': dict(entry.get('metadata', {})),
|
||
'parent_id': dataset_id,
|
||
}
|
||
return new_id
|
||
|
||
def list_datasets(self) -> list[dict]:
|
||
"""Return summary dict for every stored dataset."""
|
||
with self._datalock:
|
||
results = []
|
||
for ds_id, entry in self._stores.items():
|
||
if entry.get('type') != 'dataset':
|
||
continue
|
||
meta = entry.get('metadata', {})
|
||
results.append({
|
||
'dataset_id': ds_id,
|
||
'parent_id': entry.get('parent_id'),
|
||
'row_count': meta.get('row_count'),
|
||
'file_count': meta.get('file_count'),
|
||
'columns': list(entry.get('schema', {}).keys()),
|
||
'column_count': len(entry.get('schema', {})),
|
||
})
|
||
return results
|
||
|
||
def memory_estimate_mb(self, dataset_id: str) -> float:
|
||
"""Rough memory estimate based on schema × row_count."""
|
||
with self._datalock:
|
||
entry = self.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return 0.0
|
||
schema = entry.get('schema', {})
|
||
row_count = entry.get('metadata', {}).get('row_count', 0)
|
||
if not row_count or not schema:
|
||
return 0.0
|
||
bytes_per_row = sum(_estimate_dtype_bytes(dtype) for dtype in schema.values())
|
||
return (bytes_per_row * row_count) / (1024 * 1024)
|
||
|
||
# ── cluster-result operations ───────────────────────────────────────
|
||
|
||
def store_cluster_result(
|
||
self,
|
||
result_id: str,
|
||
labels: list[int],
|
||
n_clusters: int,
|
||
params: dict[str, Any],
|
||
parent_dataset_id: str,
|
||
) -> None:
|
||
"""Store a clustering result."""
|
||
with self._datalock:
|
||
self._stores[result_id] = {
|
||
'type': 'cluster_result',
|
||
'labels': labels,
|
||
'n_clusters': n_clusters,
|
||
'params': params,
|
||
'parent_dataset_id': parent_dataset_id,
|
||
}
|
||
|
||
def get_cluster_result(self, result_id: str) -> Optional[dict]:
|
||
"""Retrieve a stored cluster result."""
|
||
with self._datalock:
|
||
entry = self._stores.get(result_id)
|
||
if entry and entry.get('type') == 'cluster_result':
|
||
return entry
|
||
return None
|
||
|
||
def drop_cluster_result(self, result_id: str) -> bool:
|
||
"""Remove a cluster result."""
|
||
with self._datalock:
|
||
existed = result_id in self._stores and self._stores[result_id].get('type') == 'cluster_result'
|
||
if existed:
|
||
del self._stores[result_id]
|
||
return existed
|
||
|
||
# ── feature-result operations ───────────────────────────────────────
|
||
|
||
def store_feature_result(
|
||
self,
|
||
result_id: str,
|
||
features: list[dict[str, Any]],
|
||
parent_cluster_result_id: str,
|
||
) -> None:
|
||
"""Store extracted features per cluster."""
|
||
with self._datalock:
|
||
self._stores[result_id] = {
|
||
'type': 'feature_result',
|
||
'features': features,
|
||
'parent_cluster_result_id': parent_cluster_result_id,
|
||
}
|
||
|
||
def get_feature_result(self, result_id: str) -> Optional[dict]:
|
||
"""Retrieve a stored feature result."""
|
||
with self._datalock:
|
||
entry = self._stores.get(result_id)
|
||
if entry and entry.get('type') == 'feature_result':
|
||
return entry
|
||
return None
|
||
|
||
def drop_feature_result(self, result_id: str) -> bool:
|
||
"""Remove a feature result."""
|
||
with self._datalock:
|
||
existed = result_id in self._stores and self._stores[result_id].get('type') == 'feature_result'
|
||
if existed:
|
||
del self._stores[result_id]
|
||
return existed
|
||
|
||
# ── general ─────────────────────────────────────────────────────────
|
||
|
||
def drop_all(self) -> int:
|
||
"""Drop every entry and run gc. Returns count of removed entries."""
|
||
with self._datalock:
|
||
count = len(self._stores)
|
||
self._stores.clear()
|
||
if count:
|
||
gc.collect()
|
||
return count
|
||
|
||
def __repr__(self) -> str:
|
||
with self._datalock:
|
||
types = {}
|
||
for entry in self._stores.values():
|
||
t = entry.get('type', 'unknown')
|
||
types[t] = types.get(t, 0) + 1
|
||
parts = [f"{k}={v}" for k, v in sorted(types.items())]
|
||
return f"<SessionStore {'; '.join(parts)}>"
|