147 lines
5.4 KiB
Python
147 lines
5.4 KiB
Python
import importlib.util
|
|
import pathlib
|
|
from types import ModuleType
|
|
from typing import Generator, Iterable
|
|
|
|
from alembic import command as alembic_command
|
|
from alembic.config import Config as AlembicConfig
|
|
|
|
from birthday_pool_bot.repositories import RepositorySettings
|
|
|
|
|
|
class Migrator:
|
|
MIGRATION_FILENAME_TEMPLATE = (
|
|
"%%(year)d_%%(month)02d_%%(day)02d_%%(hour)02dh%%(minute)02dm%%(second)02ds_%%(slug)s"
|
|
)
|
|
|
|
def __init__(self, settings: RepositorySettings):
|
|
self._settings = settings
|
|
|
|
self._alembic_config = self._generate_alembic_config()
|
|
|
|
def _generate_alembic_config(self) -> AlembicConfig:
|
|
migrations_path = self.get_migrations_directory_path()
|
|
|
|
config = AlembicConfig()
|
|
config.set_main_option("script_location", str(migrations_path))
|
|
config.set_main_option("file_template", self.MIGRATION_FILENAME_TEMPLATE)
|
|
config.set_main_option("sqlalchemy.url", str(self._settings.dsn).replace("%", "%%"))
|
|
|
|
return config
|
|
|
|
def squash_migrations(
|
|
self,
|
|
new_migration_id: str | None = None,
|
|
**kwargs,
|
|
) -> str:
|
|
new_message = kwargs.get("new_message", None)
|
|
|
|
if new_migration_id is None:
|
|
last_migration_id = self.get_latest_migration()
|
|
new_migration_id = last_migration_id
|
|
|
|
self.rollback(migration_id="base")
|
|
self._delete_all_migrations()
|
|
new_migration_id = self.create_migration(
|
|
migration_id=new_migration_id,
|
|
message=new_message,
|
|
)
|
|
return new_migration_id
|
|
|
|
def migrate(self, migration_id: str | None = None, **kwargs) -> str | None:
|
|
alembic_revision = migration_id or "head"
|
|
alembic_command.upgrade(config=self._alembic_config, revision=alembic_revision)
|
|
|
|
return self.get_current_migration()
|
|
|
|
def rollback(self, migration_id: str | None = None, **kwargs) -> str | None:
|
|
alembic_revision = migration_id or "-1"
|
|
alembic_command.downgrade(config=self._alembic_config, revision=alembic_revision)
|
|
|
|
return self.get_current_migration()
|
|
|
|
def get_current_migration(self) -> str | None:
|
|
return None
|
|
|
|
def create_migration(
|
|
self,
|
|
migration_id: str | None = None,
|
|
**kwargs,
|
|
) -> str:
|
|
message = kwargs.get("message", None)
|
|
|
|
alembic_command.revision(
|
|
self._alembic_config,
|
|
message=message,
|
|
rev_id=migration_id,
|
|
autogenerate=True,
|
|
)
|
|
|
|
last_migration_id = self.get_latest_migration()
|
|
if last_migration_id is None:
|
|
raise RuntimeError("Migration not created")
|
|
|
|
return last_migration_id
|
|
|
|
def get_latest_migration(self) -> str | None:
|
|
last_migration_id = None
|
|
for migration_id in self.get_migrations():
|
|
last_migration_id = migration_id
|
|
return last_migration_id
|
|
|
|
def get_migrations(self) -> Generator[str, None, None]:
|
|
migrations_modules = self._get_migrations_modules()
|
|
sorted_migrations_modules = self._sort_migrations_modules(modules=migrations_modules)
|
|
for migration_module in sorted_migrations_modules:
|
|
yield migration_module.revision
|
|
|
|
def _get_migrations_modules(self) -> Iterable[ModuleType]:
|
|
migrations_path = self.get_migrations_directory_path()
|
|
versions_path = migrations_path / "versions"
|
|
|
|
for index, migration_module_path in enumerate(versions_path.glob("*.py")):
|
|
migration_module_name = f"migration_{index:08}"
|
|
migration_module_spec = importlib.util.spec_from_file_location(
|
|
name=migration_module_name,
|
|
location=migration_module_path,
|
|
)
|
|
if migration_module_spec is None or migration_module_spec.loader is None:
|
|
raise ValueError(f"Cannot load migration module '{migration_module_path}'")
|
|
migration_module = importlib.util.module_from_spec(spec=migration_module_spec)
|
|
migration_module_spec.loader.exec_module(migration_module)
|
|
yield migration_module
|
|
|
|
def _sort_migrations_modules(self, modules: Iterable[ModuleType]) -> list[ModuleType]:
|
|
next_module_map = {}
|
|
first_module = None
|
|
for module in modules:
|
|
if module.down_revision is None:
|
|
if first_module is not None:
|
|
raise ValueError((
|
|
"Found multiple first migrations: "
|
|
f"'{module.revision}' and '{first_module.revision}'"
|
|
))
|
|
first_module = module
|
|
else:
|
|
next_module_map[module.down_revision] = module
|
|
if first_module is None:
|
|
raise ValueError("Doesn't found first migration")
|
|
|
|
current_module = first_module
|
|
sorted_modules = [current_module]
|
|
while next_module_map:
|
|
current_module = next_module_map.pop(current_module.revision)
|
|
sorted_modules.append(current_module)
|
|
|
|
return sorted_modules
|
|
|
|
def _delete_all_migrations(self):
|
|
migrations_path = self.get_migrations_directory_path()
|
|
versions_path = migrations_path / "versions"
|
|
for migration_file_path in versions_path.glob("*.py"):
|
|
migration_file_path.unlink()
|
|
|
|
def get_migrations_directory_path(self) -> pathlib.Path:
|
|
directory_path = pathlib.Path(__file__).parent
|
|
return directory_path
|