fix fk migrate
This commit is contained in:
71
alice/cli.py
71
alice/cli.py
@@ -5,18 +5,17 @@ import sys
|
||||
from enum import Enum
|
||||
|
||||
import asyncclick as click
|
||||
from asyncclick import BadParameter, ClickException
|
||||
from tortoise import generate_schema_for_client, ConfigurationError, Tortoise
|
||||
from asyncclick import BadOptionUsage, BadParameter, Context, UsageError
|
||||
from tortoise import Tortoise, generate_schema_for_client
|
||||
|
||||
from alice.migrate import Migrate
|
||||
from alice.utils import get_app_connection
|
||||
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
|
||||
class Color(str, Enum):
|
||||
green = "green"
|
||||
red = "red"
|
||||
yellow = "yellow"
|
||||
|
||||
|
||||
@click.group(context_settings={"help_option_names": ["-h", "--help"]})
|
||||
@@ -37,55 +36,55 @@ class Color(str, Enum):
|
||||
)
|
||||
@click.option("--app", default="models", show_default=True, help="Tortoise-ORM app name.")
|
||||
@click.pass_context
|
||||
async def cli(ctx, config, tortoise_orm, location, app):
|
||||
async def cli(ctx: Context, config, tortoise_orm, location, app):
|
||||
ctx.ensure_object(dict)
|
||||
try:
|
||||
config_module = importlib.import_module(config)
|
||||
config_module = importlib.import_module(config, ".")
|
||||
except ModuleNotFoundError:
|
||||
raise BadParameter(param_hint=["--tortoise-orm"], message=f'No module named "{config}"')
|
||||
raise BadOptionUsage(ctx=ctx, message=f'No module named "{config}"', option_name="--config")
|
||||
config = getattr(config_module, tortoise_orm, None)
|
||||
if not config:
|
||||
raise BadParameter(
|
||||
param_hint=["--config"],
|
||||
message=f'Can\'t get "{tortoise_orm}" from module "{config_module}"',
|
||||
)
|
||||
raise BadOptionUsage(f'Can\'t get "{tortoise_orm}" from module "{config_module}"', ctx=ctx)
|
||||
|
||||
if app not in config.get("apps").keys():
|
||||
raise BadParameter(param_hint=["--app"], message=f'No app found in "{config}"')
|
||||
raise BadOptionUsage(f'No app found in "{config}"', ctx=ctx)
|
||||
|
||||
ctx.obj["config"] = config
|
||||
ctx.obj["location"] = location
|
||||
ctx.obj["app"] = app
|
||||
try:
|
||||
|
||||
if ctx.invoked_subcommand == "init":
|
||||
await Tortoise.init(config=config)
|
||||
else:
|
||||
if not os.path.isdir(location):
|
||||
raise UsageError("You must exec init first", ctx=ctx)
|
||||
await Migrate.init_with_old_models(config, app, location)
|
||||
except ConfigurationError:
|
||||
pass
|
||||
|
||||
|
||||
@cli.command(help="Generate migrate changes file.")
|
||||
@click.option("--name", default="update", show_default=True, help="Migrate name.")
|
||||
@click.pass_context
|
||||
async def migrate(ctx, name):
|
||||
async def migrate(ctx: Context, name):
|
||||
config = ctx.obj["config"]
|
||||
location = ctx.obj["location"]
|
||||
app = ctx.obj["app"]
|
||||
|
||||
ret = Migrate.migrate(name)
|
||||
if not ret:
|
||||
click.secho("No changes detected", fg=Color.green)
|
||||
else:
|
||||
Migrate.write_old_models(config, app, location)
|
||||
click.secho(f"Success migrate {ret}", fg=Color.green)
|
||||
return click.secho("No changes detected", fg=Color.yellow)
|
||||
Migrate.write_old_models(config, app, location)
|
||||
click.secho(f"Success migrate {ret}", fg=Color.green)
|
||||
|
||||
|
||||
@cli.command(help="Upgrade to latest version.")
|
||||
@click.pass_context
|
||||
async def upgrade(ctx):
|
||||
async def upgrade(ctx: Context):
|
||||
app = ctx.obj["app"]
|
||||
config = ctx.obj["config"]
|
||||
connection = get_app_connection(config, app)
|
||||
available_versions = Migrate.get_all_version_files(is_all=False)
|
||||
if not available_versions:
|
||||
return click.secho("No migrate items", fg=Color.green)
|
||||
return click.secho("No migrate items", fg=Color.yellow)
|
||||
async with connection._in_transaction() as conn:
|
||||
for file in available_versions:
|
||||
file_path = os.path.join(Migrate.migrate_location, file)
|
||||
@@ -103,16 +102,16 @@ async def upgrade(ctx):
|
||||
|
||||
@cli.command(help="Downgrade to previous version.")
|
||||
@click.pass_context
|
||||
async def downgrade(ctx):
|
||||
async def downgrade(ctx: Context):
|
||||
app = ctx.obj["app"]
|
||||
config = ctx.obj["config"]
|
||||
connection = get_app_connection(config, app)
|
||||
available_versions = Migrate.get_all_version_files()
|
||||
if not available_versions:
|
||||
return click.secho("No migrate items", fg=Color.green)
|
||||
return click.secho("No migrate items", fg=Color.yellow)
|
||||
|
||||
async with connection._in_transaction() as conn:
|
||||
for file in available_versions:
|
||||
for file in reversed(available_versions):
|
||||
file_path = os.path.join(Migrate.migrate_location, file)
|
||||
with open(file_path, "r") as f:
|
||||
content = json.load(f)
|
||||
@@ -120,7 +119,8 @@ async def downgrade(ctx):
|
||||
downgrade_query_list = content.get("downgrade")
|
||||
for downgrade_query in downgrade_query_list:
|
||||
await conn.execute_query(downgrade_query)
|
||||
|
||||
else:
|
||||
continue
|
||||
with open(file_path, "w") as f:
|
||||
content["migrate"] = False
|
||||
json.dump(content, f, indent=4)
|
||||
@@ -129,20 +129,20 @@ async def downgrade(ctx):
|
||||
|
||||
@cli.command(help="Show current available heads in migrate location.")
|
||||
@click.pass_context
|
||||
def heads(ctx):
|
||||
def heads(ctx: Context):
|
||||
for version in Migrate.get_all_version_files(is_all=False):
|
||||
click.secho(version, fg=Color.green)
|
||||
click.secho(version, fg=Color.yellow)
|
||||
|
||||
|
||||
@cli.command(help="List all migrate items.")
|
||||
@click.pass_context
|
||||
def history(ctx):
|
||||
for version in Migrate.get_all_version_files():
|
||||
click.secho(version, fg=Color.green)
|
||||
click.secho(version, fg=Color.yellow)
|
||||
|
||||
|
||||
@cli.command(
|
||||
help="Init migrate location and generate schema, you must call first before other actions."
|
||||
help="Init migrate location and generate schema, you must exec first before other cmd."
|
||||
)
|
||||
@click.option(
|
||||
"--safe",
|
||||
@@ -152,7 +152,7 @@ def history(ctx):
|
||||
show_default=True,
|
||||
)
|
||||
@click.pass_context
|
||||
async def init(ctx, safe):
|
||||
async def init(ctx: Context, safe):
|
||||
location = ctx.obj["location"]
|
||||
app = ctx.obj["app"]
|
||||
config = ctx.obj["config"]
|
||||
@@ -165,15 +165,16 @@ async def init(ctx, safe):
|
||||
os.mkdir(dirname)
|
||||
click.secho(f"Success create migrate location {dirname}", fg=Color.green)
|
||||
else:
|
||||
raise ClickException(f"Already inited app `{app}`")
|
||||
return click.secho(f'Already inited app "{app}"', fg=Color.yellow)
|
||||
|
||||
Migrate.write_old_models(config, app, location)
|
||||
|
||||
await Migrate.init_with_old_models(config, app, location)
|
||||
await generate_schema_for_client(get_app_connection(config, app), safe)
|
||||
connection = get_app_connection(config, app)
|
||||
await generate_schema_for_client(connection, safe)
|
||||
|
||||
click.secho(f"Success init for app `{app}`", fg=Color.green)
|
||||
return click.secho(f'Success init for app "{app}"', fg=Color.green)
|
||||
|
||||
|
||||
def main():
|
||||
sys.path.insert(0, ".")
|
||||
cli(_anyio_backend="asyncio")
|
||||
|
||||
@@ -101,14 +101,14 @@ class DDL:
|
||||
|
||||
fk_name = self.schema_generator._generate_fk_name(
|
||||
from_table=db_table,
|
||||
from_field=field.model_field_name,
|
||||
from_field=field.source_field or field.model_field_name + "_id",
|
||||
to_table=field.related_model._meta.db_table,
|
||||
to_field=to_field_name,
|
||||
)
|
||||
return self._ADD_FK_TEMPLATE.format(
|
||||
table_name=db_table,
|
||||
fk_name=fk_name,
|
||||
db_column=field.model_field_name,
|
||||
db_column=field.source_field or field.model_field_name + "_id",
|
||||
table=field.related_model._meta.db_table,
|
||||
field=to_field_name,
|
||||
on_delete=field.on_delete,
|
||||
@@ -118,11 +118,12 @@ class DDL:
|
||||
to_field_name = field.to_field_instance.source_field
|
||||
if not to_field_name:
|
||||
to_field_name = field.to_field_instance.model_field_name
|
||||
db_table = model._meta.db_table
|
||||
return self._DROP_FK_TEMPLATE.format(
|
||||
table_name=model._meta.db_table,
|
||||
table_name=db_table,
|
||||
fk_name=self.schema_generator._generate_fk_name(
|
||||
from_table=model._meta.db_table,
|
||||
from_field=field.model_field_name,
|
||||
from_table=db_table,
|
||||
from_field=field.source_field or field.model_field_name + "_id",
|
||||
to_table=field.related_model._meta.db_table,
|
||||
to_field=to_field_name,
|
||||
),
|
||||
|
||||
@@ -5,7 +5,7 @@ from copy import deepcopy
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Type
|
||||
|
||||
from tortoise import ForeignKeyFieldInstance, Model, Tortoise
|
||||
from tortoise import BackwardFKRelation, ForeignKeyFieldInstance, Model, Tortoise
|
||||
from tortoise.backends.mysql.schema_generator import MySQLSchemaGenerator
|
||||
from tortoise.fields import Field
|
||||
|
||||
@@ -18,6 +18,9 @@ from alice.utils import get_app_connection
|
||||
class Migrate:
|
||||
upgrade_operators: List[str] = []
|
||||
downgrade_operators: List[str] = []
|
||||
_upgrade_fk_operators: List[str] = []
|
||||
_downgrade_fk_operators: List[str] = []
|
||||
|
||||
ddl: DDL
|
||||
migrate_config: dict
|
||||
old_models = "old_models"
|
||||
@@ -73,7 +76,7 @@ class Migrate:
|
||||
filename = f"{cls._get_latest_version() + 1}_{now}_{name}.json"
|
||||
content = {
|
||||
"upgrade": cls.upgrade_operators,
|
||||
"download": cls.downgrade_operators,
|
||||
"downgrade": cls.downgrade_operators,
|
||||
"migrate": False,
|
||||
}
|
||||
with open(os.path.join(cls.migrate_location, filename), "w") as f:
|
||||
@@ -94,18 +97,26 @@ class Migrate:
|
||||
if not cls.upgrade_operators:
|
||||
return False
|
||||
|
||||
cls._merge_operators()
|
||||
|
||||
return cls._generate_diff_sql(name)
|
||||
|
||||
@classmethod
|
||||
def _add_operator(cls, operator: str, upgrade=True):
|
||||
def _add_operator(cls, operator: str, upgrade=True, fk=False):
|
||||
if upgrade:
|
||||
cls.upgrade_operators.append(operator)
|
||||
if fk:
|
||||
cls._upgrade_fk_operators.append(operator)
|
||||
else:
|
||||
cls.upgrade_operators.append(operator)
|
||||
else:
|
||||
cls.downgrade_operators.append(operator)
|
||||
if fk:
|
||||
cls._downgrade_fk_operators.append(operator)
|
||||
else:
|
||||
cls.downgrade_operators.append(operator)
|
||||
|
||||
@classmethod
|
||||
def cp_models(
|
||||
cls, model_files: List[str], old_model_file,
|
||||
cls, model_files: List[str], old_model_file,
|
||||
):
|
||||
"""
|
||||
cp currents models to old_model_files
|
||||
@@ -113,13 +124,11 @@ class Migrate:
|
||||
:param old_model_file:
|
||||
:return:
|
||||
"""
|
||||
pattern = (
|
||||
r"(ManyToManyField|ForeignKeyField|OneToOneField)\((model_name)?(\"|\')(\w+)(.+)\)"
|
||||
)
|
||||
pattern = r"(ManyToManyField|ForeignKeyField|OneToOneField)\(('|\")(\w+)."
|
||||
for i, model_file in enumerate(model_files):
|
||||
with open(model_file, "r") as f:
|
||||
content = f.read()
|
||||
ret = re.sub(pattern, rf"\1\2(\3{cls.diff_app}\5)", content)
|
||||
ret = re.sub(pattern, rf"\1(\2{cls.diff_app}.", content)
|
||||
with open(old_model_file, "w" if i == 0 else "w+a") as f:
|
||||
f.write(ret)
|
||||
|
||||
@@ -142,7 +151,7 @@ class Migrate:
|
||||
|
||||
@classmethod
|
||||
def _diff_models(
|
||||
cls, old_models: Dict[str, Type[Model]], new_models: Dict[str, Type[Model]], upgrade=True
|
||||
cls, old_models: Dict[str, Type[Model]], new_models: Dict[str, Type[Model]], upgrade=True
|
||||
):
|
||||
"""
|
||||
diff models and add operators
|
||||
@@ -184,18 +193,37 @@ class Migrate:
|
||||
new_keys = new_fields_map.keys()
|
||||
for new_key in new_keys:
|
||||
new_field = new_fields_map.get(new_key)
|
||||
if cls._exclude_field(new_field):
|
||||
continue
|
||||
if new_key not in old_keys:
|
||||
cls._add_operator(cls._add_field(new_model, new_field), upgrade)
|
||||
cls._add_operator(
|
||||
cls._add_field(new_model, new_field),
|
||||
upgrade,
|
||||
isinstance(new_field, ForeignKeyFieldInstance),
|
||||
)
|
||||
else:
|
||||
old_field = old_fields_map.get(new_key)
|
||||
if old_field.index and not new_field.index:
|
||||
cls._add_operator(cls._remove_index(old_model, old_field), upgrade)
|
||||
cls._add_operator(
|
||||
cls._remove_index(old_model, old_field),
|
||||
upgrade,
|
||||
isinstance(old_field, ForeignKeyFieldInstance),
|
||||
)
|
||||
elif new_field.index and not old_field.index:
|
||||
cls._add_operator(cls._add_index(new_model, new_field), upgrade)
|
||||
cls._add_operator(
|
||||
cls._add_index(new_model, new_field),
|
||||
upgrade,
|
||||
isinstance(new_field, ForeignKeyFieldInstance),
|
||||
)
|
||||
|
||||
for old_key in old_keys:
|
||||
if old_key not in new_keys:
|
||||
field = old_fields_map.get(old_key)
|
||||
cls._add_operator(cls._remove_field(old_model, field), upgrade)
|
||||
field = old_fields_map.get(old_key)
|
||||
if old_key not in new_keys and not cls._exclude_field(field):
|
||||
cls._add_operator(
|
||||
cls._remove_field(old_model, field),
|
||||
upgrade,
|
||||
isinstance(field, ForeignKeyFieldInstance),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _remove_index(cls, model: Type[Model], field: Field):
|
||||
@@ -205,6 +233,10 @@ class Migrate:
|
||||
def _add_index(cls, model: Type[Model], field: Field):
|
||||
return cls.ddl.add_index(model, [field.model_field_name], field.unique)
|
||||
|
||||
@classmethod
|
||||
def _exclude_field(cls, field: Field):
|
||||
return isinstance(field, BackwardFKRelation)
|
||||
|
||||
@classmethod
|
||||
def _add_field(cls, model: Type[Model], field: Field):
|
||||
if isinstance(field, ForeignKeyFieldInstance):
|
||||
@@ -217,3 +249,28 @@ class Migrate:
|
||||
if isinstance(field, ForeignKeyFieldInstance):
|
||||
return cls.ddl.drop_fk(model, field)
|
||||
return cls.ddl.drop_column(model, field.model_field_name)
|
||||
|
||||
@classmethod
|
||||
def _add_fk(cls, model: Type[Model], field: Field):
|
||||
return cls.ddl.add_fk(model, field)
|
||||
|
||||
@classmethod
|
||||
def _remove_fk(cls, model: Type[Model], field: Field):
|
||||
return cls.ddl.drop_fk(model, field)
|
||||
|
||||
@classmethod
|
||||
def _merge_operators(cls):
|
||||
"""
|
||||
fk must be last when add,first when drop
|
||||
:return:
|
||||
"""
|
||||
for _upgrade_fk_operator in cls._upgrade_fk_operators:
|
||||
if "ADD" in _upgrade_fk_operator:
|
||||
cls.upgrade_operators.append(_upgrade_fk_operator)
|
||||
else:
|
||||
cls.upgrade_operators.insert(0, _upgrade_fk_operator)
|
||||
for _downgrade_fk_operator in cls._downgrade_fk_operators:
|
||||
if "ADD" in _downgrade_fk_operator:
|
||||
cls.downgrade_operators.append(_downgrade_fk_operator)
|
||||
else:
|
||||
cls.downgrade_operators.insert(0, _downgrade_fk_operator)
|
||||
|
||||
Reference in New Issue
Block a user