144 lines
4.6 KiB
Python
144 lines
4.6 KiB
Python
"""Monthly partition maintenance for `events` (PARTITION BY RANGE(received_at)).
|
|
|
|
Idempotent: creates current + N future month partitions and drops old ones beyond
|
|
the retention window. Names are derived from year/month, so they cannot inject SQL.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from datetime import UTC, date, datetime
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker
|
|
|
|
from ..logging_config import get_logger
|
|
from ..settings import get_settings
|
|
|
|
log = get_logger("monlet.partitioning")
|
|
|
|
|
|
def _month_start(d: date) -> date:
|
|
return date(d.year, d.month, 1)
|
|
|
|
|
|
def _next_month(d: date) -> date:
|
|
year = d.year + (1 if d.month == 12 else 0)
|
|
month = 1 if d.month == 12 else d.month + 1
|
|
return date(year, month, 1)
|
|
|
|
|
|
def _shift_months(d: date, n: int) -> date:
|
|
y, m = d.year, d.month
|
|
total = y * 12 + (m - 1) + n
|
|
return date(total // 12, (total % 12) + 1, 1)
|
|
|
|
|
|
def _partition_name(d: date) -> str:
|
|
return f"events_p{d.year:04d}_{d.month:02d}"
|
|
|
|
|
|
async def ensure_event_partitions(engine_or_sm: AsyncEngine | async_sessionmaker) -> None:
|
|
"""Create current + N future month partitions for events. Idempotent."""
|
|
s = get_settings()
|
|
today = _month_start(datetime.now(UTC).date())
|
|
months = [_shift_months(today, i) for i in range(0, max(s.events_future_partitions, 0) + 1)]
|
|
|
|
statements: list[str] = []
|
|
for start in months:
|
|
end = _next_month(start)
|
|
# Explicit UTC qualifier — Postgres otherwise interprets bare date literals via
|
|
# the session TimeZone, which would shift partition bounds on non-UTC servers.
|
|
statements.append(
|
|
f"CREATE TABLE IF NOT EXISTS {_partition_name(start)} "
|
|
f"PARTITION OF events FOR VALUES FROM ('{start.isoformat()} UTC') "
|
|
f"TO ('{end.isoformat()} UTC');"
|
|
)
|
|
|
|
async with _connect(engine_or_sm) as conn:
|
|
for stmt in statements:
|
|
await conn.execute(text(stmt))
|
|
|
|
|
|
async def drop_old_event_partitions(engine_or_sm: AsyncEngine | async_sessionmaker) -> None:
|
|
"""Drop month partitions older than events_retention_months.
|
|
|
|
Each DROP runs in its own transaction so the ACCESS EXCLUSIVE lock on parent
|
|
`events` is held only for one partition at a time and ingest is not blocked
|
|
by a long maintenance batch.
|
|
"""
|
|
s = get_settings()
|
|
if s.events_retention_months <= 0:
|
|
return
|
|
cutoff = _shift_months(_month_start(datetime.now(UTC).date()), -s.events_retention_months)
|
|
|
|
engine = _resolve_engine(engine_or_sm)
|
|
async with engine.connect() as conn:
|
|
rows = (
|
|
await conn.execute(
|
|
text(
|
|
"""
|
|
SELECT c.relname
|
|
FROM pg_inherits i
|
|
JOIN pg_class c ON c.oid = i.inhrelid
|
|
JOIN pg_class p ON p.oid = i.inhparent
|
|
WHERE p.relname = 'events'
|
|
"""
|
|
)
|
|
)
|
|
).all()
|
|
|
|
to_drop: list[str] = []
|
|
for (name,) in rows:
|
|
if not name.startswith("events_p"):
|
|
continue
|
|
try:
|
|
y = int(name[8:12])
|
|
m = int(name[13:15])
|
|
except ValueError:
|
|
continue
|
|
if date(y, m, 1) < cutoff:
|
|
to_drop.append(name)
|
|
|
|
for name in to_drop:
|
|
async with engine.begin() as conn:
|
|
await conn.execute(text(f"DROP TABLE IF EXISTS {name};"))
|
|
log.info("event_partition_dropped", name=name)
|
|
|
|
|
|
def _resolve_engine(engine_or_sm: AsyncEngine | async_sessionmaker) -> AsyncEngine:
|
|
if isinstance(engine_or_sm, async_sessionmaker):
|
|
return engine_or_sm.kw["bind"]
|
|
return engine_or_sm
|
|
|
|
|
|
def _connect(engine_or_sm):
|
|
"""Return an async ctx for a transactional connection from either engine or sessionmaker."""
|
|
return _resolve_engine(engine_or_sm).begin()
|
|
|
|
|
|
async def run_partition_maintenance(sm: async_sessionmaker, stop_event: asyncio.Event) -> None:
|
|
"""Independent worker — runs even when MONLET_ENABLE_DETECTOR=false."""
|
|
interval = max(get_settings().partition_maintenance_interval_sec, 60)
|
|
while not stop_event.is_set():
|
|
try:
|
|
await ensure_event_partitions(sm)
|
|
await drop_old_event_partitions(sm)
|
|
except Exception as exc:
|
|
log.warning("partition_maintenance_failed", error=str(exc))
|
|
try:
|
|
await asyncio.wait_for(stop_event.wait(), timeout=interval)
|
|
except TimeoutError:
|
|
continue
|
|
|
|
|
|
# Re-export helpers for tests.
|
|
__all__ = [
|
|
"run_partition_maintenance",
|
|
"ensure_event_partitions",
|
|
"drop_old_event_partitions",
|
|
"_partition_name",
|
|
"_shift_months",
|
|
"_month_start",
|
|
]
|