redis-om-python/redis_om/model/migrations/migrator.py
2021-10-22 06:33:05 -07:00

156 lines
4.7 KiB
Python

import hashlib
import logging
from dataclasses import dataclass
from enum import Enum
from typing import Optional, Union
from redis import ResponseError, Redis
from aioredis import ResponseError as AResponseError, Redis as ARedis
from redis_om.model.model import model_registry
log = logging.getLogger(__name__)
import importlib # noqa: E402
import pkgutil # noqa: E402
class MigrationError(Exception):
pass
def import_submodules(root_module_name: str):
"""Import all submodules of a module, recursively."""
# TODO: Call this without specifying a module name, to import everything?
root_module = importlib.import_module(root_module_name)
if not hasattr(root_module, "__path__"):
raise MigrationError(
"The root module must be a Python package. "
f"You specified: {root_module_name}"
)
for loader, module_name, is_pkg in pkgutil.walk_packages(
root_module.__path__, root_module.__name__ + "." # type: ignore
):
importlib.import_module(module_name)
def schema_hash_key(index_name):
return f"{index_name}:hash"
async def create_index(redis: Union[Redis, ARedis], index_name, schema, current_hash):
try:
await redis.execute_command(f"ft.info {index_name}")
except (ResponseError, AResponseError):
await redis.execute_command(f"ft.create {index_name} {schema}")
await redis.set(schema_hash_key(index_name), current_hash)
else:
log.info("Index already exists, skipping. Index hash: %s", index_name)
class MigrationAction(Enum):
CREATE = 2
DROP = 1
@dataclass
class IndexMigration:
model_name: str
index_name: str
schema: str
hash: str
action: MigrationAction
redis: Union[Redis, ARedis]
previous_hash: Optional[str] = None
async def run(self):
if self.action is MigrationAction.CREATE:
await self.create()
elif self.action is MigrationAction.DROP:
await self.drop()
async def create(self):
try:
await create_index(self.redis, self.index_name, self.schema, self.hash)
except ResponseError:
log.info("Index already exists: %s", self.index_name)
async def drop(self):
try:
await self.redis.execute_command(f"FT.DROPINDEX {self.index_name}")
except ResponseError:
log.info("Index does not exist: %s", self.index_name)
class Migrator:
def __init__(self, redis: Union[Redis, ARedis], module=None):
self.module = module
self.migrations = []
self.redis = redis
async def run(self):
# Try to load any modules found under the given path or module name.
if self.module:
import_submodules(self.module)
for name, cls in model_registry.items():
hash_key = schema_hash_key(cls.Meta.index_name)
try:
schema = cls.redisearch_schema()
except NotImplementedError:
log.info("Skipping migrations for %s", name)
continue
current_hash = hashlib.sha1(schema.encode("utf-8")).hexdigest() # nosec
try:
await self.redis.execute_command("ft.info", cls.Meta.index_name)
except (ResponseError, AResponseError):
self.migrations.append(
IndexMigration(
name,
cls.Meta.index_name,
schema,
current_hash,
MigrationAction.CREATE,
self.redis
)
)
continue
stored_hash = self.redis.get(hash_key)
schema_out_of_date = current_hash != stored_hash
if schema_out_of_date:
# TODO: Switch out schema with an alias to avoid downtime -- separate migration?
self.migrations.append(
IndexMigration(
name,
cls.Meta.index_name,
schema,
current_hash,
MigrationAction.DROP,
self.redis,
stored_hash
)
)
self.migrations.append(
IndexMigration(
name,
cls.Meta.index_name,
schema,
current_hash,
MigrationAction.CREATE,
self.redis,
stored_hash
)
)
# TODO: Migration history
# TODO: Dry run with output
for migration in self.migrations:
await migration.run()