diff --git a/pkg/core/bootutils/deps.py b/pkg/core/bootutils/deps.py index 1a439af8..b2508b22 100644 --- a/pkg/core/bootutils/deps.py +++ b/pkg/core/bootutils/deps.py @@ -40,6 +40,7 @@ required_deps = { 'sqlmodel': 'sqlmodel', 'telegramify_markdown': 'telegramify-markdown', 'slack_sdk': 'slack_sdk', + 'asyncpg': 'asyncpg', } diff --git a/pkg/persistence/databases/postgresql.py b/pkg/persistence/databases/postgresql.py new file mode 100644 index 00000000..f63d8f61 --- /dev/null +++ b/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) diff --git a/pkg/persistence/databases/sqlite.py b/pkg/persistence/databases/sqlite.py index c1337459..ed096040 100644 --- a/pkg/persistence/databases/sqlite.py +++ b/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) diff --git a/pkg/persistence/mgr.py b/pkg/persistence/mgr.py index 9b926733..bab2cde5 100644 --- a/pkg/persistence/mgr.py +++ b/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() diff --git a/pkg/persistence/migrations/dbm006_plugin_install_source.py b/pkg/persistence/migrations/dbm006_plugin_install_source.py index 37f74929..7ab0afd3 100644 --- a/pkg/persistence/migrations/dbm006_plugin_install_source.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index e8d8be83..a8ada81e 100644 --- a/pyproject.toml +++ b/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", diff --git a/templates/config.yaml b/templates/config.yaml index 8c0ac685..00de2a9b 100644 --- a/templates/config.yaml +++ b/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' \ No newline at end of file