"""WESP-compatible model aliases with enterprise_id context for ported lab code.""" from __future__ import annotations import uuid from contextlib import contextmanager from contextvars import ContextVar, Token from datetime import UTC, datetime from typing import Any, Generic, Iterator, TypeVar from sqlalchemy import delete, func, select from app.modules.zootech.lab_models import ( ZootechLabAnimalProfile, ZootechLabComponentNutrientValue, ZootechLabProfileNorm, ZootechLabRacionNormyInfo, ZootechLabRacionNormyMoskwa, ZootechLabRacionNormyMoskwaMeta, ZootechLabRacionNormyPiter, ZootechLabRacionNormyPiterMeta, ZootechLabRationCalcIndicator, ZootechLabRationCalcTotal, ZootechLabRationCompoundLine, ZootechLabRationLine, ZootechLabRecipeRation, ) from app.modules.zootech.models import ZootechComponent, ZootechIngredient, ZootechRecipe from app.modules.zootech.wesp_bridge_db import db _enterprise_id: ContextVar[str | None] = ContextVar("wesp_enterprise_id", default=None) WESP_SUPPRESS_SYNC_ENQUEUE = "wesp_suppress_sync_enqueue" M = TypeVar("M") def default_uuid() -> str: return str(uuid.uuid4()) class TimestampMixin: """Legacy mixin marker for ported lab model modules.""" class SoftDeleteMixin: """Legacy mixin marker; soft_delete lives on Zootech lab rows.""" def soft_delete(self, deleted_by_user: str = "system", **_: Any) -> None: if getattr(self, "is_deleted", False): return self.is_deleted = True if hasattr(self, "deleted_at"): self.deleted_at = datetime.now(UTC) if hasattr(self, "deleted_by"): self.deleted_by = deleted_by_user def get_enterprise_id() -> str: ent = _enterprise_id.get() if not ent: raise RuntimeError("wesp_bridge_models: enterprise_id not set — use wesp_enterprise()") return ent def set_enterprise_id(enterprise_id: str) -> Token: return _enterprise_id.set(enterprise_id) def reset_enterprise_id(token: Token) -> None: _enterprise_id.reset(token) @contextmanager def wesp_enterprise(enterprise_id: str) -> Iterator[None]: token = set_enterprise_id(enterprise_id) try: yield finally: reset_enterprise_id(token) def _model_has_enterprise(model: type) -> bool: return hasattr(model, "enterprise_id") class _BridgeQuery(Generic[M]): __slots__ = ("_model", "_stmt") def __init__(self, model: type[M], stmt=None) -> None: self._model = model self._stmt = stmt def _base(self): if self._stmt is not None: return self._stmt stmt = select(self._model) if _model_has_enterprise(self._model): stmt = stmt.where(self._model.enterprise_id == get_enterprise_id()) # type: ignore[attr-defined] return stmt def filter(self, *criteria) -> _BridgeQuery[M]: stmt = self._base() for crit in criteria: stmt = stmt.where(crit) return _BridgeQuery(self._model, stmt) def filter_by(self, **kwargs: Any) -> _BridgeQuery[M]: stmt = self._base() for key, value in kwargs.items(): stmt = stmt.where(getattr(self._model, key) == value) return _BridgeQuery(self._model, stmt) def order_by(self, *clauses) -> _BridgeQuery[M]: base = self._stmt if self._stmt is not None else self._base() return _BridgeQuery(self._model, base.order_by(*clauses)) def all(self) -> list[M]: return list(db.session.scalars(self._base()).all()) def first(self) -> M | None: return db.session.scalar(self._base().limit(1)) def one_or_none(self) -> M | None: return self.first() def get(self, ident: str) -> M | None: if not hasattr(self._model, "id"): return None stmt = select(self._model).where(self._model.id == ident) # type: ignore[attr-defined] if _model_has_enterprise(self._model): stmt = stmt.where(self._model.enterprise_id == get_enterprise_id()) # type: ignore[attr-defined] return db.session.scalar(stmt) def count(self) -> int: subq = self._base().subquery() return int(db.session.scalar(select(func.count()).select_from(subq)) or 0) def delete(self, synchronize_session: bool = False) -> int: del synchronize_session rows = self.all() for row in rows: db.session.delete(row) return len(rows) def _strip_legacy_kwargs(kwargs: dict[str, Any], model: type) -> dict[str, Any]: out = dict(kwargs) for key in ("sync_timestamp", "sync_status"): out.pop(key, None) if model in (ZootechRecipe, ZootechComponent, ZootechIngredient): out.pop("created_by", None) out.pop("updated_by", None) return out def _enable_bridge(model: type[M]) -> type[M]: model.query = _BridgeQuery(model) # type: ignore[attr-defined] if not _model_has_enterprise(model): return model original_init = model.__init__ def _patched_init(self, **kwargs: Any) -> None: kwargs = _strip_legacy_kwargs(kwargs, model) if "enterprise_id" not in kwargs: kwargs["enterprise_id"] = get_enterprise_id() original_init(self, **kwargs) model.__init__ = _patched_init # type: ignore[method-assign] return model Recipe = _enable_bridge(ZootechRecipe) Component = _enable_bridge(ZootechComponent) Ingredient = _enable_bridge(ZootechIngredient) LabAnimalProfile = _enable_bridge(ZootechLabAnimalProfile) LabProfileNorm = _enable_bridge(ZootechLabProfileNorm) LabComponentNutrientValue = _enable_bridge(ZootechLabComponentNutrientValue) LabRecipeRation = _enable_bridge(ZootechLabRecipeRation) LabRationLine = _enable_bridge(ZootechLabRationLine) LabRationCalcTotal = _enable_bridge(ZootechLabRationCalcTotal) LabRationCalcIndicator = _enable_bridge(ZootechLabRationCalcIndicator) LabRationCompoundLine = _enable_bridge(ZootechLabRationCompoundLine) LabRacionNormyMoskwa = _enable_bridge(ZootechLabRacionNormyMoskwa) LabRacionNormyMoskwaMeta = _enable_bridge(ZootechLabRacionNormyMoskwaMeta) LabRacionNormyPiter = _enable_bridge(ZootechLabRacionNormyPiter) LabRacionNormyPiterMeta = _enable_bridge(ZootechLabRacionNormyPiterMeta) LabRacionNormyInfo = _enable_bridge(ZootechLabRacionNormyInfo) def _enqueue_recipe_children_sync(recipe_id: str) -> None: del recipe_id __all__ = [ "Recipe", "Component", "Ingredient", "LabAnimalProfile", "LabProfileNorm", "LabComponentNutrientValue", "LabRecipeRation", "LabRationLine", "LabRationCalcTotal", "LabRationCalcIndicator", "LabRationCompoundLine", "LabRacionNormyMoskwa", "LabRacionNormyMoskwaMeta", "LabRacionNormyPiter", "LabRacionNormyPiterMeta", "LabRacionNormyInfo", "default_uuid", "TimestampMixin", "SoftDeleteMixin", "WESP_SUPPRESS_SYNC_ENQUEUE", "get_enterprise_id", "set_enterprise_id", "wesp_enterprise", "_enqueue_recipe_children_sync", ]