infrahub-server 1.4.12__py3-none-any.whl → 1.5.0__py3-none-any.whl

infrahub_testcontainers/container.py

@@ -152,6 +152,8 @@ class InfrahubDockerCompose(DockerCompose):
                     "INFRAHUB_TESTING_TASKMGR_BACKGROUND_SVC_REPLICAS": "1",
                     "PREFECT_MESSAGING_BROKER": "prefect_redis.messaging",
                     "PREFECT_MESSAGING_CACHE": "prefect_redis.messaging",
+                    "PREFECT_SERVER_EVENTS_CAUSAL_ORDERING": "prefect_redis.ordering",
+                    "PREFECT_SERVER_CONCURRENCY_LEASE_STORAGE": "prefect_redis.lease_storage",
                     "PREFECT__SERVER_WEBSERVER_ONLY": "true",
                     "PREFECT_API_DATABASE_MIGRATE_ON_START": "false",
                     "PREFECT_API_BLOCKS_REGISTER_ON_START": "false",
@@ -237,7 +239,9 @@ class InfrahubDockerCompose(DockerCompose):
             for service_name, service_data in INFRAHUB_SERVICES.items()
         }
 
-    def database_create_backup(self, backup_name: str = "neo4j_database.backup", dest_dir: Path | None = None) -> None:
+    def database_create_backup(
+        self, backup_name: str = "neo4j_database.backup", dest_dir: Path | None = None, compress: bool = False
+    ) -> None:
         assert self.use_neo4j_enterprise
 
         self.exec_in_container(
@@ -245,7 +249,7 @@ class InfrahubDockerCompose(DockerCompose):
                 "neo4j-admin",
                 "database",
                 "backup",
-                "--compress=false",
+                f"--compress={'true' if compress else 'false'}",
                 "--to-path",
                 str(self.internal_backup_dir),
             ],
@@ -513,3 +517,111 @@ class InfrahubDockerCompose(DockerCompose):
                 )
             self.start()
             print("Database restored successfully")
+
+    def task_manager_create_backup(self, backup_name: str = "prefect.dump", dest_dir: Path | None = None) -> Path:
+        """Create a backup of the task manager PostgreSQL database using ``pg_dump``.
+
+        Args:
+            backup_name: Name of the archive file to create. Defaults to ``prefect.dump``.
+            dest_dir: Optional host directory where the backup should be copied after it is
+                produced. When omitted, the backup remains in ``external_backup_dir``.
+
+        Returns:
+            Path to the backup archive on the host filesystem.
+
+        Raises:
+            FileNotFoundError: If the pg_dump command completes but no archive is produced.
+        """
+
+        service_name = "task-manager-db"
+
+        try:
+            self.get_container(service_name=service_name)
+        except ContainerIsNotRunning:
+            self.start_container(service_name=service_name)
+
+        self.external_backup_dir.mkdir(parents=True, exist_ok=True)
+
+        internal_backup_path = self.internal_backup_dir / backup_name
+        dump_command = [
+            "pg_dump",
+            "--format=custom",
+            "--blobs",
+            "--no-owner",
+            "--no-privileges",
+            "--dbname=postgresql://postgres:postgres@localhost:5432/prefect",
+            f"--file={internal_backup_path}",
+        ]
+        self.exec_in_container(command=dump_command, service_name=service_name)
+
+        source_path = self.external_backup_dir / backup_name
+        if not source_path.exists():
+            raise FileNotFoundError(f"Backup file {source_path} was not created")
+
+        final_path = source_path
+        if dest_dir:
+            dest_dir.mkdir(parents=True, exist_ok=True)
+            if dest_dir.resolve() != self.external_backup_dir.resolve():
+                final_path = dest_dir / backup_name
+                shutil.copy(source_path, final_path)
+
+        return final_path
+
+    def task_manager_restore_backup(self, backup_file: Path) -> None:
+        """Restore the task manager PostgreSQL database from a ``pg_restore`` archive.
+
+        Args:
+            backup_file: Path to the backup archive on the host filesystem.
+
+        Raises:
+            FileNotFoundError: If the provided backup archive does not exist.
+        """
+
+        if not backup_file.exists():
+            raise FileNotFoundError(f"Backup file {backup_file} does not exist")
+
+        service_name = "task-manager-db"
+
+        try:
+            self.get_container(service_name=service_name)
+        except ContainerIsNotRunning:
+            self.start_container(service_name=service_name)
+
+        self.external_backup_dir.mkdir(parents=True, exist_ok=True)
+        target_path = self.external_backup_dir / backup_file.name
+        shutil.copy(backup_file, target_path)
+
+        admin_dsn = "postgresql://postgres:postgres@localhost:5432/postgres"
+        prefect_dsn = "postgresql://postgres:postgres@localhost:5432/prefect"
+        internal_backup_path = self.internal_backup_dir / backup_file.name
+
+        terminate_sessions_command = [
+            "psql",
+            f"--dbname={admin_dsn}",
+            "--command",
+            "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = 'prefect';",
+        ]
+        drop_database_command = [
+            "psql",
+            f"--dbname={admin_dsn}",
+            "--command",
+            "DROP DATABASE IF EXISTS prefect WITH (FORCE);",
+        ]
+        create_database_command = [
+            "psql",
+            f"--dbname={admin_dsn}",
+            "--command",
+            "CREATE DATABASE prefect OWNER postgres;",
+        ]
+        restore_command = [
+            "pg_restore",
+            "--no-owner",
+            "--role=postgres",
+            f"--dbname={prefect_dsn}",
+            str(internal_backup_path),
+        ]
+
+        self.exec_in_container(command=terminate_sessions_command, service_name=service_name)
+        self.exec_in_container(command=drop_database_command, service_name=service_name)
+        self.exec_in_container(command=create_database_command, service_name=service_name)
+        self.exec_in_container(command=restore_command, service_name=service_name)

infrahub_testcontainers/docker-compose-cluster.test.yml

@@ -184,6 +184,8 @@ services:
 
       PREFECT_MESSAGING_BROKER:
       PREFECT_MESSAGING_CACHE:
+      PREFECT_SERVER_EVENTS_CAUSAL_ORDERING:
+      PREFECT_SERVER_CONCURRENCY_LEASE_STORAGE:
       PREFECT__SERVER_WEBSERVER_ONLY:
       PREFECT_API_DATABASE_MIGRATE_ON_START:
       PREFECT_API_BLOCKS_REGISTER_ON_START:
@@ -225,6 +227,8 @@ services:
       INFRAHUB_CACHE_ADDRESS: ${INFRAHUB_TESTING_CACHE_ADDRESS}
       PREFECT_MESSAGING_BROKER: prefect_redis.messaging
       PREFECT_MESSAGING_CACHE: prefect_redis.messaging
+      PREFECT_SERVER_EVENTS_CAUSAL_ORDERING: prefect_redis.ordering
+      PREFECT_SERVER_CONCURRENCY_LEASE_STORAGE: prefect_redis.lease_storage
       PREFECT_REDIS_MESSAGING_HOST: "${INFRAHUB_TESTING_CACHE_ADDRESS:-cache}"
       PREFECT_REDIS_MESSAGING_DB: "1"
       PREFECT_REDIS_MESSAGING_CONSUMER_MIN_IDLE_TIME: "30"
@@ -244,6 +248,7 @@ services:
       - POSTGRES_DB=prefect
     volumes:
       - workflow_db:/var/lib/postgresql/data
+      - "./${INFRAHUB_TESTING_LOCAL_DB_BACKUP_DIRECTORY}:${INFRAHUB_TESTING_INTERNAL_DB_BACKUP_DIRECTORY}"
     healthcheck:
       test: ["CMD-SHELL", "pg_isready"]
       interval: 10s

infrahub_testcontainers/docker-compose.test.yml

@@ -95,6 +95,8 @@ services:
 
       PREFECT_MESSAGING_BROKER:
       PREFECT_MESSAGING_CACHE:
+      PREFECT_SERVER_EVENTS_CAUSAL_ORDERING:
+      PREFECT_SERVER_CONCURRENCY_LEASE_STORAGE:
       PREFECT__SERVER_WEBSERVER_ONLY:
       PREFECT_API_DATABASE_MIGRATE_ON_START:
       PREFECT_API_BLOCKS_REGISTER_ON_START:
@@ -136,6 +138,8 @@ services:
       INFRAHUB_CACHE_ADDRESS: ${INFRAHUB_TESTING_CACHE_ADDRESS}
       PREFECT_MESSAGING_BROKER: prefect_redis.messaging
       PREFECT_MESSAGING_CACHE: prefect_redis.messaging
+      PREFECT_SERVER_EVENTS_CAUSAL_ORDERING: prefect_redis.ordering
+      PREFECT_SERVER_CONCURRENCY_LEASE_STORAGE: prefect_redis.lease_storage
       PREFECT_REDIS_MESSAGING_HOST: "${INFRAHUB_TESTING_CACHE_ADDRESS:-cache}"
       PREFECT_REDIS_MESSAGING_DB: "1"
       PREFECT_REDIS_MESSAGING_CONSUMER_MIN_IDLE_TIME: "30"
@@ -155,6 +159,7 @@ services:
       - POSTGRES_DB=prefect
     volumes:
       - workflow_db:/var/lib/postgresql/data
+      - "./${INFRAHUB_TESTING_LOCAL_DB_BACKUP_DIRECTORY}:${INFRAHUB_TESTING_INTERNAL_DB_BACKUP_DIRECTORY}"
     healthcheck:
       test: ["CMD-SHELL", "pg_isready"]
       interval: 10s

infrahub_testcontainers/models.py

@@ -1,4 +1,4 @@
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from enum import Enum
 from typing import Any
 
@@ -27,7 +27,7 @@ class InfrahubResultContext(BaseModel):
 
 class InfrahubActiveMeasurementItem(BaseModel):
     definition: MeasurementDefinition
-    start_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    start_time: datetime = Field(default_factory=lambda: datetime.now(UTC))
     context: dict[str, Any] = Field(default_factory=dict)
 
 

infrahub_testcontainers/performance_test.py

@@ -1,6 +1,6 @@
 import hashlib
 import json
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from types import TracebackType
 from typing import Any
 
@@ -35,7 +35,7 @@ class InfrahubPerformanceTest:
         self.env_vars = {}
         self.project_name = ""
         self.test_info = {}
-        self.start_time = datetime.now(timezone.utc)
+        self.start_time = datetime.now(UTC)
         self.end_time: datetime | None = None
         self.results_url = results_url
         self.scraper_endpoint = ""
@@ -57,7 +57,7 @@ class InfrahubPerformanceTest:
 
     def finalize(self, session: pytest.Session) -> None:
         if self.initialized:
-            self.end_time = datetime.now(timezone.utc)
+            self.end_time = datetime.now(UTC)
             self.extract_test_session_information(session)
             self.send_results()
 
@@ -129,7 +129,7 @@ class InfrahubPerformanceTest:
         if not exc_type and self.active_measurements:
             self.add_measurement(
                 definition=self.active_measurements.definition,
-                value=(datetime.now(timezone.utc) - self.active_measurements.start_time).total_seconds() * 1000,
+                value=(datetime.now(UTC) - self.active_measurements.start_time).total_seconds() * 1000,
                 context=self.active_measurements.context,
             )
 

infrahub/core/convert_object_type/conversion.py

@@ -1,134 +0,0 @@
-from typing import Any
-
-from pydantic import BaseModel
-
-from infrahub.core.attribute import BaseAttribute
-from infrahub.core.branch import Branch
-from infrahub.core.constants import RelationshipCardinality
-from infrahub.core.manager import NodeManager
-from infrahub.core.node import Node
-from infrahub.core.node.create import create_node
-from infrahub.core.query.relationship import GetAllPeersIds
-from infrahub.core.query.resource_manager import PoolChangeReserved
-from infrahub.core.relationship import RelationshipManager
-from infrahub.core.schema import NodeSchema
-from infrahub.database import InfrahubDatabase
-
-
-class InputDataForDestField(BaseModel):  # Only one of these fields can be not None
-    attribute_value: Any | None = None
-    peer_id: str | None = None
-    peers_ids: list[str] | None = None
-
-    @property
-    def value(self) -> Any:
-        fields = [self.attribute_value, self.peer_id, self.peers_ids]
-        set_fields = [f for f in fields if f is not None]
-        if len(set_fields) != 1:
-            raise ValueError("Exactly one of attribute_value, peer_id, or peers_ids must be set")
-        return set_fields[0]
-
-
-class InputForDestField(BaseModel):  # Only one of these fields can be not None
-    source_field: str | None = None
-    data: InputDataForDestField | None = None
-
-    @property
-    def value(self) -> Any:
-        if self.source_field is not None and self.data is not None:
-            raise ValueError("Only one of source_field or data can be set")
-        if self.source_field is None and self.data is None:
-            raise ValueError("Either source_field or data must be set")
-        return self.source_field if self.source_field is not None else self.data
-
-
-async def get_out_rels_peers_ids(node: Node, db: InfrahubDatabase) -> list[str]:
-    all_peers: list[Node] = []
-    for name in node._relationships:
-        relm: RelationshipManager = getattr(node, name)
-        peers = await relm.get_peers(db=db)
-        all_peers.extend(peers.values())
-    return [peer.id for peer in all_peers]
-
-
-async def build_data_new_node(db: InfrahubDatabase, mapping: dict[str, InputForDestField], node: Node) -> dict:
-    """Value of a given field on the target kind to convert is either an input source attribute/relationship of the source node,
-    or a raw value."""
-
-    data = {}
-    for dest_field_name, input_for_dest_field in mapping.items():
-        value = input_for_dest_field.value
-        if isinstance(value, str):  # source_field
-            item = getattr(node, value)
-            if isinstance(item, BaseAttribute):
-                data[dest_field_name] = item.value
-            elif isinstance(item, RelationshipManager):
-                if item.schema.cardinality == RelationshipCardinality.ONE:
-                    peer = await item.get_peer(db=db)
-                    if peer is not None:
-                        data[dest_field_name] = {"id": peer.id}
-                    # else, relationship is optional, and if the target relationship is mandatory an error will be raised during creation
-                elif item.schema.cardinality == RelationshipCardinality.MANY:
-                    data[dest_field_name] = [{"id": peer.id} for _, peer in (await item.get_peers(db=db)).items()]
-                else:
-                    raise ValueError(f"Unknown cardinality {item.schema.cardinality=}")
-        else:  # user input data
-            data[dest_field_name] = value.value
-    return data
-
-
-async def get_unidirectional_rels_peers_ids(node: Node, branch: Branch, db: InfrahubDatabase) -> list[str]:
-    """
-    Returns peers ids of nodes connected to input `node` through an incoming unidirectional relationship.
-    """
-
-    out_rels_identifier = [rel.identifier for rel in node.get_schema().relationships]
-    query = await GetAllPeersIds.init(db=db, node_id=node.id, branch=branch, exclude_identifiers=out_rels_identifier)
-    await query.execute(db=db)
-    return query.get_peers_uuids()
-
-
-async def convert_object_type(
-    node: Node, target_schema: NodeSchema, mapping: dict[str, InputForDestField], branch: Branch, db: InfrahubDatabase
-) -> Node:
-    """Delete the node and return the new created one. If creation fails, the node is not deleted, and raise an error.
-    An extra check is performed on input node peers relationships to make sure they are still valid."""
-
-    node_schema = node.get_schema()
-    if not isinstance(node_schema, NodeSchema):
-        raise ValueError(f"Only a node with a NodeSchema can be converted, got {type(node_schema)}")
-
-    async with db.start_transaction() as dbt:
-        deleted_node_out_rels_peer_ids = await get_out_rels_peers_ids(node=node, db=dbt)
-        deleted_node_unidir_rels_peer_ids = await get_unidirectional_rels_peers_ids(node=node, db=dbt, branch=branch)
-
-        # Delete the node, so we delete relationships with peers as well, which might temporarily break cardinality constraints
-        # but they should be restored when creating the new node.
-        deleted_nodes = await NodeManager.delete(db=dbt, branch=branch, nodes=[node], cascade_delete=False)
-        if len(deleted_nodes) != 1:
-            raise ValueError(f"Deleted {len(deleted_nodes)} nodes instead of 1")
-
-        data_new_node = await build_data_new_node(dbt, mapping, node)
-        node_created = await create_node(
-            data=data_new_node,
-            db=dbt,
-            branch=branch,
-            schema=target_schema,
-        )
-
-        # Make sure relationships with constraints are not broken by retrieving them
-        peers_ids = deleted_node_out_rels_peer_ids + deleted_node_unidir_rels_peer_ids
-        peers = await NodeManager.get_many(ids=peers_ids, db=dbt, prefetch_relationships=True, branch=branch)
-        for peer in peers.values():
-            peer.validate_relationships()
-
-        # If the node had some value reserved in any Pools / Resource Manager, we need to change the identifier of the reservation(s)
-        query = await PoolChangeReserved.init(
-            db=dbt,
-            existing_identifier=node.get_id(),
-            new_identifier=node_created.get_id(),
-            branch=branch,
-        )
-        await query.execute(db=dbt)
-
-        return node_created