artemis/core/data/schema/base.py

124 lines
4.0 KiB
Python
Raw Normal View History

import json
import logging
from random import randrange
from typing import Any, Optional, Dict, List
2023-03-18 06:12:58 +00:00
from sqlalchemy.engine import Row
from sqlalchemy.engine.cursor import CursorResult
from sqlalchemy.engine.base import Connection
from sqlalchemy.sql import text, func, select
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy import MetaData, Table, Column
2024-05-22 03:05:22 +00:00
from sqlalchemy.types import Integer, String, TIMESTAMP, JSON, INTEGER, TEXT
from sqlalchemy.schema import ForeignKey
from sqlalchemy.dialects.mysql import insert
from core.config import CoreConfig
metadata = MetaData()
event_log = Table(
"event_log",
metadata,
Column("id", Integer, primary_key=True, nullable=False),
Column("system", String(255), nullable=False),
Column("type", String(255), nullable=False),
Column("severity", Integer, nullable=False),
2024-05-22 03:05:22 +00:00
Column("user", INTEGER, ForeignKey("aime_user.id", ondelete="cascade", onupdate="cascade")),
Column("arcade", INTEGER, ForeignKey("arcade.id", ondelete="cascade", onupdate="cascade")),
Column("machine", INTEGER, ForeignKey("machine.id", ondelete="cascade", onupdate="cascade")),
Column("ip", TEXT(39)),
2023-03-04 00:56:12 +00:00
Column("message", String(1000), nullable=False),
Column("details", JSON, nullable=False),
Column("when_logged", TIMESTAMP, nullable=False, server_default=func.now()),
2023-03-09 16:38:58 +00:00
mysql_charset="utf8mb4",
)
2023-03-09 16:38:58 +00:00
class BaseData:
def __init__(self, cfg: CoreConfig, conn: Connection) -> None:
self.config = cfg
self.conn = conn
self.logger = logging.getLogger("database")
2023-03-09 16:38:58 +00:00
2024-01-09 18:57:59 +00:00
async def execute(self, sql: str, opts: Dict[str, Any] = {}) -> Optional[CursorResult]:
res = None
try:
2024-01-09 18:57:59 +00:00
self.logger.debug(f"SQL Execute: {''.join(str(sql).splitlines())}")
res = self.conn.execute(text(sql), opts)
except SQLAlchemyError as e:
self.logger.error(f"SQLAlchemy error {e}")
return None
2023-03-09 16:38:58 +00:00
except UnicodeEncodeError as e:
self.logger.error(f"UnicodeEncodeError error {e}")
return None
2023-07-16 20:58:34 +00:00
except Exception:
try:
res = self.conn.execute(sql, opts)
except SQLAlchemyError as e:
self.logger.error(f"SQLAlchemy error {e}")
return None
2023-03-09 16:38:58 +00:00
except UnicodeEncodeError as e:
self.logger.error(f"UnicodeEncodeError error {e}")
return None
2023-07-16 20:58:34 +00:00
except Exception:
self.logger.error(f"Unknown error")
raise
return res
2023-03-09 16:38:58 +00:00
def generate_id(self) -> int:
"""
Generate a random 5-7 digit id
"""
return randrange(10000, 9999999)
2024-01-09 18:57:59 +00:00
async def log_event(
2024-05-22 03:05:22 +00:00
self, system: str, type: str, severity: int, message: str, details: Dict = {}, user: int = None, arcade: int = None, machine: int = None, ip: str = None
2023-03-09 16:38:58 +00:00
) -> Optional[int]:
sql = event_log.insert().values(
system=system,
type=type,
severity=severity,
2024-05-22 03:05:22 +00:00
user=user,
arcade=arcade,
machine=machine,
ip=ip,
2023-03-09 16:38:58 +00:00
message=message,
details=json.dumps(details),
)
2024-01-09 19:42:17 +00:00
result = await self.execute(sql)
if result is None:
2023-03-09 16:38:58 +00:00
self.logger.error(
f"{__name__}: Failed to insert event into event log! system = {system}, type = {type}, severity = {severity}, message = {message}"
)
return None
return result.lastrowid
2023-03-09 16:38:58 +00:00
2024-01-09 18:57:59 +00:00
async def get_event_log(self, entries: int = 100) -> Optional[List[Dict]]:
sql = event_log.select().limit(entries).all()
2024-01-09 19:42:17 +00:00
result = await self.execute(sql)
2023-03-09 16:38:58 +00:00
if result is None:
return None
return result.fetchall()
2023-03-09 16:38:58 +00:00
def fix_bools(self, data: Dict) -> Dict:
2023-03-09 16:38:58 +00:00
for k, v in data.items():
if k == "userName" or k == "teamName":
continue
if type(v) == str and v.lower() == "true":
data[k] = True
elif type(v) == str and v.lower() == "false":
data[k] = False
2023-03-09 16:38:58 +00:00
return data