pub/aerich
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #32 from psbleep/close-tortoise-connections

Browse Source 
Close database connections
This commit is contained in:
long2ice 2020-07-26 16:52:08 +08:00 committed by GitHub
commit 18cb75f555
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 18 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
aerich/cli.py
View File

@ -1,3 +1,4 @@
import functools
import json import json
import os import os
import sys import sys
@ -27,6 +28,16 @@ class Color(str, Enum):
parser = ConfigParser() parser = ConfigParser()
def close_db(func):
@functools.wraps(func)
async def close_db_inner(*args, **kwargs):
result = await func(*args, **kwargs)
await Tortoise.close_connections()
return result
return close_db_inner
@click.group(context_settings={"help_option_names": ["-h", "--help"]}) @click.group(context_settings={"help_option_names": ["-h", "--help"]})
@click.version_option(__version__, "-V", "--version") @click.version_option(__version__, "-V", "--version")
@click.option( @click.option(
@ -70,6 +81,7 @@ async def cli(ctx: Context, config, app, name):
@cli.command(help="Generate migrate changes file.") @cli.command(help="Generate migrate changes file.")
@click.option("--name", default="update", show_default=True, help="Migrate name.") @click.option("--name", default="update", show_default=True, help="Migrate name.")
@click.pass_context @click.pass_context
@close_db
async def migrate(ctx: Context, name): async def migrate(ctx: Context, name):
config = ctx.obj["config"] config = ctx.obj["config"]
location = ctx.obj["location"] location = ctx.obj["location"]
@ -84,6 +96,7 @@ async def migrate(ctx: Context, name):
@cli.command(help="Upgrade to latest version.") @cli.command(help="Upgrade to latest version.")
@click.pass_context @click.pass_context
@close_db
async def upgrade(ctx: Context): async def upgrade(ctx: Context):
config = ctx.obj["config"] config = ctx.obj["config"]
app = ctx.obj["app"] app = ctx.obj["app"]
@ -110,6 +123,7 @@ async def upgrade(ctx: Context):
@cli.command(help="Downgrade to previous version.") @cli.command(help="Downgrade to previous version.")
@click.pass_context @click.pass_context
@close_db
async def downgrade(ctx: Context): async def downgrade(ctx: Context):
app = ctx.obj["app"] app = ctx.obj["app"]
config = ctx.obj["config"] config = ctx.obj["config"]
@ -132,6 +146,7 @@ async def downgrade(ctx: Context):
@cli.command(help="Show current available heads in migrate location.") @cli.command(help="Show current available heads in migrate location.")
@click.pass_context @click.pass_context
@close_db
async def heads(ctx: Context): async def heads(ctx: Context):
app = ctx.obj["app"] app = ctx.obj["app"]
versions = Migrate.get_all_version_files() versions = Migrate.get_all_version_files()
@ -146,7 +161,8 @@ async def heads(ctx: Context):
@cli.command(help="List all migrate items.") @cli.command(help="List all migrate items.")
@click.pass_context @click.pass_context
def history(ctx): @close_db
async def history(ctx: Context):
versions = Migrate.get_all_version_files() versions = Migrate.get_all_version_files()
for version in versions: for version in versions:
click.secho(version, fg=Color.green) click.secho(version, fg=Color.green)
@ -196,6 +212,7 @@ async def init(
show_default=True, show_default=True,
) )
@click.pass_context @click.pass_context
@close_db
async def init_db(ctx: Context, safe): async def init_db(ctx: Context, safe):
config = ctx.obj["config"] config = ctx.obj["config"]
location = ctx.obj["location"] location = ctx.obj["location"]