Optimize Database connection and ApiCaller

This commit is contained in:
TheophileDiot 2022-10-25 11:39:30 +02:00
parent 0edef7c520
commit 9140dc3244
9 changed files with 186 additions and 313 deletions

View File

@ -1 +1,3 @@
.git
.git
__pycache__
env

6
.gitignore vendored
View File

@ -1,7 +1,5 @@
site/
.idea/
.vscode/
**/__pycache__/
ui/env/
docs/env/
env/
__pycache__
env

View File

@ -1,3 +1,4 @@
from time import sleep
from traceback import format_exc
from subprocess import run, DEVNULL, STDOUT
from glob import glob
@ -135,10 +136,18 @@ class Config(ApiCaller, ConfigCaller):
self.__services = services
self.__configs = configs
self.__config = self.__get_full_env()
if self.__db is None:
self.__db = Database(
self.__logger, sqlalchemy_string=self.__config.get("DATABASE_URI", None)
)
while not self.__db.is_initialized():
self.__logger.warning(
"Database is not initialized, retrying in 5 seconds ...",
)
sleep(5)
self._set_apis(self.__get_apis())
# write configs

View File

@ -12,15 +12,10 @@ sys_path.append("/opt/bunkerweb/utils")
sys_path.append("/opt/bunkerweb/api")
sys_path.append("/opt/bunkerweb/db")
from docker import DockerClient
from docker.errors import DockerException
from kubernetes import client as kube_client
from logger import setup_logger
from SwarmController import SwarmController
from IngressController import IngressController
from DockerController import DockerController
from Database import Database
# Get variables
logger = setup_logger("Autoconf", getenv("LOG_LEVEL", "INFO"))
@ -49,50 +44,6 @@ try:
if proc.returncode != 0:
_exit(1)
db = None
if "DATABASE_URI" in environ:
db = Database(logger)
elif kubernetes:
corev1 = kube_client.CoreV1Api()
for pod in corev1.list_pod_for_all_namespaces(watch=False).items:
if (
pod.metadata.annotations != None
and "bunkerweb.io/INSTANCE" in pod.metadata.annotations
):
for pod_env in pod.spec.containers[0].env:
if pod_env.name == "DATABASE_URI":
db = Database(
logger,
pod_env.value or getenv("DATABASE_URI", "5000"),
)
break
else:
try:
docker_client = DockerClient(base_url="tcp://docker-proxy:2375")
except DockerException:
docker_client = DockerClient(
base_url=getenv("DOCKER_HOST", "unix:///var/run/docker.sock")
)
apis = []
for instance in docker_client.containers.list(
filters={"label": "bunkerweb.INSTANCE"}
):
for var in instance.attrs["Config"]["Env"]:
if var.startswith("DATABASE_URI="):
db = Database(logger, var.replace("DATABASE_URI=", "", 1))
break
if db is None:
logger.error("No database found, exiting ...")
_exit(1)
while not db.is_initialized():
logger.warning(
"Database is not initialized, retrying in 5 seconds ...",
)
sleep(5)
# Instantiate the controller
if swarm:
logger.info("Swarm mode detected")

View File

@ -20,8 +20,6 @@ sys_path.append("/opt/bunkerweb/api")
sys_path.append("/opt/bunkerweb/db")
from docker import DockerClient
from docker.errors import DockerException
from kubernetes import client as kube_client
from logger import setup_logger
from Database import Database
@ -174,69 +172,20 @@ if __name__ == "__main__":
if config_files.get("LOG_LEVEL", logger.level) != logger.level:
logger = setup_logger("Generator", config_files["LOG_LEVEL"])
bw_integration = None
if config_files.get("SWARM_MODE", "no") == "yes":
bw_integration = "Swarm"
elif config_files.get("KUBERNETES_MODE", "no") == "yes":
bw_integration = "Local"
if config_files.get("KUBERNETES_MODE", "no") == "yes":
bw_integration = "Kubernetes"
elif config_files.get("AUTOCONF_MODE", "no") == "yes":
bw_integration = "Autoconf"
elif args.method != "autoconf" and exists("/opt/bunkerweb/INTEGRATION"):
with open("/opt/bunkerweb/INTEGRATION", "r") as f:
bw_integration = f.read().strip()
elif (
config_files.get("SWARM_MODE", "no") == "yes"
or config_files.get("AUTOCONF_MODE", "no") == "yes"
):
bw_integration = "Cluster"
db = None
if bw_integration == "Linux":
db = Database(logger, config_files["DATABASE_URI"])
elif bw_integration in ("Docker", "Swarm", "Autoconf"):
try:
docker_client = DockerClient(base_url="tcp://docker-proxy:2375")
except DockerException:
docker_client = DockerClient(
base_url=getenv("DOCKER_HOST", "unix:///var/run/docker.sock")
)
apis = []
for instance in docker_client.containers.list(
filters={"label": "bunkerweb.INSTANCE"}
):
api = None
for var in instance.attrs["Config"]["Env"]:
if db is None and var.startswith("DATABASE_URI="):
db = Database(logger, var.replace("DATABASE_URI=", "", 1))
elif var.startswith("API_HTTP_PORT="):
api = API(
f"http://{instance.name}:{var.replace('API_HTTP_PORT=', '', 1)}"
)
if api:
apis.append(api)
else:
apis.append(
API(
f"http://{instance.name}:{getenv('API_HTTP_PORT', '5000')}"
)
)
api_caller = ApiCaller(apis=apis)
elif bw_integration == "Kubernetes":
corev1 = kube_client.CoreV1Api()
for pod in corev1.list_pod_for_all_namespaces(watch=False).items:
if (
pod.metadata.annotations != None
and "bunkerweb.io/INSTANCE" in pod.metadata.annotations
):
for pod_env in pod.spec.containers[0].env:
if pod_env.name == "DATABASE_URI":
db = Database(
logger,
pod_env.value or getenv("DATABASE_URI", "5000"),
)
break
if db is None:
db = Database(logger)
db = Database(
logger,
sqlalchemy_string=getenv("DATABASE_URI", None),
bw_integration=bw_integration,
)
if args.init:
ret, err = db.init_tables(
@ -274,29 +223,7 @@ if __name__ == "__main__":
with open("/opt/bunkerweb/VERSION", "r") as f:
bw_version = f.read().strip()
bw_integration = None
if (
getenv("SWARM_MODE", config_files.get("SWARM_MODE", "no"))
== "yes"
):
bw_integration = "Swarm"
elif (
getenv(
"KUBERNETES_MODE", config_files.get("KUBERNETES_MODE", "no")
)
== "yes"
):
bw_integration = "Kubernetes"
elif (
getenv("AUTOCONF_MODE", config_files.get("AUTOCONF_MODE", "no"))
== "yes"
):
bw_integration = "Autoconf"
elif exists("/opt/bunkerweb/INTEGRATION"):
with open("/opt/bunkerweb/INTEGRATION", "r") as f:
bw_integration = f.read().strip()
if bw_integration == "Linux":
if bw_integration == "Local":
err = db.save_config(config_files, args.method)
if not err:
@ -305,9 +232,18 @@ if __name__ == "__main__":
err = None
err1 = None
err2 = db.initialize_db(
version=bw_version, integration=bw_integration
)
integration = "Linux"
if config_files.get("KUBERNETES_MODE", "no") == "yes":
integration = "Kubernetes"
elif config_files.get("SWARM_MODE", "no") == "yes":
integration = "Swarm"
elif config_files.get("AUTOCONF_MODE", "no") == "yes":
integration = "Autoconf"
elif exists("/opt/bunkerweb/INTEGRATION"):
with open("/opt/bunkerweb/INTEGRATION", "r") as f:
integration = f.read().strip()
err2 = db.initialize_db(version=bw_version, integration=integration)
if err or err1 or err2:
logger.error(
@ -325,44 +261,41 @@ if __name__ == "__main__":
config = config_files
elif args.method != "autoconf":
bw_integration = "Docker"
try:
docker_client = DockerClient(base_url="tcp://docker-proxy:2375")
except DockerException:
docker_client = DockerClient(
base_url=getenv("DOCKER_HOST", "unix:///var/run/docker.sock")
)
bw_integration = "Cluster"
docker_client = DockerClient(
base_url=getenv("DOCKER_HOST", "unix:///var/run/docker.sock")
)
tmp_config = {}
custom_confs = []
apis = []
db = None
for instance in docker_client.containers.list(
filters={"label": "bunkerweb.INSTANCE"}
):
api = None
for var in instance.attrs["Config"]["Env"]:
if custom_confs_rx.match(var.split("=", 1)[0]):
splitted = var.split("=", 1)
splitted = var.split("=", 1)
if custom_confs_rx.match(splitted[0]):
custom_confs.append(
{
"value": var.pop(0),
"value": splitted[1],
"exploded": custom_confs_rx.search(
"=".join(var)
splitted[0]
).groups(),
}
)
else:
tmp_config[var.split("=", 1)[0]] = var.split("=", 1)[1]
tmp_config[splitted[0]] = splitted[1]
if var.startswith("DATABASE_URI="):
db = Database(logger, var.replace("DATABASE_URI=", "", 1))
elif var.startswith("API_HTTP_PORT="):
api = API(
f"http://{instance.name}:{var.replace('API_HTTP_PORT=', '', 1)}"
)
if splitted[0] == "DATABASE_URI":
db = Database(
logger,
sqlalchemy_string=splitted[1],
)
elif splitted[0] == "API_HTTP_PORT":
api = API(f"http://{instance.name}:{splitted[1]}")
if api:
apis.append(api)
@ -407,63 +340,33 @@ if __name__ == "__main__":
config = config_files
else:
db = None
if getenv("KUBERNETES_MODE", "no") == "yes":
corev1 = kube_client.CoreV1Api()
for pod in corev1.list_pod_for_all_namespaces(watch=False).items:
if (
pod.metadata.annotations != None
and "bunkerweb.io/INSTANCE" in pod.metadata.annotations
and "DATABASE_URI" in pod.spec.containers[0].env
):
db = Database(
logger, pod.spec.containers[0].env["DATABASE_URI"]
)
break
elif getenv("AUTOCONF_MODE", getenv("SWARM_MODE", "no")) == "yes":
try:
docker_client = DockerClient(base_url="tcp://docker-proxy:2375")
except DockerException:
docker_client = DockerClient(
base_url=getenv("DOCKER_HOST", "unix:///var/run/docker.sock")
)
for instance in docker_client.containers.list(
filters={"label": "bunkerweb.INSTANCE"}
):
for var in instance.attrs["Config"]["Env"]:
if var.startswith("DATABASE_URI="):
db = Database(logger, var.replace("DATABASE_URI=", "", 1))
break
if db:
break
if db is None:
db = Database(logger)
db = Database(
logger,
bw_integration="Kubernetes"
if getenv("KUBERNETES_MODE", "no") == "yes"
else "Cluster",
)
config = db.get_config()
bw_integration = None
if config.get("SWARM_MODE", "no") == "yes":
bw_integration = "Swarm"
elif config.get("KUBERNETES_MODE", "no") == "yes":
bw_integration = "Local"
if config.get("KUBERNETES_MODE", "no") == "yes":
bw_integration = "Kubernetes"
elif config.get("AUTOCONF_MODE", "no") == "yes":
bw_integration = "Autoconf"
elif args.method != "autoconf" and exists("/opt/bunkerweb/INTEGRATION"):
with open("/opt/bunkerweb/INTEGRATION", "r") as f:
bw_integration = f.read().strip()
elif (
config.get("SWARM_MODE", "no") == "yes"
or config.get("AUTOCONF_MODE", "no") == "yes"
):
bw_integration = "Cluster"
logger = setup_logger("Generator", config.get("LOG_LEVEL", "INFO"))
if bw_integration == "Docker":
if args.method != "autoconf" and bw_integration == "Cluster":
while not api_caller._send_to_apis("GET", "/ping"):
logger.warning(
"Waiting for BunkerWeb's temporary nginx to start, retrying in 5 seconds ...",
)
sleep(5)
elif bw_integration == "Linux":
elif bw_integration == "Local":
retries = 0
while not exists("/opt/bunkerweb/tmp/nginx.pid"):
if retries == 5:
@ -489,7 +392,7 @@ if __name__ == "__main__":
elif isdir(file):
rmtree(file, ignore_errors=False)
if bw_integration in ("Docker", "Linux"):
if args.method != "autoconf":
logger.info(
"Generating custom configs from Database ...",
)
@ -517,14 +420,14 @@ if __name__ == "__main__":
)
templator.render()
if bw_integration == "Docker":
if args.method != "autoconf" and bw_integration == "Cluster":
ret = api_caller._send_to_apis("POST", "/reload")
if not ret:
logger.error(
"reload failed",
)
sys_exit(1)
elif bw_integration == "Linux":
elif bw_integration == "Local":
cmd = "/usr/sbin/nginx -s reload"
proc = run(cmd.split(" "), stdin=DEVNULL, stderr=STDOUT)
if proc.returncode != 0:

View File

@ -1,7 +1,7 @@
from contextlib import contextmanager
from copy import deepcopy
from logging import INFO, WARNING, Logger, getLogger
from os import _exit, environ, path
from os import _exit, environ, getenv, path
from re import search
from typing import Any, Dict, List, Optional, Tuple
from sqlalchemy import create_engine, inspect
@ -14,13 +14,52 @@ from model import *
class Database:
def __init__(self, logger: Logger, sqlalchemy_string: str = None) -> None:
def __init__(
self,
logger: Logger,
sqlalchemy_string: str = None,
bw_integration: str = "Local",
) -> None:
"""Initialize the database"""
self.__logger = logger
getLogger("sqlalchemy.engine").setLevel(
logger.level if logger.level != INFO else WARNING
)
if sqlalchemy_string is None and bw_integration != "Local":
if bw_integration == "Kubernetes":
from kubernetes import client as kube_client
corev1 = kube_client.CoreV1Api()
for pod in corev1.list_pod_for_all_namespaces(watch=False).items:
if (
pod.metadata.annotations != None
and "bunkerweb.io/INSTANCE" in pod.metadata.annotations
):
for pod_env in pod.spec.containers[0].env:
if pod_env.name == "DATABASE_URI":
sqlalchemy_string = pod_env.value
break
if sqlalchemy_string:
break
else:
from docker import DockerClient
docker_client = DockerClient(
base_url=getenv("DOCKER_HOST", "unix:///var/run/docker.sock")
)
for instance in docker_client.containers.list(
filters={"label": "bunkerweb.INSTANCE"}
):
for var in instance.attrs["Config"]["Env"]:
if var.startswith("DATABASE_URI="):
sqlalchemy_string = var.replace("DATABASE_URI=", "", 1)
break
if sqlalchemy_string:
break
if not sqlalchemy_string:
sqlalchemy_string = environ.get(
"DATABASE_URI", "sqlite:////data/db.sqlite3"

View File

@ -25,8 +25,13 @@ class JobScheduler(ApiCaller):
lock=None,
apis=[],
logger: Logger = setup_logger("Scheduler", environ.get("LOG_LEVEL", "INFO")),
auto: bool = False,
):
super().__init__(apis)
if auto is True:
self.auto_setup()
self.__logger = logger
self.__env = env
with open("/tmp/autoconf.env", "w") as f:

View File

@ -13,10 +13,7 @@ sys_path.append("/opt/bunkerweb/utils")
sys_path.append("/opt/bunkerweb/api")
sys_path.append("/opt/bunkerweb/db")
from docker import DockerClient
from docker.errors import DockerException
from dotenv import dotenv_values
from kubernetes import client as kube_client
from logger import setup_logger
from Database import Database
@ -26,7 +23,7 @@ from API import API
run = True
scheduler = None
reloading = False
logger = setup_logger("Scheduler", environ.get("LOG_LEVEL", "INFO"))
logger = setup_logger("Scheduler", getenv("LOG_LEVEL", "INFO"))
def handle_stop(signum, frame):
@ -105,8 +102,6 @@ if __name__ == "__main__":
logger.info("Scheduler started ...")
bw_integration = "Linux"
if args.variables:
logger.info(f"Variables : {args.variables}")
@ -114,42 +109,13 @@ if __name__ == "__main__":
env = dotenv_values(args.variables)
else:
# Read from database
db = None
if "DATABASE_URI" in environ:
db = Database(logger)
elif getenv("KUBERNETES_MODE", "no") == "yes":
bw_integration = "Kubernetes"
corev1 = kube_client.CoreV1Api()
for pod in corev1.list_pod_for_all_namespaces(watch=False).items:
if (
pod.metadata.annotations != None
and "bunkerweb.io/INSTANCE" in pod.metadata.annotations
and "DATABASE_URI" in pod.spec.containers[0].env
):
db = Database(
logger, pod.spec.containers[0].env["DATABASE_URI"]
)
break
else:
bw_integration = "Docker"
try:
docker_client = DockerClient(base_url="tcp://docker-proxy:2375")
except DockerException:
docker_client = DockerClient(
base_url=getenv("DOCKER_HOST", "unix:///var/run/docker.sock")
)
for instance in docker_client.containers.list(
filters={"label": "bunkerweb.INSTANCE"}
):
for var in instance.attrs["Config"]["Env"]:
if var.startswith("DATABASE_URI="):
db = Database(logger, var.replace("DATABASE_URI=", "", 1))
break
if db is None:
logger.error("No database found, exiting ...")
stop(1)
db = Database(
logger,
sqlalchemy_string=getenv("DATABASE_URI", None),
bw_integration="Kubernetes"
if getenv("KUBERNETES_MODE", "no") == "yes"
else "Cluster",
)
while not db.is_initialized():
logger.warning(
@ -180,67 +146,12 @@ if __name__ == "__main__":
logger.info("Executing job scheduler ...")
while True:
apis = []
if not args.variables:
if bw_integration == "Docker":
try:
docker_client = DockerClient(base_url="tcp://docker-proxy:2375")
except DockerException:
docker_client = DockerClient(
base_url=getenv(
"DOCKER_HOST", "unix:///var/run/docker.sock"
)
)
for instance in docker_client.containers.list(
filters={"label": "bunkerweb.INSTANCE"}
):
api = None
for var in instance.attrs["Config"]["Env"]:
if var.startswith("API_HTTP_PORT="):
api = API(
f"http://{instance.name}:{var.replace('API_HTTP_PORT=', '', 1)}"
)
break
if api:
apis.append(api)
else:
apis.append(
API(
f"http://{instance.name}:{getenv('API_HTTP_PORT', '5000')}"
)
)
elif bw_integration == "Kubernetes":
corev1 = kube_client.CoreV1Api()
for pod in corev1.list_pod_for_all_namespaces(watch=False).items:
api = None
if (
pod.metadata.annotations != None
and "bunkerweb.io/INSTANCE" in pod.metadata.annotations
):
for pod_env in instance.spec.containers[0].env:
if pod_env.name == "API_HTTP_PORT":
api = API(
f"http://{pod.status.pod_ip}:{pod_env.value or getenv('API_HTTP_PORT', '5000')}"
)
break
if api:
apis.append(api)
else:
apis.append(
API(
f"http://{pod.status.pod_ip}:{env.get('API_HTTP_PORT', getenv('API_HTTP_PORT', '5000'))}"
)
)
# Instantiate scheduler
scheduler = JobScheduler(
env=env,
apis=apis,
apis=[],
logger=logger,
auto=not args.variables,
)
# Only run jobs once

View File

@ -1,8 +1,10 @@
from io import BytesIO
from os import environ
from os import environ, getenv
from os.path import exists
from tarfile import open as taropen
from logger import setup_logger
from API import API
class ApiCaller:
@ -10,6 +12,59 @@ class ApiCaller:
self.__apis = apis
self.__logger = setup_logger("Api", environ.get("LOG_LEVEL", "INFO"))
def auto_setup(self, bw_integration: str = None):
if bw_integration is None and getenv("KUBERNETES_MODE", "no") == "yes":
bw_integration = "Kubernetes"
if bw_integration == "Kubernetes":
from kubernetes import client as kube_client
corev1 = kube_client.CoreV1Api()
for pod in corev1.list_pod_for_all_namespaces(watch=False).items:
if (
pod.metadata.annotations != None
and "bunkerweb.io/INSTANCE" in pod.metadata.annotations
):
api = None
for pod_env in pod.spec.containers[0].env:
if pod_env.name == "API_HTTP_PORT":
api = API(
f"http://{pod.status.pod_ip}:{pod_env.value or '5000'}"
)
break
if api:
self.__apis.append(api)
else:
self.__apis.append(
API(
f"http://{pod.status.pod_ip}:{getenv('API_HTTP_PORT', '5000')}"
)
)
else:
from docker import DockerClient
docker_client = DockerClient(
base_url=getenv("DOCKER_HOST", "unix:///var/run/docker.sock")
)
for instance in docker_client.containers.list(
filters={"label": "bunkerweb.INSTANCE"}
):
api = None
for var in instance.attrs["Config"]["Env"]:
if var.startswith("API_HTTP_PORT="):
api = API(
f"http://{instance.name}:{var.replace('API_HTTP_PORT=', '', 1)}"
)
break
if api:
self.__apis.append(api)
else:
self.__apis.append(
API(f"http://{instance.name}:{getenv('API_HTTP_PORT', '5000')}")
)
def _set_apis(self, apis):
self.__apis = apis