1from plain import models
  2from plain.models.db import DatabaseError
  3from plain.models.registry import ModelsRegistry
  4from plain.utils.functional import classproperty
  5from plain.utils.timezone import now
  6
  7from .exceptions import MigrationSchemaMissing
  8
  9
 10class MigrationRecorder:
 11    """
 12    Deal with storing migration records in the database.
 13
 14    Because this table is actually itself used for dealing with model
 15    creation, it's the one thing we can't do normally via migrations.
 16    We manually handle table creation/schema updating (using schema backend)
 17    and then have a floating model to do queries with.
 18
 19    If a migration is unapplied its row is removed from the table. Having
 20    a row in the table always means a migration is applied.
 21    """
 22
 23    _migration_class = None
 24
 25    @classproperty
 26    def Migration(cls):
 27        """
 28        Lazy load to avoid PackageRegistryNotReady if installed packages import
 29        MigrationRecorder.
 30        """
 31        if cls._migration_class is None:
 32            _models_registry = ModelsRegistry()
 33            _models_registry.ready = True
 34
 35            class Migration(models.Model):
 36                app = models.CharField(max_length=255)
 37                name = models.CharField(max_length=255)
 38                applied = models.DateTimeField(default=now)
 39
 40                class Meta:
 41                    models_registry = _models_registry
 42                    package_label = "migrations"
 43                    db_table = "plainmigrations"
 44
 45                def __str__(self):
 46                    return f"Migration {self.name} for {self.app}"
 47
 48            cls._migration_class = Migration
 49        return cls._migration_class
 50
 51    def __init__(self, connection):
 52        self.connection = connection
 53
 54    @property
 55    def migration_qs(self):
 56        return self.Migration.objects.all()
 57
 58    def has_table(self):
 59        """Return True if the plainmigrations table exists."""
 60        with self.connection.cursor() as cursor:
 61            tables = self.connection.introspection.table_names(cursor)
 62        return self.Migration._meta.db_table in tables
 63
 64    def ensure_schema(self):
 65        """Ensure the table exists and has the correct schema."""
 66        # If the table's there, that's fine - we've never changed its schema
 67        # in the codebase.
 68        if self.has_table():
 69            return
 70        # Make the table
 71        try:
 72            with self.connection.schema_editor() as editor:
 73                editor.create_model(self.Migration)
 74        except DatabaseError as exc:
 75            raise MigrationSchemaMissing(
 76                f"Unable to create the plainmigrations table ({exc})"
 77            )
 78
 79    def applied_migrations(self):
 80        """
 81        Return a dict mapping (package_name, migration_name) to Migration instances
 82        for all applied migrations.
 83        """
 84        if self.has_table():
 85            return {
 86                (migration.app, migration.name): migration
 87                for migration in self.migration_qs
 88            }
 89        else:
 90            # If the plainmigrations table doesn't exist, then no migrations
 91            # are applied.
 92            return {}
 93
 94    def record_applied(self, app, name):
 95        """Record that a migration was applied."""
 96        self.ensure_schema()
 97        self.migration_qs.create(app=app, name=name)
 98
 99    def record_unapplied(self, app, name):
100        """Record that a migration was unapplied."""
101        self.ensure_schema()
102        self.migration_qs.filter(app=app, name=name).delete()
103
104    def flush(self):
105        """Delete all migration records. Useful for testing migrations."""
106        self.migration_qs.all().delete()