JQC_multi_tenant/app/tenancy/quota.py

"""
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.reported_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 flask import g, current_app

from app.utils.time_utils import now_eastern

logger = logging.getLogger(__name__)


def _mt_enabled():
    return current_app.config.get('MULTI_TENANT_ENABLED', False)


def _month_window():
    """[start, end) of the current month, in the timezone the rows are stamped in.

    now_eastern(), not datetime.now(): every timestamp in the tenant DB is
    written by now_eastern() (rule 2). On a UTC server the two differ by 4-5
    hours, so a plain now() puts the month boundary in the wrong place and the
    first hours of each month count the wrong rows — a discrepancy that only
    appears on the 1st and is gone before anyone investigates it.
    """
    now = now_eastern()
    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():
    """Issues filed this month.

    `reported_at`, NOT `created_at` — the issues table has no created_at
    column. (IssueComment does, in the same module, which is how the wrong name
    got here.) Referencing a missing column raises AttributeError while the
    query is built, and every caller wraps this in a try/except, so the failure
    was invisible: the plan page silently showed 0 for EVERY axis and the
    issues quota was never evaluated at all.
    """
    from app.models.issue import Issue
    start, end = _month_window()
    return (Issue.query
            .filter(Issue.reported_at >= start,
                    Issue.reported_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)