patching command framework infrahub db patch plan/apply/revert (#6311)

* WIP patch for removing dup kind-migrated nodes following a merge

* WIP patch framework

* patch commands and framework

* add restore method

* update tests for reverting a patch

* generate cli docs

* update docs

* remove actual patch from this branch

* consolidate file names for patch plans

* make elementId function use dynamic

* one more docstring

* refactor to track deleted IDs in a file and add failure tests

* delete edges first during revert

* move call order around

* small changes for PR feedback

* do vertex/edge adds in transactions

* generate docs again

* fix for reverting deleted edge linked to deleted vertex

Author: ajtmccarty

backend/infrahub/cli/db.py

@@ -54,12 +54,14 @@
 from infrahub.services.adapters.workflow.local import WorkflowLocalExecution
 
 from .constants import ERROR_BADGE, FAILED_BADGE, SUCCESS_BADGE
+from .patch import patch_app
 
 if TYPE_CHECKING:
     from infrahub.cli.context import CliContext
     from infrahub.database import InfrahubDatabase
 
 app = AsyncTyper()
+app.add_typer(patch_app, name="patch")
 
 PERMISSIONS_AVAILABLE = ["read", "write", "admin"]
 

backend/infrahub/cli/patch.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import importlib
+import inspect
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import typer
+from infrahub_sdk.async_typer import AsyncTyper
+from rich import print as rprint
+
+from infrahub import config
+from infrahub.patch.edge_adder import PatchPlanEdgeAdder
+from infrahub.patch.edge_deleter import PatchPlanEdgeDeleter
+from infrahub.patch.edge_updater import PatchPlanEdgeUpdater
+from infrahub.patch.plan_reader import PatchPlanReader
+from infrahub.patch.plan_writer import PatchPlanWriter
+from infrahub.patch.queries.base import PatchQuery
+from infrahub.patch.runner import (
+    PatchPlanEdgeDbIdTranslator,
+    PatchRunner,
+)
+from infrahub.patch.vertex_adder import PatchPlanVertexAdder
+from infrahub.patch.vertex_deleter import PatchPlanVertexDeleter
+from infrahub.patch.vertex_updater import PatchPlanVertexUpdater
+
+from .constants import ERROR_BADGE, SUCCESS_BADGE
+
+if TYPE_CHECKING:
+    from infrahub.cli.context import CliContext
+    from infrahub.database import InfrahubDatabase
+
+
+patch_app = AsyncTyper(help="Commands for planning, applying, and reverting database patches")
+
+
+def get_patch_runner(db: InfrahubDatabase) -> PatchRunner:
+    return PatchRunner(
+        plan_writer=PatchPlanWriter(),
+        plan_reader=PatchPlanReader(),
+        edge_db_id_translator=PatchPlanEdgeDbIdTranslator(),
+        vertex_adder=PatchPlanVertexAdder(db=db),
+        vertex_deleter=PatchPlanVertexDeleter(db=db),
+        vertex_updater=PatchPlanVertexUpdater(db=db),
+        edge_adder=PatchPlanEdgeAdder(db=db),
+        edge_deleter=PatchPlanEdgeDeleter(db=db),
+        edge_updater=PatchPlanEdgeUpdater(db=db),
+    )
+
+
+@patch_app.command(name="plan")
+async def plan_patch_cmd(
+    ctx: typer.Context,
+    patch_path: str = typer.Argument(
+        help="Path to the file containing the PatchQuery instance to run. Use Python-style dot paths, such as infrahub.cli.patch.queries.base"
+    ),
+    patch_plans_dir: Path = typer.Option(Path("infrahub-patches"), help="Path to patch plans directory"),  # noqa: B008
+    apply: bool = typer.Option(False, help="Apply the patch immediately after creating it"),
+    config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
+) -> None:
+    """Create a plan for a given patch and save it in the patch plans directory to be applied/reverted"""
+    logging.getLogger("infrahub").setLevel(logging.WARNING)
+    logging.getLogger("neo4j").setLevel(logging.ERROR)
+    logging.getLogger("prefect").setLevel(logging.ERROR)
+
+    patch_module = importlib.import_module(patch_path)
+    patch_query_class = None
+    patch_query_class_count = 0
+    for _, cls in inspect.getmembers(patch_module, inspect.isclass):
+        if issubclass(cls, PatchQuery) and cls is not PatchQuery:
+            patch_query_class = cls
+            patch_query_class_count += 1
+
+    patch_query_path = f"{PatchQuery.__module__}.{PatchQuery.__name__}"
+    if patch_query_class is None:
+        rprint(f"{ERROR_BADGE} No subclass of {patch_query_path} found in {patch_path}")
+        raise typer.Exit(1)
+    if patch_query_class_count > 1:
+        rprint(
+            f"{ERROR_BADGE} Multiple subclasses of {patch_query_path} found in {patch_path}. Please only define one per file."
+        )
+        raise typer.Exit(1)
+
+    config.load_and_exit(config_file_name=config_file)
+
+    context: CliContext = ctx.obj
+    dbdriver = await context.init_db(retry=1)
+
+    patch_query_instance = patch_query_class(db=dbdriver)
+    async with dbdriver.start_session() as db:
+        patch_runner = get_patch_runner(db=db)
+        patch_plan_dir = await patch_runner.prepare_plan(patch_query_instance, directory=Path(patch_plans_dir))
+        rprint(f"{SUCCESS_BADGE} Patch plan created at {patch_plan_dir}")
+        if apply:
+            await patch_runner.apply(patch_plan_directory=patch_plan_dir)
+            rprint(f"{SUCCESS_BADGE} Patch plan successfully applied")
+
+    await dbdriver.close()
+
+
+@patch_app.command(name="apply")
+async def apply_patch_cmd(
+    ctx: typer.Context,
+    patch_plan_dir: Path = typer.Argument(help="Path to the directory containing a patch plan"),
+    config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
+) -> None:
+    """Apply a given patch plan"""
+    logging.getLogger("infrahub").setLevel(logging.WARNING)
+    logging.getLogger("neo4j").setLevel(logging.ERROR)
+    logging.getLogger("prefect").setLevel(logging.ERROR)
+
+    config.load_and_exit(config_file_name=config_file)
+
+    context: CliContext = ctx.obj
+    dbdriver = await context.init_db(retry=1)
+
+    if not patch_plan_dir.exists() or not patch_plan_dir.is_dir():
+        rprint(f"{ERROR_BADGE} patch_plan_dir must be an existing directory")
+        raise typer.Exit(1)
+
+    async with dbdriver.start_session() as db:
+        patch_runner = get_patch_runner(db=db)
+        await patch_runner.apply(patch_plan_directory=patch_plan_dir)
+        rprint(f"{SUCCESS_BADGE} Patch plan successfully applied")
+
+    await dbdriver.close()
+
+
+@patch_app.command(name="revert")
+async def revert_patch_cmd(
+    ctx: typer.Context,
+    patch_plan_dir: Path = typer.Argument(help="Path to the directory containing a patch plan"),
+    config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
+) -> None:
+    """Revert a given patch plan"""
+    logging.getLogger("infrahub").setLevel(logging.WARNING)
+    logging.getLogger("neo4j").setLevel(logging.ERROR)
+    logging.getLogger("prefect").setLevel(logging.ERROR)
+    config.load_and_exit(config_file_name=config_file)
+
+    context: CliContext = ctx.obj
+    db = await context.init_db(retry=1)
+
+    if not patch_plan_dir.exists() or not patch_plan_dir.is_dir():
+        rprint(f"{ERROR_BADGE} patch_plan_dir must be an existing directory")
+        raise typer.Exit(1)
+
+    patch_runner = get_patch_runner(db=db)
+    await patch_runner.revert(patch_plan_directory=patch_plan_dir)
+    rprint(f"{SUCCESS_BADGE} Patch plan successfully reverted")
+
+    await db.close()

backend/infrahub/patch/__init__.py

@@ -0,0 +1,13 @@
+from enum import Enum
+
+
+class PatchPlanFilename(str, Enum):
+    VERTICES_TO_ADD = "vertices_to_add.json"
+    VERTICES_TO_UPDATE = "vertices_to_update.json"
+    VERTICES_TO_DELETE = "vertices_to_delete.json"
+    EDGES_TO_ADD = "edges_to_add.json"
+    EDGES_TO_UPDATE = "edges_to_update.json"
+    EDGES_TO_DELETE = "edges_to_delete.json"
+    ADDED_DB_IDS = "added_db_ids.json"
+    DELETED_DB_IDS = "deleted_db_ids.json"
+    REVERTED_DELETED_DB_IDS = "reverted_deleted_db_ids.json"

backend/infrahub/patch/constants.py

@@ -0,0 +1,64 @@
+from collections import defaultdict
+from dataclasses import asdict
+from typing import AsyncGenerator
+
+from infrahub.core.query import QueryType
+from infrahub.database import InfrahubDatabase
+
+from .models import EdgeToAdd
+
+
+class PatchPlanEdgeAdder:
+    def __init__(self, db: InfrahubDatabase, batch_size_limit: int = 1000) -> None:
+        self.db = db
+        self.batch_size_limit = batch_size_limit
+
+    async def _run_add_query(self, edge_type: str, edges_to_add: list[EdgeToAdd]) -> dict[str, str]:
+        query = """
+UNWIND $edges_to_add AS edge_to_add
+MATCH (a) WHERE %(id_func_name)s(a) = edge_to_add.from_id
+MATCH (b) WHERE %(id_func_name)s(b) = edge_to_add.to_id
+CREATE (a)-[e:%(edge_type)s]->(b)
+SET e = edge_to_add.after_props
+RETURN edge_to_add.identifier AS abstract_id, %(id_func_name)s(e) AS db_id
+        """ % {
+            "edge_type": edge_type,
+            "id_func_name": self.db.get_id_function_name(),
+        }
+        edges_to_add_dicts = [asdict(v) for v in edges_to_add]
+        # use transaction to make sure we record the results before committing them
+        try:
+            txn_db = self.db.start_transaction()
+            async with txn_db as txn:
+                results = await txn.execute_query(
+                    query=query, params={"edges_to_add": edges_to_add_dicts}, type=QueryType.WRITE
+                )
+            abstract_to_concrete_id_map: dict[str, str] = {}
+            for result in results:
+                abstract_id = result.get("abstract_id")
+                concrete_id = result.get("db_id")
+                abstract_to_concrete_id_map[abstract_id] = concrete_id
+        finally:
+            await txn_db.close()
+        return abstract_to_concrete_id_map
+
+    async def execute(
+        self,
+        edges_to_add: list[EdgeToAdd],
+    ) -> AsyncGenerator[dict[str, str], None]:
+        """
+        Create edges_to_add on the database.
+        Returns a generator that yields dictionaries mapping EdgeToAdd.identifier to the database-level ID of the newly created edge.
+        """
+        edges_map_queue: dict[str, list[EdgeToAdd]] = defaultdict(list)
+        for edge_to_add in edges_to_add:
+            edges_map_queue[edge_to_add.edge_type].append(edge_to_add)
+            if len(edges_map_queue[edge_to_add.edge_type]) > self.batch_size_limit:
+                yield await self._run_add_query(
+                    edge_type=edge_to_add.edge_type,
+                    edges_to_add=edges_map_queue[edge_to_add.edge_type],
+                )
+                edges_map_queue[edge_to_add.edge_type] = []
+
+        for edge_type, edges_group in edges_map_queue.items():
+            yield await self._run_add_query(edge_type=edge_type, edges_to_add=edges_group)

backend/infrahub/patch/edge_adder.py

@@ -0,0 +1,33 @@
+from typing import AsyncGenerator
+
+from infrahub.core.query import QueryType
+from infrahub.database import InfrahubDatabase
+
+from .models import EdgeToDelete
+
+
+class PatchPlanEdgeDeleter:
+    def __init__(self, db: InfrahubDatabase, batch_size_limit: int = 1000) -> None:
+        self.db = db
+        self.batch_size_limit = batch_size_limit
+
+    async def _run_delete_query(self, ids_to_delete: list[str]) -> set[str]:
+        query = """
+MATCH ()-[e]-()
+WHERE %(id_func_name)s(e) IN $ids_to_delete
+DELETE e
+RETURN %(id_func_name)s(e) AS deleted_id
+        """ % {"id_func_name": self.db.get_id_function_name()}
+        results = await self.db.execute_query(
+            query=query, params={"ids_to_delete": ids_to_delete}, type=QueryType.WRITE
+        )
+        deleted_ids: set[str] = set()
+        for result in results:
+            deleted_id = result.get("deleted_id")
+            deleted_ids.add(deleted_id)
+        return deleted_ids
+
+    async def execute(self, edges_to_delete: list[EdgeToDelete]) -> AsyncGenerator[set[str], None]:
+        for i in range(0, len(edges_to_delete), self.batch_size_limit):
+            ids_to_delete = [e.db_id for e in edges_to_delete[i : i + self.batch_size_limit]]
+            yield await self._run_delete_query(ids_to_delete=ids_to_delete)

backend/infrahub/patch/edge_deleter.py

@@ -0,0 +1,28 @@
+from dataclasses import asdict
+
+from infrahub.core.query import QueryType
+from infrahub.database import InfrahubDatabase
+
+from .models import EdgeToUpdate
+
+
+class PatchPlanEdgeUpdater:
+    def __init__(self, db: InfrahubDatabase, batch_size_limit: int = 1000) -> None:
+        self.db = db
+        self.batch_size_limit = batch_size_limit
+
+    async def _run_update_query(self, edges_to_update: list[EdgeToUpdate]) -> None:
+        query = """
+UNWIND $edges_to_update AS edge_to_update
+MATCH ()-[e]-()
+WHERE %(id_func_name)s(e) = edge_to_update.db_id
+SET e = edge_to_update.after_props
+        """ % {"id_func_name": self.db.get_id_function_name()}
+        await self.db.execute_query(
+            query=query, params={"edges_to_update": [asdict(e) for e in edges_to_update]}, type=QueryType.WRITE
+        )
+
+    async def execute(self, edges_to_update: list[EdgeToUpdate]) -> None:
+        for i in range(0, len(edges_to_update), self.batch_size_limit):
+            vertices_slice = edges_to_update[i : i + self.batch_size_limit]
+            await self._run_update_query(edges_to_update=vertices_slice)