Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use aerich for database migrations #1110

Merged
merged 7 commits into from
Oct 2, 2024
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Draft: use aerich for database migrations
bigherc18 committed Sep 29, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
commit 0ec6868f220be53acdaed80e0550dbc5dab88c01
45 changes: 42 additions & 3 deletions pdm.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -80,6 +80,7 @@ dependencies = [
"tortoise-orm==0.21.6",
"uvloop~=0.20",
"web3~=7.2",
"aerich~=0.7.2",
]

[tool.pdm.resolution]
65 changes: 64 additions & 1 deletion src/dipdup/cli.py
Original file line number Diff line number Diff line change
@@ -8,12 +8,14 @@
from contextlib import AsyncExitStack
from contextlib import suppress
from dataclasses import dataclass
from dataclasses import fields
from functools import wraps
from pathlib import Path
from typing import TYPE_CHECKING
from typing import Any
from typing import TypeVar
from typing import cast
from typing import Iterable

import click
import uvloop
@@ -28,8 +30,11 @@
from dipdup.report import save_report
from dipdup.sys import fire_and_forget
from dipdup.sys import set_up_process
from aerich.cli import cli as aerich_cli
from aerich.cli import Command as AerichCommand

if TYPE_CHECKING:
from types import ModuleType
from dipdup.config import DipDupConfig

ROOT_CONFIG = 'dipdup.yaml'
@@ -42,6 +47,13 @@
'config',
}

AERICH_CMDS = {
'history',
'heads',
'migrate',
'upgrade',
'downgrade',
}

_logger = logging.getLogger(__name__)
_click_wrap_text = click.formatting.wrap_text
@@ -144,6 +156,13 @@ def _print() -> None:
class CLIContext:
config_paths: list[str]
config: 'DipDupConfig'
command: AerichCommand

# NOTE: aerich CLI access the context via `ctx.obj["command"]`
def __getitem__(self, item):
if item in {f.name for f in fields(self)}:
return getattr(self, item)
raise KeyError(f"{item} is not a valid field name.")


def _cli_wrapper(fn: WrappedCommandT) -> WrappedCommandT:
@@ -221,6 +240,18 @@ def _skip_cli_group() -> bool:
return True


# TODO: add DipDupConfig type hint and import it
async def _create_aerich_command(config) -> AerichCommand:
from tortoise.backends.base.config_generator import generate_config

db_url = config.database.connection_string
# TODO: Refactor building the app_modules dict and use here and in the tortoise_wrapper function ?
modules: dict[str, Iterable[str | ModuleType]] = {"models": [f'{config.package}.models', "aerich.models"]}
tortoise_config = generate_config(db_url=db_url, app_modules=modules)
# TODO: add location in config file (and CLI option?)
return AerichCommand(tortoise_config=tortoise_config, app="models")


@click.group(
context_settings={'max_content_width': 120},
help=WELCOME_ASCII,
@@ -316,9 +347,12 @@ async def cli(ctx: click.Context, config: list[str], env_file: list[str], c: lis
if ctx.invoked_subcommand != 'init':
raise InitializationRequiredError(f'Failed to create a project package: {e}') from e

aerich_command = await _create_aerich_command(_config)

ctx.obj = CLIContext(
config_paths=config,
config=_config,
command=aerich_command,
)


@@ -335,6 +369,15 @@ async def run(ctx: click.Context) -> None:
config: DipDupConfig = ctx.obj.config
config.initialize()

aerich_command: AerichCommand = ctx.obj.command
# TODO: add safe option to the CLI and/or config
# TODO: make the logs DEBUG
try:
_logger.info("Trying to initializing database migrations")
await aerich_command.init_db(safe=True)
except FileExistsError:
_logger.info("Database migrations already initialized")

dipdup = DipDup(config)
await dipdup.run()

@@ -400,9 +443,13 @@ async def migrate(ctx: click.Context, dry_run: bool) -> None:
unsafe=True,
)
config.initialize()

aerich_command = await _create_aerich_command(config)

ctx.obj = CLIContext(
config_paths=ctx.parent.params['config'],
config=config,
command=aerich_command,
)
await _cli_unwrapper(init)(
ctx=ctx,
@@ -562,7 +609,16 @@ async def hasura_configure(ctx: click.Context, force: bool) -> None:
@_cli_wrapper
async def schema(ctx: click.Context) -> None:
"""Commands to manage database schema."""
pass
if ctx.invoked_subcommand in AERICH_CMDS:
aerich_command: AerichCommand = ctx.obj.command
await aerich_command.init()


schema.add_command(aerich_cli.commands["history"])
schema.add_command(aerich_cli.commands["heads"])
schema.add_command(aerich_cli.commands["migrate"])
schema.add_command(aerich_cli.commands["upgrade"])
schema.add_command(aerich_cli.commands["downgrade"])


@schema.command(name='approve')
@@ -701,6 +757,10 @@ async def schema_init(ctx: click.Context) -> None:
config.database.schema_name,
)

aerich_command: AerichCommand = ctx.obj.command
# TODO: add safe option to the CLI
await aerich_command.init_db(safe=True)

_logger.info('Schema initialized')


@@ -789,9 +849,12 @@ async def new(
_logger.info('Initializing project')
config = DipDupConfig.load([Path(answers['package'])])
config.initialize()

aerich_command = await _create_aerich_command(config)
ctx.obj = CLIContext(
config_paths=[Path(answers['package']).joinpath(ROOT_CONFIG).as_posix()],
config=config,
command=aerich_command,
)
await _cli_unwrapper(init)(
ctx=ctx,
2 changes: 1 addition & 1 deletion src/dipdup/database.py
Original file line number Diff line number Diff line change
@@ -73,7 +73,7 @@ async def tortoise_wrapper(
if models:
if not models.endswith('.models'):
models += '.models'
model_modules['models'] = [models]
model_modules['models'] = [models, "aerich.models"]

# NOTE: Must be called before entering Tortoise context
decimal_precision = decimal_precision or guess_decimal_precision(models)
2 changes: 2 additions & 0 deletions src/dipdup/dipdup.py
Original file line number Diff line number Diff line change
@@ -726,6 +726,7 @@ async def _initialize_schema(self) -> None:
self._schema = await Schema.get_or_none(name=schema_name)

# NOTE: Call with existing Schema too to create new tables if missing
# TODO: Check if it doesn't conflict with aerich migrations
try:
await generate_schema(
conn,
@@ -740,6 +741,7 @@ async def _initialize_schema(self) -> None:

schema_hash = get_schema_hash(conn)

# TODO: Advise to run `dipdup schema migrate` before `dipdup schema approve`
if self._schema is None:
await self._ctx.fire_hook('on_reindex')