Add the database primary key mode config (#953)

* Add the database primary key mode config

* Update auto to autoincrement

backend/plugin/code_generator/crud/crud_code.py

@@ -3,6 +3,7 @@ from collections.abc import Sequence
 from sqlalchemy import Row, RowMapping, text
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from backend.common.enums import DataBaseType
 from backend.core.conf import settings
 
 
@@ -18,7 +19,7 @@ class CRUDGen:
         :param table_schema: 数据库 schema 名称
         :return:
         """
-        if settings.DATABASE_TYPE == 'mysql':
+        if DataBaseType.mysql == settings.DATABASE_TYPE:
             sql = """
             SELECT table_name AS table_name, table_comment AS table_comment
             FROM information_schema.tables
@@ -48,7 +49,7 @@ class CRUDGen:
         :param table_name: 表名
         :return:
         """
-        if settings.DATABASE_TYPE == 'mysql':
+        if DataBaseType.mysql == settings.DATABASE_TYPE:
             sql = """
             SELECT table_name AS table_name, table_comment AS table_comment
             FROM information_schema.tables
@@ -79,7 +80,7 @@ class CRUDGen:
         :param table_name: 表名
         :return:
         """
-        if settings.DATABASE_TYPE == 'mysql':
+        if DataBaseType.mysql == settings.DATABASE_TYPE:
             sql = """
             SELECT column_name AS column_name,
             CASE WHEN column_key = 'PRI' THEN 1 ELSE 0 END AS is_pk,

backend/plugin/code_generator/service/column_service.py

@@ -2,6 +2,7 @@ from collections.abc import Sequence
 
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from backend.common.enums import DataBaseType
 from backend.common.exception import errors
 from backend.core.conf import settings
 from backend.plugin.code_generator.crud.crud_column import gen_column_dao
@@ -32,7 +33,7 @@ class GenColumnService:
     @staticmethod
     async def get_types() -> list[str]:
         """获取所有列类型"""
-        if settings.DATABASE_TYPE == 'mysql':
+        if DataBaseType.mysql == settings.DATABASE_TYPE:
             types = GenMySQLColumnType.get_member_keys()
         else:
             types = GenPostgreSQLColumnType.get_member_keys()

backend/plugin/code_generator/utils/type_conversion.py

@@ -1,5 +1,6 @@
 from functools import lru_cache
 
+from backend.common.enums import DataBaseType
 from backend.core.conf import settings
 from backend.plugin.code_generator.enums import GenMySQLColumnType, GenPostgreSQLColumnType
 
@@ -12,7 +13,7 @@ def sql_type_to_sqlalchemy(typing: str) -> str:
     :param typing: SQL 类型字符串
     :return:
     """
-    if settings.DATABASE_TYPE == 'mysql':
+    if DataBaseType.mysql == settings.DATABASE_TYPE:
         if typing in GenMySQLColumnType.get_member_keys():
             return typing
     else:
@@ -30,7 +31,7 @@ def sql_type_to_pydantic(typing: str) -> str:
     :return:
     """
     try:
-        if settings.DATABASE_TYPE == 'mysql':
+        if DataBaseType.mysql == settings.DATABASE_TYPE:
             return GenMySQLColumnType[typing].value
         if typing == 'CHARACTER VARYING':  # postgresql 中 DDL VARCHAR 的别名
             return 'str'

backend/plugin/tools.py

@@ -79,16 +79,16 @@ async def get_plugin_sql(plugin: str, db_type: DataBaseType, pk_type: PrimaryKey
     """
     if db_type == DataBaseType.mysql:
         mysql_dir = PLUGIN_DIR / plugin / 'sql' / 'mysql'
-        if pk_type == PrimaryKeyType.autoincrement:
-            sql_file = mysql_dir / 'init.sql'
-        else:
-            sql_file = mysql_dir / 'init_snowflake.sql'
+        sql_file = (
+            mysql_dir / 'init.sql' if pk_type == PrimaryKeyType.autoincrement else mysql_dir / 'init_snowflake.sql'
+        )
     else:
         postgresql_dir = PLUGIN_DIR / plugin / 'sql' / 'postgresql'
-        if pk_type == PrimaryKeyType.autoincrement:
-            sql_file = postgresql_dir / 'init.sql'
-        else:
-            sql_file = postgresql_dir / 'init_snowflake.sql'
+        sql_file = (
+            postgresql_dir / 'init.sql'
+            if pk_type == PrimaryKeyType.autoincrement
+            else postgresql_dir / 'init_snowflake.sql'
+        )
 
     path = anyio.Path(sql_file)
     if not await path.exists():