chore: 优化一些内容

Author: tianxiu2b2t

.gitignore

@@ -5,10 +5,6 @@ logs
 bmclapi
 config
 cache
+database
 
-start.sh
-
-**/__pycache__
-
-# pipy plugins
-tianxiu2b2t
+**/__pycache__

core/__init__.py

@@ -15,6 +15,7 @@
 from .config import API_VERSION, VERSION, cfg
 from .logger import logger
 from .utils import runtime
+from .database import init as init_database
 from . import web
 import platform
 
@@ -52,9 +53,15 @@ def load_clusters():
             logger.terror("core.initialize.cluster.missing", err=e)
             continue
 
+async def load_database():
+    await init_database(
+        cfg.get('database', {})
+    )
+
 async def load_config():
     load_storages()
     load_clusters()
+    await load_database()
 
     if clusters.count == 0 or clusters.storages.count == 0:
         logger.terror("core.initialize.missing", clusters=clusters.count, storages=clusters.storages.count)

core/cluster.py

@@ -4,6 +4,7 @@
 import datetime
 import hmac
 from pathlib import Path
+import sys
 import tempfile
 import time
 from typing import Any, Optional
@@ -20,6 +21,7 @@
 from .logger import logger
 from .config import API_VERSION, ROOT_PATH, cfg, USER_AGENT
 from .storage import CheckStorage, StorageManager
+from .database import get_db
 
 class TokenManager:
     def __init__(
@@ -132,12 +134,13 @@ class Cluster:
     def __init__(
         self,
         id: str,
-        secret: str
+        secret: str,
+        manager: 'ClusterManager'
     ):
         self._token = TokenManager(id, secret)
         self._last_modified = 0
         self.sio = socketio.AsyncClient(
-            handle_sigint=False
+            handle_sigint=False,
         )
         self._keepalive_lock = utils.CustomLock(locked=True)
         self._storage_wait = utils.CustomLock(locked=True)
@@ -149,6 +152,7 @@ def __init__(
         self._failed_keepalive = 0
         self._retry_times = 0
         self._task_group = None
+        self._manager = manager
 
     @property
     def id(self):
@@ -167,14 +171,17 @@ async def setup(
         @self.sio.on("warden-error") # type: ignore
         async def _(message: Any):
             logger.twarning("cluster.warden", id=self.id, msg=message)
+            await get_db().insert_cluster_info(self.id, "server-push", "warden-error", message)
 
         @self.sio.on("exception") # type: ignore
         async def _(message: Any):
             logger.terror("cluster.exception", id=self.id, msg=message)
+            await get_db().insert_cluster_info(self.id, "server-push", "exception", message)
 
         @self.sio.on("message") # type: ignore
         async def _(message: Any):
             logger.tinfo("cluster.message", id=self.id, msg=message)
+            await get_db().insert_cluster_info(self.id, "server-push", "message", message)
 
         @self.sio.on("connect") # type: ignore
         async def _():
@@ -183,6 +190,7 @@ async def _():
                 return
             self._enabled = False
             logger.tinfo("cluster.reconnect", id=self.id)
+            await get_db().insert_cluster_info(self.id, "socketio", "reconnect")
             await self.enable()
 
 
@@ -225,6 +233,10 @@ async def keepalive(self):
 
                     self.counter.hits -= current_counter.hits
                     self.counter.bytes -= current_counter.bytes
+
+                    # insert into db
+                    await get_db().upsert_cluster_counter(self.id, current_counter.hits, current_counter.bytes)
+
                     resp_time = time.time() - datetime.datetime.fromisoformat(res.ack).timestamp()
                     logger.tsuccess("cluster.keepalive", id=self.id, hits=units.format_number(current_counter.hits), bytes=units.format_number(current_counter.bytes), delay=F"{resp_time * 1000:.4f}")
             except:
@@ -252,14 +264,19 @@ async def request_cert(self):
         return Certificate(CertificateType.CLUSTER, str(cert), str(key))
 
     async def connect(self):
-        await self.sio.connect(
-            cfg.base_url,
-            transports=['websocket'],
-            headers={
-                "User-Agent": USER_AGENT,
-            },
-            auth=self._token.get_socketio_token
-        )
+        try:
+            await self.sio.connect(
+                cfg.base_url,
+                transports=['websocket'],
+                headers={
+                    "User-Agent": USER_AGENT,
+                },
+                auth=self._token.get_socketio_token
+            )
+        except:
+            logger.debug_traceback()
+            await anyio.sleep(5)
+            await self.connect()
 
     async def emit(self, event: str, data: Any = None, timeout: Optional[int] = None) -> SocketEmitResult:
         err, ack = None, None
@@ -269,8 +286,13 @@ async def callback(data: Any):
                 err, ack = data
             else:
                 err, ack = data, None
+            if event not in ("keep-alive", ):
+                await get_db().insert_cluster_info(self.id, "server", event, err)
             fut.set()
 
+        if event not in ("keep-alive", ):
+            await get_db().insert_cluster_info(self.id, "client", event)
+
         await self.sio.emit(event, data, callback=callback)
 
         fut = anyio.Event()
@@ -352,8 +374,8 @@ async def enable(self):
                     "noFastEnable": False,
                     "byoc": utils.get_certificate_type() != CertificateType.CLUSTER,
                     "flavor": {
-                        "runtime": "python",
-                        "storage": "local"
+                        "runtime": f"python/{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
+                        "storage": self._manager.storages.get_storage_type.type
                     }
                 }, 300)
                 if res.err is not None:
@@ -546,7 +568,7 @@ def add_cluster(
         id: str,
         secret: str
     ):
-        self._clusters[id] = Cluster(id, secret)
+        self._clusters[id] = Cluster(id, secret, self)
 
     def get_cluster(
         self,

core/config.py

@@ -28,12 +28,12 @@ def _get_keys(self, key: str):
     def get(self, key: str, default = None) -> Any:
         val = os.getenv(key) or self._get_value(self._data, self._get_keys(key))
         if val is None:
-            print(f"[Config] Key '{key}' is not set?")
             if key in DEFAULT_CONFIG:
                 self.set(key, DEFAULT_CONFIG[key])
                 val = DEFAULT_CONFIG[key]
                 self.save()
             else:
+                print(f"[Config] Key '{key}' is not set?")
                 val = default
         return val
 

core/logger.py

@@ -80,24 +80,6 @@ def tdebug_traceback(self, key: str, *args, **kwargs):
 
 logger = Loglogger()
 
-class LoguruHandler(logging.Handler):  # pragma: no cover
-    """logging 与 loguru 之间的桥梁，将 logging 的日志转发到 loguru。"""
-
-    def emit(self, record: logging.LogRecord):
-        try:
-            level = Logger.level(record.levelname).name
-        except ValueError:
-            level = record.levelno
-
-        frame, depth = inspect.currentframe(), 0
-        while frame and (depth == 0 or frame.f_code.co_filename == logging.__file__):
-            frame = frame.f_back
-            depth += 1
-
-        logger._log_with_args(
-            level,
-            record.getMessage(),
-        )
 
 def _log(*values):
     data = []
@@ -118,7 +100,10 @@ def emit(self, record):
         except ValueError:
             level = record.levelno
 
-        print(level, record.getMessage())
+        logger.log.opt(depth=6).log(
+            level,
+            record.getMessage(),
+        )
 
 # 配置拦截处理器
 logging.basicConfig(handlers=[InterceptHandler()], level=logging.DEBUG)

core/storage/__init__.py

@@ -1,4 +1,5 @@
 from collections import deque
+from dataclasses import dataclass
 from typing import Type
 
 import anyio.abc
@@ -22,6 +23,26 @@
     "s3": S3Storage
 }
 
+@dataclass
+class StorageTypeCount:
+    file: int = 0
+    alist: int = 0
+    webdav: int = 0
+    s3: int = 0
+
+    @property
+    def type(self):
+        res = []
+        if self.file > 0:
+            res.append("file")
+        if self.alist > 0:
+            res.append("alist")
+        if self.webdav > 0:
+            res.append("webdav")
+        if self.s3 > 0:
+            res.append("s3")
+        return "+".join(res)
+
 class StorageManager:
     def __init__(
         self
@@ -30,6 +51,20 @@ def __init__(
         self._weight_storages: deque[Storage] = deque()
         self._online_storages: deque[Storage] = deque()
         self._status = False
+    
+    @property
+    def get_storage_type(self):
+        res = StorageTypeCount()
+        for storage in self._storages:
+            if isinstance(storage, LocalStorage):
+                res.file += 1
+            elif isinstance(storage, AlistStorage):
+                res.alist += 1
+            elif isinstance(storage, WebDavStorage):
+                res.webdav += 1
+            elif isinstance(storage, S3Storage):
+                res.s3 += 1
+        return res
 
 
     def add_storage(

core/storage/alist.py

@@ -21,6 +21,12 @@ def __init__(
 
     def __repr__(self):
         return f"<AlistResponse code={self.code} data={self.data} message={self.message}>"
+    
+    def raise_for_status(self):
+        if self.code == 200:
+            return
+        raise Exception(f"Status: {self.code}, message: {self.message}")
+        
 
 class AlistStorage(abc.Storage):
     type = "alist"
@@ -64,6 +70,8 @@ async def _fetch_token(self):
                 }
             ) as resp:
                 data = AlistResponse(await resp.json())
+                data.raise_for_status()
+
                 self._token = data.data["token"]
 
                 assert self._task_group is not None
@@ -155,6 +163,7 @@ async def upload(self, path: str, tmp_file: _TemporaryFileWrapper):
                 data=tmp_file.file
             ) as resp:
                 data = AlistResponse(await resp.json())
+                data.raise_for_status()
                 return True
 
     async def get_response_file(self, hash: str) -> abc.ResponseFile: