52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
from typing import Type
|
|
|
|
from tortoise import Model
|
|
from tortoise.backends.asyncpg.schema_generator import AsyncpgSchemaGenerator
|
|
|
|
from aerich.ddl import BaseDDL
|
|
|
|
|
|
class PostgresDDL(BaseDDL):
|
|
schema_generator_cls = AsyncpgSchemaGenerator
|
|
DIALECT = AsyncpgSchemaGenerator.DIALECT
|
|
_ADD_INDEX_TEMPLATE = 'CREATE {unique}INDEX "{index_name}" ON "{table_name}" ({column_names})'
|
|
_DROP_INDEX_TEMPLATE = 'DROP INDEX "{index_name}"'
|
|
_ALTER_NULL_TEMPLATE = 'ALTER TABLE "{table_name}" ALTER COLUMN "{column}" {set_drop} NOT NULL'
|
|
_MODIFY_COLUMN_TEMPLATE = (
|
|
'ALTER TABLE "{table_name}" ALTER COLUMN "{column}" TYPE {datatype}{using}'
|
|
)
|
|
_SET_COMMENT_TEMPLATE = 'COMMENT ON COLUMN "{table_name}"."{column}" IS {comment}'
|
|
_DROP_FK_TEMPLATE = 'ALTER TABLE "{table_name}" DROP CONSTRAINT "{fk_name}"'
|
|
|
|
def alter_column_null(self, model: "Type[Model]", field_describe: dict):
|
|
db_table = model._meta.db_table
|
|
return self._ALTER_NULL_TEMPLATE.format(
|
|
table_name=db_table,
|
|
column=field_describe.get("db_column"),
|
|
set_drop="DROP" if field_describe.get("nullable") else "SET",
|
|
)
|
|
|
|
def modify_column(self, model: "Type[Model]", field_describe: dict, is_pk: bool = False):
|
|
db_table = model._meta.db_table
|
|
db_field_types = field_describe.get("db_field_types")
|
|
db_column = field_describe.get("db_column")
|
|
datatype = db_field_types.get(self.DIALECT) or db_field_types.get("")
|
|
return self._MODIFY_COLUMN_TEMPLATE.format(
|
|
table_name=db_table,
|
|
column=db_column,
|
|
datatype=datatype,
|
|
using=f' USING "{db_column}"::{datatype}',
|
|
)
|
|
|
|
def set_comment(self, model: "Type[Model]", field_describe: dict):
|
|
db_table = model._meta.db_table
|
|
return self._SET_COMMENT_TEMPLATE.format(
|
|
table_name=db_table,
|
|
column=field_describe.get("db_column") or field_describe.get("raw_field"),
|
|
comment=(
|
|
"'{}'".format(field_describe.get("description"))
|
|
if field_describe.get("description")
|
|
else "NULL"
|
|
),
|
|
)
|