From 5cff5d81ab87fd44feaeb57a64712e2569efac93 Mon Sep 17 00:00:00 2001
From: Aaron McCarty <aaron@opsmill.com>
Date: Wed, 18 Sep 2024 15:28:54 -0700
Subject: [PATCH 1/5] IFC-699 use elementId() instead of ID() for neo4j queries

---
 backend/infrahub/core/protocols_base.py      |  1 +
 backend/infrahub/core/query/branch.py        | 21 ++++++++-----
 backend/infrahub/core/query/diff.py          | 31 ++++++++++----------
 backend/infrahub/core/query/ipam.py          |  8 +++--
 backend/infrahub/core/query/relationship.py  |  5 +++-
 backend/infrahub/core/query/standard_node.py | 11 +++----
 backend/infrahub/core/utils.py               | 10 +++----
 backend/infrahub/database/__init__.py        |  5 ++++
 8 files changed, 54 insertions(+), 38 deletions(-)

diff --git a/backend/infrahub/core/protocols_base.py b/backend/infrahub/core/protocols_base.py
index 7336ece96d..ea8cd38d00 100644
--- a/backend/infrahub/core/protocols_base.py
+++ b/backend/infrahub/core/protocols_base.py
@@ -60,6 +60,7 @@ async def run_query(
     def render_list_comprehension(self, items: str, item_name: str) -> str: ...
     def render_list_comprehension_with_list(self, items: str, item_names: list[str]) -> str: ...
     def render_uuid_generation(self, node_label: str, node_attr: str) -> str: ...
+    def get_id_function_name(self) -> str: ...
 
 
 @runtime_checkable
diff --git a/backend/infrahub/core/query/branch.py b/backend/infrahub/core/query/branch.py
index 21d3916502..15716e130a 100644
--- a/backend/infrahub/core/query/branch.py
+++ b/backend/infrahub/core/query/branch.py
@@ -24,11 +24,13 @@ def __init__(self, node_id: int, **kwargs: Any):
     async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
         query = """
         MATCH (root:Root)
-        MATCH (d) WHERE ID(d) = $node_id
+        MATCH (d) WHERE %(id_func)s(d) = $node_id
         WITH root,d
         CREATE (d)-[r:IS_PART_OF { branch: $branch, branch_level: $branch_level, from: $now, to: null, status: $status }]->(root)
-        RETURN ID(r)
-        """
+        RETURN %(id_func)s(r)
+        """ % {
+            "id_func": db.get_id_function_name(),
+        }
 
         self.params["node_id"] = element_id_to_id(self.node_id)
         self.params["now"] = self.at.to_string()
@@ -100,10 +102,12 @@ def __init__(self, ids: list[str], **kwargs: Any) -> None:
     async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
         query = """
         MATCH ()-[r]->()
-        WHERE ID(r) IN $ids
+        WHERE %(id_func)s(r) IN $ids
         SET r.from = $at
         SET r.conflict = NULL
-        """
+        """ % {
+            "id_func": db.get_id_function_name(),
+        }
 
         self.add_to_query(query=query)
 
@@ -126,13 +130,13 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
         if config.SETTINGS.database.db_type == config.DatabaseType.MEMGRAPH:
             query = """
             MATCH p = (s)-[r]-(d)
-            WHERE ID(r) IN $ids
+            WHERE %(id_func)s(r) IN $ids
             DELETE r
             """
         else:
             query = """
             MATCH p = (s)-[r]-(d)
-            WHERE ID(r) IN $ids
+            WHERE %(id_func)s(r) IN $ids
             DELETE r
             WITH *
             UNWIND nodes(p) AS n
@@ -140,6 +144,9 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
             WHERE NOT exists((n)--())
             DELETE n
             """
+        query %= {
+            "id_func": db.get_id_function_name(),
+        }
 
         self.add_to_query(query=query)
 
diff --git a/backend/infrahub/core/query/diff.py b/backend/infrahub/core/query/diff.py
index 23ea3180b4..8048479407 100644
--- a/backend/infrahub/core/query/diff.py
+++ b/backend/infrahub/core/query/diff.py
@@ -584,13 +584,13 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
                 OPTIONAL MATCH path = (
                     (:Root)<-[r_root:IS_PART_OF]-(n:Node)-[r_node]-(inner_p)-[inner_diff_rel]->(inner_q)
                 )
-                WHERE ID(inner_p) = ID(p) AND ID(inner_diff_rel) = ID(diff_rel) AND ID(inner_q) = ID(q)
+                WHERE %(id_func)s(inner_p) = %(id_func)s(p) AND %(id_func)s(inner_diff_rel) = %(id_func)s(diff_rel) AND %(id_func)s(inner_q) = %(id_func)s(q)
                 AND any(l in labels(inner_p) WHERE l in ["Attribute", "Relationship"])
                 AND type(inner_diff_rel) IN ["IS_VISIBLE", "IS_PROTECTED", "HAS_SOURCE", "HAS_OWNER", "HAS_VALUE"]
                 AND any(l in labels(inner_q) WHERE l in ["Boolean", "Node", "AttributeValue"])
                 AND type(r_node) IN ["HAS_ATTRIBUTE", "IS_RELATED"]
                 AND %(n_node_where)s
-                AND [ID(n), type(r_node)] <> [ID(inner_q), type(inner_diff_rel)]
+                AND [%(id_func)s(n), type(r_node)] <> [%(id_func)s(inner_q), type(inner_diff_rel)]
                 AND ALL(
                     r in [r_root, r_node]
                     WHERE r.from <= $to_time AND r.branch IN $branch_names
@@ -600,8 +600,8 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
                 AND (r_node.status = "deleted" OR r_root.status = "active")
                 WITH path AS diff_rel_path, diff_rel, r_root, n, r_node, p
                 ORDER BY
-                    ID(n) DESC,
-                    ID(p) DESC,
+                    %(id_func)s(n) DESC,
+                    %(id_func)s(p) DESC,
                     r_node.branch = diff_rel.branch DESC,
                     r_root.branch = diff_rel.branch DESC,
                     r_node.from DESC,
@@ -619,9 +619,9 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
                 // get base branch version of the diff path, if it exists
                 WITH diff_rel_path, diff_rel, r_root, n, r_node, p
                 OPTIONAL MATCH latest_base_path = (:Root)<-[r_root2]-(n2)-[r_node2]-(inner_p2)-[base_diff_rel]->(base_prop)
-                WHERE ID(r_root2) = ID(r_root) AND ID(n2) = ID(n) AND ID(r_node2) = ID(r_node)  AND ID(inner_p2) = ID(p)
+                WHERE %(id_func)s(r_root2) = %(id_func)s(r_root) AND %(id_func)s(n2) = %(id_func)s(n) AND %(id_func)s(r_node2) = %(id_func)s(r_node)  AND %(id_func)s(inner_p2) = %(id_func)s(p)
                 AND any(r in relationships(diff_rel_path) WHERE r.branch = $branch_name)
-                AND ID(n2) <> ID(base_prop)
+                AND %(id_func)s(n2) <> %(id_func)s(base_prop)
                 AND type(base_diff_rel) = type(diff_rel)
                 AND all(
                     r in relationships(latest_base_path)
@@ -638,9 +638,9 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
                 OPTIONAL MATCH base_peer_path = (
                    (:Root)<-[r_root3]-(n3)-[r_node3]-(inner_p3:Relationship)-[base_r_peer:IS_RELATED]-(base_peer:Node)
                 )
-                WHERE ID(r_root3) = ID(r_root) AND ID(n3) = ID(n) AND ID(r_node3) = ID(r_node) AND ID(inner_p3) = ID(p)
+                WHERE %(id_func)s(r_root3) = %(id_func)s(r_root) AND %(id_func)s(n3) = %(id_func)s(n) AND %(id_func)s(r_node3) = %(id_func)s(r_node) AND %(id_func)s(inner_p3) = %(id_func)s(p)
                 AND type(diff_rel) <> "IS_RELATED"
-                AND [ID(n3), type(r_node3)] <> [ID(base_peer), type(base_r_peer)]
+                AND [%(id_func)s(n3), type(r_node3)] <> [%(id_func)s(base_peer), type(base_r_peer)]
                 AND base_r_peer.from <= $to_time
                 AND base_r_peer.branch IN $branch_names
                 // exclude paths where an active edge is below a deleted edge
@@ -663,7 +663,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
                 OPTIONAL MATCH path = (
                     (:Root)<-[r_root:IS_PART_OF]-(inner_p)-[inner_diff_rel]-(inner_q)-[r_prop]-(prop)
                 )
-                WHERE ID(inner_p) = ID(p) AND ID(inner_diff_rel) = ID(diff_rel) AND ID(inner_q) = ID(q)
+                WHERE %(id_func)s(inner_p) = %(id_func)s(p) AND %(id_func)s(inner_diff_rel) = %(id_func)s(diff_rel) AND %(id_func)s(inner_q) = %(id_func)s(q)
                 AND "Node" IN labels(inner_p)
                 AND type(inner_diff_rel) IN ["HAS_ATTRIBUTE", "IS_RELATED"]
                 AND any(l in labels(inner_q) WHERE l in ["Attribute", "Relationship"])
@@ -673,13 +673,13 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
                     r in [r_root, r_prop]
                     WHERE r.from <= $to_time AND r.branch IN $branch_names
                 )
-                AND [ID(inner_p), type(inner_diff_rel)] <> [ID(prop), type(r_prop)]
+                AND [%(id_func)s(inner_p), type(inner_diff_rel)] <> [%(id_func)s(prop), type(r_prop)]
                 // exclude paths where an active edge is below a deleted edge
                 AND (inner_diff_rel.status = "active" OR (r_prop.status = "deleted" AND inner_diff_rel.branch = r_prop.branch))
                 AND (inner_diff_rel.status = "deleted" OR r_root.status = "active")
                 WITH path, prop, r_prop, r_root
                 ORDER BY
-                    ID(prop),
+                    %(id_func)s(prop),
                     r_prop.branch = diff_rel.branch DESC,
                     r_root.branch = diff_rel.branch DESC,
                     r_prop.from DESC,
@@ -697,7 +697,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
                 OPTIONAL MATCH path = (
                     (inner_q:Root)<-[inner_diff_rel:IS_PART_OF]-(inner_p:Node)-[r_node]-(node)-[r_prop]-(prop)
                 )
-                WHERE ID(inner_p) = ID(p) AND ID(inner_diff_rel) = ID(diff_rel) AND ID(inner_q) = ID(q)
+                WHERE %(id_func)s(inner_p) = %(id_func)s(p) AND %(id_func)s(inner_diff_rel) = %(id_func)s(diff_rel) AND %(id_func)s(inner_q) = %(id_func)s(q)
                 AND type(r_node) IN ["HAS_ATTRIBUTE", "IS_RELATED"]
                 AND any(l in labels(node) WHERE l in ["Attribute", "Relationship"])
                 AND type(r_prop) IN ["IS_VISIBLE", "IS_PROTECTED", "HAS_SOURCE", "HAS_OWNER", "HAS_VALUE", "IS_RELATED"]
@@ -706,7 +706,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
                     r in [r_node, r_prop]
                     WHERE r.from <= $to_time AND r.branch IN $branch_names
                 )
-                AND [ID(inner_p), type(r_node)] <> [ID(prop), type(r_prop)]
+                AND [%(id_func)s(inner_p), type(r_node)] <> [%(id_func)s(prop), type(r_prop)]
                 // exclude paths where an active edge is below a deleted edge
                 AND (inner_diff_rel.status = "active" OR
                     (
@@ -717,8 +717,8 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
                 AND (r_prop.status = "deleted" OR r_node.status = "active")
                 WITH path, node, prop, r_prop, r_node
                 ORDER BY
-                    ID(node),
-                    ID(prop),
+                    %(id_func)s(node),
+                    %(id_func)s(prop),
                     r_prop.branch = diff_rel.branch DESC,
                     r_node.branch = diff_rel.branch DESC,
                     r_prop.from DESC,
@@ -730,6 +730,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
             WITH p, q, diff_rel, full_diff_paths + latest_paths AS full_diff_paths
         """ % {
             "diff_rel_filter": diff_rel_filter,
+            "id_func": db.get_id_function_name(),
             "p_node_where": p_node_where,
             "n_node_where": n_node_where,
         }
diff --git a/backend/infrahub/core/query/ipam.py b/backend/infrahub/core/query/ipam.py
index 0a0ac44f8d..1c69b61c85 100644
--- a/backend/infrahub/core/query/ipam.py
+++ b/backend/infrahub/core/query/ipam.py
@@ -274,8 +274,8 @@ def rel_filter(rel_name: str) -> str:
             -[r_attr:HAS_ATTRIBUTE]->(attr:Attribute)
             -[r_attr_val:HAS_VALUE]->(av:AttributeValue)
         )
-        WHERE ID(r_1) = ID(r_rel1)
-        AND ID(r_2) = ID(r_rel2)
+        WHERE %(id_func)s(r_1) = %(id_func)s(r_rel1)
+        AND %(id_func)s(r_2) = %(id_func)s(r_rel2)
         AND ({rel_filter("r_attr")})
         AND ({rel_filter("r_attr_val")})
         AND attr.name IN ["prefix", "address"]
@@ -301,7 +301,9 @@ def rel_filter(rel_name: str) -> str:
             deepest_branch_details[1] AS branch,
             head(collect(is_active)) AS is_latest_active
         WHERE is_latest_active = TRUE
-        """
+        """ % {
+            "id_func": db.get_id_function_name(),
+        }
         self.return_labels = ["pfx", "child", "av", "branch_level", "branch"]
         self.add_to_query(query)
 
diff --git a/backend/infrahub/core/query/relationship.py b/backend/infrahub/core/query/relationship.py
index f8f0223f7c..fab5ce5927 100644
--- a/backend/infrahub/core/query/relationship.py
+++ b/backend/infrahub/core/query/relationship.py
@@ -391,7 +391,10 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
         self.return_labels = ["s", "d", "rl"]
 
         for prop_name, prop in self.data.properties.items():
-            self.add_to_query("MATCH (prop_%s) WHERE ID(prop_%s) = $prop_%s_id" % (prop_name, prop_name, prop_name))
+            self.add_to_query(
+                "MATCH (prop_%(prop_name)s) WHERE %(id_func)s(prop_%(prop_name)s) = $prop_%(prop_name)s_id"
+                % {"prop_name": prop_name, "id_func": db.get_id_function_name()}
+            )
             self.params[f"prop_{prop_name}_id"] = element_id_to_id(prop.prop_db_id)
             self.return_labels.append(f"prop_{prop_name}")
 
diff --git a/backend/infrahub/core/query/standard_node.py b/backend/infrahub/core/query/standard_node.py
index 3c59f9d215..2ab42297e5 100644
--- a/backend/infrahub/core/query/standard_node.py
+++ b/backend/infrahub/core/query/standard_node.py
@@ -109,13 +109,10 @@ def __init__(self, node_id: str, node_type: str, **kwargs: Any) -> None:
         super().__init__(**kwargs)
 
     async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
-        query = (
-            """
-            MATCH (n:%s)
-            WHERE ID(n) = $node_id OR n.uuid = $node_id
-            """
-            % self.node_type
-        )
+        query = """
+            MATCH (n:%(node_type)s)
+            WHERE %(id_func)s(n) = $node_id OR n.uuid = $node_id
+            """ % {"node_type": self.node_type, "id_func": db.get_id_function_name()}
 
         self.params["node_id"] = self.node_id
         self.add_to_query(query)
diff --git a/backend/infrahub/core/utils.py b/backend/infrahub/core/utils.py
index fd136b4c97..43f81a26b2 100644
--- a/backend/infrahub/core/utils.py
+++ b/backend/infrahub/core/utils.py
@@ -27,12 +27,12 @@ async def add_relationship(
     status=RelationshipStatus.ACTIVE,
 ):
     create_rel_query = """
-    MATCH (s) WHERE ID(s) = $src_node_id
-    MATCH (d) WHERE ID(d) = $dst_node_id
+    MATCH (s) WHERE %(id_func)s(s) = $src_node_id
+    MATCH (d) WHERE %(id_func)s(d) = $dst_node_id
     WITH s,d
-    CREATE (s)-[r:%s { branch: $branch, branch_level: $branch_level, from: $at, to: null, status: $status }]->(d)
-    RETURN ID(r)
-    """ % str(rel_type).upper()
+    CREATE (s)-[r:%(rel_type)s { branch: $branch, branch_level: $branch_level, from: $at, to: null, status: $status }]->(d)
+    RETURN %(id_func)s(r)
+    """ % {"id_func": db.get_id_function_name(), "rel_type": str(rel_type).upper()}
 
     at = Timestamp(at)
 
diff --git a/backend/infrahub/database/__init__.py b/backend/infrahub/database/__init__.py
index fb69d7d1b0..347a616ff9 100644
--- a/backend/infrahub/database/__init__.py
+++ b/backend/infrahub/database/__init__.py
@@ -336,6 +336,11 @@ def render_uuid_generation(self, node_label: str, node_attr: str, index: int = 1
             """
         return generate_uuid_query
 
+    def get_id_function_name(self) -> str:
+        if self.db_type == DatabaseType.NEO4J:
+            return "elementId"
+        return "ID"
+
 
 async def create_database(driver: AsyncDriver, database_name: str) -> None:
     default_db = driver.session()

From ef2e5a420ec37dcb21f8d011548e392db1193fa6 Mon Sep 17 00:00:00 2001
From: Aaron McCarty <aaron@opsmill.com>
Date: Wed, 18 Sep 2024 17:25:01 -0700
Subject: [PATCH 2/5] update id function in more places, update
 element_id_to_id

---
 backend/infrahub/core/query/branch.py        |  2 +-
 backend/infrahub/core/query/node.py          |  2 +-
 backend/infrahub/core/query/standard_node.py |  2 +-
 backend/infrahub/core/utils.py               | 32 ++++++++------------
 4 files changed, 15 insertions(+), 23 deletions(-)

diff --git a/backend/infrahub/core/query/branch.py b/backend/infrahub/core/query/branch.py
index 15716e130a..a357d29e71 100644
--- a/backend/infrahub/core/query/branch.py
+++ b/backend/infrahub/core/query/branch.py
@@ -113,7 +113,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
 
         self.params["at"] = self.at.to_string()
         self.params["ids"] = [element_id_to_id(id) for id in self.ids]
-        self.return_labels = ["ID(r)"]
+        self.return_labels = [f"{db.get_id_function_name}(r)"]
 
 
 class RebaseBranchDeleteRelationshipQuery(Query):
diff --git a/backend/infrahub/core/query/node.py b/backend/infrahub/core/query/node.py
index 1e6ae51565..904700a434 100644
--- a/backend/infrahub/core/query/node.py
+++ b/backend/infrahub/core/query/node.py
@@ -794,7 +794,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
         self.order_by = []
         self.params["node_kind"] = self.schema.kind
 
-        self.return_labels = ["n.uuid", "rb.branch", "ID(rb) as rb_id"]
+        self.return_labels = ["n.uuid", "rb.branch", f"{db.get_id_function_name()}(rb) as rb_id"]
         where_clause_elements = []
 
         branch_filter, branch_params = self.branch.get_query_filter_path(
diff --git a/backend/infrahub/core/query/standard_node.py b/backend/infrahub/core/query/standard_node.py
index 2ab42297e5..e36d4d5ac0 100644
--- a/backend/infrahub/core/query/standard_node.py
+++ b/backend/infrahub/core/query/standard_node.py
@@ -158,4 +158,4 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
         self.add_to_query(query)
 
         self.return_labels = ["n"]
-        self.order_by = ["ID(n)"]
+        self.order_by = [f"{db.get_id_function_name()}(n)"]
diff --git a/backend/infrahub/core/utils.py b/backend/infrahub/core/utils.py
index 43f81a26b2..45b16bb7eb 100644
--- a/backend/infrahub/core/utils.py
+++ b/backend/infrahub/core/utils.py
@@ -65,18 +65,16 @@ async def update_relationships_to(ids: list[str], db: InfrahubDatabase, to: Time
     if not ids:
         return None
 
-    list_matches = [f"id(r) = {element_id_to_id(id)}" for id in ids]
-
     to = Timestamp(to)
 
-    query = f"""
+    query = """
     MATCH ()-[r]->()
-    WHERE {' or '.join(list_matches)}
+    WHERE %(id_func)s(r) IN $ids
     SET r.to = $to
-    RETURN ID(r)
-    """
+    RETURN %(id_func)s(r)
+    """ % {"id_func": db.get_id_function_name()}
 
-    params = {"to": to.to_string()}
+    params = {"to": to.to_string(), "ids": [element_id_to_id(_id) for _id in ids]}
 
     return await db.execute_query(query=query, params=params, name="update_relationships_to")
 
@@ -98,13 +96,10 @@ async def get_paths_between_nodes(
         relationships_str = ":" + "|".join(relationships)
 
     query = """
-    MATCH p = (s)-[%s*%s]-(d)
-    WHERE ID(s) = $source_id AND ID(d) = $destination_id
+    MATCH p = (s)-[%(rel)s*%(length_limit)s]-(d)
+    WHERE %(id_func)s(s) = $source_id AND %(id_func)s(d) = $destination_id
     RETURN p
-    """ % (
-        relationships_str.upper(),
-        length_limit,
-    )
+    """ % {"rel": relationships_str.upper(), "length_limit": length_limit, "id_func": db.get_id_function_name()}
 
     if print_query:
         print(query)
@@ -170,14 +165,11 @@ async def delete_all_nodes(db: InfrahubDatabase):
     return await db.execute_query(query=query, params=params, name="delete_all_nodes")
 
 
-def element_id_to_id(element_id: Union[str, int]) -> int:
-    if isinstance(element_id, int):
-        return element_id
-
-    if isinstance(element_id, str) and ":" not in element_id:
+def element_id_to_id(element_id: str | int) -> str | int:
+    try:
         return int(element_id)
-
-    return int(element_id.split(":")[2])
+    except ValueError:
+        return element_id
 
 
 def extract_field_filters(field_name: str, filters: dict) -> dict[str, Any]:

From 7349c70f456fe06f7672c4496bdfbf1c1dc428fc Mon Sep 17 00:00:00 2001
From: Aaron McCarty <aaron@opsmill.com>
Date: Wed, 18 Sep 2024 17:57:32 -0700
Subject: [PATCH 3/5] actually call function

---
 backend/infrahub/core/query/branch.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/backend/infrahub/core/query/branch.py b/backend/infrahub/core/query/branch.py
index a357d29e71..f79495e076 100644
--- a/backend/infrahub/core/query/branch.py
+++ b/backend/infrahub/core/query/branch.py
@@ -113,7 +113,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
 
         self.params["at"] = self.at.to_string()
         self.params["ids"] = [element_id_to_id(id) for id in self.ids]
-        self.return_labels = [f"{db.get_id_function_name}(r)"]
+        self.return_labels = [f"{db.get_id_function_name()}(r)"]
 
 
 class RebaseBranchDeleteRelationshipQuery(Query):

From aafaef6aaca353cb4206affac13636f604ba3b8d Mon Sep 17 00:00:00 2001
From: Aaron McCarty <aaron@opsmill.com>
Date: Wed, 18 Sep 2024 18:21:19 -0700
Subject: [PATCH 4/5] move element_id_to_id somewhere a little more appropriate

---
 backend/infrahub/core/query/branch.py       |  7 +++----
 backend/infrahub/core/query/relationship.py |  4 ++--
 backend/infrahub/core/utils.py              | 17 +++++------------
 backend/infrahub/database/__init__.py       |  8 ++++++++
 backend/tests/unit/core/test_core_utils.py  |  7 -------
 5 files changed, 18 insertions(+), 25 deletions(-)

diff --git a/backend/infrahub/core/query/branch.py b/backend/infrahub/core/query/branch.py
index f79495e076..4b47b0e301 100644
--- a/backend/infrahub/core/query/branch.py
+++ b/backend/infrahub/core/query/branch.py
@@ -5,7 +5,6 @@
 from infrahub import config
 from infrahub.core.constants import RelationshipStatus
 from infrahub.core.query import Query, QueryType
-from infrahub.core.utils import element_id_to_id
 
 if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
@@ -32,7 +31,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
             "id_func": db.get_id_function_name(),
         }
 
-        self.params["node_id"] = element_id_to_id(self.node_id)
+        self.params["node_id"] = db.to_database_id(self.node_id)
         self.params["now"] = self.at.to_string()
         self.params["branch"] = self.branch.name
         self.params["branch_level"] = self.branch.hierarchy_level
@@ -112,7 +111,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
         self.add_to_query(query=query)
 
         self.params["at"] = self.at.to_string()
-        self.params["ids"] = [element_id_to_id(id) for id in self.ids]
+        self.params["ids"] = [db.to_database_id(id) for id in self.ids]
         self.return_labels = [f"{db.get_id_function_name()}(r)"]
 
 
@@ -150,4 +149,4 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
 
         self.add_to_query(query=query)
 
-        self.params["ids"] = [element_id_to_id(id) for id in self.ids]
+        self.params["ids"] = [db.to_database_id(id) for id in self.ids]
diff --git a/backend/infrahub/core/query/relationship.py b/backend/infrahub/core/query/relationship.py
index fab5ce5927..7cd75fef6f 100644
--- a/backend/infrahub/core/query/relationship.py
+++ b/backend/infrahub/core/query/relationship.py
@@ -11,7 +11,7 @@
 from infrahub.core.query import Query, QueryType
 from infrahub.core.query.subquery import build_subquery_filter, build_subquery_order
 from infrahub.core.timestamp import Timestamp
-from infrahub.core.utils import element_id_to_id, extract_field_filters
+from infrahub.core.utils import extract_field_filters
 
 if TYPE_CHECKING:
     from uuid import UUID
@@ -395,7 +395,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs):
                 "MATCH (prop_%(prop_name)s) WHERE %(id_func)s(prop_%(prop_name)s) = $prop_%(prop_name)s_id"
                 % {"prop_name": prop_name, "id_func": db.get_id_function_name()}
             )
-            self.params[f"prop_{prop_name}_id"] = element_id_to_id(prop.prop_db_id)
+            self.params[f"prop_{prop_name}_id"] = db.to_database_id(prop.prop_db_id)
             self.return_labels.append(f"prop_{prop_name}")
 
         self.params["rel_prop"] = self.get_relationship_properties_dict(status=RelationshipStatus.DELETED)
diff --git a/backend/infrahub/core/utils.py b/backend/infrahub/core/utils.py
index 45b16bb7eb..55d6fe9ee7 100644
--- a/backend/infrahub/core/utils.py
+++ b/backend/infrahub/core/utils.py
@@ -37,8 +37,8 @@ async def add_relationship(
     at = Timestamp(at)
 
     params = {
-        "src_node_id": element_id_to_id(src_node_id),
-        "dst_node_id": element_id_to_id(dst_node_id),
+        "src_node_id": db.to_database_id(src_node_id),
+        "dst_node_id": db.to_database_id(dst_node_id),
         "at": at.to_string(),
         "branch": branch_name or registry.default_branch,
         "branch_level": branch_level or 1,
@@ -74,7 +74,7 @@ async def update_relationships_to(ids: list[str], db: InfrahubDatabase, to: Time
     RETURN %(id_func)s(r)
     """ % {"id_func": db.get_id_function_name()}
 
-    params = {"to": to.to_string(), "ids": [element_id_to_id(_id) for _id in ids]}
+    params = {"to": to.to_string(), "ids": [db.to_database_id(_id) for _id in ids]}
 
     return await db.execute_query(query=query, params=params, name="update_relationships_to")
 
@@ -105,8 +105,8 @@ async def get_paths_between_nodes(
         print(query)
 
     params = {
-        "source_id": element_id_to_id(source_id),
-        "destination_id": element_id_to_id(destination_id),
+        "source_id": db.to_database_id(source_id),
+        "destination_id": db.to_database_id(destination_id),
     }
 
     return await db.execute_query(query=query, params=params, name="get_paths_between_nodes")
@@ -165,13 +165,6 @@ async def delete_all_nodes(db: InfrahubDatabase):
     return await db.execute_query(query=query, params=params, name="delete_all_nodes")
 
 
-def element_id_to_id(element_id: str | int) -> str | int:
-    try:
-        return int(element_id)
-    except ValueError:
-        return element_id
-
-
 def extract_field_filters(field_name: str, filters: dict) -> dict[str, Any]:
     """Extract the filters for a given field (attribute or relationship) from a filters dict."""
     return {
diff --git a/backend/infrahub/database/__init__.py b/backend/infrahub/database/__init__.py
index 347a616ff9..b5e59977cd 100644
--- a/backend/infrahub/database/__init__.py
+++ b/backend/infrahub/database/__init__.py
@@ -341,6 +341,14 @@ def get_id_function_name(self) -> str:
             return "elementId"
         return "ID"
 
+    def to_database_id(self, db_id: str | int) -> str | int:
+        if self.db_type == DatabaseType.NEO4J:
+            return db_id
+        try:
+            return int(db_id)
+        except ValueError:
+            return db_id
+
 
 async def create_database(driver: AsyncDriver, database_name: str) -> None:
     default_db = driver.session()
diff --git a/backend/tests/unit/core/test_core_utils.py b/backend/tests/unit/core/test_core_utils.py
index 086447dcd4..d33892dbd8 100644
--- a/backend/tests/unit/core/test_core_utils.py
+++ b/backend/tests/unit/core/test_core_utils.py
@@ -4,7 +4,6 @@
 from infrahub.core.utils import (
     count_relationships,
     delete_all_nodes,
-    element_id_to_id,
     get_paths_between_nodes,
     parse_node_kind,
 )
@@ -15,12 +14,6 @@ async def test_delete_all_nodes(db: InfrahubDatabase):
     assert await delete_all_nodes(db) == []
 
 
-def test_element_id_to_id():
-    assert element_id_to_id("4:c0814fa2-df5b-4d66-ba5f-9a01817f16fb:167") == 167
-    assert element_id_to_id("198") == 198
-    assert element_id_to_id(167) == 167
-
-
 def test_parse_node_kind():
     assert parse_node_kind(kind="TestMyModel") == NodeKind(namespace="Test", name="MyModel")
     assert parse_node_kind(kind="Test3Myname1234") == NodeKind(namespace="Test3", name="Myname1234")

From a73675c15fc64b665c6794314d7eeee81d76f652 Mon Sep 17 00:00:00 2001
From: Aaron McCarty <aaron@opsmill.com>
Date: Wed, 18 Sep 2024 18:38:38 -0700
Subject: [PATCH 5/5] add new method to database protocol

---
 backend/infrahub/core/protocols_base.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/backend/infrahub/core/protocols_base.py b/backend/infrahub/core/protocols_base.py
index ea8cd38d00..31fc37f277 100644
--- a/backend/infrahub/core/protocols_base.py
+++ b/backend/infrahub/core/protocols_base.py
@@ -61,6 +61,7 @@ def render_list_comprehension(self, items: str, item_name: str) -> str: ...
     def render_list_comprehension_with_list(self, items: str, item_names: list[str]) -> str: ...
     def render_uuid_generation(self, node_label: str, node_attr: str) -> str: ...
     def get_id_function_name(self) -> str: ...
+    def to_database_id(self, db_id: str | int) -> str | int: ...
 
 
 @runtime_checkable