71 lines
1.9 KiB
Python
71 lines
1.9 KiB
Python
"""db_xp.py – minimal user helper"""
|
||
|
||
from __future__ import annotations
|
||
from datetime import datetime, timezone
|
||
from typing import Dict, Any
|
||
|
||
from sqlalchemy import (
|
||
Table, MetaData, select, func, insert as sa_insert, text, inspect
|
||
)
|
||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||
|
||
from backend.core.db import SessionLocal, engine
|
||
|
||
_NOW = lambda: datetime.now(timezone.utc)
|
||
metadata = MetaData()
|
||
|
||
_IS_PG = engine.url.get_backend_name().startswith("postgres")
|
||
|
||
def _get_users_table() -> Table:
|
||
return Table("users", metadata, autoload_with=engine)
|
||
|
||
def _get_column_info():
|
||
try:
|
||
insp = inspect(engine)
|
||
cols = {c["name"]: c for c in insp.get_columns("users")}
|
||
return cols
|
||
except Exception:
|
||
return {}
|
||
|
||
def _insert_ignore(**vals):
|
||
users = _get_users_table()
|
||
if _IS_PG:
|
||
return (
|
||
pg_insert(users)
|
||
.values(**vals)
|
||
.on_conflict_do_nothing(index_elements=["ip"])
|
||
)
|
||
return sa_insert(users).values(**vals).prefix_with("OR IGNORE")
|
||
|
||
def ensure_user(ip: str) -> None:
|
||
cols = _get_column_info()
|
||
has_data = (
|
||
"data" in cols and
|
||
not cols["data"].get("nullable", True) # NOT NULL
|
||
)
|
||
|
||
vals = dict(
|
||
ip=ip,
|
||
first_visit=_NOW(),
|
||
ban_status=False,
|
||
soft_banned=False,
|
||
)
|
||
if has_data:
|
||
vals["data"] = {}
|
||
|
||
stmt = _insert_ignore(**vals)
|
||
with SessionLocal.begin() as s:
|
||
s.execute(stmt)
|
||
|
||
def is_ip_banned(ip: str) -> bool:
|
||
users = _get_users_table()
|
||
with SessionLocal() as s:
|
||
return bool(s.scalar(select(users.c.ban_status).where(users.c.ip == ip)))
|
||
|
||
def get_status(ip: str) -> Dict[str, Any]:
|
||
ensure_user(ip)
|
||
users = _get_users_table()
|
||
with SessionLocal() as s:
|
||
soft = s.scalar(select(users.c.soft_banned).where(users.c.ip == ip))
|
||
return {"soft_banned": bool(soft)}
|