def test_migrations(app_before_init_db):
"""The database migrations complete successfully."""
app, conf_obj = app_before_init_db
dbp.connect()
if conf_obj.database.engine == "PostgresqlDatabase":
dbp.execute_sql("DROP SCHEMA public CASCADE;")
dbp.execute_sql("CREATE SCHEMA public;")
dbp.execute_sql("GRANT ALL ON SCHEMA public TO public;")
router = Router(dbp, migrate_dir="migrations", ignore=["basemodel"])
router.run()
applied_migrations = list(router.done)
applied_migrations.reverse()
# Shut up a warning in rollback that we can't do anything about.
logging.getLogger("peewee_migrate").warn = Mock()
# Make sure new rollbacks work. The existing ones are what they are.
for m in applied_migrations:
if m == "029_message_read":
break
router.rollback(m)
dbp.close()
def rollback(name):
"""
Rolls back migration
"""
migrator = Router(db, logger=LOGGER)
migrator.migrate_dir = os.path.join(base_dir, 'application', 'migrations')
migrator.rollback(name)
def rollback(self):
from peewee_migrate import Router
with self._db() as db, db.atomic():
router = Router(db,
migrate_table=self.migrate_table,
migrate_dir=self.migrate_dir)
if router.done:
router.rollback(router.done[-1])
def rollback(self, _):
"""Rollbacks a latest migration."""
from peewee_migrate import Router
with self._db('cardinal') as db, db.atomic():
router = Router(db,
migrate_table=self.migrate_table,
migrate_dir=self.migrate_dir)
if router.done:
router.rollback(router.done[-1])
def rollback():
try:
router = Router(models.database, ignore=[models.MyModel])
if not router.done:
logging.warning("没有migrate记录，无法回滚")
sys.exit(-1)
name = router.done[-1]
router.rollback(name)
clean(name)
except:
logging.exception("rollback失败")
sys.exit(-1)
def _rollbackmigration(migration_name):
router = Router(DB)
router.rollback(migration_name)
router = Router(database,
migrate_dir='../migrations'
if os.getcwd().endswith('scripts') else 'migrations',
ignore=['basemodel'])
parser = argparse.ArgumentParser(
description='Apply or manage database migrations.')
parser.add_argument('-c',
'--create',
metavar='NAME',
help='Creates a new migration')
parser.add_argument('-a',
'--auto',
metavar='NAME',
help='Creates a new migration (automatic)')
parser.add_argument('-r',
'--rollback',
metavar='NAME',
help='Rolls back a migration')
args = parser.parse_args()
if args.create:
router.create(args.create)
elif args.auto:
router.create(args.auto, 'app')
elif args.rollback:
router.rollback(args.rollback)
else:
router.run()
class Manager:
def __init__(self, app: typing.Union[str, apistar.App, apistar.ASyncApp]):
"""
Create a database manager from API Star application.
:param app: API Star application.
:return: Manager instance.
"""
self.app = self.get_app(app) if isinstance(app, str) else app
self.component = self.get_database_component(self.app)
self.router = Router(self.component.database)
@classmethod
def get_app(cls, path: str) -> typing.Union[apistar.App, apistar.ASyncApp]:
"""
Get database manager from API Star app path.
:param path: API Star app path in format <package>.<module>:<variable>
:return: Database manager.
"""
try:
m, c = path.rsplit(":", 1)
module = importlib.import_module(m)
app = getattr(module, c)
except ValueError:
raise ImportError(
"Wrong path path, it should be: <package>.<module>:<variable>")
except (AttributeError, ImportError):
raise ImportError("API Star path not found '{}'".format(path))
return app
@classmethod
def get_database_component(
cls, app: typing.Union[apistar.App, apistar.ASyncApp]
) -> PeeweeDatabaseComponent:
"""
Get database component from API Star application.
:param app: API Star application.
:return: Database component
"""
for component in app.injector.components:
if isinstance(component, PeeweeDatabaseComponent):
return component
raise ValueError(
"No 'PeeweeDatabaseComponent' found in API Star application")
def upgrade(self, target: typing.Optional[str] = None, fake: bool = False):
"""
Apply migrations sequentially to a target.
:param target: Target migration name.
:param fake: Fake migrations.
"""
self.router.run(target, fake)
def downgrade(self, target: typing.Optional[str] = None):
"""
Rollback migrations sequentially to a target.
:param target: Target migration name.
"""
self.router.rollback(target)
def merge(self, name: typing.Optional[str] = None):
"""
Merge all migrations into a single one.
:param name: Migration name.
"""
self.router.merge(name)
def create(self, name: str, module: typing.Optional[str] = None):
"""
Create a new migration with given name. If module is provided, the migration will be generated automatically.
:param name: Migration name.
:param module: Module to infer changes.
"""
self.router.create(name, auto=module)
@property
def models(self) -> typing.List[peewee.Model]:
"""
Returns a list of registered models.
:return: List of registered models.
"""
return list(Model.register) # noqa
def create_tables(self):
"""
Create tables in database.
"""
self.component.database.create_tables(self.models)
def drop_tables(self):
"""
Drop tables from database.
"""
self.component.database.drop_tables(self.models)
def __repr__(self):
migrations = "\n".join(
[f" [x] {migration}" for migration in self.router.done])
migrations += "\n".join(
[f" [ ] {migration}" for migration in self.router.diff])
if len(Model.register) == 0:
models = "No models found."
elif len(Model.register) == 1:
models = f"{len(Model.register)} model found:\n"
else:
models = f"{len(Model.register)} models found:\n"
models += "\n".join([
f" - {model.__module__}.{model.__name__}"
for model in Model.register
])
return f"Migrations\n----------\n{migrations}\n\nModels\n------\n{models}"
