From 919a296988d8babfb2f85f3c5cb75033b06616fc Mon Sep 17 00:00:00 2001
From: KimigaiiWuyi <444835641@qq.com>
Date: Mon, 25 Sep 2023 23:52:24 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20=E6=96=B0=E5=A2=9E`on=5Fcore=5Fstar?=
 =?UTF-8?q?t`=E8=A3=85=E9=A5=B0=E5=99=A8=E5=92=8C=E6=8B=86=E5=88=86`exec?=
 =?UTF-8?q?=5Flist`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 gsuid_core/core.py                       | 23 +++++++++++++++++++++++
 gsuid_core/server.py                     |  8 ++++++++
 gsuid_core/utils/database/base_models.py |  2 +-
 gsuid_core/utils/database/startup.py     | 18 ++++++++++++++++++
 4 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 gsuid_core/utils/database/startup.py

diff --git a/gsuid_core/core.py b/gsuid_core/core.py
index 78ee7ac..8bfa6a7 100644
--- a/gsuid_core/core.py
+++ b/gsuid_core/core.py
@@ -19,7 +19,9 @@ from gsuid_core.config import core_config  # noqa: E402
 from gsuid_core.data_store import image_res  # noqa: E402
 from gsuid_core.handler import handle_event  # noqa: E402
 from gsuid_core.models import MessageReceive  # noqa: E402
+from gsuid_core.server import core_start_def  # noqa: E402
 from gsuid_core.webconsole.mount_app import site  # noqa: E402
+from gsuid_core.utils.database.startup import exec_list  # noqa: E402
 from gsuid_core.aps import start_scheduler, shutdown_scheduler  # noqa: E402
 from gsuid_core.utils.plugins_config.models import (  # noqa: E402
     GsListStrConfig,
@@ -40,6 +42,21 @@ app = FastAPI()
 HOST = core_config.get_config('HOST')
 PORT = int(core_config.get_config('PORT'))
 
+exec_list.extend(
+    [
+        'ALTER TABLE GsBind ADD COLUMN group_id TEXT',
+        'ALTER TABLE GsBind ADD COLUMN sr_uid TEXT',
+        'ALTER TABLE GsUser ADD COLUMN sr_uid TEXT',
+        'ALTER TABLE GsUser ADD COLUMN sr_region TEXT',
+        'ALTER TABLE GsUser ADD COLUMN fp TEXT',
+        'ALTER TABLE GsUser ADD COLUMN device_id TEXT',
+        'ALTER TABLE GsUser ADD COLUMN sr_sign_switch TEXT DEFAULT "off"',
+        'ALTER TABLE GsUser ADD COLUMN sr_push_switch TEXT DEFAULT "off"',
+        'ALTER TABLE GsUser ADD COLUMN draw_switch TEXT DEFAULT "off"',
+        'ALTER TABLE GsCache ADD COLUMN sr_uid TEXT',
+    ]
+)
+
 
 @app.websocket('/ws/{bot_id}')
 async def websocket_endpoint(websocket: WebSocket, bot_id: str):
@@ -66,6 +83,12 @@ async def startup_event():
         from gsuid_core.webconsole.__init__ import start_check
 
         await start_check()
+        try:
+            _task = [_def() for _def in core_start_def]
+            asyncio.gather(*_task)
+        except Exception as e:
+            logger.exception(e)
+
     except ImportError:
         logger.warning('未加载GenshinUID...网页控制台启动失败...')
     await start_scheduler()
diff --git a/gsuid_core/server.py b/gsuid_core/server.py
index 473b94f..fee93d5 100644
--- a/gsuid_core/server.py
+++ b/gsuid_core/server.py
@@ -9,6 +9,14 @@ from fastapi import WebSocket
 from gsuid_core.bot import _Bot
 from gsuid_core.logger import logger
 
+core_start_def = set()
+
+
+def on_core_start(func: Callable):
+    if func not in core_start_def:
+        core_start_def.add(func)
+    return func
+
 
 class GsServer:
     _instance = None
diff --git a/gsuid_core/utils/database/base_models.py b/gsuid_core/utils/database/base_models.py
index 12f85a1..f4a3fde 100644
--- a/gsuid_core/utils/database/base_models.py
+++ b/gsuid_core/utils/database/base_models.py
@@ -523,7 +523,7 @@ class User(BaseModel):
     ) -> List[T_User]:
         sql = select(cls).where(cls.cookie is not None, cls.cookie != '')
         result = await session.execute(sql)
-        data: List[T_User] = result.scalars().all()
+        data = result.scalars().all()
         return data
 
     @classmethod
diff --git a/gsuid_core/utils/database/startup.py b/gsuid_core/utils/database/startup.py
new file mode 100644
index 0000000..3964459
--- /dev/null
+++ b/gsuid_core/utils/database/startup.py
@@ -0,0 +1,18 @@
+from sqlalchemy.sql import text
+
+from gsuid_core.server import on_core_start
+
+from .base_models import async_maker
+
+exec_list = []
+
+
+@on_core_start
+async def sr_adapter():
+    async with async_maker() as session:
+        for _t in exec_list:
+            try:
+                await session.execute(text(_t))
+                await session.commit()
+            except:  # noqa: E722
+                pass