99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
"""
|
|
db.py — MySQL/MariaDB connection for NetAdmin (PyMySQL).
|
|
All tables live in the shared 'netadmin' database.
|
|
Rows returned as dicts. Interface compatible with sqlite3 usage patterns.
|
|
"""
|
|
|
|
import logging
|
|
from contextlib import contextmanager
|
|
from typing import Any
|
|
|
|
import pymysql
|
|
import pymysql.cursors
|
|
|
|
from config import cfg
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class _ConnWrapper:
|
|
"""
|
|
Thin wrapper around a PyMySQL connection that adds .execute() / .fetchone()
|
|
/ .fetchall() methods directly on the connection object — matching the sqlite3
|
|
interface used throughout the codebase.
|
|
"""
|
|
|
|
def __init__(self, conn):
|
|
self._conn = conn
|
|
self._cur = conn.cursor()
|
|
|
|
# ── Cursor proxy ─────────────────────────────────────────────────────
|
|
def execute(self, sql: str, params=None):
|
|
self._cur.execute(sql, params or ())
|
|
return self._cur
|
|
|
|
def fetchone(self):
|
|
return self._cur.fetchone()
|
|
|
|
def fetchall(self):
|
|
return self._cur.fetchall()
|
|
|
|
# ── Transaction proxy ─────────────────────────────────────────────────
|
|
def commit(self):
|
|
self._conn.commit()
|
|
|
|
def rollback(self):
|
|
self._conn.rollback()
|
|
|
|
def close(self):
|
|
try:
|
|
self._cur.close()
|
|
except Exception:
|
|
pass
|
|
self._conn.close()
|
|
|
|
|
|
def _connect() -> _ConnWrapper:
|
|
db = cfg.database
|
|
raw = pymysql.connect(
|
|
host = db.host,
|
|
port = db.port,
|
|
user = db.user,
|
|
password = db.password,
|
|
database = db.name,
|
|
charset = "utf8mb4",
|
|
cursorclass = pymysql.cursors.DictCursor,
|
|
autocommit = False,
|
|
)
|
|
return _ConnWrapper(raw)
|
|
|
|
|
|
@contextmanager
|
|
def get_db():
|
|
"""Yield a connection wrapper. Commits on success, rolls back on exception."""
|
|
conn = _connect()
|
|
try:
|
|
yield conn
|
|
conn.commit()
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ── Datetime helper ────────────────────────────────────────────────────────
|
|
|
|
def _dt(v) -> str | None:
|
|
"""Convert datetime/date → ISO string. Pass str/None through unchanged.
|
|
|
|
MariaDB returns DATETIME columns as datetime.datetime objects via PyMySQL.
|
|
Pydantic models expect str for these fields.
|
|
"""
|
|
import datetime as _datetime
|
|
if v is None:
|
|
return None
|
|
if isinstance(v, (_datetime.datetime, _datetime.date)):
|
|
return v.isoformat()
|
|
return str(v) if v else None
|