From f15cbaf9e0269fca88741fae135bc7952fafd444 Mon Sep 17 00:00:00 2001 From: long2ice Date: Wed, 29 Dec 2021 21:36:23 +0800 Subject: [PATCH] Support migration for specified index. (#203) --- CHANGELOG.md | 4 ++++ README.md | 2 +- aerich/cli.py | 7 +++--- aerich/coder.py | 31 +++++++++++++++++++++++++++ aerich/ddl/__init__.py | 6 ++++++ aerich/migrate.py | 42 ++++++++++++++++++++++++++++++------ aerich/models.py | 4 +++- aerich/version.py | 2 +- poetry.lock | 48 +++++++++++++++++++++++------------------- pyproject.toml | 4 ++-- 10 files changed, 113 insertions(+), 37 deletions(-) create mode 100644 aerich/coder.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a26a8ed..c9022b6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## 0.6 +### 0.6.2 + +- Support migration for specified index. (#203) + ### 0.6.1 - Fix `pyproject.toml` not existing error. (#217) diff --git a/README.md b/README.md index 439c05e..a8e969a 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ ## Introduction -Aerich is a database migrations tool for Tortoise-ORM, which is like alembic for SQLAlchemy, or like Django ORM with +Aerich is a database migrations tool for TortoiseORM, which is like alembic for SQLAlchemy, or like Django ORM with it\'s own migration solution. ## Install diff --git a/aerich/cli.py b/aerich/cli.py index c3856b6..65308f4 100644 --- a/aerich/cli.py +++ b/aerich/cli.py @@ -10,12 +10,11 @@ from tomlkit.exceptions import NonExistentKey from tortoise import Tortoise +from aerich import Command +from aerich.enums import Color from aerich.exceptions import DowngradeError from aerich.utils import add_src_path, get_tortoise_config - -from . import Command -from .enums import Color -from .version import __version__ +from aerich.version import __version__ CONFIG_DEFAULT_VALUES = { "src_folder": ".", diff --git a/aerich/coder.py b/aerich/coder.py new file mode 100644 index 0000000..bc41d72 --- /dev/null +++ b/aerich/coder.py @@ -0,0 +1,31 @@ +import base64 +import json +import pickle # nosec: B301 + +from tortoise.indexes import Index + + +class JsonEncoder(json.JSONEncoder): + def default(self, obj): + if isinstance(obj, Index): + return { + "type": "index", + "val": base64.b64encode(pickle.dumps(obj)).decode(), + } # nosec: B301 + else: + return super().default(obj) + + +def object_hook(obj): + _type = obj.get("type") + if not _type: + return obj + return pickle.loads(base64.b64decode(obj["val"])) # nosec: B301 + + +def encoder(obj: dict): + return json.dumps(obj, cls=JsonEncoder) + + +def decoder(obj: str): + return json.loads(obj, object_hook=object_hook) diff --git a/aerich/ddl/__init__.py b/aerich/ddl/__init__.py index 2818955..cb01073 100644 --- a/aerich/ddl/__init__.py +++ b/aerich/ddl/__init__.py @@ -195,6 +195,12 @@ def drop_index(self, model: "Type[Model]", field_names: List[str], unique=False) table_name=model._meta.db_table, ) + def drop_index_by_name(self, model: "Type[Model]", index_name: str): + return self._DROP_INDEX_TEMPLATE.format( + index_name=index_name, + table_name=model._meta.db_table, + ) + def add_fk(self, model: "Type[Model]", field_describe: dict, reference_table_describe: dict): db_table = model._meta.db_table diff --git a/aerich/migrate.py b/aerich/migrate.py index ab2f361..6085efc 100644 --- a/aerich/migrate.py +++ b/aerich/migrate.py @@ -1,12 +1,14 @@ import os from datetime import datetime +from hashlib import md5 from pathlib import Path -from typing import Dict, List, Optional, Tuple, Type +from typing import Dict, List, Optional, Tuple, Type, Union import click from dictdiffer import diff from tortoise import BaseDBAsyncClient, Model, Tortoise from tortoise.exceptions import OperationalError +from tortoise.indexes import Index from aerich.ddl import BaseDDL from aerich.models import MAX_VERSION_LENGTH, Aerich @@ -32,7 +34,7 @@ class Migrate: ddl: BaseDDL _last_version_content: Optional[dict] = None app: str - migrate_location: str + migrate_location: Path dialect: str _db_version: Optional[str] = None @@ -157,6 +159,18 @@ def _add_operator(cls, operator: str, upgrade=True, fk_m2m_index=False): else: cls.downgrade_operators.append(operator) + @classmethod + def _handle_indexes(cls, model: Type[Model], indexes: List[Union[Tuple[str], Index]]): + ret = [] + for index in indexes: + if isinstance(index, Index): + index.__hash__ = lambda self: md5( # nosec: B303 + self.index_name(cls.ddl.schema_generator, model).encode() + + self.__class__.__name__.encode() + ).hexdigest() + ret.append(index) + return ret + @classmethod def diff_models(cls, old_models: Dict[str, dict], new_models: Dict[str, dict], upgrade=True): """ @@ -192,8 +206,18 @@ def diff_models(cls, old_models: Dict[str, dict], new_models: Dict[str, dict], u new_unique_together = set( map(lambda x: tuple(x), new_model_describe.get("unique_together")) ) - old_indexes = set(map(lambda x: tuple(x), old_model_describe.get("indexes", []))) - new_indexes = set(map(lambda x: tuple(x), new_model_describe.get("indexes", []))) + old_indexes = set( + map( + lambda x: x if isinstance(x, Index) else tuple(x), + cls._handle_indexes(model, old_model_describe.get("indexes", [])), + ) + ) + new_indexes = set( + map( + lambda x: x if isinstance(x, Index) else tuple(x), + cls._handle_indexes(model, new_model_describe.get("indexes", [])), + ) + ) old_pk_field = old_model_describe.get("pk_field") new_pk_field = new_model_describe.get("pk_field") # pk field @@ -463,12 +487,18 @@ def _resolve_fk_fields_name(cls, model: Type[Model], fields_name: Tuple[str]): return ret @classmethod - def _drop_index(cls, model: Type[Model], fields_name: Tuple[str], unique=False): + def _drop_index(cls, model: Type[Model], fields_name: Union[Tuple[str], Index], unique=False): + if isinstance(fields_name, Index): + return cls.ddl.drop_index_by_name( + model, fields_name.index_name(cls.ddl.schema_generator, model) + ) fields_name = cls._resolve_fk_fields_name(model, fields_name) return cls.ddl.drop_index(model, fields_name, unique) @classmethod - def _add_index(cls, model: Type[Model], fields_name: Tuple[str], unique=False): + def _add_index(cls, model: Type[Model], fields_name: Union[Tuple[str], Index], unique=False): + if isinstance(fields_name, Index): + return fields_name.get_sql(cls.ddl.schema_generator, model, False) fields_name = cls._resolve_fk_fields_name(model, fields_name) return cls.ddl.add_index(model, fields_name, unique) diff --git a/aerich/models.py b/aerich/models.py index a689f03..1e0514f 100644 --- a/aerich/models.py +++ b/aerich/models.py @@ -1,12 +1,14 @@ from tortoise import Model, fields +from aerich.coder import decoder, encoder + MAX_VERSION_LENGTH = 255 class Aerich(Model): version = fields.CharField(max_length=MAX_VERSION_LENGTH) app = fields.CharField(max_length=20) - content = fields.JSONField() + content = fields.JSONField(encoder=encoder, decoder=decoder) class Meta: ordering = ["-id"] diff --git a/aerich/version.py b/aerich/version.py index 43c4ab0..22049ab 100644 --- a/aerich/version.py +++ b/aerich/version.py @@ -1 +1 @@ -__version__ = "0.6.1" +__version__ = "0.6.2" diff --git a/poetry.lock b/poetry.lock index 966e735..96672e2 100644 --- a/poetry.lock +++ b/poetry.lock @@ -43,17 +43,17 @@ python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" [[package]] name = "attrs" -version = "21.2.0" +version = "21.4.0" description = "Classes Without Boilerplate" category = "dev" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" [package.extras] -dev = ["coverage[toml] (>=5.0.2)", "hypothesis", "pympler", "pytest (>=4.3.0)", "six", "mypy", "pytest-mypy-plugins", "zope.interface", "furo", "sphinx", "sphinx-notfound-page", "pre-commit"] +dev = ["coverage[toml] (>=5.0.2)", "hypothesis", "pympler", "pytest (>=4.3.0)", "six", "mypy", "pytest-mypy-plugins", "zope.interface", "furo", "sphinx", "sphinx-notfound-page", "pre-commit", "cloudpickle"] docs = ["furo", "sphinx", "zope.interface", "sphinx-notfound-page"] -tests = ["coverage[toml] (>=5.0.2)", "hypothesis", "pympler", "pytest (>=4.3.0)", "six", "mypy", "pytest-mypy-plugins", "zope.interface"] -tests_no_zope = ["coverage[toml] (>=5.0.2)", "hypothesis", "pympler", "pytest (>=4.3.0)", "six", "mypy", "pytest-mypy-plugins"] +tests = ["coverage[toml] (>=5.0.2)", "hypothesis", "pympler", "pytest (>=4.3.0)", "six", "mypy", "pytest-mypy-plugins", "zope.interface", "cloudpickle"] +tests_no_zope = ["coverage[toml] (>=5.0.2)", "hypothesis", "pympler", "pytest (>=4.3.0)", "six", "mypy", "pytest-mypy-plugins", "cloudpickle"] [[package]] name = "bandit" @@ -310,11 +310,11 @@ python-versions = ">=2.6" [[package]] name = "platformdirs" -version = "2.4.0" +version = "2.4.1" description = "A small Python module for determining appropriate platform-specific dirs, e.g. a \"user data dir\"." category = "dev" optional = false -python-versions = ">=3.6" +python-versions = ">=3.7" [package.extras] docs = ["Sphinx (>=4)", "furo (>=2021.7.5b38)", "proselint (>=0.10.2)", "sphinx-autodoc-typehints (>=1.12)"] @@ -528,23 +528,30 @@ python-versions = ">=3.6,<4.0" [[package]] name = "tortoise-orm" -version = "0.18.0" +version = "0.18.1" description = "Easy async ORM for python, built with relations in mind" category = "main" optional = false -python-versions = ">=3.7,<4.0" +python-versions = "^3.7" +develop = false [package.dependencies] -aiosqlite = ">=0.16.0,<0.18.0" -iso8601 = ">=0.1.13,<0.2.0" -pypika-tortoise = ">=0.1.2,<0.2.0" +aiosqlite = ">=0.16.0, <0.18.0" +iso8601 = "^0.1.13" +pypika-tortoise = "^0.1.2" pytz = "*" [package.extras] +accel = ["ciso8601", "uvloop", "orjson"] +asyncpg = ["asyncpg"] aiomysql = ["aiomysql"] asyncmy = ["asyncmy"] -asyncpg = ["asyncpg"] -accel = ["ciso8601 (>=2.1.2,<3.0.0)", "orjson", "uvloop (>=0.16.0,<0.17.0)"] + +[package.source] +type = "git" +url = "https://github.com/tortoise/tortoise-orm.git" +reference = "develop" +resolved_reference = "af2aab8d4f52cb0d1b986865195c00b74def6b36" [[package]] name = "typed-ast" @@ -581,7 +588,7 @@ asyncpg = ["asyncpg"] [metadata] lock-version = "1.1" python-versions = "^3.7" -content-hash = "972c0d0a1eb0c6a697ccbce550af36e1bb43cb0d74b7819f80846f5599bb8c2b" +content-hash = "bb4dffaf7ba754eae8b9ed38d3931eacfe520f0290fef3e2ff31846e14b8a8aa" [metadata.files] aiosqlite = [ @@ -636,8 +643,8 @@ atomicwrites = [ {file = "atomicwrites-1.4.0.tar.gz", hash = "sha256:ae70396ad1a434f9c7046fd2dd196fc04b12f9e91ffb859164193be8b6168a7a"}, ] attrs = [ - {file = "attrs-21.2.0-py2.py3-none-any.whl", hash = "sha256:149e90d6d8ac20db7a955ad60cf0e6881a3f20d37096140088356da6c716b0b1"}, - {file = "attrs-21.2.0.tar.gz", hash = "sha256:ef6aaac3ca6cd92904cdd0d83f629a15f18053ec84e6432106f7a4d04ae4f5fb"}, + {file = "attrs-21.4.0-py2.py3-none-any.whl", hash = "sha256:2d27e3784d7a565d36ab851fe94887c5eccd6a463168875832a1be79c82828b4"}, + {file = "attrs-21.4.0.tar.gz", hash = "sha256:626ba8234211db98e869df76230a137c4c40a12d72445c45d5f5b716f076e2fd"}, ] bandit = [ {file = "bandit-1.7.1-py3-none-any.whl", hash = "sha256:f5acd838e59c038a159b5c621cf0f8270b279e884eadd7b782d7491c02add0d4"}, @@ -790,8 +797,8 @@ pbr = [ {file = "pbr-5.8.0.tar.gz", hash = "sha256:672d8ebee84921862110f23fcec2acea191ef58543d34dfe9ef3d9f13c31cddf"}, ] platformdirs = [ - {file = "platformdirs-2.4.0-py3-none-any.whl", hash = "sha256:8868bbe3c3c80d42f20156f22e7131d2fb321f5bc86a2a345375c6481a67021d"}, - {file = "platformdirs-2.4.0.tar.gz", hash = "sha256:367a5e80b3d04d2428ffa76d33f124cf11e8fff2acdaa9b43d545f5c7d661ef2"}, + {file = "platformdirs-2.4.1-py3-none-any.whl", hash = "sha256:1d7385c7db91728b83efd0ca99a5afb296cab9d0ed8313a45ed8ba17967ecfca"}, + {file = "platformdirs-2.4.1.tar.gz", hash = "sha256:440633ddfebcc36264232365d7840a970e75e1018d15b4327d11f91909045fda"}, ] pluggy = [ {file = "pluggy-1.0.0-py2.py3-none-any.whl", hash = "sha256:74134bbf457f031a36d68416e1509f34bd5ccc019f0bcc952c7b909d06b37bd3"}, @@ -900,10 +907,7 @@ tomlkit = [ {file = "tomlkit-0.8.0-py3-none-any.whl", hash = "sha256:b824e3466f1d475b2b5f1c392954c6cb7ea04d64354ff7300dc7c14257dc85db"}, {file = "tomlkit-0.8.0.tar.gz", hash = "sha256:29e84a855712dfe0e88a48f6d05c21118dbafb283bb2eed614d46f80deb8e9a1"}, ] -tortoise-orm = [ - {file = "tortoise-orm-0.18.0.tar.gz", hash = "sha256:52c33b712279ba8b04f5a1bed3010ade4961f0200480df5362cf674486a9016a"}, - {file = "tortoise_orm-0.18.0-py3-none-any.whl", hash = "sha256:1f5422a94b2d58f61b91bf8993b678533554abdebdadf783ba0c484122fb63c1"}, -] +tortoise-orm = [] typed-ast = [ {file = "typed_ast-1.5.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:5d8314c92414ce7481eee7ad42b353943679cf6f30237b5ecbf7d835519e1212"}, {file = "typed_ast-1.5.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:b53ae5de5500529c76225d18eeb060efbcec90ad5e030713fe8dab0fb4531631"}, diff --git a/pyproject.toml b/pyproject.toml index a6bf9c2..9f1a3d8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "aerich" -version = "0.6.1" +version = "0.6.2" description = "A database migrations tool for Tortoise ORM." authors = ["long2ice "] license = "Apache-2.0" @@ -16,7 +16,7 @@ include = ["CHANGELOG.md", "LICENSE", "README.md"] [tool.poetry.dependencies] python = "^3.7" -tortoise-orm = "*" +tortoise-orm = { git = "https://github.com/tortoise/tortoise-orm.git", branch = "develop" } click = "*" asyncpg = { version = "*", optional = true } asyncmy = { version = "*", optional = true }