51 lines
2.1 KiB
Python
51 lines
2.1 KiB
Python
from typing import Type
|
|
|
|
from tortoise import Model
|
|
from tortoise.backends.asyncpg.schema_generator import AsyncpgSchemaGenerator
|
|
from tortoise.fields import Field
|
|
|
|
from aerich.ddl import BaseDDL
|
|
|
|
|
|
class PostgresDDL(BaseDDL):
|
|
schema_generator_cls = AsyncpgSchemaGenerator
|
|
DIALECT = AsyncpgSchemaGenerator.DIALECT
|
|
_ALTER_DEFAULT_TEMPLATE = 'ALTER TABLE "{table_name}" ALTER COLUMN "{column}" {default}'
|
|
_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}'
|
|
_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_default(self, model: "Type[Model]", field_object: Field):
|
|
db_table = model._meta.db_table
|
|
default = self._get_default(model, field_object)
|
|
return self._ALTER_DEFAULT_TEMPLATE.format(
|
|
table_name=db_table,
|
|
column=field_object.model_field_name,
|
|
default="SET" + default if default else "DROP DEFAULT",
|
|
)
|
|
|
|
def alter_column_null(self, model: "Type[Model]", field_object: Field):
|
|
db_table = model._meta.db_table
|
|
return self._ALTER_NULL_TEMPLATE.format(
|
|
table_name=db_table,
|
|
column=field_object.model_field_name,
|
|
set_drop="DROP" if field_object.null else "SET",
|
|
)
|
|
|
|
def modify_column(self, model: "Type[Model]", field_object: Field):
|
|
db_table = model._meta.db_table
|
|
return self._MODIFY_COLUMN_TEMPLATE.format(
|
|
table_name=db_table,
|
|
column=field_object.model_field_name,
|
|
datatype=field_object.get_for_dialect(self.DIALECT, "SQL_TYPE"),
|
|
)
|
|
|
|
def set_comment(self, model: "Type[Model]", field_object: Field):
|
|
db_table = model._meta.db_table
|
|
return self._SET_COMMENT_TEMPLATE.format(
|
|
table_name=db_table,
|
|
column=field_object.model_field_name,
|
|
comment="'{}'".format(field_object.description) if field_object.description else "NULL",
|
|
)
|