plain.postgres 0.84.0__py3-none-any.whl

plain/postgres/migrations/exceptions.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import Any
+
+from plain.postgres.db import DatabaseError
+
+
+class AmbiguityError(Exception):
+    """More than one migration matches a name prefix."""
+
+    pass
+
+
+class BadMigrationError(Exception):
+    """There's a bad migration (unreadable/bad format/etc.)."""
+
+    pass
+
+
+class CircularDependencyError(Exception):
+    """There's an impossible-to-resolve circular dependency."""
+
+    pass
+
+
+class InconsistentMigrationHistory(Exception):
+    """An applied migration has some of its dependencies not applied."""
+
+    pass
+
+
+class InvalidBasesError(ValueError):
+    """A model's base classes can't be resolved."""
+
+    pass
+
+
+class NodeNotFoundError(LookupError):
+    """An attempt on a node is made that is not available in the graph."""
+
+    def __init__(self, message: str, node: Any, origin: Any = None) -> None:
+        self.message = message
+        self.origin = origin
+        self.node = node
+
+    def __str__(self) -> str:
+        return self.message
+
+    def __repr__(self) -> str:
+        return f"NodeNotFoundError({self.node!r})"
+
+
+class MigrationSchemaMissing(DatabaseError):
+    pass

plain/postgres/migrations/executor.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, Any
+
+from ..transaction import atomic
+from .loader import MigrationLoader
+from .migration import Migration
+from .recorder import MigrationRecorder
+from .state import ProjectState
+
+if TYPE_CHECKING:
+    from plain.postgres.connection import DatabaseConnection
+
+
+class MigrationExecutor:
+    """
+    End-to-end migration execution - load migrations and run them up or down
+    to a specified set of targets.
+    """
+
+    def __init__(
+        self,
+        connection: DatabaseConnection,
+        progress_callback: Callable[..., Any] | None = None,
+    ) -> None:
+        self.connection = connection
+        self.loader = MigrationLoader(self.connection)
+        self.recorder = MigrationRecorder(self.connection)
+        self.progress_callback = progress_callback
+
+    def migration_plan(
+        self, targets: list[tuple[str, str]], clean_start: bool = False
+    ) -> list[Migration]:
+        """
+        Given a set of targets, return a list of Migration instances.
+        """
+        plan = []
+        if clean_start:
+            applied = {}
+        else:
+            applied_source = self.loader.applied_migrations or {}
+            applied = dict(applied_source)
+        for target in targets:
+            for migration in self.loader.graph.forwards_plan(target):
+                if migration not in applied:
+                    plan.append(self.loader.graph.nodes[migration])
+                    applied[migration] = self.loader.graph.nodes[migration]
+        return plan
+
+    def _create_project_state(
+        self, with_applied_migrations: bool = False
+    ) -> ProjectState:
+        """
+        Create a project state including all the applications without
+        migrations and applied migrations if with_applied_migrations=True.
+        """
+        state = ProjectState(real_packages=self.loader.unmigrated_packages)
+        if with_applied_migrations:
+            # Create the forwards plan Plain would follow on an empty database
+            full_plan = self.migration_plan(
+                self.loader.graph.leaf_nodes(), clean_start=True
+            )
+            applied_source = self.loader.applied_migrations or {}
+            applied_migrations = {
+                self.loader.graph.nodes[key]
+                for key in applied_source
+                if key in self.loader.graph.nodes
+            }
+            for migration in full_plan:
+                if migration in applied_migrations:
+                    migration.mutate_state(state, preserve=False)
+        return state
+
+    def migrate(
+        self,
+        targets: list[tuple[str, str]],
+        plan: list[Migration] | None = None,
+        state: ProjectState | None = None,
+        fake: bool = False,
+        atomic_batch: bool = False,
+    ) -> ProjectState:
+        """
+        Migrate the database up to the given targets.
+
+        Plain first needs to create all project states before a migration is
+        (un)applied and in a second step run all the database operations.
+
+        atomic_batch: Whether to run all migrations in a single transaction.
+        """
+        # The plain_migrations table must be present to record applied
+        # migrations, but don't create it if there are no migrations to apply.
+        if plan == []:
+            if not self.recorder.has_table():
+                return self._create_project_state(with_applied_migrations=False)
+        else:
+            self.recorder.ensure_schema()
+
+        if plan is None:
+            plan = self.migration_plan(targets)
+        # Create the forwards plan Plain would follow on an empty database
+        full_plan = self.migration_plan(
+            self.loader.graph.leaf_nodes(), clean_start=True
+        )
+
+        if not plan:
+            if state is None:
+                # The resulting state should include applied migrations.
+                state = self._create_project_state(with_applied_migrations=True)
+        else:
+            if state is None:
+                # The resulting state should still include applied migrations.
+                state = self._create_project_state(with_applied_migrations=True)
+
+            migrations_to_run = set(plan)
+
+            # Choose context manager based on atomic_batch
+            batch_context = atomic if (atomic_batch and len(plan) > 1) else nullcontext
+
+            with batch_context():
+                for migration in full_plan:
+                    if not migrations_to_run:
+                        # We remove every migration that we applied from these sets so
+                        # that we can bail out once the last migration has been applied
+                        # and don't always run until the very end of the migration
+                        # process.
+                        break
+                    if migration in migrations_to_run:
+                        if "models_registry" not in state.__dict__:
+                            state.models_registry  # Render all -- performance critical
+                        state = self.apply_migration(state, migration, fake=fake)
+                        migrations_to_run.remove(migration)
+
+        self.check_replacements()
+
+        assert state is not None
+        return state
+
+    def apply_migration(
+        self, state: ProjectState, migration: Migration, fake: bool = False
+    ) -> ProjectState:
+        """Run a migration forwards."""
+        migration_recorded = False
+        if self.progress_callback:
+            self.progress_callback("apply_start", migration=migration, fake=fake)
+        if not fake:
+            # Alright, do it normally
+            with self.connection.schema_editor(
+                atomic=migration.atomic
+            ) as schema_editor:
+                state = migration.apply(
+                    state, schema_editor, operation_callback=self.progress_callback
+                )
+                if not schema_editor.deferred_sql:
+                    self.record_migration(migration)
+                    migration_recorded = True
+        if not migration_recorded:
+            self.record_migration(migration)
+        # Report progress
+        if self.progress_callback:
+            self.progress_callback("apply_success", migration=migration, fake=fake)
+        return state
+
+    def record_migration(self, migration: Migration) -> None:
+        # For replacement migrations, record individual statuses
+        if migration.replaces:
+            for package_label, name in migration.replaces:
+                self.recorder.record_applied(package_label, name)
+        else:
+            self.recorder.record_applied(migration.package_label, migration.name)
+
+    def check_replacements(self) -> None:
+        """
+        Mark replacement migrations applied if their replaced set all are.
+
+        Do this unconditionally on every migrate, rather than just when
+        migrations are applied or unapplied, to correctly handle the case
+        when a new squash migration is pushed to a deployment that already had
+        all its replaced migrations applied. In this case no new migration will
+        be applied, but the applied state of the squashed migration must be
+        maintained.
+        """
+        applied = self.recorder.applied_migrations()
+        for key, migration in self.loader.replacements.items():
+            all_applied = all(m in applied for m in migration.replaces)
+            if all_applied and key not in applied:
+                self.recorder.record_applied(*key)

plain/postgres/migrations/graph.py

@@ -0,0 +1,364 @@
+from __future__ import annotations
+
+from functools import total_ordering
+from typing import TYPE_CHECKING, Any, cast
+
+from plain.postgres.migrations.state import ProjectState
+
+from .exceptions import CircularDependencyError, NodeNotFoundError
+
+if TYPE_CHECKING:
+    from plain.postgres.migrations.migration import Migration
+
+
+@total_ordering
+class Node:
+    """
+    A single node in the migration graph. Contains direct links to adjacent
+    nodes in either direction.
+    """
+
+    def __init__(self, key: tuple[str, str]):
+        self.key = key
+        self.children: set[Node] = set()
+        self.parents: set[Node] = set()
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Node):
+            return self.key == other.key
+        return self.key == other
+
+    def __lt__(self, other: object) -> bool:
+        if isinstance(other, Node):
+            return self.key < other.key
+        if isinstance(other, tuple):
+            return self.key < cast(tuple[str, str], other)
+        return NotImplemented
+
+    def __hash__(self) -> int:
+        return hash(self.key)
+
+    def __getitem__(self, item: int) -> str:
+        return self.key[item]
+
+    def __str__(self) -> str:
+        return str(self.key)
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}: ({self.key[0]!r}, {self.key[1]!r})>"
+
+    def add_child(self, child: Node) -> None:
+        self.children.add(child)
+
+    def add_parent(self, parent: Node) -> None:
+        self.parents.add(parent)
+
+
+class DummyNode(Node):
+    """
+    A node that doesn't correspond to a migration file on disk.
+    (A squashed migration that was removed, for example.)
+
+    After the migration graph is processed, all dummy nodes should be removed.
+    If there are any left, a nonexistent dependency error is raised.
+    """
+
+    def __init__(
+        self,
+        key: tuple[str, str],
+        origin: Migration | tuple[str, str] | None,
+        error_message: str,
+    ):
+        super().__init__(key)
+        self.origin = origin
+        self.error_message = error_message
+
+    def raise_error(self) -> None:
+        raise NodeNotFoundError(self.error_message, self.key, origin=self.origin)
+
+
+class MigrationGraph:
+    """
+    Represent the digraph of all migrations in a project.
+
+    Each migration is a node, and each dependency is an edge. There are
+    no implicit dependencies between numbered migrations - the numbering is
+    merely a convention to aid file listing. Every new numbered migration
+    has a declared dependency to the previous number, meaning that VCS
+    branch merges can be detected and resolved.
+
+    Migrations files can be marked as replacing another set of migrations -
+    this is to support the "squash" feature. The graph handler isn't responsible
+    for these; instead, the code to load them in here should examine the
+    migration files and if the replaced migrations are all either unapplied
+    or not present, it should ignore the replaced ones, load in just the
+    replacing migration, and repoint any dependencies that pointed to the
+    replaced migrations to point to the replacing one.
+
+    A node should be a tuple: (app_path, migration_name). The tree special-cases
+    things within an app - namely, root nodes and leaf nodes ignore dependencies
+    to other packages.
+    """
+
+    def __init__(self):
+        self.node_map: dict[tuple[str, str], Node] = {}
+        self.nodes: dict[tuple[str, str], Migration | None] = {}
+
+    def add_node(self, key: tuple[str, str], migration: Migration) -> None:
+        assert key not in self.node_map
+        node = Node(key)
+        self.node_map[key] = node
+        self.nodes[key] = migration
+
+    def add_dummy_node(
+        self,
+        key: tuple[str, str],
+        origin: Migration | tuple[str, str] | None,
+        error_message: str,
+    ) -> None:
+        node = DummyNode(key, origin, error_message)
+        self.node_map[key] = node
+        self.nodes[key] = None
+
+    def add_dependency(
+        self,
+        migration: Migration | tuple[str, str] | None,
+        child: tuple[str, str],
+        parent: tuple[str, str],
+        skip_validation: bool = False,
+    ) -> None:
+        """
+        This may create dummy nodes if they don't yet exist. If
+        `skip_validation=True`, validate_consistency() should be called
+        afterward.
+        """
+        if child not in self.nodes:
+            error_message = (
+                f"Migration {migration} dependencies reference nonexistent"
+                f" child node {child!r}"
+            )
+            self.add_dummy_node(child, migration, error_message)
+        if parent not in self.nodes:
+            error_message = (
+                f"Migration {migration} dependencies reference nonexistent"
+                f" parent node {parent!r}"
+            )
+            self.add_dummy_node(parent, migration, error_message)
+        self.node_map[child].add_parent(self.node_map[parent])
+        self.node_map[parent].add_child(self.node_map[child])
+        if not skip_validation:
+            self.validate_consistency()
+
+    def remove_replaced_nodes(
+        self, replacement: tuple[str, str], replaced: list[tuple[str, str]]
+    ) -> None:
+        """
+        Remove each of the `replaced` nodes (when they exist). Any
+        dependencies that were referencing them are changed to reference the
+        `replacement` node instead.
+        """
+        # Cast list of replaced keys to set to speed up lookup later.
+        replaced_set: set[tuple[str, str]] = set(replaced)
+        try:
+            replacement_node = self.node_map[replacement]
+        except KeyError as err:
+            raise NodeNotFoundError(
+                f"Unable to find replacement node {replacement!r}. It was either never added"
+                " to the migration graph, or has been removed.",
+                replacement,
+            ) from err
+        for replaced_key in replaced_set:
+            self.nodes.pop(replaced_key, None)
+            replaced_node = self.node_map.pop(replaced_key, None)
+            if replaced_node:
+                for child in replaced_node.children:
+                    child.parents.remove(replaced_node)
+                    # We don't want to create dependencies between the replaced
+                    # node and the replacement node as this would lead to
+                    # self-referencing on the replacement node at a later iteration.
+                    if child.key not in replaced_set:
+                        replacement_node.add_child(child)
+                        child.add_parent(replacement_node)
+                for parent in replaced_node.parents:
+                    parent.children.remove(replaced_node)
+                    # Again, to avoid self-referencing.
+                    if parent.key not in replaced_set:
+                        replacement_node.add_parent(parent)
+                        parent.add_child(replacement_node)
+
+    def remove_replacement_node(
+        self, replacement: tuple[str, str], replaced: list[tuple[str, str]]
+    ) -> None:
+        """
+        The inverse operation to `remove_replaced_nodes`. Almost. Remove the
+        replacement node `replacement` and remap its child nodes to `replaced`
+        - the list of nodes it would have replaced. Don't remap its parent
+        nodes as they are expected to be correct already.
+        """
+        self.nodes.pop(replacement, None)
+        try:
+            replacement_node = self.node_map.pop(replacement)
+        except KeyError as err:
+            raise NodeNotFoundError(
+                f"Unable to remove replacement node {replacement!r}. It was either never added"
+                " to the migration graph, or has been removed already.",
+                replacement,
+            ) from err
+        replaced_nodes: set[Node] = set()
+        replaced_nodes_parents: set[Node] = set()
+        for key in replaced:
+            replaced_node = self.node_map.get(key)
+            if replaced_node:
+                replaced_nodes.add(replaced_node)
+                replaced_nodes_parents |= replaced_node.parents
+        # We're only interested in the latest replaced node, so filter out
+        # replaced nodes that are parents of other replaced nodes.
+        replaced_nodes -= replaced_nodes_parents
+        for child in replacement_node.children:
+            child.parents.remove(replacement_node)
+            for replaced_node in replaced_nodes:
+                replaced_node.add_child(child)
+                child.add_parent(replaced_node)
+        for parent in replacement_node.parents:
+            parent.children.remove(replacement_node)
+            # NOTE: There is no need to remap parent dependencies as we can
+            # assume the replaced nodes already have the correct ancestry.
+
+    def validate_consistency(self) -> None:
+        """Ensure there are no dummy nodes remaining in the graph."""
+        [n.raise_error() for n in self.node_map.values() if isinstance(n, DummyNode)]
+
+    def forwards_plan(self, target: tuple[str, str]) -> list[tuple[str, str]]:
+        """
+        Given a node, return a list of which previous nodes (dependencies) must
+        be applied, ending with the node itself. This is the list you would
+        follow if applying the migrations to a database.
+        """
+        if target not in self.nodes:
+            raise NodeNotFoundError(f"Node {target!r} not a valid node", target)
+        return self.iterative_dfs(self.node_map[target])
+
+    def iterative_dfs(
+        self, start: Node, forwards: bool = True
+    ) -> list[tuple[str, str]]:
+        """Iterative depth-first search for finding dependencies."""
+        visited: list[tuple[str, str]] = []
+        visited_set: set[Node] = set()
+        stack: list[tuple[Node, bool]] = [(start, False)]
+        while stack:
+            node, processed = stack.pop()
+            if node in visited_set:
+                pass
+            elif processed:
+                visited_set.add(node)
+                visited.append(node.key)
+            else:
+                stack.append((node, True))
+                stack += [
+                    (n, False)
+                    for n in sorted(node.parents if forwards else node.children)
+                ]
+        return visited
+
+    def root_nodes(self, app: str | None = None) -> list[tuple[str, str]]:
+        """
+        Return all root nodes - that is, nodes with no dependencies inside
+        their app. These are the starting point for an app.
+        """
+        roots: set[tuple[str, str]] = set()
+        for node in self.nodes:
+            if all(key[0] != node[0] for key in self.node_map[node].parents) and (
+                not app or app == node[0]
+            ):
+                roots.add(node)
+        return sorted(roots)
+
+    def leaf_nodes(self, app: str | None = None) -> list[tuple[str, str]]:
+        """
+        Return all leaf nodes - that is, nodes with no dependents in their app.
+        These are the "most current" version of an app's schema.
+        Having more than one per app is technically an error, but one that
+        gets handled further up, in the interactive command - it's usually the
+        result of a VCS merge and needs some user input.
+        """
+        leaves: set[tuple[str, str]] = set()
+        for node in self.nodes:
+            if all(key[0] != node[0] for key in self.node_map[node].children) and (
+                not app or app == node[0]
+            ):
+                leaves.add(node)
+        return sorted(leaves)
+
+    def ensure_not_cyclic(self) -> None:
+        # Algo from GvR:
+        # https://neopythonic.blogspot.com/2009/01/detecting-cycles-in-directed-graph.html
+        todo: set[tuple[str, str]] = set(self.nodes)
+        while todo:
+            node = todo.pop()
+            stack: list[tuple[str, str]] = [node]
+            while stack:
+                top = stack[-1]
+                for child in self.node_map[top].children:
+                    # Use child.key instead of child to speed up the frequent
+                    # hashing.
+                    node = child.key
+                    if node in stack:
+                        cycle = stack[stack.index(node) :]
+                        raise CircularDependencyError(
+                            ", ".join("{}.{}".format(*n) for n in cycle)
+                        )
+                    if node in todo:
+                        stack.append(node)
+                        todo.remove(node)
+                        break
+                else:
+                    node = stack.pop()
+
+    def __str__(self) -> str:
+        return "Graph: {} nodes, {} edges".format(*self._nodes_and_edges())
+
+    def __repr__(self) -> str:
+        nodes, edges = self._nodes_and_edges()
+        return f"<{self.__class__.__name__}: nodes={nodes}, edges={edges}>"
+
+    def _nodes_and_edges(self) -> tuple[int, int]:
+        return len(self.nodes), sum(
+            len(node.parents) for node in self.node_map.values()
+        )
+
+    def _generate_plan(
+        self, nodes: list[tuple[str, str]], at_end: bool
+    ) -> list[tuple[str, str]]:
+        plan: list[tuple[str, str]] = []
+        for node in nodes:
+            for migration in self.forwards_plan(node):
+                if migration not in plan and (at_end or migration not in nodes):
+                    plan.append(migration)
+        return plan
+
+    def make_state(
+        self,
+        nodes: tuple[str, str] | list[tuple[str, str]] | None = None,
+        at_end: bool = True,
+        real_packages: Any = None,
+    ) -> ProjectState:
+        """
+        Given a migration node or nodes, return a complete ProjectState for it.
+        If at_end is False, return the state before the migration has run.
+        If nodes is not provided, return the overall most current project state.
+        """
+        if nodes is None:
+            nodes = list(self.leaf_nodes())
+        if not nodes:
+            return ProjectState()
+        if not isinstance(nodes[0], tuple):
+            nodes = cast(list[tuple[str, str]], [nodes])
+        assert isinstance(nodes, list)  # Type narrowing after checks above
+        plan = self._generate_plan(nodes, at_end)
+        project_state = ProjectState(real_packages=real_packages)
+        for node in plan:
+            project_state = self.nodes[node].mutate_state(project_state, preserve=False)  # type: ignore[attr-defined]
+        return project_state
+
+    def __contains__(self, node: tuple[str, str]) -> bool:
+        return node in self.nodes