Add basic migrations, query expression resolver

redis_developer/orm/migrations/migrator.py

@@ -0,0 +1,116 @@
+import hashlib
+import logging
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional
+
+from redis import ResponseError
+
+from redis_developer.connections import get_redis_connection
+from redis_developer.orm.model import model_registry
+
+redis = get_redis_connection()
+log = logging.getLogger(__name__)
+
+
+import importlib
+import pkgutil
+
+
+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)
+    for loader, module_name, is_pkg in pkgutil.walk_packages(
+        root_module.__path__, root_module.__name__ + '.'):
+        importlib.import_module(module_name)
+
+
+def schema_hash_key(index_name):
+    return f"{index_name}:hash"
+
+
+def create_index(index_name, schema, current_hash):
+    redis.execute_command(f"ft.create {index_name} "
+                          f"{schema}")
+    redis.set(schema_hash_key(index_name), current_hash)
+
+
+class MigrationAction(Enum):
+    CREATE = 2
+    DROP = 1
+
+
+@dataclass            
+class IndexMigration:
+    model_name: str
+    index_name: str
+    schema: str
+    hash: str
+    action: MigrationAction
+    previous_hash: Optional[str] = None
+    
+    def run(self):
+        if self.action is MigrationAction.CREATE:
+            self.create()
+        elif self.action is MigrationAction.DROP:
+            self.drop()
+    
+    def create(self):
+        return create_index(self.index_name, self.schema, self.hash)
+    
+    def drop(self):
+        redis.execute_command(f"FT.DROPINDEX {self.index_name}")
+
+
+class Migrator:
+    def __init__(self, module=None):
+        # Try to load any modules found under the given path or module name.
+        if module:
+            import_submodules(module)
+
+        self.migrations = []
+
+        for name, cls in model_registry.items():
+            hash_key = schema_hash_key(cls.Meta.index_name)
+            try:
+                schema = cls.schema()
+            except NotImplementedError:
+                log.info("Skipping migrations for %s", name)
+                continue
+            current_hash = hashlib.sha1(schema.encode("utf-8")).hexdigest()
+
+            try:
+                redis.execute_command("ft.info", cls.Meta.index_name)
+            except ResponseError:
+                self.migrations.append(
+                    IndexMigration(name, cls.Meta.index_name, schema, current_hash,
+                                   MigrationAction.CREATE))
+
+            stored_hash = 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, stored_hash))
+                self.migrations.append(
+                     IndexMigration(name, cls.Meta.index_name, schema, current_hash,
+                                   MigrationAction.CREATE, stored_hash))
+    
+    @property
+    def valid_migrations(self):
+        return self.missing_indexes.keys() + self.out_of_date_indexes.keys()
+        
+    def validate_migration(self, model_class_name):
+        if model_class_name not in self.valid_migrations:
+            migrations = ", ".join(self.valid_migrations)
+            raise RuntimeError(f"No migration found for {model_class_name}."
+                               f"Valid migrations are: {migrations}")
+
+    def run(self):
+        # TODO: Migration history
+        # TODO: Dry run with output
+        for migration in self.migrations:
+            migration.run()