Adaptation to postgresql code generation

wu-clan · wu-clan · commit a4b392f04b13 · 2024-12-24T21:57:13.000+08:00
diff --git a/backend/app/generator/crud/crud_gen.py b/backend/app/generator/crud/crud_gen.py
@@ -5,47 +5,100 @@
 from sqlalchemy import Row, text
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from backend.core.conf import settings
+
 
 class CRUDGen:
     @staticmethod
     async def get_all_tables(db: AsyncSession, table_schema: str) -> Sequence[str]:
-        stmt = text(
-            'select table_name as table_name '
-            'from information_schema.tables '
-            'where table_name not like "sys_gen_%" '
-            'and table_schema = :table_schema;'
-        ).bindparams(table_schema=table_schema)
+        if settings.DATABASE_TYPE == 'mysql':
+            sql = """
+            SELECT table_name AS table_name FROM information_schema.tables 
+            WHERE table_name NOT LIKE 'sys_gen_%' 
+            AND table_schema = :table_schema;
+            """
+        else:
+            sql = """
+            SELECT table_name AS table_name FROM information_schema.tables 
+            WHERE table_name NOT LIKE 'sys_gen_%' 
+            AND table_catalog = :table_schema
+            AND table_schema = 'public'; -- schema 通常是 'public'
+            """
+        stmt = text(sql).bindparams(table_schema=table_schema)
         result = await db.execute(stmt)
         return result.scalars().all()
 
     @staticmethod
     async def get_table(db: AsyncSession, table_name: str) -> Row[tuple]:
-        stmt = text(
-            'select table_name as table_name, table_comment as table_comment '
-            'from information_schema.tables '
-            'where table_name not like "sys_gen_%" '
-            'and table_name = :table_name;'
-        ).bindparams(table_name=table_name)
+        if settings.DATABASE_TYPE == 'mysql':
+            sql = """
+            SELECT table_name AS table_name, table_comment AS table_comment FROM information_schema.tables 
+            WHERE table_name NOT LIKE 'sys_gen_%' 
+            AND table_name = :table_name;
+            """
+        else:
+            sql = """
+            SELECT t.tablename AS table_name,
+            pg_catalog.obj_description(t.tablename::regclass, 'pg_class') AS table_comment
+            FROM pg_tables t
+            WHERE t.tablename NOT LIKE 'sys_gen_%'
+            AND t.tablename = :table_name
+            AND t.schemaname = 'public'; -- schema 通常是 'public'
+            """
+        stmt = text(sql).bindparams(table_name=table_name)
         result = await db.execute(stmt)
         return result.fetchone()
 
     @staticmethod
     async def get_all_columns(db: AsyncSession, table_schema: str, table_name: str) -> Sequence[Row[tuple]]:
-        stmt = text(
-            'select column_name AS column_name, '
-            'case when column_key = "PRI" then 1 else 0 end as is_pk, '
-            'case when is_nullable = "NO" or column_key = "PRI" then 0 else 1 end as is_nullable, '
-            'ordinal_position as sort, '
-            'column_comment as column_comment, '
-            'column_type as column_type '
-            'from information_schema.columns '
-            'where table_schema = :table_schema '
-            'and table_name = :table_name '
-            'and column_name != "id" '
-            'and column_name != "created_time" '
-            'and column_name != "updated_time" '
-            'order by sort;'
-        ).bindparams(table_schema=table_schema, table_name=table_name)
+        if settings.DATABASE_TYPE == 'mysql':
+            sql = """
+            SELECT column_name AS column_name, 
+            CASE WHEN column_key = 'PRI' THEN 1 ELSE 0 END AS is_pk, 
+            CASE WHEN is_nullable = 'NO' OR column_key = 'PRI' THEN 0 ELSE 1 END AS is_nullable, 
+            ordinal_position AS sort, column_comment AS column_comment, 
+            column_type AS column_type FROM information_schema.columns 
+            WHERE table_schema = :table_schema 
+            AND table_name = :table_name 
+            AND column_name <> 'id' 
+            AND column_name <> 'created_time' 
+            AND column_name <> 'updated_time' 
+            ORDER BY sort;
+            """
+            stmt = text(sql).bindparams(table_schema=table_schema, table_name=table_name)
+        else:
+            sql = """
+            SELECT a.attname AS column_name,
+            CASE WHEN EXISTS (
+            SELECT 1
+            FROM pg_constraint c
+            WHERE c.conrelid = t.oid
+            AND c.contype = 'p'
+            AND a.attnum = ANY(c.conkey)
+            ) THEN 1 ELSE 0 END AS is_pk,
+            CASE WHEN a.attnotnull OR EXISTS (
+            SELECT 1
+            FROM pg_constraint c
+            WHERE c.conrelid = t.oid
+            AND c.contype = 'p'
+            AND a.attnum = ANY(c.conkey)
+            ) THEN 0 ELSE 1 END AS is_nullable,
+            a.attnum AS sort,
+            col_description(t.oid, a.attnum) AS column_comment,
+            pg_catalog.format_type(a.atttypid, a.atttypmod) AS column_type
+            FROM pg_attribute a
+            JOIN pg_class t ON a.attrelid = t.oid
+            JOIN pg_namespace n ON n.oid = t.relnamespace
+            WHERE n.nspname = 'public'  -- 根据你的实际情况修改 schema 名称，通常是 'public'
+            AND t.relname = :table_name
+            AND a.attnum > 0
+            AND NOT a.attisdropped
+            AND a.attname <> 'id'
+            AND a.attname <> 'created_time'
+            AND a.attname <> 'updated_time'
+            ORDER BY sort;
+            """
+            stmt = text(sql).bindparams(table_name=table_name)
         result = await db.execute(stmt)
         return result.fetchall()
 
diff --git a/backend/templates/py/api.jinja b/backend/templates/py/api.jinja
@@ -39,7 +39,7 @@ async def get_pagination_{{ table_name_en }}(db: CurrentSession) -> ResponseMode
     '',
     summary='创建{{ table_simple_name_zh }}',
     dependencies=[
-        Depends(RequestPermission('{{ permission_sign }}:add')),
+        Depends(RequestPermission('{{ permission }}:add')),
         DependsRBAC,
     ],
 )
@@ -52,7 +52,7 @@ async def create_{{ table_name_en }}(obj: Create{{ schema_name }}Param) -> Respo
     '/{pk}',
     summary='更新{{ table_simple_name_zh }}',
     dependencies=[
-        Depends(RequestPermission('{{ permission_sign }}:edit')),
+        Depends(RequestPermission('{{ permission }}:edit')),
         DependsRBAC,
     ],
 )
@@ -67,7 +67,7 @@ async def update_{{ table_name_en }}(pk: Annotated[int, Path(...)], obj: Update{
     '',
     summary='（批量）删除{{ table_simple_name_zh }}',
     dependencies=[
-        Depends(RequestPermission('{{ permission_sign }}:del')),
+        Depends(RequestPermission('{{ permission }}:del')),
         DependsRBAC,
     ],
 )
diff --git a/backend/templates/py/model.jinja b/backend/templates/py/model.jinja
@@ -5,12 +5,16 @@ from uuid import UUID
 
 import sqlalchemy as sa
 
-from backend.common.model import {% if have_datetime_column %}Base{% else %}MappedBase{% endif %}, id_key
+from backend.common.model import {% if default_datetime_column %}Base{% else %}MappedBase{% endif %}, id_key
+{% if database_type == 'mysql' -%}
 from sqlalchemy.dialects import mysql
+{% else -%}
+from sqlalchemy.dialects import postgresql
+{% endif -%}
 from sqlalchemy.orm import Mapped, mapped_column
 
 
-class {{ table_name_class }}({% if have_datetime_column %}Base{% else %}MappedBase{% endif %}):
+class {{ table_name_class }}({% if default_datetime_column %}Base{% else %}MappedBase{% endif %}):
     """{{ table_name_zh }}"""
 
     __tablename__ = '{{ table_name_en }}'
@@ -23,9 +27,14 @@ class {{ table_name_class }}({% if have_datetime_column %}Base{% else %}MappedBa
     {%- endif %} = mapped_column(
     {%- if model.type in ['NVARCHAR', 'String', 'Unicode', 'VARCHAR'] -%}
         sa.String({{ model.length }})
-    {%- elif model.type in ['BIT', 'ENUM', 'LONGBLOB', 'LONGTEXT', 'MEDIUMBLOB', 'MEDIUMINT', 'MEDIUMTEXT', 'SET',
-    'TINYBLOB', 'TINYINT', 'TINYTEXT', 'YEAR'] -%}
+    {%- elif database_type == 'mysql' and model.type in ['BIT', 'ENUM', 'LONGBLOB', 'LONGTEXT', 'MEDIUMBLOB',
+    'MEDIUMINT', 'MEDIUMTEXT', 'SET', 'TINYBLOB', 'TINYINT', 'TINYTEXT', 'YEAR'] -%}
         mysql.{{ model.type }}()
+    {%- elif database_type == 'postgresql' and model.type in [
+    'ARRAY', 'BIT', 'BYTEA', 'CIDR', 'CITEXT', 'DATEMULTIRANGE', 'DATERANGE', 'DOMAIN', 'ENUM', 'HSTORE', 'INET',
+    'INT4MULTIRANGE', 'INT4RANGE', 'INT8MULTIRANGE', 'INT8RANGE', 'INTERVAL', 'JSONB', 'JSONPATH', 'MACADDR',
+    'MACADDR8', 'MONEY', 'NUMMULTIRANGE', 'NUMRANGE', 'OID', 'REGCLASS', 'REGCONFIG', 'TSMULTIRANGE', 'TSQUERY',
+    'TSRANGE', 'TSTZMULTIRANGE', 'TSTZRANGE', 'TSVECTOR'] -%}
     {%- else -%}
         sa.{{ model.type }}()
     {%- endif -%}, default=
diff --git a/backend/templates/py/schema.jinja b/backend/templates/py/schema.jinja
@@ -26,7 +26,7 @@ class Get{{ schema_name }}ListDetails({{ schema_name }}SchemaBase):
     model_config = ConfigDict(from_attributes=True)
 
     id: int
-    {% if have_datetime_column %}
+    {% if default_datetime_column %}
     created_time: datetime
     updated_time: datetime | None = None
     {% endif %}
diff --git a/backend/utils/gen_template.py b/backend/utils/gen_template.py
@@ -7,6 +7,7 @@
 
 from backend.app.generator.conf import generator_settings
 from backend.app.generator.model import GenBusiness, GenModel
+from backend.core.conf import settings
 from backend.core.path_conf import JINJA2_TEMPLATE_DIR
 
 
@@ -95,8 +96,9 @@ def get_vars(business: GenBusiness, models: Sequence[GenModel]) -> dict:
             'table_simple_name_zh': business.table_simple_name_zh,
             'table_comment': business.table_comment,
             'schema_name': to_pascal(business.schema_name),
-            'have_datetime_column': business.default_datetime_column,
-            'permission_sign': str(business.table_name_en.replace('_', ':')),
+            'default_datetime_column': business.default_datetime_column,
+            'permission': str(business.table_name_en.replace('_', ':')),
+            'database_type': settings.DATABASE_TYPE,
             'models': models,
         }
 
diff --git a/backend/utils/type_conversion.py b/backend/utils/type_conversion.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-from backend.common.enums import GenModelMySQLColumnType
+from backend.common.enums import GenModelMySQLColumnType, GenModelPostgreSQLColumnType
+from backend.core.conf import settings
 
 
 def sql_type_to_sqlalchemy(typing: str) -> str:
@@ -10,8 +11,12 @@ def sql_type_to_sqlalchemy(typing: str) -> str:
     :param typing:
     :return:
     """
-    if typing in GenModelMySQLColumnType.get_member_keys():
-        return typing
+    if settings.DATABASE_TYPE == 'mysql':
+        if typing in GenModelMySQLColumnType.get_member_keys():
+            return typing
+    else:
+        if typing in GenModelPostgreSQLColumnType.get_member_keys():
+            return typing
     return 'String'
 
 
@@ -23,6 +28,11 @@ def sql_type_to_pydantic(typing: str) -> str:
     :return:
     """
     try:
-        return GenModelMySQLColumnType[typing].value
+        if settings.DATABASE_TYPE == 'mysql':
+            return GenModelMySQLColumnType[typing].value
+        else:
+            if typing == 'CHARACTER VARYING':  # postgresql 中 DDL VARCHAR 的别名
+                return 'str'
+            return GenModelPostgreSQLColumnType[typing].value
     except KeyError:
         return 'str'
