diff --git a/pkg/persistence/databases/postgresql.py b/pkg/persistence/databases/postgresql.py new file mode 100644 index 00000000..19b3706b --- /dev/null +++ b/pkg/persistence/databases/postgresql.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import pip +import sqlalchemy.ext.asyncio as sqlalchemy_asyncio +import sys + +from .. import database + + +@database.manager_class('postgresql') +class PostgreSQLDatabaseManager(database.BaseDatabaseManager): + """PostgreSQL database manager""" + + async def initialize(self) -> None: + + # default to PostgreSQL with asyncpg driver + try: + __import__("asyncpg") + except ImportError: + print('以下依赖包未安装,将自动安装,请完成后重启程序:') + print( + 'The dependence package asyncpg is missing, it will be installed automatically, please restart the program after completion:' + ) + pip.main(['install', "asyncpg"]) + print('已自动安装缺失的依赖包 asyncpg ,请重启程序。') + print('The missing dependence asyncpg have been installed automatically, please restart the program.') + sys.exit(0) + + engine_url = self.ap.instance_config.data['system'].get('database', {}).get('engine_url', 'postgresql+asyncpg://root:***@127.0.0.1:5432/postgres') + 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..7b219e4c 100644 --- a/pkg/persistence/databases/sqlite.py +++ b/pkg/persistence/databases/sqlite.py @@ -10,5 +10,5 @@ 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}') + engine_url = self.ap.instance_config.data['system'].get('database', {}).get('engine_url', 'sqlite+aiosqlite:///data/langbot.db') + self.engine = sqlalchemy_asyncio.create_async_engine(engine_url) diff --git a/pkg/persistence/mgr.py b/pkg/persistence/mgr.py index 9b926733..06ccc859 100644 --- a/pkg/persistence/mgr.py +++ b/pkg/persistence/mgr.py @@ -38,9 +38,12 @@ class PersistenceManager: async def initialize(self): self.ap.logger.info('Initializing database...') + database_type = self.ap.instance_config.data['system'].get('database', {}).get('type', 'sqlite') 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/templates/config.yaml b/templates/config.yaml index 8c0ac685..dbb5cc46 100644 --- a/templates/config.yaml +++ b/templates/config.yaml @@ -20,6 +20,9 @@ system: jwt: expire: 604800 secret: '' +database: + type: sqlite + engine_url: 'sqlite+aiosqlite:///data/langbot.db' vdb: use: chroma qdrant: @@ -30,4 +33,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