287 lines
8.6 KiB
Python
287 lines
8.6 KiB
Python
"""Timezone-aware datetime utilities.
|
|
|
|
Central source for all local-time operations in the project.
|
|
Reads the configured timezone from clock_config table (default: Asia/Shanghai).
|
|
All backend modules should use these functions instead of bare datetime.now().
|
|
|
|
Two time domains:
|
|
- **Local time** (configurable offset) — for date boundaries, display timestamps,
|
|
user-facing schedule logic. Use now(), today(), strftime(), strptime().
|
|
- **Unix timestamps** (UTC, always) — for interval math, expiry, cooldowns.
|
|
Use ts_now(), ts_from_local(), ts_to_local().
|
|
|
|
The timezone is cached per-request via Flask app config to avoid DB lookups.
|
|
"""
|
|
import time as _time
|
|
import threading
|
|
from datetime import datetime, timedelta, timezone, date
|
|
|
|
import pytz
|
|
|
|
# Thread-local "current user" context (set by the scheduler while it processes
|
|
# a specific user's job; request threads resolve the user from the session).
|
|
_TLS = threading.local()
|
|
|
|
|
|
def set_thread_user(user_id):
|
|
"""Set the current user for this thread (scheduler per-user iteration)."""
|
|
_TLS.user_id = user_id
|
|
try:
|
|
del _TLS.tz
|
|
except AttributeError:
|
|
pass
|
|
|
|
|
|
def clear_thread_user():
|
|
try:
|
|
del _TLS.user_id
|
|
del _TLS.tz
|
|
except AttributeError:
|
|
pass
|
|
|
|
|
|
def get_thread_user():
|
|
return getattr(_TLS, "user_id", None)
|
|
|
|
|
|
def _ctx_user_id():
|
|
"""Resolve the current user id: request auth context first, then thread-local."""
|
|
try:
|
|
from flask import has_request_context, g
|
|
if has_request_context():
|
|
uid = getattr(g, "user_id", None)
|
|
if uid is not None:
|
|
return int(uid)
|
|
except Exception:
|
|
pass
|
|
return getattr(_TLS, "user_id", None) or 0
|
|
|
|
# ─── Timezone loading ───────────────────────
|
|
|
|
_DEFAULT_TZ_STR = "Asia/Shanghai"
|
|
_DEFAULT_UTC_OFFSET = timezone(timedelta(hours=8))
|
|
|
|
|
|
def _load_tz() -> "pytz.BaseTzInfo":
|
|
"""Load timezone from the current user's clock_config. Falls back to Asia/Shanghai."""
|
|
try:
|
|
from flask import current_app
|
|
db_path = current_app.config.get("DB_PATH", "data.db")
|
|
from backend.models import get_db
|
|
db = get_db(db_path)
|
|
row = db.execute(
|
|
"SELECT timezone FROM clock_config WHERE is_primary=1 AND user_id=?",
|
|
(_ctx_user_id(),)
|
|
).fetchone()
|
|
db.close()
|
|
if row and row["timezone"]:
|
|
tz_str = row["timezone"].strip()
|
|
return pytz.timezone(tz_str)
|
|
except Exception:
|
|
pass
|
|
return pytz.timezone(_DEFAULT_TZ_STR)
|
|
|
|
|
|
def get_tz() -> "pytz.BaseTzInfo":
|
|
"""Return the configured timezone object.
|
|
|
|
Cached per-request (Flask ``g``) and per-thread (scheduler TLS), NOT per
|
|
process: a timezone change made on one Gunicorn worker must be visible to
|
|
the next request on ANY worker immediately (a process-wide cache would
|
|
serve stale timezones to other workers).
|
|
"""
|
|
try:
|
|
from flask import has_request_context, g
|
|
if has_request_context():
|
|
cached = getattr(g, "_tz_cache", None)
|
|
if cached is not None:
|
|
return cached
|
|
tz = _load_tz()
|
|
g._tz_cache = tz
|
|
return tz
|
|
except Exception:
|
|
pass
|
|
# No request context (scheduler threads): cache per thread-user so each
|
|
# user's jobs use their own timezone without a DB hit per call.
|
|
tz = getattr(_TLS, "tz", None)
|
|
if tz is not None:
|
|
return tz
|
|
tz = _load_tz()
|
|
_TLS.tz = tz
|
|
return tz
|
|
|
|
|
|
def clear_tz_cache():
|
|
"""Force reload timezone on next call (called after timezone config changes)."""
|
|
try:
|
|
from flask import g
|
|
g.__dict__.pop("_tz_cache", None)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
del _TLS.tz
|
|
except AttributeError:
|
|
pass
|
|
|
|
|
|
def set_tz(tz_str: str) -> bool:
|
|
"""Set the timezone in clock_config DB and clear cache.
|
|
|
|
Args:
|
|
tz_str: IANA timezone name (e.g. "Asia/Shanghai", "America/New_York")
|
|
or UTC offset string (e.g. "UTC+8", "UTC-5")
|
|
|
|
Returns:
|
|
True on success, False on invalid timezone.
|
|
"""
|
|
# Normalize "UTC+X" / "UTC-X" to pytz-compatible
|
|
normalized = tz_str.strip()
|
|
if normalized.upper().startswith("UTC") or normalized.upper().startswith("UT"):
|
|
try:
|
|
cleaned = normalized.upper().replace("UTC", "").replace("UT", "").strip()
|
|
sign = 1
|
|
if cleaned.startswith("+"):
|
|
cleaned = cleaned[1:]
|
|
elif cleaned.startswith("-"):
|
|
sign = -1
|
|
cleaned = cleaned[1:]
|
|
hours = int(cleaned)
|
|
# Convert to an IANA-like name: pytz has FixedOffset but easiest is
|
|
# just use a fixed-offset timezone
|
|
import pytz
|
|
pytz.FixedOffset(sign * hours * 60) # FixedOffset expects minutes
|
|
except (ValueError, IndexError):
|
|
return False
|
|
else:
|
|
# IANA name — validate by trying to load
|
|
try:
|
|
import pytz
|
|
pytz.timezone(normalized)
|
|
except Exception:
|
|
return False
|
|
|
|
# Save to DB
|
|
try:
|
|
from flask import current_app
|
|
db_path = current_app.config.get("DB_PATH", "data.db")
|
|
from backend.models import get_db
|
|
db = get_db(db_path)
|
|
uid = _ctx_user_id()
|
|
db.execute(
|
|
"UPDATE clock_config SET timezone=? WHERE is_primary=1 AND user_id=?",
|
|
(normalized, uid)
|
|
)
|
|
if db.total_changes == 0:
|
|
db.execute(
|
|
"INSERT INTO clock_config (user_id, city, timezone, is_primary) VALUES (?, '', ?, 1)",
|
|
(uid, normalized)
|
|
)
|
|
db.commit()
|
|
db.close()
|
|
clear_tz_cache()
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def get_tz_str() -> str:
|
|
"""Return the configured timezone as a string."""
|
|
try:
|
|
from flask import current_app
|
|
db_path = current_app.config.get("DB_PATH", "data.db")
|
|
from backend.models import get_db
|
|
db = get_db(db_path)
|
|
row = db.execute(
|
|
"SELECT timezone FROM clock_config WHERE is_primary=1 AND user_id=?",
|
|
(_ctx_user_id(),)
|
|
).fetchone()
|
|
db.close()
|
|
if row and row["timezone"]:
|
|
return row["timezone"]
|
|
except Exception:
|
|
pass
|
|
return _DEFAULT_TZ_STR
|
|
|
|
|
|
# ─── Local time functions ───────────────────
|
|
# Use these for date boundaries, display timestamps, schedule logic.
|
|
|
|
def now() -> datetime:
|
|
"""Return current datetime in the configured timezone.
|
|
|
|
This is the primary replacement for bare `datetime.now()`.
|
|
"""
|
|
return datetime.now(get_tz())
|
|
|
|
|
|
def today() -> date:
|
|
"""Return current date in the configured timezone."""
|
|
return now().date()
|
|
|
|
|
|
def today_str(fmt: str = "%Y%m%d") -> str:
|
|
"""Return today's date formatted as string.
|
|
|
|
Common format: "%Y%m%d" for DB comparisons, "%Y-%m-%d" for displays.
|
|
"""
|
|
return now().strftime(fmt)
|
|
|
|
|
|
def strftime(fmt: str = "%Y-%m-%d %H:%M:%S") -> str:
|
|
"""Return current datetime formatted as string in local timezone."""
|
|
return now().strftime(fmt)
|
|
|
|
|
|
def strptime(date_string: str, fmt: str) -> datetime:
|
|
"""Parse a string into a timezone-aware datetime.
|
|
|
|
The result is in the configured timezone. If the input string
|
|
has no timezone info, it's assumed to be in local time.
|
|
"""
|
|
dt = datetime.strptime(date_string, fmt)
|
|
return get_tz().localize(dt)
|
|
|
|
|
|
def isoformat() -> str:
|
|
"""Return current time in ISO 8601 format with timezone offset."""
|
|
return now().isoformat()
|
|
|
|
|
|
def localize(dt: datetime) -> datetime:
|
|
"""Convert a naive datetime to timezone-aware in configured timezone.
|
|
|
|
If dt is already timezone-aware, it's converted to local timezone.
|
|
"""
|
|
tz = get_tz()
|
|
if dt.tzinfo is None:
|
|
return tz.localize(dt)
|
|
return dt.astimezone(tz)
|
|
|
|
|
|
# ─── Unix timestamp utilities ───────────────
|
|
# These are timezone-independent — always UTC for epoch math.
|
|
|
|
def ts_now() -> float:
|
|
"""Return current Unix timestamp (seconds since epoch, UTC).
|
|
|
|
Equivalent to time.time(). Use for expiry, cooldowns, interval math.
|
|
"""
|
|
return _time.time()
|
|
|
|
|
|
def ts_from_local(dt: datetime = None) -> float:
|
|
"""Convert a local-timezone datetime to Unix timestamp.
|
|
|
|
If dt is None, returns current timestamp (same as ts_now()).
|
|
"""
|
|
if dt is None:
|
|
return _time.time()
|
|
if dt.tzinfo is None:
|
|
dt = localize(dt)
|
|
return dt.timestamp()
|
|
|
|
|
|
def ts_to_local(ts: float) -> datetime:
|
|
"""Convert a Unix timestamp to local timezone datetime."""
|
|
return datetime.fromtimestamp(ts, get_tz())
|