From 56e125a6fc09c60d4c2401faa28b199927ebc499 Mon Sep 17 00:00:00 2001
From: Chang Liu <liuc223@gmail.com>
Date: Mon, 10 Jul 2023 11:44:15 -0400
Subject: [PATCH] Add recursive rel support for networkx

---
 tools/python_api/src_py/query_result.py | 16 +++++++++++++++-
 tools/python_api/src_py/types.py        |  1 +
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/tools/python_api/src_py/query_result.py b/tools/python_api/src_py/query_result.py
index 5954e43d13..133484d555 100644
--- a/tools/python_api/src_py/query_result.py
+++ b/tools/python_api/src_py/query_result.py
@@ -278,6 +278,20 @@ def encode_node_id(node, table_primary_key_dict):
                     rels[(_src["table"], _src["offset"], _dst["table"],
                           _dst["offset"])] = row[i]
 
+                elif column_type == Type.RECURSIVE_REL.value:
+                    for node in row[i]['_nodes']:
+                        _id = node["_id"]
+                        nodes[(_id["table"], _id["offset"])] = node
+                        table_to_label_dict[_id["table"]] = node["_label"]
+                    for rel in row[i]['_rels']:
+                        for key in rel:
+                            if rel[key] is None:
+                                del rel[key]
+                        _src = rel["_src"]
+                        _dst = rel["_dst"]
+                        rels[(_src["table"], _src["offset"], _dst["table"],
+                              _dst["offset"])] = rel
+
         # Add nodes
         for node in nodes.values():
             _id = node["_id"]
@@ -313,7 +327,7 @@ def _get_properties_to_extract(self):
         for i in range(len(column_names)):
             column_name = column_names[i]
             column_type = column_types[i]
-            if column_type in [Type.NODE.value, Type.REL.value]:
+            if column_type in [Type.NODE.value, Type.REL.value, Type.RECURSIVE_REL.value]:
                 properties_to_extract[i] = (column_type, column_name)
         return properties_to_extract
 
diff --git a/tools/python_api/src_py/types.py b/tools/python_api/src_py/types.py
index a259a9d390..5e6ea38a1d 100644
--- a/tools/python_api/src_py/types.py
+++ b/tools/python_api/src_py/types.py
@@ -18,3 +18,4 @@ class Type(Enum):
     REL = "REL"
     NODE_ID = "NODE_ID"
     STRUCT = "STRUCT"
+    RECURSIVE_REL = "RECURSIVE_REL"