infrahub-server 1.6.0b0__tar.gz → 1.6.2__tar.gz

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: infrahub-server
-Version: 1.6.0b0
+Version: 1.6.2
 Summary: Infrahub is taking a new approach to Infrastructure Management by providing a new generation of datastore to organize and control all the data that defines how an infrastructure should run.
 Project-URL: Homepage, https://opsmill.com
 Project-URL: Repository, https://github.com/opsmill/infrahub
@@ -21,7 +21,6 @@ Requires-Dist: bcrypt<4.2,>=4.1
 Requires-Dist: boto3==1.34.129
 Requires-Dist: cachetools-async==0.0.5
 Requires-Dist: click==8.1.7
-Requires-Dist: copier==9.8.0
 Requires-Dist: deepdiff==8.6.1
 Requires-Dist: dulwich==0.22.7
 Requires-Dist: email-validator<2.2,>=2.1
@@ -60,7 +59,7 @@ Requires-Dist: tomli>=1.1.0; python_version <= '3.11'
 Requires-Dist: typer==0.19.2
 Requires-Dist: ujson<6,>=5
 Requires-Dist: uvicorn[standard]<0.33,>=0.32
-Requires-Dist: whenever==0.7.3
+Requires-Dist: whenever==0.9.3
 Description-Content-Type: text/markdown
 
 <h1 align="center">

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/api/oauth2.py

@@ -12,10 +12,12 @@ from opentelemetry import trace
 from infrahub import config, models
 from infrahub.api.dependencies import get_db
 from infrahub.auth import (
+    SSOStateCache,
     get_groups_from_provider,
     signin_sso_account,
     validate_auth_response,
 )
+from infrahub.auth_pkce import compute_code_challenge, generate_code_verifier
 from infrahub.exceptions import ProcessingError
 from infrahub.log import get_logger
 from infrahub.message_bus.types import KVTTL
@@ -42,6 +44,7 @@ async def authorize(request: Request, provider_name: str, final_url: str | None
     with trace.get_tracer(__name__).start_as_current_span("sso_oauth2_client_configuration") as span:
         span.set_attribute("provider_name", provider_name)
         span.set_attribute("scopes", provider.scopes)
+        span.set_attribute("pkce_enabled", provider.pkce_enabled)
 
         client = AsyncOAuth2Client(
             client_id=provider.client_id,
@@ -52,14 +55,32 @@ async def authorize(request: Request, provider_name: str, final_url: str | None
     redirect_uri = _get_redirect_url(request=request, provider_name=provider_name)
     final_url = final_url or config.SETTINGS.main.public_url or str(request.base_url)
 
+    # Generate PKCE parameters if enabled
+    code_verifier = None
+    pkce_params: dict[str, str] = {}
+    if provider.pkce_enabled:
+        code_verifier = generate_code_verifier()
+        code_challenge = compute_code_challenge(code_verifier)
+        pkce_params = {
+            "code_challenge": code_challenge,
+            "code_challenge_method": "S256",
+        }
+
     authorization_uri, state = client.create_authorization_url(
-        url=provider.authorization_url, redirect_uri=redirect_uri, scope=provider.scopes, final_url=final_url
+        url=provider.authorization_url,
+        redirect_uri=redirect_uri,
+        scope=provider.scopes,
+        final_url=final_url,
+        **pkce_params,
     )
 
     service: InfrahubServices = request.app.state.service
 
+    cache_data = SSOStateCache(final_url=final_url, code_verifier=code_verifier)
     await service.cache.set(
-        key=f"security:oauth2:provider:{provider_name}:state:{state}", value=final_url, expires=KVTTL.TWO_HOURS
+        key=f"security:oauth2:provider:{provider_name}:state:{state}",
+        value=cache_data.model_dump_json(),
+        expires=KVTTL.TWO_HOURS,
     )
 
     if config.SETTINGS.dev.frontend_redirect_sso:
@@ -82,13 +103,15 @@ async def token(
     service: InfrahubServices = request.app.state.service
 
     cache_key = f"security:oauth2:provider:{provider_name}:state:{state}"
-    stored_final_url = await service.cache.get(key=cache_key)
+    cached_data = await service.cache.get(key=cache_key)
     await service.cache.delete(key=cache_key)
 
-    if not stored_final_url:
+    if not cached_data:
         raise ProcessingError(message="Invalid 'state' parameter")
 
-    token_data = {
+    sso_state = SSOStateCache.model_validate_json(cached_data)
+
+    token_data: dict[str, str | None] = {
         "code": code,
         "client_id": provider.client_id,
         "client_secret": provider.client_secret,
@@ -96,6 +119,10 @@ async def token(
         "grant_type": "authorization_code",
     }
 
+    # Add code_verifier if PKCE was used
+    if sso_state.code_verifier:
+        token_data["code_verifier"] = sso_state.code_verifier
+
     token_response = await service.http.post(provider.token_url, data=token_data)
     validate_auth_response(response=token_response, provider_type="OAuth 2.0")
 
@@ -139,5 +166,5 @@ async def token(
     )
 
     return models.UserTokenWithUrl(
-        access_token=user_token.access_token, refresh_token=user_token.refresh_token, final_url=stored_final_url
+        access_token=user_token.access_token, refresh_token=user_token.refresh_token, final_url=sso_state.final_url
     )

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/api/oidc.py

@@ -14,10 +14,12 @@ from pydantic import BaseModel, HttpUrl
 from infrahub import config, models
 from infrahub.api.dependencies import get_db
 from infrahub.auth import (
+    SSOStateCache,
     get_groups_from_provider,
     signin_sso_account,
     validate_auth_response,
 )
+from infrahub.auth_pkce import compute_code_challenge, generate_code_verifier
 from infrahub.exceptions import ProcessingError
 from infrahub.log import get_logger
 from infrahub.message_bus.types import KVTTL
@@ -58,6 +60,10 @@ class OIDCDiscoveryConfig(BaseModel):
     tls_client_certificate_bound_access_tokens: bool | None = None
     mtls_endpoint_aliases: dict[str, HttpUrl] | None = None
 
+    @property
+    def supports_pkce(self) -> bool:
+        return "S256" in (self.code_challenge_methods_supported or [])
+
 
 def _get_redirect_url(request: Request, provider_name: str) -> str:
     """Return public redirect URL."""
@@ -74,10 +80,14 @@ async def authorize(request: Request, provider_name: str, final_url: str | None
     validate_auth_response(response=response, provider_type="OIDC")
     oidc_config = OIDCDiscoveryConfig(**response.json())
 
+    pkce_supported = oidc_config.supports_pkce
+
     with trace.get_tracer(__name__).start_as_current_span("sso_oauth2_client_configuration") as span:
         span.set_attribute("provider_name", provider_name)
         span.set_attribute("scopes", provider.scopes)
         span.set_attribute("discovery_url", provider.discovery_url)
+        span.set_attribute("pkce_enabled", provider.pkce_enabled)
+        span.set_attribute("pkce_supported", pkce_supported)
 
         client = AsyncOAuth2Client(
             client_id=provider.client_id,
@@ -88,12 +98,26 @@ async def authorize(request: Request, provider_name: str, final_url: str | None
     redirect_uri = _get_redirect_url(request=request, provider_name=provider_name)
     final_url = final_url or config.SETTINGS.main.public_url or str(request.base_url)
 
+    # Generate PKCE parameters if enabled and supported by provider
+    code_verifier = None
+    pkce_params: dict[str, str] = {}
+    if provider.pkce_enabled and pkce_supported:
+        code_verifier = generate_code_verifier()
+        code_challenge = compute_code_challenge(code_verifier)
+        pkce_params = {
+            "code_challenge": code_challenge,
+            "code_challenge_method": "S256",
+        }
+
     authorization_uri, state = client.create_authorization_url(
-        url=str(oidc_config.authorization_endpoint), redirect_uri=redirect_uri, scope=provider.scopes
+        url=str(oidc_config.authorization_endpoint), redirect_uri=redirect_uri, scope=provider.scopes, **pkce_params
     )
 
+    cache_data = SSOStateCache(final_url=final_url, code_verifier=code_verifier)
     await service.cache.set(
-        key=f"security:oidc:provider:{provider_name}:state:{state}", value=final_url, expires=KVTTL.TWO_HOURS
+        key=f"security:oidc:provider:{provider_name}:state:{state}",
+        value=cache_data.model_dump_json(),
+        expires=KVTTL.TWO_HOURS,
     )
 
     if config.SETTINGS.dev.frontend_redirect_sso:
@@ -116,13 +140,15 @@ async def token(
     service: InfrahubServices = request.app.state.service
 
     cache_key = f"security:oidc:provider:{provider_name}:state:{state}"
-    stored_final_url = await service.cache.get(key=cache_key)
+    cached_data = await service.cache.get(key=cache_key)
     await service.cache.delete(key=cache_key)
 
-    if not stored_final_url:
+    if not cached_data:
         raise ProcessingError(message="Invalid 'state' parameter")
 
-    token_data = {
+    sso_state = SSOStateCache.model_validate_json(cached_data)
+
+    token_data: dict[str, str | None] = {
         "code": code,
         "client_id": provider.client_id,
         "client_secret": provider.client_secret,
@@ -130,6 +156,10 @@ async def token(
         "grant_type": "authorization_code",
     }
 
+    # Add code_verifier if PKCE was used
+    if sso_state.code_verifier:
+        token_data["code_verifier"] = sso_state.code_verifier
+
     discovery_response = await service.http.get(url=provider.discovery_url)
     validate_auth_response(response=discovery_response, provider_type="OIDC")
 
@@ -183,7 +213,7 @@ async def token(
     )
 
     return models.UserTokenWithUrl(
-        access_token=user_token.access_token, refresh_token=user_token.refresh_token, final_url=stored_final_url
+        access_token=user_token.access_token, refresh_token=user_token.refresh_token, final_url=sso_state.final_url
     )
 
 

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/auth.py

@@ -51,6 +51,17 @@ class AccountSession(BaseModel):
         return self.auth_type == AuthType.JWT
 
 
+class SSOStateCache(BaseModel):
+    """Cache data stored during OAuth2/OIDC authorization flow.
+
+    This model is used to store state information between the authorization
+    request and the token exchange, including PKCE code_verifier when enabled.
+    """
+
+    final_url: str
+    code_verifier: str | None = None
+
+
 async def validate_active_account(db: InfrahubDatabase, account_id: str) -> None:
     account = await NodeManager.get_one(db=db, kind=CoreGenericAccount, id=account_id, raise_on_error=True)
     if account.status.value != AccountStatus.ACTIVE.value:

infrahub_server-1.6.2/backend/infrahub/auth_pkce.py

@@ -0,0 +1,41 @@
+"""PKCE (RFC 7636) utilities for OAuth2/OIDC authentication.
+
+This module provides functions to generate code verifiers and compute
+code challenges for the Proof Key for Code Exchange (PKCE) extension.
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import secrets
+
+
+def generate_code_verifier() -> str:
+    """Generate a cryptographically random code verifier.
+
+    The code verifier is a high-entropy cryptographic random string using
+    the unreserved characters [A-Z] / [a-z] / [0-9] / "-" / "." / "_" / "~",
+    with a minimum length of 43 characters and a maximum length of 128
+    characters.
+
+    Returns:
+        A 43-character URL-safe string (256 bits of entropy).
+    """
+    return secrets.token_urlsafe(32)
+
+
+def compute_code_challenge(code_verifier: str) -> str:
+    """Compute S256 code challenge from verifier.
+
+    Implements the S256 code challenge method as defined in RFC 7636:
+    code_challenge = BASE64URL(SHA256(code_verifier))
+
+    Args:
+        code_verifier: The code verifier string.
+
+    Returns:
+        Base64URL-encoded SHA256 hash without padding.
+    """
+    digest = hashlib.sha256(code_verifier.encode("ascii")).digest()
+    return base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/config.py

@@ -327,7 +327,7 @@ class DevelopmentSettings(BaseSettings):
         description="Allow enterprise configuration in development mode, this will not enable the features just allow the configuration.",
     )
     git_credential_helper: str = Field(
-        default="/usr/local/bin/infrahub-git-credential",
+        default="infrahub-git-credential",
         description="Location of git credential helper",
     )
 
@@ -585,11 +585,14 @@ class SecurityOIDCBaseSettings(BaseSettings):
     icon: str = Field(default="mdi:account-key")
     display_label: str = Field(default="Single Sign on")
     userinfo_method: UserInfoMethod = Field(default=UserInfoMethod.GET)
+    pkce_enabled: bool = Field(
+        default=True, description="Enable PKCE (RFC 7636) with S256 method for authorization code flow"
+    )
 
 
 class SecurityOIDCSettings(SecurityOIDCBaseSettings):
     client_id: str = Field(..., description="Client ID of the application created in the auth provider")
-    client_secret: str = Field(..., description="Client secret as defined in auth provider")
+    client_secret: str | None = Field(default=None, description="Client secret as defined in auth provider")
     discovery_url: str = Field(..., description="The OIDC discovery URL xyz/.well-known/openid-configuration")
     scopes: list[str] = Field(default_factory=_default_scopes)
 
@@ -637,13 +640,16 @@ class SecurityOAuth2BaseSettings(BaseSettings):
 
     icon: str = Field(default="mdi:account-key")
     userinfo_method: UserInfoMethod = Field(default=UserInfoMethod.GET)
+    pkce_enabled: bool = Field(
+        default=True, description="Enable PKCE (RFC 7636) with S256 method for authorization code flow"
+    )
 
 
 class SecurityOAuth2Settings(SecurityOAuth2BaseSettings):
     """Common base for Oauth2 providers"""
 
     client_id: str = Field(..., description="Client ID of the application created in the auth provider")
-    client_secret: str = Field(..., description="Client secret as defined in auth provider")
+    client_secret: str | None = Field(default=None, description="Client secret as defined in auth provider")
     authorization_url: str = Field(...)
     token_url: str = Field(...)
     userinfo_url: str = Field(...)

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/core/branch/models.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import re
 from typing import TYPE_CHECKING, Any, Optional, Self, Union, cast
 
-from neo4j.graph import Node as Neo4jNode
 from pydantic import Field, field_validator
 
 from infrahub.core.branch.enums import BranchStatus
@@ -24,6 +23,8 @@ from infrahub.core.timestamp import Timestamp
 from infrahub.exceptions import BranchNotFoundError, InitializationError, ValidationError
 
 if TYPE_CHECKING:
+    from neo4j.graph import Node as Neo4jNode
+
     from infrahub.database import InfrahubDatabase
 
 
@@ -168,7 +169,7 @@ class Branch(StandardNode):
         )
         await query.execute(db=db)
 
-        return [cls.from_db(node=cast(Neo4jNode, result.get("n"))) for result in query.get_results()]
+        return [cls.from_db(node=cast("Neo4jNode", result.get("n"))) for result in query.get_results()]
 
     @classmethod
     async def get_list_count(

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/core/branch/tasks.py

@@ -157,7 +157,12 @@ async def rebase_branch(branch: str, context: InfrahubContext, send_events: bool
             responses = await schema_validate_migrations(
                 message=SchemaValidateMigrationData(branch=obj, schema_branch=candidate_schema, constraints=constraints)
             )
-            error_messages = [violation.message for response in responses for violation in response.violations]
+            error_messages = [
+                f"{violation.message} for constraint {response.constraint_name} {response.schema_path.field_name} {response.schema_path.property_name} and node {violation.node_id} {violation.node_kind}"  # noqa: E501
+                for response in responses
+                for violation in response.violations
+            ]
+
             if error_messages:
                 raise ValidationError(",\n".join(error_messages))
 

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/core/changelog/models.py

@@ -290,7 +290,7 @@ class NodeChangelog(BaseModel):
                     name=relationship.schema.name
                 )
             relationship_container = cast(
-                RelationshipCardinalityManyChangelog, self.relationships[relationship.schema.name]
+                "RelationshipCardinalityManyChangelog", self.relationships[relationship.schema.name]
             )
 
             relationship_container.add_new_peer(relationship=relationship)
@@ -311,7 +311,7 @@ class NodeChangelog(BaseModel):
                     name=relationship.schema.name
                 )
             relationship_container = cast(
-                RelationshipCardinalityManyChangelog, self.relationships[relationship.schema.name]
+                "RelationshipCardinalityManyChangelog", self.relationships[relationship.schema.name]
             )
             relationship_container.remove_peer(
                 peer_id=relationship.get_peer_id(), peer_kind=relationship.get_peer_kind()

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/core/constants/__init__.py

@@ -391,3 +391,4 @@ DEFAULT_REL_IDENTIFIER_LENGTH = 128
 OBJECT_TEMPLATE_RELATIONSHIP_NAME = "object_template"
 OBJECT_TEMPLATE_NAME_ATTR = "template_name"
 PROFILE_NODE_RELATIONSHIP_IDENTIFIER = "node__profile"
+PROFILE_TEMPLATE_RELATIONSHIP_IDENTIFIER = "template__profile"

infrahub_server-1.6.2/backend/infrahub/core/graph/__init__.py

@@ -0,0 +1 @@
+GRAPH_VERSION = 48

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/core/integrity/object_conflict/conflict_recorder.py

@@ -24,7 +24,7 @@ class ObjectConflictValidatorRecorder:
             )
         except NodeNotFoundError:
             return []
-        proposed_change = cast(CoreProposedChange, proposed_change)
+        proposed_change = cast("CoreProposedChange", proposed_change)
         validator = await self.get_or_create_validator(proposed_change)
         await self.initialize_validator(validator)
 

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/core/manager.py

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, Iterable, Literal, TypeVar, overload
 
 from infrahub_sdk.utils import deep_merge_dict, is_valid_uuid
 
-from infrahub.core.constants import RelationshipCardinality, RelationshipDirection
+from infrahub.core.constants import InfrahubKind, RelationshipCardinality, RelationshipDirection
 from infrahub.core.node import Node
 from infrahub.core.node.delete_validator import NodeDeleteValidator
 from infrahub.core.query.node import (
@@ -188,17 +188,22 @@ class NodeManager:
         await query.execute(db=db)
         node_ids = query.get_node_ids()
 
-        # if display_label or hfid has been requested we need to ensure we are querying the right fields
-        if fields and "display_label" in fields:
+        if (
+            fields
+            and "identifier" in fields
+            and node_schema.kind
+            in [
+                InfrahubKind.BASEPERMISSION,
+                InfrahubKind.GLOBALPERMISSION,
+                InfrahubKind.OBJECTPERMISSION,
+            ]
+        ):
+            # This is a workaround to ensure we are querying the right fields for permissions
+            # The identifier for permissions needs the same fields as the display label
             schema_branch = db.schema.get_schema_branch(name=branch.name)
             display_label_fields = schema_branch.generate_fields_for_display_label(name=node_schema.kind)
             if display_label_fields:
-                fields = deep_merge_dict(dicta=fields, dictb=display_label_fields)
-
-        if fields and "hfid" in fields and node_schema.human_friendly_id:
-            hfid_fields = node_schema.generate_fields_for_hfid()
-            if hfid_fields:
-                fields = deep_merge_dict(dicta=fields, dictb=hfid_fields)
+                deep_merge_dict(dicta=fields, dictb=display_label_fields)
 
         response = await cls.get_many(
             ids=node_ids,
@@ -300,6 +305,8 @@ class NodeManager:
         branch: Branch | str | None = None,
         branch_agnostic: bool = False,
         fetch_peers: bool = False,
+        include_source: bool = False,
+        include_owner: bool = False,
     ) -> list[Relationship]:
         branch = await registry.get_branch(branch=branch, db=db)
         at = Timestamp(at)
@@ -324,24 +331,31 @@ class NodeManager:
         if not peers_info:
             return []
 
-        # if display_label has been requested we need to ensure we are querying the right fields
-        if fields and "display_label" in fields:
-            peer_schema = schema.get_peer_schema(db=db, branch=branch)
-            schema_branch = db.schema.get_schema_branch(name=branch.name)
-            display_label_fields = schema_branch.generate_fields_for_display_label(name=peer_schema.kind)
-            if display_label_fields:
-                fields = deep_merge_dict(dicta=fields, dictb=display_label_fields)
-
-        if fields and "hfid" in fields:
+        if fields and "identifier" in fields:
+            # This is a workaround to ensure we are querying the right fields for permissions
+            # The identifier for permissions needs the same fields as the display label
             peer_schema = schema.get_peer_schema(db=db, branch=branch)
-            hfid_fields = peer_schema.generate_fields_for_hfid()
-            if hfid_fields:
-                fields = deep_merge_dict(dicta=fields, dictb=hfid_fields)
+            if peer_schema.kind in [
+                InfrahubKind.BASEPERMISSION,
+                InfrahubKind.GLOBALPERMISSION,
+                InfrahubKind.OBJECTPERMISSION,
+            ]:
+                schema_branch = db.schema.get_schema_branch(name=branch.name)
+                display_label_fields = schema_branch.generate_fields_for_display_label(name=peer_schema.kind)
+                if display_label_fields:
+                    deep_merge_dict(dicta=fields, dictb=display_label_fields)
 
         if fetch_peers:
             peer_ids = [peer.peer_id for peer in peers_info]
             peer_nodes = await cls.get_many(
-                db=db, ids=peer_ids, fields=fields, at=at, branch=branch, branch_agnostic=branch_agnostic
+                db=db,
+                ids=peer_ids,
+                fields=fields,
+                at=at,
+                branch=branch,
+                branch_agnostic=branch_agnostic,
+                include_source=include_source,
+                include_owner=include_owner,
             )
 
         results = []
@@ -420,15 +434,6 @@ class NodeManager:
         if not peers_ids:
             return {}
 
-        hierarchy_schema = node_schema.get_hierarchy_schema(db=db, branch=branch)
-
-        # if display_label has been requested we need to ensure we are querying the right fields
-        if fields and "display_label" in fields:
-            schema_branch = db.schema.get_schema_branch(name=branch.name)
-            display_label_fields = schema_branch.generate_fields_for_display_label(name=hierarchy_schema.kind)
-            if display_label_fields:
-                fields = deep_merge_dict(dicta=fields, dictb=display_label_fields)
-
         return await cls.get_many(
             db=db, ids=peers_ids, fields=fields, at=at, branch=branch, include_owner=True, include_source=True
         )

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/core/migrations/graph/__init__.py

@@ -48,6 +48,8 @@ from .m043_create_hfid_display_label_in_db import Migration043
 from .m044_backfill_hfid_display_label_in_db import Migration044
 from .m045_backfill_hfid_display_label_in_db_profile_template import Migration045
 from .m046_fill_agnostic_hfid_display_labels import Migration046
+from .m047_backfill_or_null_display_label import Migration047
+from .m048_undelete_rel_props import Migration048
 
 if TYPE_CHECKING:
     from ..shared import MigrationTypes
@@ -100,6 +102,8 @@ MIGRATIONS: list[type[MigrationTypes]] = [
     Migration044,
     Migration045,
     Migration046,
+    Migration047,
+    Migration048,
 ]
 
 

{infrahub_server-1.6.0b0 → infrahub_server-1.6.2}/backend/infrahub/core/migrations/graph/m041_deleted_dup_edges.py

@@ -68,33 +68,49 @@ DELETE added_e
         self.add_to_query(query)
 
 
-class DeleteDuplicateEdgesForMigratedKindNodes(Query):
-    name = "delete_duplicate_edges_for_migrated_kind_nodes_query"
+class DeleteDuplicatedRelationshipEdges(Query):
+    name = "delete_duplicated_relationship_edges_query"
     type = QueryType.WRITE
     insert_return = False
 
+    def __init__(self, migrated_kind_nodes_only: bool = True, **kwargs: Any):
+        self.migrated_kind_nodes_only = migrated_kind_nodes_only
+        super().__init__(**kwargs)
+
     async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:  # noqa: ARG002
-        query = """
+        if not self.migrated_kind_nodes_only:
+            relationship_filter_query = """
+MATCH (rel:Relationship)
+            """
+
+        else:
+            relationship_filter_query = """
 // ------------
 // get UUIDs for migrated kind/inheritance nodes
 // ------------
 MATCH (n:Node)
 WITH n.uuid AS node_uuid, count(*) AS num_nodes_with_uuid
 WHERE num_nodes_with_uuid > 1
-CALL (node_uuid) {
-    // ------------
-    // find any Relationships for these nodes
-    // ------------
-    MATCH (n:Node {uuid: node_uuid})-[:IS_RELATED]-(rel:Relationship)
-    WITH DISTINCT rel
-    MATCH (rel)-[e]->(peer)
+// ------------
+// find any Relationships for these nodes
+// ------------
+MATCH (n:Node {uuid: node_uuid})-[:IS_RELATED]-(rel:Relationship)
+WITH DISTINCT rel
+
+            """
+        self.add_to_query(relationship_filter_query)
+
+        query = """
+CALL (rel) {
+    MATCH (rel)-[e]-(peer)
     WITH
+        elementId(rel) AS rel_element_id,
         type(e) AS e_type,
         e.branch AS e_branch,
         e.from AS e_from,
         e.to AS e_to,
         e.status AS e_status,
-        e.peer AS e_peer,
+        elementId(peer) AS peer_element_id,
         CASE
             WHEN startNode(e) = rel THEN "out" ELSE "in"
         END AS direction,
@@ -142,7 +158,9 @@ class Migration041(ArbitraryMigration):
         rprint("done")
 
         rprint("Deleting duplicate edges for migrated kind/inheritance nodes", end="...")
-        delete_duplicate_edges_query = await DeleteDuplicateEdgesForMigratedKindNodes.init(db=db)
+        delete_duplicate_edges_query = await DeleteDuplicatedRelationshipEdges.init(
+            db=db, migrated_kind_nodes_only=True
+        )
         await delete_duplicate_edges_query.execute(db=db)
         rprint("done")