Merge pull request #1647 from 0xzmz/master

add database connect config

pkg/core/bootutils/deps.py

@@ -40,6 +40,7 @@ required_deps = {
     'sqlmodel': 'sqlmodel',
     'telegramify_markdown': 'telegramify-markdown',
     'slack_sdk': 'slack_sdk',
+    'asyncpg': 'asyncpg',
 }
 
 

pkg/persistence/databases/postgresql.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
+
+from .. import database
+
+
+@database.manager_class('postgresql')
+class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
+    """PostgreSQL database manager"""
+
+    async def initialize(self) -> None:
+        postgresql_config = self.ap.instance_config.data.get('database', {}).get('postgresql', {})
+
+        host = postgresql_config.get('host', '127.0.0.1')
+        port = postgresql_config.get('port', 5432)
+        user = postgresql_config.get('user', 'postgres')
+        password = postgresql_config.get('password', 'postgres')
+        database = postgresql_config.get('database', 'postgres')
+        engine_url = f'postgresql+asyncpg://{user}:{password}@{host}:{port}/{database}'
+        self.engine = sqlalchemy_asyncio.create_async_engine(engine_url)

pkg/persistence/databases/sqlite.py

@@ -10,5 +10,6 @@ class SQLiteDatabaseManager(database.BaseDatabaseManager):
     """SQLite database manager"""
 
     async def initialize(self) -> None:
-        sqlite_path = 'data/langbot.db'
-        self.engine = sqlalchemy_asyncio.create_async_engine(f'sqlite+aiosqlite:///{sqlite_path}')
+        db_file_path = self.ap.instance_config.data.get('database', {}).get('sqlite', {}).get('path', 'data/langbot.db')
+        engine_url = f'sqlite+aiosqlite:///{db_file_path}'
+        self.engine = sqlalchemy_asyncio.create_async_engine(engine_url)

pkg/persistence/mgr.py

@@ -36,11 +36,13 @@ class PersistenceManager:
         self.meta = base.Base.metadata
 
     async def initialize(self):
-        self.ap.logger.info('Initializing database...')
-
+        database_type = self.ap.instance_config.data.get('database', {}).get('use', 'sqlite')
+        self.ap.logger.info(f'Initializing database type: {database_type}...')
         for manager in database.preregistered_managers:
-            self.db = manager(self.ap)
-            await self.db.initialize()
+            if manager.name == database_type:
+                self.db = manager(self.ap)
+                await self.db.initialize()
+                break
 
         await self.create_tables()
 

pkg/persistence/migrations/dbm006_plugin_install_source.py

@@ -9,9 +9,21 @@ class DBMigratePluginInstallSource(migration.DBMigration):
     async def upgrade(self):
         """升级"""
         # 查询表结构获取所有列名（异步执行 SQL）
-        result = await self.ap.persistence_mgr.execute_async(sqlalchemy.text('PRAGMA table_info(plugin_settings);'))
-        # fetchall() 是同步方法，无需 await
-        columns = [row[1] for row in result.fetchall()]
+
+        columns = []
+
+        if self.ap.persistence_mgr.db.name == 'postgresql':
+            result = await self.ap.persistence_mgr.execute_async(
+                sqlalchemy.text(
+                    "SELECT column_name FROM information_schema.columns WHERE table_name = 'plugin_settings';"
+                )
+            )
+            all_result = result.fetchall()
+            columns = [row[0] for row in all_result]
+        else:
+            result = await self.ap.persistence_mgr.execute_async(sqlalchemy.text('PRAGMA table_info(plugin_settings);'))
+            all_result = result.fetchall()
+            columns = [row[1] for row in all_result]
 
         # 检查并添加 install_source 列
         if 'install_source' not in columns:

pyproject.toml

@@ -19,7 +19,7 @@ dependencies = [
     "dashscope>=1.23.2",
     "dingtalk-stream>=0.24.0",
     "discord-py>=2.5.2",
-    "pynacl>=1.5.0",  # Required for Discord voice support
+    "pynacl>=1.5.0", # Required for Discord voice support
     "gewechat-client>=0.1.5",
     "lark-oapi>=1.4.15",
     "mcp>=1.8.1",
@@ -63,6 +63,7 @@ dependencies = [
     "chromadb>=0.4.24",
     "qdrant-client (>=1.15.1,<2.0.0)",
     "langbot-plugin==0.1.1b6",
+    "asyncpg>=0.30.0",
 ]
 keywords = [
     "bot",

templates/config.yaml

@@ -20,6 +20,16 @@ system:
     jwt:
         expire: 604800
         secret: ''
+database:
+    use: sqlite
+    sqlite:
+        path: 'data/langbot.db'
+    postgresql:
+        host: '127.0.0.1'
+        port: 5432
+        user: 'postgres'
+        password: 'postgres'
+        database: 'postgres'
 vdb:
     use: chroma
     qdrant:
@@ -30,4 +40,4 @@ vdb:
 plugin:
     runtime_ws_url: 'ws://langbot_plugin_runtime:5400/control/ws'
     enable_marketplace: true
-    cloud_service_url: 'https://space.langbot.app'
+    cloud_service_url: 'https://space.langbot.app'