"""Unit of Work — wraps a SQLAlchemy session for explicit transaction control. Usage in routers:: with UnitOfWork(db) as uow: service_a(db, ...) service_b(db, ...) uow.commit() # single commit for the entire operation If an exception propagates, ``__exit__`` issues a rollback automatically. Services should **never** call ``db.commit()``; they use ``db.add()`` / ``db.flush()`` to stage work and let the caller decide when to commit. **Documented exceptions** (services that may commit internally): - Import services (atomic_import, sigma_import, etc.) — self-contained sync ops. - Background jobs (campaign_scheduler, intel_service, stale_detection, mitre_sync) — self-contained operations. - Self-contained batch ops (e.g. detection_rule_service.auto_associate_rules, snapshot_service.create_snapshot, campaign_service.generate_campaign_from_*, osint_enrichment_service.enrich_technique_with_cves). """ from __future__ import annotations from sqlalchemy.orm import Session class UnitOfWork: """Lightweight transaction wrapper around an existing SQLAlchemy session.""" def __init__(self, session: Session) -> None: self._session = session # -- context manager ----------------------------------------------------- def __enter__(self) -> "UnitOfWork": return self def __exit__(self, exc_type, exc_val, exc_tb) -> None: if exc_type is not None: self.rollback() # -- public API ---------------------------------------------------------- def commit(self) -> None: """Flush pending changes and commit the transaction.""" self._session.commit() def rollback(self) -> None: """Roll back the current transaction.""" self._session.rollback() def flush(self) -> None: """Flush pending changes without committing (useful for getting IDs).""" self._session.flush()