diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d8756ee35..39eba06e3 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,6 +14,7 @@ Changelog Added ^^^^^ - ``Q.__bool__()`` so ``Q`` objects with no filters/children (including nested empty ``Q`` children) are falsy. +- ``db_collation`` argument on fields to set a column collation, emitted as a ``COLLATE`` clause in the generated schema. (#686) 1.1.8 ----- diff --git a/tests/schema/models_collation.py b/tests/schema/models_collation.py new file mode 100644 index 000000000..6bf731df7 --- /dev/null +++ b/tests/schema/models_collation.py @@ -0,0 +1,12 @@ +""" +This example demonstrates SQL Schema generation for fields that set a db_collation. +""" + +from tortoise import fields +from tortoise.models import Model + + +class Account(Model): + name = fields.CharField(max_length=50, db_collation="NOCASE") + bio = fields.TextField(db_collation="NOCASE") + plain = fields.CharField(max_length=20) diff --git a/tests/schema/test_generate_schema.py b/tests/schema/test_generate_schema.py index 63098bce1..9147dc693 100644 --- a/tests/schema/test_generate_schema.py +++ b/tests/schema/test_generate_schema.py @@ -409,6 +409,25 @@ async def test_schema_no_db_constraint(): await _teardown_tortoise() +@pytest.mark.asyncio +async def test_schema_db_collation(): + await _reset_tortoise() + try: + await _init_for_sqlite("tests.schema.models_collation") + sql = get_schema_sql(connections.get("default"), safe=False) + assert ( + sql.strip() + == """CREATE TABLE "account" ( + "id" INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, + "name" VARCHAR(50) COLLATE NOCASE NOT NULL, + "bio" TEXT COLLATE NOCASE NOT NULL, + "plain" VARCHAR(20) NOT NULL +);""" + ) + finally: + await _teardown_tortoise() + + @pytest.mark.asyncio async def test_schema(): await _reset_tortoise() diff --git a/tortoise/backends/base/schema_generator.py b/tortoise/backends/base/schema_generator.py index d1b9725b5..dee0e317e 100644 --- a/tortoise/backends/base/schema_generator.py +++ b/tortoise/backends/base/schema_generator.py @@ -25,7 +25,7 @@ class BaseSchemaGenerator(SchemaQuotingMixin): DIALECT = "sql" TABLE_CREATE_TEMPLATE = "CREATE TABLE {exists}{table_name} ({fields}){extra}{comment};" - FIELD_TEMPLATE = '"{name}" {type}{nullable}{unique}{primary}{default}{comment}' + FIELD_TEMPLATE = '"{name}" {type}{collate}{nullable}{unique}{primary}{default}{comment}' INDEX_CREATE_TEMPLATE = ( 'CREATE {index_type}INDEX {exists}"{index_name}" ON {table_name} ({fields}){extra};' ) @@ -52,12 +52,14 @@ def _create_string( is_primary_key: bool, comment: str, default: str, + collation: str = "", ) -> str: # children can override this function to customize their sql queries return self.FIELD_TEMPLATE.format( name=db_column, type=field_type, + collate=collation, nullable=nullable, unique="" if is_primary_key else unique, comment=comment if self.client.capabilities.inline_comment else "", @@ -103,6 +105,11 @@ def _column_comment_generator(self, table: str, column: str, comment: str) -> st # needs to be implemented for each supported client raise NotImplementedError() # pragma: nocoverage + def _column_collation_generator(self, collation: str) -> str: + # The collation name is a bare identifier for most dialects. Backends that + # need it quoted (e.g. Postgres) override this. + return f" COLLATE {collation}" + def _post_table_hook(self) -> str: # This method provides a mechanism where you can perform a set of # operation on the database table after it's initialized. This method @@ -248,6 +255,8 @@ def _get_field_sql_and_related_table( nullable = " NOT NULL" if not field_object.null else "" unique = " UNIQUE" if field_object.unique else "" field_type = field_object.get_for_dialect(self.DIALECT, "SQL_TYPE") + db_collation = getattr(field_object, "db_collation", None) + collation = self._column_collation_generator(db_collation) if db_collation else "" qualified_table_name = self._qualify_table_name(table_name, schema) field_creation_string, related_table_name = "", "" @@ -270,6 +279,7 @@ def _get_field_sql_and_related_table( is_primary_key=field_object.pk, comment="", default=default, + collation=collation, ) + self._create_fk_string( constraint_name=self._get_fk_name( table_name, @@ -292,6 +302,7 @@ def _get_field_sql_and_related_table( is_primary_key=field_object.pk, comment=comment, default=default, + collation=collation, ) return field_creation_string, related_table_name diff --git a/tortoise/backends/base_postgres/schema_generator.py b/tortoise/backends/base_postgres/schema_generator.py index fdbe29207..a1405dc88 100644 --- a/tortoise/backends/base_postgres/schema_generator.py +++ b/tortoise/backends/base_postgres/schema_generator.py @@ -51,6 +51,10 @@ def _column_comment_generator(self, table: str, column: str, comment: str) -> st self.comments_array.append(comment) return "" + def _column_collation_generator(self, collation: str) -> str: + # Postgres collation names are identifiers and need double quoting. + return f' COLLATE "{collation}"' + def _post_table_hook(self) -> str: val = "\n".join(self.comments_array) self.comments_array = [] diff --git a/tortoise/backends/mssql/schema_generator.py b/tortoise/backends/mssql/schema_generator.py index 1a56124b9..99df78fcd 100644 --- a/tortoise/backends/mssql/schema_generator.py +++ b/tortoise/backends/mssql/schema_generator.py @@ -15,7 +15,7 @@ class MSSQLSchemaGenerator(MSSQLQuotingMixin, BaseSchemaGenerator): DIALECT = "mssql" TABLE_CREATE_TEMPLATE = "CREATE TABLE {table_name} ({fields}){extra};" - FIELD_TEMPLATE = "[{name}] {type}{nullable}{unique}{primary}{default}" + FIELD_TEMPLATE = "[{name}] {type}{collate}{nullable}{unique}{primary}{default}" INDEX_CREATE_TEMPLATE = "CREATE INDEX [{index_name}] ON {table_name} ({fields});" UNIQUE_CONSTRAINT_CREATE_TEMPLATE = "CONSTRAINT [{index_name}] UNIQUE ({fields})" GENERATED_PK_TEMPLATE = "[{field_name}] {generated_sql}" @@ -109,6 +109,7 @@ def _create_string( is_primary_key: bool, comment: str, default: str, + collation: str = "", ) -> str: if nullable == "": unique = "" @@ -120,6 +121,7 @@ def _create_string( is_primary_key=is_primary_key, comment=comment, default=default, + collation=collation, ) def _get_inner_statements(self) -> list[str]: diff --git a/tortoise/backends/mysql/schema_generator.py b/tortoise/backends/mysql/schema_generator.py index ec4ca03a1..68aafb571 100644 --- a/tortoise/backends/mysql/schema_generator.py +++ b/tortoise/backends/mysql/schema_generator.py @@ -18,7 +18,7 @@ class MySQLSchemaGenerator(MySQLQuotingMixin, BaseSchemaGenerator): INDEX_CREATE_TEMPLATE = "{index_type}KEY `{index_name}` ({fields}){extra}" UNIQUE_CONSTRAINT_CREATE_TEMPLATE = "UNIQUE KEY `{index_name}` ({fields})" UNIQUE_INDEX_CREATE_TEMPLATE = UNIQUE_CONSTRAINT_CREATE_TEMPLATE - FIELD_TEMPLATE = "`{name}` {type}{nullable}{unique}{primary}{comment}{default}" + FIELD_TEMPLATE = "`{name}` {type}{collate}{nullable}{unique}{primary}{comment}{default}" GENERATED_PK_TEMPLATE = "`{field_name}` {generated_sql}{comment}" FK_TEMPLATE = ( "{constraint}FOREIGN KEY (`{db_column}`)" diff --git a/tortoise/backends/oracle/schema_generator.py b/tortoise/backends/oracle/schema_generator.py index 0203a3bbf..083048fab 100644 --- a/tortoise/backends/oracle/schema_generator.py +++ b/tortoise/backends/oracle/schema_generator.py @@ -15,7 +15,7 @@ class OracleSchemaGenerator(BaseSchemaGenerator): DIALECT = "oracle" TABLE_CREATE_TEMPLATE = "CREATE TABLE {table_name} ({fields}){extra};" - FIELD_TEMPLATE = '"{name}" {type}{default}{nullable}{unique}{primary}' + FIELD_TEMPLATE = '"{name}" {type}{collate}{default}{nullable}{unique}{primary}' TABLE_COMMENT_TEMPLATE = "COMMENT ON TABLE {table} IS '{comment}';" COLUMN_COMMENT_TEMPLATE = "COMMENT ON COLUMN {table}.\"{column}\" IS '{comment}';" INDEX_CREATE_TEMPLATE = 'CREATE INDEX "{index_name}" ON {table_name} ({fields});' diff --git a/tortoise/fields/base.py b/tortoise/fields/base.py index 81db786c0..6064ec9f9 100644 --- a/tortoise/fields/base.py +++ b/tortoise/fields/base.py @@ -123,6 +123,10 @@ class Field(Generic[VALUE], metaclass=_FieldMeta): :param description: Field description. Will also appear in ``Tortoise.describe_model()`` and as DB comments in the generated DDL. :param validators: Validators for this field. + :param db_collation: Set a database collation for the column, emitted as ``COLLATE`` in the + generated DDL. Only meaningful for text based columns and the value is passed through to + the database as given, so use a collation the target database knows (for example + ``NOCASE`` on SQLite or ``utf8mb4_unicode_ci`` on MySQL). **Class Attributes:** These attributes needs to be defined when defining an actual field type. @@ -230,6 +234,7 @@ def __init__( description: str | None = None, model: Model | None = None, validators: list[Validator | Callable] | None = None, + db_collation: str | None = None, **kwargs: Any, ) -> None: if (index := kwargs.pop("index", None)) is not None: @@ -284,6 +289,7 @@ def __init__( self.index = bool(db_index) self.model_field_name = "" self.description = description + self.db_collation = db_collation self.docstring: str | None = None self.validators: list[Validator | Callable] = validators or [] # TODO: consider making this not be set from constructor