diff --git a/backend/cli.py b/backend/cli.py index a7d47199..dbb5f280 100644 --- a/backend/cli.py +++ b/backend/cli.py @@ -18,12 +18,13 @@ from rich.prompt import IntPrompt, Prompt from rich.table import Table from rich.text import Text from sqlalchemy import text -from sqlalchemy.ext.asyncio import create_async_engine +from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession from watchfiles import PythonFilter from backend import __version__ from backend.common.enums import DataBaseType, PrimaryKeyType from backend.common.exception.errors import BaseExceptionError +from backend.common.model import MappedBase from backend.core.conf import settings from backend.core.path_conf import ( ENV_EXAMPLE_FILE_PATH, @@ -31,8 +32,13 @@ from backend.core.path_conf import ( MYSQL_SCRIPT_DIR, POSTGRESQL_SCRIPT_DIR, ) -from backend.database.db import async_db_session, create_tables, drop_tables -from backend.database.redis import redis_client +from backend.database.db import ( + async_db_session, + create_database_async_engine, + create_database_async_session, + create_database_url, +) +from backend.database.redis import RedisCli, redis_client from backend.plugin.core import get_plugin_sql, get_plugins from backend.plugin.installer import install_git_plugin, install_zip_plugin from backend.utils.console import console @@ -103,19 +109,10 @@ def setup_env_file() -> bool: return True -async def create_database_if_not_exists() -> bool: - from sqlalchemy import URL - +async def create_database(conn: AsyncConnection) -> bool: try: terminate_sql = None if DataBaseType.mysql == settings.DATABASE_TYPE: - url = URL.create( - drivername='mysql+asyncmy', - username=settings.DATABASE_USER, - password=settings.DATABASE_PASSWORD, - host=settings.DATABASE_HOST, - port=settings.DATABASE_PORT, - ) check_sql = f"SHOW DATABASES LIKE '{settings.DATABASE_SCHEMA}'" drop_sql = f'DROP DATABASE IF EXISTS `{settings.DATABASE_SCHEMA}`' create_sql = ( @@ -123,38 +120,23 @@ async def create_database_if_not_exists() -> bool: f'COLLATE {settings.DATABASE_CHARSET}_unicode_ci' ) else: - url = URL.create( - drivername='postgresql+asyncpg', - username=settings.DATABASE_USER, - password=settings.DATABASE_PASSWORD, - host=settings.DATABASE_HOST, - port=settings.DATABASE_PORT, - database='postgres', - ) check_sql = f"SELECT 1 FROM pg_database WHERE datname = '{settings.DATABASE_SCHEMA}'" + drop_sql = f'DROP DATABASE IF EXISTS {settings.DATABASE_SCHEMA}' + create_sql = f'CREATE DATABASE {settings.DATABASE_SCHEMA}' terminate_sql = ( f'SELECT pg_terminate_backend(pid) FROM pg_stat_activity ' f"WHERE datname = '{settings.DATABASE_SCHEMA}' AND pid <> pg_backend_pid()" ) - drop_sql = f'DROP DATABASE IF EXISTS {settings.DATABASE_SCHEMA}' - create_sql = f'CREATE DATABASE {settings.DATABASE_SCHEMA}' - engine = create_async_engine(url, isolation_level='AUTOCOMMIT') - - try: - async with engine.connect() as conn: - result = await conn.execute(text(check_sql)) - exists = result.fetchone() is not None - - console.print(f'重建 {settings.DATABASE_SCHEMA} 数据库...', style='white') - if exists: - if terminate_sql: - await conn.execute(text(terminate_sql)) - await conn.execute(text(drop_sql)) - await conn.execute(text(create_sql)) - console.print('数据库创建成功', style='green') - finally: - await engine.dispose() + result = await conn.execute(text(check_sql)) + exists = result.fetchone() is not None + console.print(f'重建 {settings.DATABASE_SCHEMA} 数据库...', style='white') + if exists: + if terminate_sql: + await conn.execute(text(terminate_sql)) + await conn.execute(text(drop_sql)) + await conn.execute(text(create_sql)) + console.print('数据库创建成功', style='green') except Exception as e: console.print(f'数据库创建失败: {e}', style='red') return False @@ -174,50 +156,63 @@ async def auto_init() -> None: console.print(Panel(panel_content, title=f'fba (v{__version__}) - 环境变量', border_style='cyan', padding=(1, 2))) if not setup_env_file(): raise cappa.Exit('.env 文件配置失败', code=1) + async_init_engine = create_database_async_engine(create_database_url(with_database=False)) + async_init_db_session = create_database_async_session(async_init_engine) + redis_init_client = RedisCli( + host=settings.REDIS_HOST, + port=settings.REDIS_PORT, + password=settings.REDIS_PASSWORD, + db=settings.REDIS_DATABASE, + ) + await redis_init_client.init() console.print('\n[bold cyan]步骤 2/3:[/] 数据库创建', style='bold') panel_content = Text() panel_content.append('【数据库配置】', style='bold green') panel_content.append('\n\n • 类型: ') panel_content.append(f'{settings.DATABASE_TYPE}', style='yellow') - panel_content.append('\n • 数据库：') - panel_content.append(f'{settings.DATABASE_SCHEMA}', style='yellow') panel_content.append('\n • 主机：') panel_content.append(f'{settings.DATABASE_HOST}:{settings.DATABASE_PORT}', style='yellow') + panel_content.append('\n • 数据库：') + panel_content.append(f'{settings.DATABASE_SCHEMA}', style='yellow') + panel_content.append('\n • 主键模式：') + panel_content.append(f'{settings.DATABASE_PK_MODE}', style='yellow') console.print(Panel(panel_content, title=f'fba (v{__version__}) - 数据库', border_style='cyan', padding=(1, 2))) ok = Prompt.ask('即将[red]新建/重建数据库[/red]，确认继续吗？', choices=['y', 'n'], default='n') if ok.lower() == 'y': - if not await create_database_if_not_exists(): - raise cappa.Exit('数据库创建失败', code=1) + async with async_init_engine.connect() as conn: + await conn.execution_options(isolation_level='AUTOCOMMIT') + if not await create_database(conn): + raise cappa.Exit('数据库创建失败', code=1) else: console.print('已取消数据库操作', style='yellow') console.print('\n[bold cyan]步骤 3/3:[/] 初始化数据库表和数据', style='bold') - await init() + async with async_init_db_session.begin() as db: + await init(db, redis_init_client) -async def init() -> None: +async def init(db: AsyncSession, redis: RedisCli) -> None: panel_content = Text() panel_content.append('【数据库配置】', style='bold green') panel_content.append('\n\n • 类型: ') panel_content.append(f'{settings.DATABASE_TYPE}', style='yellow') - panel_content.append('\n • 数据库：') - panel_content.append(f'{settings.DATABASE_SCHEMA}', style='yellow') panel_content.append('\n • 主机：') panel_content.append(f'{settings.DATABASE_HOST}:{settings.DATABASE_PORT}', style='yellow') + panel_content.append('\n • 数据库：') + panel_content.append(f'{settings.DATABASE_SCHEMA}', style='yellow') panel_content.append('\n • 主键模式：') - panel_content.append( - f'{settings.DATABASE_PK_MODE}', - style='yellow', - ) + panel_content.append(f'{settings.DATABASE_PK_MODE}', style='yellow') pk_details = panel_content.from_markup( '[link=https://fastapi-practices.github.io/fastapi_best_architecture_docs/backend/reference/pk.html]（了解详情）[/]' ) panel_content.append(pk_details) panel_content.append('\n\n【Redis 配置】', style='bold green') - panel_content.append('\n\n • 数据库：') + panel_content.append('\n\n • 主机：') + panel_content.append(f'{settings.REDIS_HOST}:{settings.REDIS_PORT}', style='yellow') + panel_content.append('\n • 数据库：') panel_content.append(f'{settings.REDIS_DATABASE}', style='yellow') plugins = get_plugins() panel_content.append('\n\n【已安装插件】', style='bold green') @@ -236,18 +231,25 @@ async def init() -> None: console.print('开始初始化...', style='white') try: console.print('清理 Redis 缓存', style='white') - await redis_client.delete_prefix(settings.JWT_USER_REDIS_PREFIX) - await redis_client.delete_prefix(settings.TOKEN_EXTRA_INFO_REDIS_PREFIX) - await redis_client.delete_prefix(settings.TOKEN_REDIS_PREFIX) - await redis_client.delete_prefix(settings.TOKEN_REFRESH_REDIS_PREFIX) + for prefix in [ + settings.JWT_USER_REDIS_PREFIX, + settings.TOKEN_EXTRA_INFO_REDIS_PREFIX, + settings.TOKEN_REDIS_PREFIX, + settings.TOKEN_REFRESH_REDIS_PREFIX, + ]: + await redis.delete_prefix(prefix) + console.print('重建数据库表', style='white') - await drop_tables() - await create_tables() + conn = await db.connection() + await conn.run_sync(MappedBase.metadata.drop_all) + await conn.run_sync(MappedBase.metadata.create_all) + console.print('执行 SQL 脚本', style='white') sql_scripts = await get_sql_scripts() for sql_script in sql_scripts: console.print(f'正在执行：{sql_script}', style='white') - await execute_sql_scripts(sql_script, is_init=True) + await execute_sql_scripts(db, sql_script, is_init=True) + console.print('初始化成功', style='green') console.print('\n快试试 [bold cyan]fba run[/bold cyan] 启动服务吧~') except Exception as e: @@ -354,7 +356,8 @@ async def install_plugin( sql_file = await get_plugin_sql(plugin_name, db_type, pk_type) if sql_file and not no_sql: console.print('开始自动执行插件 SQL 脚本...', style='bold cyan') - await execute_sql_scripts(sql_file) + async with async_db_session.begin() as db: + await execute_sql_scripts(db, sql_file) except Exception as e: raise cappa.Exit(e.msg if isinstance(e, BaseExceptionError) else str(e), code=1) @@ -382,14 +385,13 @@ async def get_sql_scripts() -> list[str]: return sql_scripts -async def execute_sql_scripts(sql_scripts: str, *, is_init: bool = False) -> None: - async with async_db_session.begin() as db: - try: - stmts = await parse_sql_script(sql_scripts) - for stmt in stmts: - await db.execute(text(stmt)) - except Exception as e: - raise cappa.Exit(f'SQL 脚本执行失败：{e}', code=1) +async def execute_sql_scripts(db: AsyncSession, sql_scripts: str, *, is_init: bool = False) -> None: + try: + stmts = await parse_sql_script(sql_scripts) + for stmt in stmts: + await db.execute(text(stmt)) + except Exception as e: + raise cappa.Exit(f'SQL 脚本执行失败：{e}', code=1) if not is_init: console.print('SQL 脚本已执行完成', style='bold green') @@ -464,7 +466,8 @@ class Init: if self.auto: await auto_init() else: - await init() + async with async_db_session.begin() as db: + await init(db, redis_client) @cappa.command(help='运行 API 服务', default_long=True) @@ -621,7 +624,8 @@ class FbaCli: async def __call__(self) -> None: if self.sql: - await execute_sql_scripts(self.sql) + async with async_db_session.begin() as db: + await execute_sql_scripts(db, self.sql) def main() -> None: