Update database model + Save instances to database + add the option to add logs into the database
This commit is contained in:
TheophileDiot
parent
c87c3637db
commit
e56f96d04b
|
|
@ -2,7 +2,12 @@ from contextlib import contextmanager
|
|||
from copy import deepcopy
|
||||
from datetime import datetime
|
||||
from hashlib import sha256
|
||||
from logging import INFO, WARNING, Logger, getLogger
|
||||
from logging import (
|
||||
NOTSET,
|
||||
Logger,
|
||||
_levelToName,
|
||||
_nameToLevel,
|
||||
)
|
||||
import oracledb
|
||||
from os import _exit, getenv, listdir, makedirs
|
||||
from os.path import dirname, exists
|
||||
|
|
@ -24,6 +29,8 @@ from traceback import format_exc
|
|||
|
||||
from model import (
|
||||
Base,
|
||||
Instances,
|
||||
Logs,
|
||||
Plugins,
|
||||
Settings,
|
||||
Global_values,
|
||||
|
|
@ -55,10 +62,6 @@ class Database:
|
|||
self.__sql_session = None
|
||||
self.__sql_engine = None
|
||||
|
||||
getLogger("sqlalchemy.engine").setLevel(
|
||||
logger.level if logger.level != INFO else WARNING
|
||||
)
|
||||
|
||||
if not sqlalchemy_string:
|
||||
sqlalchemy_string = getenv(
|
||||
"DATABASE_URI", "sqlite:////var/lib/bunkerweb/db.sqlite3"
|
||||
|
|
@ -77,7 +80,6 @@ class Database:
|
|||
sqlalchemy_string,
|
||||
encoding="utf-8",
|
||||
future=True,
|
||||
logging_name="sqlalchemy.engine",
|
||||
)
|
||||
except ArgumentError:
|
||||
self.__logger.error(f"Invalid database URI: {sqlalchemy_string}")
|
||||
|
|
@ -1266,3 +1268,76 @@ class Database:
|
|||
.all()
|
||||
)
|
||||
}
|
||||
|
||||
def save_log(
|
||||
self,
|
||||
log: str,
|
||||
level: Tuple[str, int],
|
||||
component: str,
|
||||
) -> str:
|
||||
"""Save log."""
|
||||
with self.__db_session() as session:
|
||||
session.add(
|
||||
Logs(
|
||||
id=int(datetime.now().timestamp()),
|
||||
message=log,
|
||||
level=str(_levelToName[_nameToLevel.get(level, NOTSET)])
|
||||
if isinstance(level, str)
|
||||
else _levelToName.get(level, "NOTSET"),
|
||||
component=component,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
session.commit()
|
||||
except BaseException:
|
||||
return format_exc()
|
||||
|
||||
return ""
|
||||
|
||||
def add_instance(self, hostname: str, port: int, server_name: str) -> str:
|
||||
"""Add instance."""
|
||||
with self.__db_session() as session:
|
||||
db_instance = (
|
||||
session.query(Instances)
|
||||
.with_entities(Instances.hostname)
|
||||
.filter_by(hostname=hostname)
|
||||
.first()
|
||||
)
|
||||
|
||||
if db_instance is not None:
|
||||
return "An instance with the same hostname already exists."
|
||||
|
||||
session.add(
|
||||
Instances(hostname=hostname, port=int(port), server_name=server_name)
|
||||
)
|
||||
|
||||
try:
|
||||
session.commit()
|
||||
except BaseException:
|
||||
return f"An error occurred while adding the instance {hostname} (port: {port}, server name: {server_name}).\n{format_exc()}"
|
||||
|
||||
return ""
|
||||
|
||||
def update_instances(self, instances: List[Dict[str, Any]]) -> str:
|
||||
"""Update instances."""
|
||||
to_put = []
|
||||
with self.__db_session() as session:
|
||||
session.query(Instances).delete()
|
||||
|
||||
for instance in instances:
|
||||
to_put.append(
|
||||
Instances(
|
||||
hostname=instance["hostname"],
|
||||
port=instance["env"].get("API_HTTP_PORT", 5000),
|
||||
server_name=instance["env"].get("API_SERVER_NAME", "bwapi"),
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
session.add_all(to_put)
|
||||
session.commit()
|
||||
except BaseException:
|
||||
return format_exc()
|
||||
|
||||
return ""
|
||||
|
|
|
|||
|
|
@ -10,7 +10,6 @@ from sqlalchemy import (
|
|||
PrimaryKeyConstraint,
|
||||
SmallInteger,
|
||||
String,
|
||||
text,
|
||||
TIMESTAMP,
|
||||
)
|
||||
from sqlalchemy.orm import declarative_base, relationship
|
||||
|
|
@ -30,7 +29,15 @@ CUSTOM_CONFIGS_TYPES_ENUM = Enum(
|
|||
"stream_http",
|
||||
name="custom_configs_types_enum",
|
||||
)
|
||||
LOG_LEVELS_ENUM = Enum("DEBUG", "INFO", "WARNING", "ERROR", name="log_levels_enum")
|
||||
LOG_LEVELS_ENUM = Enum(
|
||||
"CRITICAL",
|
||||
"ERROR",
|
||||
"WARNING",
|
||||
"INFO",
|
||||
"DEBUG",
|
||||
"NOTSET",
|
||||
name="log_levels_enum",
|
||||
)
|
||||
INTEGRATIONS_ENUM = Enum(
|
||||
"Linux",
|
||||
"Docker",
|
||||
|
|
@ -265,6 +272,14 @@ class Logs(Base):
|
|||
component = Column(String(255), nullable=False)
|
||||
|
||||
|
||||
class Instances(Base):
|
||||
__tablename__ = "instances"
|
||||
|
||||
hostname = Column(String(255), primary_key=True)
|
||||
port = Column(Integer, nullable=False)
|
||||
server_name = Column(String(255), nullable=False)
|
||||
|
||||
|
||||
class Metadata(Base):
|
||||
__tablename__ = "metadata"
|
||||
|
||||
|
|
|
|||
|
|
@ -369,6 +369,25 @@ if __name__ == "__main__":
|
|||
sys_exit(1)
|
||||
else:
|
||||
logger.info("Config successfully saved to database")
|
||||
|
||||
if apis:
|
||||
for api in apis:
|
||||
endpoint_data = api.get_endpoint().replace("http://", "").split(":")
|
||||
ret = db.add_instance(
|
||||
endpoint_data[0], endpoint_data[1], api.get_host()
|
||||
)
|
||||
|
||||
if ret:
|
||||
logger.warning(ret)
|
||||
else:
|
||||
ret = db.add_instance(
|
||||
"localhost",
|
||||
config_files.get("API_HTTP_PORT", 5000),
|
||||
config_files.get("API_SERVER_NAME", "bwapi"),
|
||||
)
|
||||
|
||||
if ret:
|
||||
logger.warning(ret)
|
||||
except SystemExit as e:
|
||||
sys_exit(e)
|
||||
except:
|
||||
|
|
|
|||
|
|
@ -1,38 +1,89 @@
|
|||
import logging
|
||||
from logging import (
|
||||
CRITICAL,
|
||||
DEBUG,
|
||||
ERROR,
|
||||
INFO,
|
||||
WARNING,
|
||||
Logger,
|
||||
_levelToName,
|
||||
_nameToLevel,
|
||||
addLevelName,
|
||||
basicConfig,
|
||||
getLogger,
|
||||
setLoggerClass,
|
||||
)
|
||||
from os import getenv
|
||||
from threading import Lock
|
||||
|
||||
logging.basicConfig(
|
||||
|
||||
class BWLogger(Logger):
|
||||
def __init__(self, name, level=INFO):
|
||||
self.name = name
|
||||
self.db_lock = Lock()
|
||||
return super(BWLogger, self).__init__(name, level)
|
||||
|
||||
def _log(
|
||||
self,
|
||||
level,
|
||||
msg,
|
||||
args,
|
||||
exc_info=None,
|
||||
extra=None,
|
||||
stack_info=False,
|
||||
stacklevel=1,
|
||||
*,
|
||||
store_db: bool = False,
|
||||
db=None,
|
||||
):
|
||||
if store_db is True and db is not None:
|
||||
with self.db_lock:
|
||||
ret = db.save_log(msg, level, self.name)
|
||||
|
||||
if ret:
|
||||
self.warning("Failed to save log to database")
|
||||
|
||||
return super(BWLogger, self)._log(
|
||||
level, msg, args, exc_info, extra, stack_info, stacklevel
|
||||
)
|
||||
|
||||
|
||||
setLoggerClass(BWLogger)
|
||||
|
||||
default_level = _nameToLevel.get(getenv("LOG_LEVEL", "INFO").upper(), INFO)
|
||||
basicConfig(
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
datefmt="[%Y-%m-%d %H:%M:%S]",
|
||||
level=logging.INFO,
|
||||
level=default_level,
|
||||
)
|
||||
|
||||
getLogger("sqlalchemy.orm.mapper.Mapper").setLevel(
|
||||
default_level if default_level != INFO else WARNING
|
||||
)
|
||||
getLogger("sqlalchemy.orm.relationships.RelationshipProperty").setLevel(
|
||||
default_level if default_level != INFO else WARNING
|
||||
)
|
||||
getLogger("sqlalchemy.orm.strategies.LazyLoader").setLevel(
|
||||
default_level if default_level != INFO else WARNING
|
||||
)
|
||||
getLogger("sqlalchemy.pool.impl.QueuePool").setLevel(
|
||||
default_level if default_level != INFO else WARNING
|
||||
)
|
||||
getLogger("sqlalchemy.engine.Engine").setLevel(
|
||||
default_level if default_level != INFO else WARNING
|
||||
)
|
||||
|
||||
# Edit the default levels of the logging module
|
||||
logging.addLevelName(logging.CRITICAL, "🚨")
|
||||
logging.addLevelName(logging.DEBUG, "🐛")
|
||||
logging.addLevelName(logging.ERROR, "❌")
|
||||
logging.addLevelName(logging.INFO, "ℹ️ ")
|
||||
logging.addLevelName(logging.WARNING, "⚠️ ")
|
||||
addLevelName(CRITICAL, "🚨")
|
||||
addLevelName(DEBUG, "🐛")
|
||||
addLevelName(ERROR, "❌")
|
||||
addLevelName(INFO, "ℹ️ ")
|
||||
addLevelName(WARNING, "⚠️ ")
|
||||
|
||||
|
||||
def setup_logger(title: str, level=logging.INFO) -> logging.Logger:
|
||||
def setup_logger(title: str, level=INFO) -> Logger:
|
||||
"""Set up local logger"""
|
||||
title = title.upper()
|
||||
logger = logging.getLogger(title)
|
||||
|
||||
if level not in (
|
||||
logging.DEBUG,
|
||||
logging.INFO,
|
||||
logging.WARNING,
|
||||
logging.ERROR,
|
||||
logging.CRITICAL,
|
||||
"DEBUG",
|
||||
"INFO",
|
||||
"WARNING",
|
||||
"ERROR",
|
||||
"CRITICAL",
|
||||
):
|
||||
level = logging.INFO
|
||||
|
||||
logger.setLevel(level)
|
||||
logger = getLogger(title)
|
||||
logger.setLevel(_nameToLevel.get(level, _levelToName.get(level, INFO)))
|
||||
|
||||
return logger
|
||||
|
|
|
|||
Loading…
Reference in New Issue