FastapiAdmin/backend/app/api/v1/module_generator/gencode/service.py

# -*- coding:utf-8 -*-

import io
import json
import os
import zipfile
from datetime import datetime
from sqlalchemy.ext.asyncio import AsyncSession
from typing import Any, List, Dict, Optional, Sequence

from app.config.setting import settings
from app.core.base_model import CamelCaseUtil
from app.core.exceptions import CustomException
from app.utils.gen_util import GenUtils
from app.utils.template_util import TemplateInitializer, TemplateUtils
from app.common.constant import GenConstant
from app.common.response import SuccessResponse
from app.api.v1.module_system.user.schema import UserOutSchema
from app.api.v1.module_system.auth.schema import AuthSchema
from .schema import GenTableCreateSchema, GenTableUpdateSchema, GenTableOutSchema, GenTableDeleteSchema, GenTableColumnCreateSchema, GenTableColumnUpdateSchema, GenTableColumnOutSchema, GenTableColumnDeleteSchema
from .param import GenTableQueryParam
from .crud import GenTableColumnCRUD, GenTableCRUD
from .model import GenTableModel, GenTableColumnModel


# 定义默认的GenConfig值
GEN_PATH = "generated_code"  # 默认生成路径


class GenTableService:
    """
    代码生成业务表服务层
    """

    @classmethod
    async def get_gen_table_list_service(
        cls, auth: AuthSchema, query_object: GenTableQueryParam, is_page: bool = False
    ):
        """
        获取代码生成业务表列表信息service

        :param auth: 认证信息
        :param query_object: 查询参数对象
        :param is_page: 是否开启分页
        :return: 代码生成业务列表信息对象
        """
        gen_table_dao = GenTableCRUD(auth=auth)
        gen_table_list_result = await gen_table_dao.get_gen_table_list(auth.db, query_object, is_page)

        return gen_table_list_result

    @classmethod
    async def get_gen_db_table_list_service(
        cls, auth: AuthSchema, query_object: GenTableQueryParam, is_page: bool = False
    ):
        """
        获取数据库列表信息service

        :param auth: 认证信息
        :param query_object: 查询参数对象
        :param is_page: 是否开启分页
        :return: 数据库列表信息对象
        """
        gen_table_dao = GenTableCRUD(auth=auth)
        gen_db_table_list_result = await gen_table_dao.get_gen_db_table_list(auth.db, query_object, is_page)

        return gen_db_table_list_result

    @classmethod
    async def get_gen_db_table_list_by_name_service(cls, auth: AuthSchema, table_names: List[str]) -> list[GenTableOutSchema]:
        """
        根据表名称组获取数据库列表信息service

        :param auth: 认证信息
        :param table_names: 表名称组
        :return: 数据库列表信息对象
        """
        gen_table_dao = GenTableCRUD(auth=auth)
        gen_db_table_list_result = await gen_table_dao.get_gen_db_table_list_by_names(auth.db, table_names)

        return [GenTableOutSchema(**gen_table) for gen_table in CamelCaseUtil.transform_result(gen_db_table_list_result)]

    @classmethod
    async def import_gen_table_service(
        cls, auth: AuthSchema, gen_table_list: List[GenTableOutSchema], current_user: UserOutSchema
    ):
        """
        导入表结构service

        :param auth: 认证信息
        :param gen_table_list: 导入表列表
        :param current_user: 当前用户信息对象
        :return: 导入结果
        """
        try:
            gen_table_dao = GenTableCRUD(auth=auth)
            gen_table_column_dao = GenTableColumnCRUD(auth=auth)
            
            for table in gen_table_list:
                table_name = table.table_name
                GenUtils.init_table(table, current_user.username)  # 使用username而不是user.user_name
                add_gen_table = await gen_table_dao.create(data=table.model_dump())
                if add_gen_table:
                    table.table_id = add_gen_table.id
                    gen_table_columns = await gen_table_column_dao.get_gen_db_table_columns_by_name(auth.db, table_name or "")
                    for column in [
                        GenTableColumnOutSchema(**gen_table_column)
                        for gen_table_column in CamelCaseUtil.transform_result(gen_table_columns)
                    ]:
                        GenUtils.init_column_field(column, table)
                        await gen_table_column_dao.create(data=column.model_dump())
            await auth.db.commit()
            return SuccessResponse(msg='导入成功')
        except Exception as e:
            try:
                await auth.db.rollback()
            except:
                pass  # 忽略回滚错误
            raise CustomException(msg=f'导入失败, {str(e)}')

    @classmethod
    async def edit_gen_table_service(cls, auth: AuthSchema, page_object: GenTableUpdateSchema) -> Dict[str, Any]:
        """
        编辑业务表信息service

        :param auth: 认证信息
        :param page_object: 编辑业务表对象
        :return: 编辑业务表校验结果
        """
        gen_table_dao = GenTableCRUD(auth=auth)
        gen_table_column_dao = GenTableColumnCRUD(auth=auth)
        
        # 检查必要字段是否存在
        if getattr(page_object, 'table_id', None) is None:
            raise CustomException(msg='业务表ID不能为空')
            
        edit_gen_table = page_object.model_dump(exclude_unset=True, by_alias=True)
        gen_table_info = await cls.get_gen_table_by_id_service(auth, page_object.table_id)
        if gen_table_info.table_id:
            try:
                # 确保options字段存在且为有效JSON
                if 'options' not in edit_gen_table or edit_gen_table['options'] is None:
                    edit_gen_table['options'] = '{}'  # 默认空对象
                else:
                    # 验证options是否为有效的JSON
                    try:
                        json.loads(edit_gen_table['options'])
                    except json.JSONDecodeError:
                        edit_gen_table['options'] = '{}'
                
                await gen_table_dao.update(id=page_object.table_id, data=edit_gen_table)
                if page_object.columns:
                    for gen_table_column in page_object.columns:
                        gen_table_column.update_by = page_object.update_by
                        gen_table_column.update_time = datetime.now()
                        if gen_table_column.column_id is not None:
                            await gen_table_column_dao.update(
                                id=gen_table_column.column_id, 
                                data=gen_table_column.model_dump(by_alias=True)
                            )
                await auth.db.commit()
                return {"is_success": True, "message": "更新成功"}
            except Exception as e:
                try:
                    await auth.db.rollback()
                except:
                    pass  # 忽略回滚错误
                raise CustomException(msg=f'更新失败: {str(e)}')
        else:
            raise CustomException(msg='业务表不存在')

    @classmethod
    async def delete_gen_table_service(cls, auth: AuthSchema, page_object: GenTableDeleteSchema) -> SuccessResponse:
        """
        删除业务表信息service

        :param auth: 认证信息
        :param page_object: 删除业务表对象
        :return: 删除业务表校验结果
        """
        gen_table_dao = GenTableCRUD(auth=auth)
        gen_table_column_dao = GenTableColumnCRUD(auth=auth)
        
        if page_object.table_ids:
            table_id_list = page_object.table_ids.split(',')
            try:
                for table_id in table_id_list:
                    await gen_table_dao.delete(ids=[int(table_id)])
                    # 删除相关的字段信息
                    # 这里需要先查询出所有相关的column_id，然后删除
                    columns = await gen_table_column_dao.get_gen_table_column_list_by_table_id(auth.db, int(table_id))
                    if columns:
                        column_ids = [column.id for column in columns]
                        await gen_table_column_dao.delete(ids=column_ids)
                await auth.db.commit()
                return SuccessResponse(msg='删除成功')
            except Exception as e:
                try:
                    await auth.db.rollback()
                except:
                    pass  # 忽略回滚错误
                raise CustomException(msg=f'删除失败: {str(e)}')
        else:
            raise CustomException(msg='传入业务表id为空')

    @classmethod
    async def get_gen_table_by_id_service(cls, auth: AuthSchema, table_id: int) -> GenTableOutSchema:
        """
        获取需要生成的业务表详细信息service

        :param auth: 认证信息
        :param table_id: 需要生成的业务表id
        :return: 需要生成的业务表id对应的信息
        """
        gen_table_dao = GenTableCRUD(auth=auth)
        gen_table = await gen_table_dao.get_gen_table_by_id(auth.db, table_id)
        if gen_table:
            result = await cls.set_table_from_options(GenTableOutSchema(**CamelCaseUtil.transform_result(gen_table)))
            return result
        else:
            raise CustomException(msg='业务表不存在')

    @classmethod
    async def get_gen_table_all_service(cls, auth: AuthSchema) -> list[GenTableOutSchema]:
        """
        获取所有业务表信息service

        :param auth: 认证信息
        :return: 所有业务表信息列表
        """
        gen_table_dao = GenTableCRUD(auth=auth)
        gen_tables = await gen_table_dao.get_gen_table_all(auth.db)
        result = []
        for table in gen_tables:
            table_info = await cls.set_table_from_options(GenTableOutSchema(**CamelCaseUtil.transform_result(table)))
            result.append(table_info)
        return result

    @classmethod
    async def create_table_service(cls, auth: AuthSchema, sql: str, current_user: UserOutSchema) -> SuccessResponse:
        """
        创建表结构service

        :param auth: 认证信息
        :param sql: 建表语句
        :param current_user: 当前用户信息对象
        :return: 创建表结构结果
        """
        gen_table_dao = GenTableCRUD(auth=auth)
        
        try:
            # 执行SQL语句创建表
            await gen_table_dao.create_table_by_sql_dao(auth.db, [sql])
            await auth.db.commit()
            return SuccessResponse(msg='创建表结构成功')
        except Exception as e:
            try:
                await auth.db.rollback()
            except:
                pass  # 忽略回滚错误
            raise CustomException(msg=f'创建表结构失败: {str(e)}')

    @classmethod
    async def preview_code_service(cls, auth: AuthSchema, table_id: int) -> dict[Any, Any]:
        """
        预览代码service

        :param auth: 认证信息
        :param table_id: 业务表id
        :return: 预览数据列表
        """
        gen_table = await cls.get_gen_table_by_id_service(auth, table_id)
        await cls.set_sub_table(auth, gen_table)
        await cls._set_pk_column(gen_table)
        env = TemplateInitializer.init_jinja2()
        context = TemplateUtils.prepare_context(gen_table)
        template_list = TemplateUtils.get_template_list(
            gen_table.tpl_category or "", 
            gen_table.tpl_web_type or ""
        )
        preview_code_result = {}
        for template in template_list:
            render_content = env.get_template(template).render(**context)
            preview_code_result[template] = render_content
        return preview_code_result

    @classmethod
    async def generate_code_service(cls, auth: AuthSchema, table_name: str) -> SuccessResponse:
        """
        生成代码至指定路径service

        :param auth: 认证信息
        :param table_name: 业务表名称
        :return: 生成代码结果
        """
        env = TemplateInitializer.init_jinja2()
        render_info = await cls.__get_gen_render_info(auth, table_name)
        for template in render_info[0]:
            try:
                render_content = env.get_template(template).render(**render_info[2])
                gen_path = cls.__get_gen_path(render_info[3], template)
                if gen_path:
                    os.makedirs(os.path.dirname(gen_path), exist_ok=True)
                    with open(gen_path, 'w', encoding='utf-8') as f:
                        f.write(render_content)
            except Exception as e:
                raise CustomException(msg=f'渲染模板失败，表名：{render_info[3].table_name}，详细错误信息：{str(e)}')

        return SuccessResponse(msg='生成代码成功')

    @classmethod
    async def batch_gen_code_service(cls, auth: AuthSchema, table_names: List[str]) -> bytes:
        """
        批量生成代码service

        :param auth: 认证信息
        :param table_names: 业务表名称组
        :return: 下载代码结果
        """
        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
            for table_name in table_names:
                env = TemplateInitializer.init_jinja2()
                render_info = await cls.__get_gen_render_info(auth, table_name)
                for template_file, output_file in zip(render_info[0], render_info[1]):
                    render_content = env.get_template(template_file).render(**render_info[2])
                    zip_file.writestr(output_file, render_content)

        zip_data = zip_buffer.getvalue()
        zip_buffer.close()
        return zip_data

    @classmethod
    async def __get_gen_render_info(cls, auth: AuthSchema, table_name: str) -> list[Any]:
        """
        获取生成代码渲染模板相关信息

        :param auth: 认证信息
        :param table_name: 业务表名称
        :return: 生成代码渲染模板相关信息
        """
            
        gen_table_dao = GenTableCRUD(auth=auth)
        gen_table = await gen_table_dao.get_gen_table_by_name(auth.db, table_name)
        if gen_table:
            gen_table_schema = GenTableOutSchema(**CamelCaseUtil.transform_result(gen_table))
            await cls.set_sub_table(auth, gen_table_schema)
            await cls._set_pk_column(gen_table_schema)
            context = TemplateUtils.prepare_context(gen_table_schema)
            template_list = TemplateUtils.get_template_list(
                gen_table_schema.tpl_category or "", 
                gen_table_schema.tpl_web_type or ""
            )
            output_files = [TemplateUtils.get_file_name([template], gen_table_schema)[0] for template in template_list]

            return [template_list, output_files, context, gen_table_schema]
        else:
            raise CustomException(msg=f'业务表 {table_name} 不存在')

    @classmethod
    def __get_gen_path(cls, gen_table: GenTableOutSchema, template: str) -> Optional[str]:
        """
        根据GenTableModel对象和模板名称生成路径

        :param gen_table: GenTableModel对象
        :param template: 模板名称
        :return: 生成的路径
        """
        try:
            gen_path = gen_table.gen_path or ""
            if gen_path == '/':
                file_name = TemplateUtils.get_file_name([template], gen_table)[0]
                return os.path.join(os.getcwd(), GEN_PATH, file_name)
            else:
                file_name = TemplateUtils.get_file_name([template], gen_table)[0]
                return os.path.join(gen_path, file_name)
        except Exception:
            return None

    @classmethod
    async def sync_db_service(cls, auth: AuthSchema, table_name: str) -> SuccessResponse:
        """
        同步数据库service

        :param auth: 认证信息
        :param table_name: 业务表名称
        :return: 同步数据库结果
        """
        gen_table_dao = GenTableCRUD(auth=auth)
        gen_table_column_dao = GenTableColumnCRUD(auth=auth)
        
        gen_table = await gen_table_dao.get_gen_table_by_name(auth.db, table_name)
        if gen_table:
            table = GenTableOutSchema(**CamelCaseUtil.transform_result(gen_table))
            table_columns = table.columns or []  # 确保不为None
            table_column_map = {column.column_name: column for column in table_columns}
            query_db_table_columns = await gen_table_column_dao.get_gen_db_table_columns_by_name(auth.db, table_name)
            db_table_columns = [
                GenTableColumnOutSchema(**column) for column in CamelCaseUtil.transform_result(query_db_table_columns)
            ]
            if not db_table_columns:
                raise CustomException('同步数据失败，原表结构不存在')
            db_table_column_names = [column.column_name for column in db_table_columns]
            try:
                for column in db_table_columns:
                    GenUtils.init_column_field(column, table)
                    if column.column_name in table_column_map:
                        prev_column = table_column_map[column.column_name]
                        column.column_id = prev_column.column_id
                        if getattr(column, 'list', False):  # 使用getattr安全访问属性
                            column.dict_type = prev_column.dict_type
                            column.query_type = prev_column.query_type
                        if (
                            prev_column.is_required != ''
                            and not column.pk
                            and (column.insert or column.edit)
                            and (column.usable_column or column.super_column)
                        ):
                            column.is_required = prev_column.is_required
                            column.html_type = prev_column.html_type
                        if column.column_id is not None:
                            await gen_table_column_dao.update(id=column.column_id, data=column.model_dump(by_alias=True))
                    else:
                        await gen_table_column_dao.create(data=column.model_dump(by_alias=True))
                del_columns = [column for column in table_columns if column.column_name not in db_table_column_names]
                if del_columns:
                    for column in del_columns:
                        if column.column_id is not None:
                            await gen_table_column_dao.delete(ids=[column.column_id])
                await auth.db.commit()
                return SuccessResponse(msg='同步成功')
            except Exception as e:
                try:
                    await auth.db.rollback()
                except:
                    pass  # 忽略回滚错误
                raise CustomException(msg=f'同步失败: {str(e)}')
        else:
            raise CustomException('业务表不存在')

    @classmethod
    async def set_sub_table(cls, auth: AuthSchema, gen_table: GenTableOutSchema) -> None:
        """
        设置主子表信息

        :param auth: 认证信息
        :param gen_table: 业务表信息
        :return:
        """
            
        if gen_table.sub_table_name:
            gen_table_dao = GenTableCRUD(auth=auth)
            sub_table = await gen_table_dao.get_gen_table_by_name(auth.db, gen_table.sub_table_name)
            if sub_table:
                gen_table.sub_table = GenTableOutSchema(**CamelCaseUtil.transform_result(sub_table))

    @classmethod
    async def _set_pk_column(cls, gen_table: GenTableOutSchema) -> None:
        """
        设置主键列信息

        :param gen_table: 业务表信息
        :return:
        """
        if gen_table.columns:
            for column in gen_table.columns:
                if column.pk:
                    gen_table.pk_column = column
                    break
        if gen_table.pk_column is None and gen_table.columns:
            gen_table.pk_column = gen_table.columns[0]
        if gen_table.tpl_category == GenConstant.TPL_SUB and gen_table.sub_table:
            if gen_table.sub_table.columns:
                for column in gen_table.sub_table.columns:
                    if column.pk:
                        gen_table.sub_table.pk_column = column
                        break
            if gen_table.sub_table.pk_column is None and gen_table.sub_table.columns:
                gen_table.sub_table.pk_column = gen_table.sub_table.columns[0]

    @classmethod
    async def set_table_from_options(cls, gen_table: GenTableOutSchema) -> GenTableOutSchema:
        """
        设置代码生成其他选项值

        :param gen_table: 生成对象
        :return: 设置后的生成对象
        """
        params_obj = json.loads(gen_table.options) if gen_table.options else None
        if params_obj:
            gen_table.tree_code = params_obj.get(GenConstant.TREE_CODE)
            gen_table.tree_parent_code = params_obj.get(GenConstant.TREE_PARENT_CODE)
            gen_table.tree_name = params_obj.get(GenConstant.TREE_NAME)
            gen_table.parent_menu_id = params_obj.get(GenConstant.PARENT_MENU_ID)
            gen_table.parent_menu_name = params_obj.get(GenConstant.PARENT_MENU_NAME)

        return gen_table

    @classmethod
    async def validate_edit(cls, edit_gen_table: GenTableUpdateSchema):
        """
        编辑保存参数校验

        :param edit_gen_table: 编辑业务表对象
        """
        if edit_gen_table.tpl_category == GenConstant.TPL_TREE:
            # 从options字段获取参数，而不是params
            if not edit_gen_table.options:
                raise CustomException(msg='树表参数不能为空')
            
            params_obj = json.loads(edit_gen_table.options)

            if GenConstant.TREE_CODE not in params_obj:
                raise CustomException(msg='树编码字段不能为空')
            elif GenConstant.TREE_PARENT_CODE not in params_obj:
                raise CustomException(msg='树父编码字段不能为空')
            elif GenConstant.TREE_NAME not in params_obj:
                raise CustomException(msg='树名称字段不能为空')
        elif edit_gen_table.tpl_category == GenConstant.TPL_SUB:
            if not edit_gen_table.sub_table_name:
                raise CustomException(msg='关联子表的表名不能为空')
            elif not edit_gen_table.sub_table_fk_name:
                raise CustomException(msg='子表关联的外键名不能为空')


class GenTableColumnService:
    """
    代码生成业务表字段服务层
    """

    @classmethod
    async def get_gen_table_column_list_by_table_id_service(cls, auth: AuthSchema, table_id: int) -> List[GenTableColumnOutSchema]:
        """
        获取业务表字段列表信息service

        :param auth: 认证信息
        :param table_id: 业务表格id
        :return: 业务表字段列表信息对象
        """
        gen_table_column_dao = GenTableColumnCRUD(auth=auth)
        gen_table_column_list_result = await gen_table_column_dao.get_gen_table_column_list_by_table_id(auth.db, table_id)

        return [
            GenTableColumnOutSchema(**gen_table_column)
            for gen_table_column in CamelCaseUtil.transform_result(gen_table_column_list_result)
        ]