apple / bt-source /panel /mod /project /node /dbutil /load_db.py
GGSheng's picture
feat: deploy to hf space | model=claude-72604b
62f6d9a verified
Raw
History Blame Contribute Delete
18.6 kB
import json
import os.path
import re
import sys
from dataclasses import dataclass, field
from typing import Tuple, Optional, List, Union
if "/www/server/panel/class" not in sys.path:
sys.path.insert(0, "/www/server/panel/class")
import public
import db
if "/www/server/panel" not in sys.path:
sys.path.insert(0, "/www/server/panel")
@dataclass
class LoadSite:
name: str
site_name: str
site_type: str
ps: str = ''
http_config: dict = field(default_factory=lambda: {
"proxy_next_upstream": "error timeout http_500 http_502 http_503 http_504",
"http_alg": "sticky_cookie",
"proxy_cache_status": False,
"cache_time": "1d",
"cache_suffix": "css,js,jpg,jpeg,gif,png,webp,woff,eot,ttf,svg,ico,css.map,js.map",
})
tcp_config: dict = field(default_factory=lambda: {
"proxy_connect_timeout": 8,
"proxy_timeout": 86400,
"host": "127.0.0.1",
"port": 80,
"type": "tcp"
})
created_at: int = 0
load_id: int = 0
site_id: int = 0
@classmethod
def bind_http_load(cls, data: dict) -> Tuple[Optional["LoadSite"], str]:
check_msg = cls.base_check(data)
if check_msg:
return None, check_msg
if not data.get('site_name', None):
return None, 'site_name is required'
if not public.is_domain(data['site_name']):
return None, 'site_name is invalid'
if not isinstance(data.get('http_config', None), dict):
return None, 'http_config is required'
else:
if "proxy_cache_status" not in dict.keys(data['http_config']): #兼容旧版本数据
data['http_config']["proxy_cache_status"] = False
data['http_config']["cache_time"] = "1d"
data['http_config']["cache_suffix"] = "css,js,jpg,jpeg,gif,png,webp,woff,eot,ttf,svg,ico,css.map,js.map"
for k in ['proxy_next_upstream', 'http_alg', "proxy_cache_status", "cache_time", "cache_suffix"]:
if k not in dict.keys(data['http_config']):
return None, 'http_config.{} is required'.format(k)
for i in data['http_config']['proxy_next_upstream'].split():
if i not in ('error', 'timeout') and not re.match(r'^http_\d{3}$', i):
return None, 'http_config.proxy_next_upstream is invalid'
if data['http_config']['http_alg'] not in ('sticky_cookie', 'round_robin', 'least_conn', 'ip_hash'):
return None, 'http_config.http_alg is invalid'
if not isinstance(data['http_config']['proxy_cache_status'], bool):
return None, 'http_config.proxy_cache_status is invalid'
if not isinstance(data['http_config']['cache_time'], str):
return None, 'http_config.cache_time is invalid'
if not re.match(r"^[0-9]+([smhd])$", data['http_config']['cache_time']):
return None, 'http_config.cache_time is invalid'
cache_suffix = data['http_config']['cache_suffix']
cache_suffix_list = []
for suffix in cache_suffix.split(","):
tmp_suffix = re.sub(r"\s", "", suffix)
if not tmp_suffix:
continue
cache_suffix_list.append(tmp_suffix)
real_cache_suffix = ",".join(cache_suffix_list)
if not real_cache_suffix:
real_cache_suffix = "css,js,jpg,jpeg,gif,png,webp,woff,eot,ttf,svg,ico,css.map,js.map"
data['http_config']['cache_suffix'] = real_cache_suffix
l = LoadSite(data.get('name'), data.get('site_name'), 'http', data.get('ps', ''),
http_config=data.get('http_config'),
created_at=data.get('created_at', 0), load_id=data.get('load_id', 0),
site_id=data.get('site_id', 0))
return l, ""
@classmethod
def base_check(cls, data) -> str:
if not data.get('name', None):
return 'name is required'
if not re.match(r'^[a-zA-Z0-9][a-zA-Z0-9_]+$', data['name']):
return '名称只能包含字母、数字、下划线且不能以数字和下划线开头'
if not len(data['name']) >= 3:
return '名称长度不能小于3个字符'
return ""
@classmethod
def bind_tcp_load(cls, data: dict) -> Tuple[Optional["LoadSite"], str]:
check_msg = cls.base_check(data)
if check_msg:
return None, check_msg
if not isinstance(data.get('tcp_config', None), dict):
return None, 'tcp_config is required'
else:
for k in ['proxy_connect_timeout', 'proxy_timeout', 'host', 'port', 'type']:
if not data['tcp_config'].get(k):
return None, 'tcp_config.{} is required'.format(k)
if data['tcp_config']['type'] not in ('tcp', 'udp'):
return None, 'tcp_config.type is invalid'
if not isinstance(data['tcp_config']['port'], int) and not 1 <= data['tcp_config']['port'] <= 65535:
return None, 'tcp_config.port is invalid'
if not public.check_ip(data['tcp_config']['host']):
return None, 'tcp_config.host is invalid'
l = LoadSite(data.get('name'), data.get('site_name'), 'tcp', ps=data.get('ps', ''),
tcp_config=data.get('tcp_config'),
created_at=data.get('created_at', 0), load_id=data.get('load_id', 0),
site_id=data.get('site_id', 0))
return l, ""
def to_dict(self) -> dict:
return {
"name": self.name,
"site_name": self.site_name,
"site_type": self.site_type,
"ps": self.ps,
"http_config": self.http_config,
"tcp_config": self.tcp_config,
"created_at": self.created_at,
"load_id": self.load_id,
"site_id": self.site_id
}
@dataclass
class HttpNode:
node_id: int
node_site_name: str
port: int
location: str = "/"
path: str = "/"
node_status: str = "online" # online, backup, down
weight: int = 1
max_fail: int = 3
fail_timeout: int = 600
ps: str = ""
created_at: int = 0
node_site_id: int = 0
id: int = 0
load_id: int = 0
@classmethod
def bind(cls, data: dict) -> Tuple[Optional["HttpNode"], str]:
if not isinstance(data.get('node_site_name', None), str):
return None, 'node_site_name is required'
if not public.is_domain(data['node_site_name']) and not public.check_ip(data['node_site_name']):
return None, 'node_site_name is invalid'
if not isinstance(data.get('port', None), int):
return None, 'port is required'
if not 1 <= data['port'] <= 65535:
return None, 'port is invalid'
if not isinstance(data.get('node_id', None), int):
return None, 'node_id is required'
if not isinstance(data.get('node_status', None), str):
return None, 'node_status is required'
if not data['node_status'] in ('online', 'backup', 'down'):
return None, 'node_status is invalid'
n = HttpNode(data.get('node_id'), data.get('node_site_name'), data.get('port'), "/",
data.get('path', "/"), data.get('node_status', "online"), data.get('weight', 1),
data.get('max_fail', 3), data.get('fail_timeout', 600), data.get('ps', ''),
data.get('created_at', 0), data.get('node_site_id', 0), data.get('id', 0),
data.get('load_id', 0)
)
return n, ""
def to_dict(self) -> dict:
return {
"node_id": self.node_id,
"node_site_name": self.node_site_name,
"port": self.port,
"location": self.location,
"path": self.path,
"node_status": self.node_status,
"weight": self.weight,
"max_fail": self.max_fail,
"fail_timeout": self.fail_timeout,
"ps": self.ps,
"created_at": self.created_at,
"node_site_id": self.node_site_id,
"id": self.id,
"load_id": self.load_id
}
@dataclass
class TcpNode:
node_id: int
host: str
port: int
id: int = 0
load_id: int = 0
node_status: str = "online" # online, backup, down
weight: int = 1
max_fail: int = 3
fail_timeout: int = 600
ps: str = ""
created_at: int = 0
@classmethod
def bind(cls, data: dict) -> Tuple[Optional["TcpNode"], str]:
if not isinstance(data.get('node_status', None), str):
return None, 'node_status is required'
if not data['node_status'] in ('online', 'backup', 'down'):
return None, 'node_status is invalid'
if not isinstance(data.get('host', None), str):
return None, 'host is required'
if not isinstance(data.get('node_id', None), int):
return None, 'node_id is required'
if not isinstance(data.get('port', None), int):
return None, 'port is required'
if not 1 <= data['port'] <= 65535:
return None, 'port is invalid'
n = TcpNode(data.get('node_id'), data.get('host'), data.get('port'), data.get('id', 0), data.get('load_id', 0),
data.get('node_status', "online"), data.get('weight', 1), data.get('max_fail', 3),
data.get('fail_timeout', 600), data.get('ps', ''), data.get('created_at', 0))
return n, ""
def to_dict(self) -> dict:
return {
"node_id": self.node_id,
"host": self.host,
"port": self.port,
"id": self.id,
"load_id": self.load_id,
"node_status": self.node_status,
"weight": self.weight,
"max_fail": self.max_fail,
"fail_timeout": self.fail_timeout,
"ps": self.ps,
"created_at": self.created_at
}
class NodeDB:
_DB_FILE = public.get_panel_path() + "/data/db/node_load_balance.db"
_DB_INIT_FILE = os.path.dirname(__file__) + "/load_balancer.sql"
def __init__(self):
sql = db.Sql()
sql._Sql__DB_FILE = self._DB_FILE
self.db = sql
def init_db(self):
sql_data = public.readFile(self._DB_INIT_FILE)
if not os.path.exists(self._DB_FILE) or os.path.getsize(self._DB_FILE) == 0:
public.writeFile(self._DB_FILE, "")
import sqlite3
conn = sqlite3.connect(self._DB_FILE)
c = conn.cursor()
c.executescript(sql_data)
conn.commit()
conn.close()
def close(self):
self.db.close()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, exc_trackback):
self.close()
def __del__(self):
self.close()
def update_load_key(self, load_id: int, load_data: dict) -> str:
if not isinstance(load_id, int):
return "load_id is required"
if not isinstance(load_data, dict):
return "load_data is required"
err = self.db.table("load_sites").where("load_id = ?", load_id).update(load_data)
if isinstance(err, str):
return err
return ""
def name_exist(self, name: str) -> bool:
return self.db.table("load_sites").where("name = ?", name).count() > 0
def load_site_name_exist(self, name: str) -> bool:
return self.db.table("load_sites").where("site_name = ?", name).count() > 0
def load_id_exist(self, load_id: int) -> bool:
return self.db.table("load_sites").where("load_id = ?", load_id).count() > 0
def loads_count(self, site_type: str, query: str = "") -> int:
if site_type == "http":
if not query:
return self.db.table("load_sites").where("site_type = ?", "http").count()
return self.db.table("load_sites").where(
"site_type = ? AND ps like ?", ("http", "%" + query + "%")).count()
else:
if not query:
return self.db.table("load_sites").where("site_type = ?", "tcp").count()
return self.db.table("load_sites").where(
"site_type = ? AND ps like ?", ("tcp", "%" + query + "%")).count()
def loads_list(self, site_type: str, offset: int, limit: int, query: str = ""):
if site_type == "all":
if query:
return self.db.table("load_sites").where("ps like ?", "%" + query + "%").limit(limit, offset).select()
return self.db.table("load_sites").limit(limit, offset).select()
if site_type == "http":
if not query:
return self.db.table("load_sites").where("site_type = ?", "http").limit(limit, offset).select()
return self.db.table("load_sites").where(
"site_type = ? AND ps like ?", ("http", "%" + query + "%")).limit(limit, offset).select()
else:
if not query:
return self.db.table("load_sites").where("site_type = ?", "tcp").limit(limit, offset).select()
return self.db.table("load_sites").where(
"site_type = ? AND ps like ?", ("tcp", "%" + query + "%")).limit(limit, offset).select()
def create_load(self, site_type: str, load: LoadSite, nodes: List[Union[HttpNode, TcpNode]]) -> str:
load_data = load.to_dict()
load_data.pop('load_id')
load_data.pop('created_at')
load_data["http_config"] = json.dumps(load.http_config)
load_data["tcp_config"] = json.dumps(load.tcp_config)
try:
err = self.db.table("load_sites").insert(load_data)
if isinstance(err, str):
return err
load.load_id = err
for node in nodes:
node_data = node.to_dict()
node_data.pop('id')
node_data.pop('created_at')
node_data['load_id'] = load.load_id
if site_type == "http" and isinstance(node, HttpNode):
err = self.db.table("http_nodes").insert(node_data)
else:
err = self.db.table("tcp_nodes").insert(node_data)
if isinstance(err, str):
return err
except Exception as e:
return "数据库操作错误:" + str(e)
return ""
def update_load(self, site_type: str, load: LoadSite, nodes: List[Union[HttpNode, TcpNode]]) -> str:
load_data = load.to_dict()
if not load.load_id:
return "load_id is required"
load_data.pop('created_at')
load_data.pop('load_id')
load_data["http_config"] = json.dumps(load.http_config)
load_data["tcp_config"] = json.dumps(load.tcp_config)
try:
err = self.db.table("load_sites").where("load_id = ?", load.load_id).update(load_data)
if isinstance(err, str):
return err
except Exception as e:
return "数据库操作错误:" + str(e)
old_nodes, err = self.get_nodes(load.load_id, site_type)
if err:
return err
old_nodes_map = {}
for old_node in old_nodes:
old_nodes_map[old_node['id']] = old_node
try:
for node in nodes:
node_data = node.to_dict()
node_data.pop('id')
node_data.pop('created_at')
node_data['load_id'] = load.load_id
if node.id in old_nodes_map:
if site_type == "http" and isinstance(node, HttpNode):
err = self.db.table("http_nodes").where("id = ?", node.id).update(node_data)
else:
err = self.db.table("tcp_nodes").where("id = ?", node.id).update(node_data)
if isinstance(err, str):
return err
old_nodes_map.pop(node.id)
else:
if site_type == "http" and isinstance(node, HttpNode):
err = self.db.table("http_nodes").insert(node_data)
else:
err = self.db.table("tcp_nodes").insert(node_data)
if isinstance(err, str):
return err
for node_id in old_nodes_map:
if site_type == "http":
err = self.db.table("http_nodes").where("id = ?", node_id).delete()
else:
err = self.db.table("tcp_nodes").where("id = ?", node_id).delete()
if isinstance(err, str):
return err
except Exception as e:
return "数据库操作错误:" + str(e)
return ""
def get_nodes(self, load_id: int, site_type: str) -> Tuple[List[dict], str]:
if site_type == "http":
nodes: List[dict] = self.db.table("http_nodes").where("load_id = ?", load_id).select()
else:
nodes: List[dict] = self.db.table("tcp_nodes").where("load_id = ?", load_id).select()
if isinstance(nodes, str):
return [], nodes
if not nodes and self.db.ERR_INFO:
return [], self.db.ERR_INFO
return nodes, ""
def get_load(self, load_id: int) -> Tuple[Optional[dict], str]:
load_data = self.db.table("load_sites").where("load_id = ?", load_id).find()
if isinstance(load_data, str):
return None, load_data
if self.db.ERR_INFO:
return None, self.db.ERR_INFO
if len(load_data) == 0:
return None, "未查询到该负载配置"
return load_data, ""
def delete(self, load_id: int) -> str:
load_data = self.db.table("load_sites").where("load_id = ?", load_id).find()
if isinstance(load_data, str):
return load_data
if self.db.ERR_INFO:
return self.db.ERR_INFO
if len(load_data) == 0:
return ""
if load_data["site_type"] == "http":
err = self.db.table("http_nodes").where("load_id = ?", load_id).delete()
else:
err = self.db.table("tcp_nodes").where("load_id = ?", load_id).delete()
if isinstance(err, str):
return err
err = self.db.table("load_sites").where("load_id = ?", load_id).delete()
if isinstance(err, str):
return err
return ""