import io import os import shutil import tempfile import zipfile from collections.abc import Sequence import anyio from anyio import open_file from pydantic.alias_generators import to_pascal from sqlalchemy import RowMapping from sqlalchemy.ext.asyncio import AsyncSession from starlette.concurrency import run_in_threadpool from backend.common.exception import errors from backend.core.conf import settings from backend.core.path_conf import BASE_PATH from backend.plugin.code_generator.crud.crud_business import code_gen_business_dao from backend.plugin.code_generator.crud.crud_code_gen import code_gen_dao from backend.plugin.code_generator.crud.crud_column import code_gen_column_dao from backend.plugin.code_generator.model import CodeGenBusiness from backend.plugin.code_generator.schema.business import CreateCodeGenBusinessParam from backend.plugin.code_generator.schema.column import CreateCodeGenColumnInternalParam from backend.plugin.code_generator.schema.code_gen import ImportParam from backend.plugin.code_generator.service.column_service import code_gen_column_service from backend.plugin.code_generator.utils.format_code import format_python_code from backend.plugin.code_generator.utils.gen_template import gen_template from backend.plugin.code_generator.utils.type_conversion import sql_type_to_pydantic from backend.utils.locks import acquire_distributed_reload_lock class CodeGenService: """代码生成服务类""" @staticmethod async def get_tables(*, db: AsyncSession, table_schema: str) -> Sequence[RowMapping]: """ 获取指定 schema 下的所有表名 :param db: 数据库会话 :param table_schema: 数据库 schema 名称 :return: """ return await code_gen_dao.get_all_tables(db, table_schema) @staticmethod async def import_business_and_model(*, db: AsyncSession, obj: ImportParam) -> None: """ 导入业务和模型列数据 :param db: 数据库会话 :param obj: 导入参数对象 :return: """ if settings.ENVIRONMENT != 'dev': raise errors.ForbiddenError(msg='禁止在非开发环境下导入代码生成业务') table_info = await code_gen_dao.get_table(db, obj.table_schema, obj.table_name) if not table_info: raise errors.NotFoundError(msg='数据库表不存在') business_info = await code_gen_business_dao.get_by_name(db, obj.table_name) if business_info: raise errors.ConflictError(msg='已存在相同数据库表业务') table_name = table_info['table_name'] doc_comment = ( table_info['table_comment'][:-1] if table_info['table_comment'][-1] == '表' else table_info['table_comment'] or table_name.split('_')[-1] ) new_business = CodeGenBusiness( **CreateCodeGenBusinessParam( app_name=obj.app, table_name=table_name, doc_comment=doc_comment, table_comment=table_info['table_comment'], class_name=to_pascal(table_name), schema_name=to_pascal(table_name), filename=table_name, tag=doc_comment, ).model_dump(), ) db.add(new_business) await db.flush() column_info = await code_gen_dao.get_all_columns(db, obj.table_schema, table_name) code_gen_columns = [] for column in column_info: column_type = column['column_type'].split('(')[0].upper() pd_type = sql_type_to_pydantic(column_type) code_gen_columns.append( CreateCodeGenColumnInternalParam( name=column['column_name'], comment=column['column_comment'], type=column_type, sort=column['sort'], length=column['column_type'].split('(')[1][:-1] if pd_type == 'str' and '(' in column['column_type'] else 0, is_pk=column['is_pk'], is_nullable=column['is_nullable'], code_gen_business_id=new_business.id, pd_type=pd_type, ), ) await code_gen_column_dao.bulk_create(db, code_gen_columns) @staticmethod async def _render_tpl_code(*, db: AsyncSession, business: CodeGenBusiness) -> dict[str, str]: """ 渲染模板代码 :param db: 数据库会话 :param business: 业务对象 :return: """ gen_models = await code_gen_column_service.get_columns(db=db, business_id=business.id) if not gen_models: raise errors.NotFoundError(msg='代码生成模型表为空') gen_vars = gen_template.get_vars(business, gen_models) template_mapping = gen_template.get_template_path_mapping(business) rendered_codes = {} for template_path, output_path in template_mapping.items(): code = await gen_template.get_template(template_path).render_async(**gen_vars) if output_path.endswith('.py'): code = await format_python_code(code) rendered_codes[output_path] = code return rendered_codes @staticmethod async def _inject_app_router(*, app_name: str, write: bool = True) -> str | None: """ 注入应用路由 :param app_name: :param write: 是否写入文件 :return: """ app_root_router = BASE_PATH / 'app' / 'router.py' async with await open_file(app_root_router, 'r', encoding='utf-8') as f: content = await f.read() import_line = f'from backend.app.{app_name}.api.router import v1 as {app_name}_v1' include_line = f'router.include_router({app_name}_v1)' has_import = import_line in content # type: ignore has_include = include_line in content # type: ignore if has_import and has_include: return None if not has_import: content = f'{import_line}\n{content}' if not has_include: content = f'{content}\n{include_line}' content = await format_python_code(content) if write: async with await open_file(app_root_router, 'w', encoding='utf-8') as f: await f.write(content) return content async def preview(self, *, db: AsyncSession, pk: int) -> dict[str, bytes]: """ 预览生成的代码 :param db: 数据库会话 :param pk: 业务 ID :return: """ business = await code_gen_business_dao.get(db, pk) if not business: raise errors.NotFoundError(msg='业务不存在') codes = {} backend_path = 'fastapi-best-architecture/backend/app/' init_files = gen_template.get_init_files(business) for filepath, content in init_files.items(): codes[f'{backend_path}{filepath}'] = content.encode('utf-8') rendered_codes = await self._render_tpl_code(db=db, business=business) for filepath, code in rendered_codes.items(): codes[f'{backend_path}{filepath}'] = code.encode('utf-8') app_router_content = await self._inject_app_router(app_name=business.app_name, write=False) if app_router_content: codes[f'{backend_path}router.py'] = app_router_content.encode('utf-8') return codes @staticmethod async def get_generate_path(*, db: AsyncSession, pk: int) -> list[str]: """ 获取代码生成路径 :param db: 数据库会话 :param pk: 业务 ID :return: """ business = await code_gen_business_dao.get(db, pk) if not business: raise errors.NotFoundError(msg='业务不存在') gen_path = business.gen_path or '/backend/app' paths = [] init_files = gen_template.get_init_files(business) paths.extend(os.path.join(gen_path, *filepath.split('/')) for filepath in init_files.keys()) template_mapping = gen_template.get_template_path_mapping(business) paths.extend(os.path.join(gen_path, *filepath.split('/')) for filepath in template_mapping.values()) return paths async def generate(self, *, db: AsyncSession, pk: int) -> str: """ 生成代码文件 :param db: 数据库会话 :param pk: 业务 ID :return: """ if settings.ENVIRONMENT != 'dev': raise errors.ForbiddenError(msg='禁止在非开发环境下生成代码') business = await code_gen_business_dao.get(db, pk) if not business: raise errors.NotFoundError(msg='业务不存在') gen_path = business.gen_path or str(BASE_PATH / 'app') async with acquire_distributed_reload_lock(): with tempfile.TemporaryDirectory() as tmp_dir: all_files = {} init_files = gen_template.get_init_files(business) all_files.update(init_files) rendered_codes = await self._render_tpl_code(db=db, business=business) all_files.update(rendered_codes) for filepath, content in all_files.items(): full_path = os.path.join(tmp_dir, *filepath.split('/')) code_folder = anyio.Path(full_path).parent await code_folder.mkdir(parents=True, exist_ok=True) async with await open_file(full_path, 'w', encoding='utf-8') as f: await f.write(content) for item in os.listdir(tmp_dir): src = os.path.join(tmp_dir, item) dst = os.path.join(gen_path, item) src_path = anyio.Path(src) if await src_path.is_dir(): await run_in_threadpool(shutil.copytree, src, dst, dirs_exist_ok=True) else: await run_in_threadpool(shutil.copy2, src, dst) await self._inject_app_router(app_name=business.app_name) return gen_path async def download(self, *, db: AsyncSession, pk: int) -> io.BytesIO: """ 下载生成的代码 :param db: 数据库会话 :param pk: 业务 ID :return: """ business = await code_gen_business_dao.get(db, pk) if not business: raise errors.NotFoundError(msg='业务不存在') all_files = {} init_files = gen_template.get_init_files(business) all_files.update(init_files) rendered_codes = await self._render_tpl_code(db=db, business=business) all_files.update(rendered_codes) app_router_content = await self._inject_app_router(app_name=business.app_name, write=False) if app_router_content: all_files['router.py'] = app_router_content bio = io.BytesIO() with zipfile.ZipFile(bio, 'w') as zf: for filepath, content in all_files.items(): zf.writestr(filepath, content) bio.seek(0) return bio code_gen_service: CodeGenService = CodeGenService()