Jun 27 MT-5

app/tenancy/quota.py

@@ -0,0 +1,156 @@
+"""
+app/tenancy/quota.py
+---------------------
+Live quota helpers for MT-5.
+
+Counts current-month rows directly in the tenant DB — no counter table, always
+accurate. Plan limits are read from g.tenant (loaded once by the resolver) so
+no second control-DB round-trip is needed.
+
+Quota axes:
+  inspections  → Inspection.inspection_date in current month, status='completed'
+  issues       → Issue.created_at in current month
+  users        → User.active == True  (total, not monthly)
+  facilities   → Facility.active == True  (total, not monthly)
+
+Returns None when MULTI_TENANT_ENABLED is False or g.tenant is absent —
+all quota checks pass, single-tenant behaviour unchanged.
+"""
+
+import logging
+from datetime import datetime
+
+from flask import g, current_app
+
+logger = logging.getLogger(__name__)
+
+
+def _mt_enabled():
+    return current_app.config.get('MULTI_TENANT_ENABLED', False)
+
+
+def _month_window():
+    now = datetime.now()
+    start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
+    if now.month == 12:
+        end = now.replace(year=now.year + 1, month=1, day=1,
+                          hour=0, minute=0, second=0, microsecond=0)
+    else:
+        end = now.replace(month=now.month + 1, day=1,
+                          hour=0, minute=0, second=0, microsecond=0)
+    return start, end
+
+
+def count_inspections_this_month():
+    from app.models.inspection import Inspection
+    start, end = _month_window()
+    return (Inspection.query
+            .filter(Inspection.status == 'completed',
+                    Inspection.inspection_date >= start,
+                    Inspection.inspection_date < end)
+            .count())
+
+
+def count_issues_this_month():
+    from app.models.issue import Issue
+    start, end = _month_window()
+    return (Issue.query
+            .filter(Issue.created_at >= start,
+                    Issue.created_at < end)
+            .count())
+
+
+def count_active_users():
+    from app.models.user import User
+    return User.query.filter_by(active=True).count()
+
+
+def count_active_facilities():
+    from app.models.facility import Facility
+    return Facility.query.filter_by(active=True).count()
+
+
+# ── Public interface ──────────────────────────────────────────────────────────
+
+class QuotaStatus:
+    __slots__ = ('axis', 'limit', 'current', 'exceeded')
+
+    def __init__(self, axis, limit, current):
+        self.axis = axis
+        self.limit = limit
+        self.current = current
+        self.exceeded = (limit is not None and current >= limit)
+
+    def __repr__(self):
+        return (f'<QuotaStatus {self.axis} {self.current}/{self.limit} '
+                f'exceeded={self.exceeded}>')
+
+
+def check_quota(axis: str):
+    """
+    Check one quota axis. Returns QuotaStatus or None (MT disabled / no tenant).
+    Reads plan limits from g.tenant — no extra DB query.
+    """
+    if not _mt_enabled():
+        return None
+    tenant = getattr(g, 'tenant', None)
+    if tenant is None:
+        return None
+
+    _limit_attrs = {
+        'inspections': 'max_inspections_month',
+        'issues':      'max_issues_month',
+        'users':       'max_users',
+        'facilities':  'max_facilities',
+    }
+    _counter_fns = {
+        'inspections': count_inspections_this_month,
+        'issues':      count_issues_this_month,
+        'users':       count_active_users,
+        'facilities':  count_active_facilities,
+    }
+    if axis not in _limit_attrs:
+        logger.warning('quota.check_quota: unknown axis %r', axis)
+        return None
+
+    limit = getattr(tenant, _limit_attrs[axis], None)
+    if limit is None:
+        # Unlimited — still return a status so caller can display current usage
+        try:
+            current = _counter_fns[axis]()
+        except Exception as exc:
+            logger.error('quota.check_quota count failed axis=%s err=%s', axis, exc)
+            return None
+        return QuotaStatus(axis=axis, limit=None, current=current)
+
+    try:
+        current = _counter_fns[axis]()
+    except Exception as exc:
+        logger.error('quota.check_quota count failed axis=%s err=%s', axis, exc)
+        return None
+
+    return QuotaStatus(axis=axis, limit=limit, current=current)
+
+
+def check_feature(feature_key: str) -> bool:
+    """
+    Return True if the current tenant's plan allows feature_key.
+    Reads from g.tenant — no extra DB query.
+    Returns True when MT is disabled.
+    """
+    if not _mt_enabled():
+        return True
+    tenant = getattr(g, 'tenant', None)
+    if tenant is None:
+        return True
+
+    _gate_attrs = {
+        'mobile_api':        'allow_mobile_api',
+        'scheduled_reports': 'allow_scheduled_reports',
+        'branding':          'allow_branding',
+        'custom_domain':     'allow_custom_domain',
+    }
+    attr = _gate_attrs.get(feature_key)
+    if attr is None:
+        return True   # unknown feature → fail open
+    return getattr(tenant, attr, True)