infrahub-server 1.5.0b0__py3-none-any.whl → 1.5.0b2__py3-none-any.whl

infrahub/cli/db.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import importlib
 import logging
 import os
 from collections import defaultdict
@@ -16,11 +15,15 @@ from infrahub_sdk.async_typer import AsyncTyper
 from prefect.testing.utilities import prefect_test_harness
 from rich import print as rprint
 from rich.console import Console
-from rich.logging import RichHandler
 from rich.table import Table
 
 from infrahub import config
+from infrahub.auth import AccountSession, AuthType
+from infrahub.context import InfrahubContext
 from infrahub.core import registry
+from infrahub.core.branch import Branch
+from infrahub.core.branch.tasks import rebase_branch
+from infrahub.core.constants import GLOBAL_BRANCH_NAME
 from infrahub.core.graph import GRAPH_VERSION
 from infrahub.core.graph.constraints import ConstraintManagerBase, ConstraintManagerMemgraph, ConstraintManagerNeo4j
 from infrahub.core.graph.index import node_indexes, rel_indexes
@@ -32,28 +35,24 @@ from infrahub.core.graph.schema import (
     GraphRelationshipIsPartOf,
     GraphRelationshipProperties,
 )
-from infrahub.core.initialization import (
-    first_time_initialization,
-    get_root_node,
-    initialization,
-    initialize_registry,
-)
+from infrahub.core.initialization import get_root_node, initialize_registry
+from infrahub.core.migrations.exceptions import MigrationFailureError
 from infrahub.core.migrations.graph import get_graph_migrations, get_migration_by_number
 from infrahub.core.migrations.schema.models import SchemaApplyMigrationData
 from infrahub.core.migrations.schema.tasks import schema_apply_migrations
 from infrahub.core.schema import SchemaRoot, core_models, internal_schema
 from infrahub.core.schema.definitions.deprecated import deprecated_models
 from infrahub.core.schema.manager import SchemaManager
-from infrahub.core.utils import delete_all_nodes
 from infrahub.core.validators.models.validate_migration import SchemaValidateMigrationData
 from infrahub.core.validators.tasks import schema_validate_migrations
 from infrahub.database import DatabaseType
 from infrahub.database.memgraph import IndexManagerMemgraph
 from infrahub.database.neo4j import IndexManagerNeo4j
-from infrahub.log import get_logger
+from infrahub.exceptions import ValidationError
 
 from .constants import ERROR_BADGE, FAILED_BADGE, SUCCESS_BADGE
 from .db_commands.check_inheritance import check_inheritance
+from .db_commands.clean_duplicate_schema_fields import clean_duplicate_schema_fields
 from .patch import patch_app
 
 
@@ -64,7 +63,7 @@ def get_timestamp_string() -> str:
 
 if TYPE_CHECKING:
     from infrahub.cli.context import CliContext
-    from infrahub.core.migrations.shared import ArbitraryMigration, GraphMigration, InternalSchemaMigration
+    from infrahub.core.migrations.shared import MigrationTypes
     from infrahub.database import InfrahubDatabase
     from infrahub.database.index import IndexManagerBase
 
@@ -93,75 +92,43 @@ def callback() -> None:
     """
 
 
-@app.command()
-async def init(
+@app.command(name="migrate")
+async def migrate_cmd(
     ctx: typer.Context,
-    config_file: str = typer.Option(
-        "infrahub.toml", envvar="INFRAHUB_CONFIG", help="Location of the configuration file to use for Infrahub"
-    ),
+    check: bool = typer.Option(False, help="Check the state of the database without applying the migrations."),
+    config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
+    migration_number: int | None = typer.Option(None, help="Apply a specific migration by number"),
 ) -> None:
-    """Erase the content of the database and initialize it with the core schema."""
-
-    log = get_logger()
-
-    # --------------------------------------------------
-    # CLEANUP
-    #  - For now we delete everything in the database
-    #   TODO, if possible try to implement this in an idempotent way
-    # --------------------------------------------------
-
+    """Check the current format of the internal graph and apply the necessary migrations"""
+    logging.getLogger("infrahub").setLevel(logging.WARNING)
     logging.getLogger("neo4j").setLevel(logging.ERROR)
-    config.load_and_exit(config_file_name=config_file)
-
-    context: CliContext = ctx.obj
-    dbdriver = await context.init_db(retry=1)
-    async with dbdriver.start_transaction() as db:
-        log.info("Delete All Nodes")
-        await delete_all_nodes(db=db)
-        await first_time_initialization(db=db)
-
-    await dbdriver.close()
-
-
-@app.command()
-async def load_test_data(
-    ctx: typer.Context,
-    config_file: str = typer.Option(
-        "infrahub.toml", envvar="INFRAHUB_CONFIG", help="Location of the configuration file to use for Infrahub"
-    ),
-    dataset: str = "dataset01",
-) -> None:
-    """Load test data into the database from the `test_data` directory."""
+    logging.getLogger("prefect").setLevel(logging.ERROR)
 
-    logging.getLogger("neo4j").setLevel(logging.ERROR)
     config.load_and_exit(config_file_name=config_file)
 
     context: CliContext = ctx.obj
     dbdriver = await context.init_db(retry=1)
 
-    async with dbdriver.start_session() as db:
-        await initialization(db=db)
-
-        log_level = "DEBUG"
+    root_node = await get_root_node(db=dbdriver)
+    migrations = await detect_migration_to_run(
+        current_graph_version=root_node.graph_version, migration_number=migration_number
+    )
 
-        FORMAT = "%(message)s"
-        logging.basicConfig(level=log_level, format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
-        logging.getLogger("infrahub")
+    if check or not migrations:
+        return
 
-        dataset_module = importlib.import_module(f"infrahub.test_data.{dataset}")
-        await dataset_module.load_data(db=db)
+    await migrate_database(db=dbdriver, migrations=migrations, initialize=True)
 
     await dbdriver.close()
 
 
-@app.command(name="migrate")
-async def migrate_cmd(
+@app.command(name="check-inheritance")
+async def check_inheritance_cmd(
     ctx: typer.Context,
-    check: bool = typer.Option(False, help="Check the state of the database without applying the migrations."),
+    fix: bool = typer.Option(False, help="Fix the inheritance of any invalid nodes."),
     config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
-    migration_number: int | None = typer.Option(None, help="Apply a specific migration by number"),
 ) -> None:
-    """Check the current format of the internal graph and apply the necessary migrations"""
+    """Check the database for any vertices with incorrect inheritance"""
     logging.getLogger("infrahub").setLevel(logging.WARNING)
     logging.getLogger("neo4j").setLevel(logging.ERROR)
     logging.getLogger("prefect").setLevel(logging.ERROR)
@@ -170,19 +137,22 @@ async def migrate_cmd(
 
     context: CliContext = ctx.obj
     dbdriver = await context.init_db(retry=1)
+    await initialize_registry(db=dbdriver)
 
-    await migrate_database(db=dbdriver, initialize=True, check=check, migration_number=migration_number)
+    success = await check_inheritance(db=dbdriver, fix=fix)
+    if not success:
+        raise typer.Exit(code=1)
 
     await dbdriver.close()
 
 
-@app.command(name="check-inheritance")
-async def check_inheritance_cmd(
+@app.command(name="check-duplicate-schema-fields")
+async def check_duplicate_schema_fields_cmd(
     ctx: typer.Context,
-    fix: bool = typer.Option(False, help="Fix the inheritance of any invalid nodes."),
+    fix: bool = typer.Option(False, help="Fix the duplicate schema fields on the default branch."),
     config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
 ) -> None:
-    """Check the database for any vertices with incorrect inheritance"""
+    """Check for any duplicate schema attributes or relationships on the default branch"""
     logging.getLogger("infrahub").setLevel(logging.WARNING)
     logging.getLogger("neo4j").setLevel(logging.ERROR)
     logging.getLogger("prefect").setLevel(logging.ERROR)
@@ -191,9 +161,8 @@ async def check_inheritance_cmd(
 
     context: CliContext = ctx.obj
     dbdriver = await context.init_db(retry=1)
-    await initialize_registry(db=dbdriver)
 
-    success = await check_inheritance(db=dbdriver, fix=fix)
+    success = await clean_duplicate_schema_fields(db=dbdriver, fix=fix)
     if not success:
         raise typer.Exit(code=1)
 
@@ -311,8 +280,38 @@ async def index(
     await dbdriver.close()
 
 
+async def detect_migration_to_run(
+    current_graph_version: int, migration_number: int | str | None = None
+) -> Sequence[MigrationTypes]:
+    """Return a sequence of migrations to apply to upgrade the database."""
+    rprint("Checking current state of the database")
+    migrations: list[MigrationTypes] = []
+
+    if migration_number:
+        migration = get_migration_by_number(migration_number)
+        migrations.append(migration)
+        if current_graph_version > migration.minimum_version:
+            rprint(
+                f"Migration {migration_number} already applied. To apply again, run the command without the --check flag."
+            )
+            return []
+        rprint(
+            f"Migration {migration_number} needs to be applied. Run `infrahub db migrate` to apply all outstanding migrations."
+        )
+    else:
+        migrations.extend(await get_graph_migrations(current_graph_version=current_graph_version))
+        if not migrations:
+            rprint(f"Database up-to-date (v{current_graph_version}), no migration to execute.")
+            return []
+
+    rprint(
+        f"Database needs to be updated (v{current_graph_version} -> v{GRAPH_VERSION}), {len(migrations)} migrations pending"
+    )
+    return migrations
+
+
 async def migrate_database(
-    db: InfrahubDatabase, initialize: bool = False, check: bool = False, migration_number: int | str | None = None
+    db: InfrahubDatabase, migrations: Sequence[MigrationTypes], initialize: bool = False
 ) -> bool:
     """Apply the latest migrations to the database, this function will print the status directly in the console.
 
@@ -320,40 +319,16 @@ async def migrate_database(
 
     Args:
         db: The database object.
-        check: If True, the function will only check the status of the database and not apply the migrations. Defaults to False.
-        migration_number: If provided, the function will only apply the migration with the given number. Defaults to None.
+        migrations: Sequence of migrations to apply.
+        initialize: Whether to initialize the registry before running migrations.
     """
-    rprint("Checking current state of the Database")
+    if not migrations:
+        return True
 
     if initialize:
         await initialize_registry(db=db)
 
     root_node = await get_root_node(db=db)
-    if migration_number:
-        migration = get_migration_by_number(migration_number)
-        migrations: Sequence[GraphMigration | InternalSchemaMigration | ArbitraryMigration] = [migration]
-        if check:
-            if root_node.graph_version > migration.minimum_version:
-                rprint(
-                    f"Migration {migration_number} already applied. To apply again, run the command without the --check flag."
-                )
-                return True
-            rprint(
-                f"Migration {migration_number} needs to be applied. Run `infrahub db migrate` to apply all outstanding migrations."
-            )
-            return False
-    else:
-        migrations = await get_graph_migrations(root=root_node)
-        if not migrations:
-            rprint(f"Database up-to-date (v{root_node.graph_version}), no migration to execute.")
-            return True
-
-        rprint(
-            f"Database needs to be updated (v{root_node.graph_version} -> v{GRAPH_VERSION}), {len(migrations)} migrations pending"
-        )
-
-    if check:
-        return True
 
     for migration in migrations:
         execution_result = await migration.execute(db=db)
@@ -378,6 +353,36 @@ async def migrate_database(
     return True
 
 
+async def trigger_rebase_branches(db: InfrahubDatabase) -> None:
+    """Trigger rebase of non-default branches, also triggering migrations in the process."""
+    branches = [b for b in await Branch.get_list(db=db) if b.name not in [registry.default_branch, GLOBAL_BRANCH_NAME]]
+    if not branches:
+        return
+
+    rprint(f"Planning rebase and migrations for {len(branches)} branches: {', '.join([b.name for b in branches])}")
+
+    for branch in branches:
+        if branch.graph_version == GRAPH_VERSION:
+            rprint(
+                f"Ignoring branch rebase and migrations for '{branch.name}' (ID: {branch.uuid}), it is already up-to-date"
+            )
+            continue
+
+        rprint(f"Rebasing branch '{branch.name}' (ID: {branch.uuid})...", end="")
+        try:
+            await registry.schema.load_schema(db=db, branch=branch)
+            await rebase_branch(
+                branch=branch.name,
+                context=InfrahubContext.init(
+                    branch=branch, account=AccountSession(auth_type=AuthType.NONE, authenticated=False, account_id="")
+                ),
+                send_events=False,
+            )
+            rprint(SUCCESS_BADGE)
+        except (ValidationError, MigrationFailureError):
+            rprint(FAILED_BADGE)
+
+
 async def initialize_internal_schema() -> None:
     registry.schema = SchemaManager()
     schema = SchemaRoot(**internal_schema)

infrahub/cli/db_commands/clean_duplicate_schema_fields.py

@@ -0,0 +1,212 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+from rich import print as rprint
+from rich.console import Console
+from rich.table import Table
+
+from infrahub.cli.constants import FAILED_BADGE, SUCCESS_BADGE
+from infrahub.core.query import Query, QueryType
+from infrahub.database import InfrahubDatabase
+
+
+class SchemaFieldType(str, Enum):
+    ATTRIBUTE = "attribute"
+    RELATIONSHIP = "relationship"
+
+
+@dataclass
+class SchemaFieldDetails:
+    schema_kind: str
+    schema_uuid: str
+    field_type: SchemaFieldType
+    field_name: str
+
+
+class DuplicateSchemaFields(Query):
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        query = """
+MATCH (root:Root)
+LIMIT 1
+WITH root.default_branch AS default_branch
+MATCH (field:SchemaAttribute|SchemaRelationship)
+CALL (default_branch, field) {
+    MATCH (field)-[is_part_of:IS_PART_OF]->(:Root)
+    WHERE is_part_of.branch = default_branch
+    ORDER BY is_part_of.from DESC
+    RETURN is_part_of
+    LIMIT 1
+}
+WITH default_branch, field, CASE
+    WHEN is_part_of.status = "active" AND is_part_of.to IS NULL THEN is_part_of.from
+    ELSE NULL
+END AS active_from
+WHERE active_from IS NOT NULL
+WITH default_branch, field, active_from, "SchemaAttribute" IN labels(field) AS is_attribute
+CALL (field, default_branch) {
+    MATCH (field)-[r1:HAS_ATTRIBUTE]->(:Attribute {name: "name"})-[r2:HAS_VALUE]->(name_value:AttributeValue)
+    WHERE r1.branch = default_branch AND r2.branch = default_branch
+    AND r1.status = "active" AND r2.status = "active"
+    AND r1.to IS NULL AND r2.to IS NULL
+    ORDER BY r1.from DESC, r1.status ASC, r2.from DESC, r2.status ASC
+    LIMIT 1
+    RETURN name_value.value AS field_name
+}
+CALL (field, default_branch) {
+    MATCH (field)-[r1:IS_RELATED]-(rel:Relationship)-[r2:IS_RELATED]-(peer:SchemaNode|SchemaGeneric)
+    WHERE rel.name IN ["schema__node__relationships", "schema__node__attributes"]
+    AND r1.branch = default_branch AND r2.branch = default_branch
+    AND r1.status = "active" AND r2.status = "active"
+    AND r1.to IS NULL AND r2.to IS NULL
+    ORDER BY r1.from DESC, r1.status ASC, r2.from DESC, r2.status ASC
+    LIMIT 1
+    RETURN peer AS schema_vertex
+}
+WITH default_branch, field, field_name, is_attribute, active_from, schema_vertex
+ORDER BY active_from DESC
+WITH default_branch, field_name, is_attribute, schema_vertex, collect(field) AS fields_reverse_chron
+WHERE size(fields_reverse_chron) > 1
+        """
+        self.add_to_query(query)
+
+
+class GetDuplicateSchemaFields(DuplicateSchemaFields):
+    """
+    Get the kind, field type, and field name for any duplicated attributes or relationships on a given schema
+    on the default branch
+    """
+
+    name = "get_duplicate_schema_fields"
+    type = QueryType.READ
+    insert_return = False
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+        await super().query_init(db=db, **kwargs)
+        query = """
+CALL (schema_vertex, default_branch) {
+    MATCH (schema_vertex)-[r1:HAS_ATTRIBUTE]->(:Attribute {name: "namespace"})-[r2:HAS_VALUE]->(name_value:AttributeValue)
+    WHERE r1.branch = default_branch AND r2.branch = default_branch
+    ORDER BY r1.from DESC, r1.status ASC, r2.from DESC, r2.status ASC
+    LIMIT 1
+    RETURN name_value.value AS schema_namespace
+}
+CALL (schema_vertex, default_branch) {
+    MATCH (schema_vertex)-[r1:HAS_ATTRIBUTE]->(:Attribute {name: "name"})-[r2:HAS_VALUE]->(name_value:AttributeValue)
+    WHERE r1.branch = default_branch AND r2.branch = default_branch
+    ORDER BY r1.from DESC, r1.status ASC, r2.from DESC, r2.status ASC
+    LIMIT 1
+    RETURN name_value.value AS schema_name
+}
+RETURN schema_namespace + schema_name AS schema_kind, schema_vertex.uuid AS schema_uuid, field_name, is_attribute
+ORDER BY schema_kind ASC, is_attribute DESC, field_name ASC
+        """
+        self.return_labels = ["schema_kind", "schema_uuid", "field_name", "is_attribute"]
+        self.add_to_query(query)
+
+    def get_schema_field_details(self) -> list[SchemaFieldDetails]:
+        schema_field_details: list[SchemaFieldDetails] = []
+        for result in self.results:
+            schema_kind = result.get_as_type(label="schema_kind", return_type=str)
+            schema_uuid = result.get_as_type(label="schema_uuid", return_type=str)
+            field_name = result.get_as_type(label="field_name", return_type=str)
+            is_attribute = result.get_as_type(label="is_attribute", return_type=bool)
+            schema_field_details.append(
+                SchemaFieldDetails(
+                    schema_kind=schema_kind,
+                    schema_uuid=schema_uuid,
+                    field_name=field_name,
+                    field_type=SchemaFieldType.ATTRIBUTE if is_attribute else SchemaFieldType.RELATIONSHIP,
+                )
+            )
+        return schema_field_details
+
+
+class FixDuplicateSchemaFields(DuplicateSchemaFields):
+    """
+    Fix the duplicate schema fields by hard deleting the earlier duplicate(s)
+    """
+
+    name = "fix_duplicate_schema_fields"
+    type = QueryType.WRITE
+    insert_return = False
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+        await super().query_init(db=db, **kwargs)
+        query = """
+WITH default_branch, tail(fields_reverse_chron) AS fields_to_delete
+UNWIND fields_to_delete AS field_to_delete
+CALL (field_to_delete, default_branch) {
+    MATCH (field_to_delete)-[r:IS_PART_OF {branch: default_branch}]-()
+    DELETE r
+    WITH field_to_delete
+    MATCH (field_to_delete)-[:IS_RELATED {branch: default_branch}]-(rel:Relationship)
+    WITH DISTINCT field_to_delete, rel
+    MATCH (rel)-[r {branch: default_branch}]-()
+    DELETE r
+    WITH field_to_delete, rel
+    OPTIONAL MATCH (rel)
+    WHERE NOT exists((rel)--())
+    DELETE rel
+    WITH DISTINCT field_to_delete
+    MATCH (field_to_delete)-[:HAS_ATTRIBUTE {branch: default_branch}]->(attr:Attribute)
+    MATCH (attr)-[r {branch: default_branch}]-()
+    DELETE r
+    WITH field_to_delete, attr
+    OPTIONAL MATCH (attr)
+    WHERE NOT exists((attr)--())
+    DELETE attr
+    WITH DISTINCT field_to_delete
+    OPTIONAL MATCH (field_to_delete)
+    WHERE NOT exists((field_to_delete)--())
+    DELETE field_to_delete
+}
+        """
+        self.add_to_query(query)
+
+
+def display_duplicate_schema_fields(duplicate_schema_fields: list[SchemaFieldDetails]) -> None:
+    console = Console()
+
+    table = Table(title="Duplicate Schema Fields on Default Branch")
+
+    table.add_column("Schema Kind")
+    table.add_column("Schema UUID")
+    table.add_column("Field Name")
+    table.add_column("Field Type")
+
+    for duplicate_schema_field in duplicate_schema_fields:
+        table.add_row(
+            duplicate_schema_field.schema_kind,
+            duplicate_schema_field.schema_uuid,
+            duplicate_schema_field.field_name,
+            duplicate_schema_field.field_type.value,
+        )
+
+    console.print(table)
+
+
+async def clean_duplicate_schema_fields(db: InfrahubDatabase, fix: bool = False) -> bool:
+    """
+    Identify any attributes or relationships that are duplicated in a schema on the default branch
+    If fix is True, runs cypher queries to hard delete the earlier duplicate
+    """
+
+    duplicate_schema_fields_query = await GetDuplicateSchemaFields.init(db=db)
+    await duplicate_schema_fields_query.execute(db=db)
+    duplicate_schema_fields = duplicate_schema_fields_query.get_schema_field_details()
+
+    if not duplicate_schema_fields:
+        rprint(f"{SUCCESS_BADGE} No duplicate schema fields found")
+        return True
+
+    display_duplicate_schema_fields(duplicate_schema_fields)
+
+    if not fix:
+        rprint(f"{FAILED_BADGE} Use the --fix flag to fix the duplicate schema fields")
+        return False
+
+    fix_duplicate_schema_fields_query = await FixDuplicateSchemaFields.init(db=db)
+    await fix_duplicate_schema_fields_query.execute(db=db)
+    rprint(f"{SUCCESS_BADGE} Duplicate schema fields deleted from the default branch")
+    return True

infrahub/cli/dev.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import importlib
+import logging
+from pathlib import Path  # noqa: TC003
+from typing import TYPE_CHECKING
+
+import typer
+from graphql import parse, print_ast, print_schema
+from infrahub_sdk.async_typer import AsyncTyper
+from rich.logging import RichHandler
+
+from infrahub import config
+from infrahub.core.initialization import (
+    first_time_initialization,
+    initialization,
+)
+from infrahub.core.schema import SchemaRoot, core_models, internal_schema
+from infrahub.core.schema.schema_branch import SchemaBranch
+from infrahub.core.utils import delete_all_nodes
+from infrahub.graphql.manager import GraphQLSchemaManager
+from infrahub.graphql.schema_sort import sort_schema_ast
+from infrahub.log import get_logger
+
+if TYPE_CHECKING:
+    from infrahub.cli.context import CliContext
+
+app = AsyncTyper()
+
+
+@app.command(name="export-graphql-schema")
+async def export_graphql_schema(
+    ctx: typer.Context,  # noqa: ARG001
+    config_file: str = typer.Option("infrahub.toml", envvar="INFRAHUB_CONFIG"),
+    out: Path = typer.Option("schema.graphql"),  # noqa: B008
+) -> None:
+    """Export the Core GraphQL schema to a file."""
+
+    config.load_and_exit(config_file_name=config_file)
+
+    schema = SchemaRoot(**internal_schema)
+    full_schema = schema.merge(schema=SchemaRoot(**core_models))
+
+    schema_branch = SchemaBranch(cache={}, name="default")
+    schema_branch.load_schema(schema=full_schema)
+
+    schema_branch.process()
+
+    gqlm = GraphQLSchemaManager(schema=schema_branch)
+    gql_schema = gqlm.generate()
+
+    schema_str = print_schema(gql_schema)
+    schema_ast = parse(schema_str)
+    sorted_schema_ast = sort_schema_ast(schema_ast)
+    sorted_schema_str = print_ast(sorted_schema_ast)
+
+    out.write_text(sorted_schema_str)
+
+
+@app.command(name="db-init")
+async def database_init(
+    ctx: typer.Context,
+    config_file: str = typer.Option(
+        "infrahub.toml", envvar="INFRAHUB_CONFIG", help="Location of the configuration file to use for Infrahub"
+    ),
+) -> None:
+    """Erase the content of the database and initialize it with the core schema."""
+
+    log = get_logger()
+
+    # --------------------------------------------------
+    # CLEANUP
+    #  - For now we delete everything in the database
+    #   TODO, if possible try to implement this in an idempotent way
+    # --------------------------------------------------
+
+    logging.getLogger("neo4j").setLevel(logging.ERROR)
+    config.load_and_exit(config_file_name=config_file)
+
+    context: CliContext = ctx.obj
+    dbdriver = await context.init_db(retry=1)
+    async with dbdriver.start_transaction() as db:
+        log.info("Delete All Nodes")
+        await delete_all_nodes(db=db)
+        await first_time_initialization(db=db)
+
+    await dbdriver.close()
+
+
+@app.command(name="load-test-data")
+async def load_test_data(
+    ctx: typer.Context,
+    config_file: str = typer.Option(
+        "infrahub.toml", envvar="INFRAHUB_CONFIG", help="Location of the configuration file to use for Infrahub"
+    ),
+    dataset: str = "dataset01",
+) -> None:
+    """Load test data into the database from the `test_data` directory."""
+
+    logging.getLogger("neo4j").setLevel(logging.ERROR)
+    config.load_and_exit(config_file_name=config_file)
+
+    context: CliContext = ctx.obj
+    dbdriver = await context.init_db(retry=1)
+
+    async with dbdriver.start_session() as db:
+        await initialization(db=db)
+
+        log_level = "DEBUG"
+
+        FORMAT = "%(message)s"
+        logging.basicConfig(level=log_level, format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
+        logging.getLogger("infrahub")
+
+        dataset_module = importlib.import_module(f"infrahub.test_data.{dataset}")
+        await dataset_module.load_data(db=db)
+
+    await dbdriver.close()