Merge pull request #20 from arangoml/query-options

Introduce AQL query options to create_networkx_graph
arangoml · Oct 7, 2021 · c62bd09 · c62bd09
2 parents 75e37d9 + b00e778
commit c62bd09
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -43,6 +43,9 @@ attributes = { 'vertexCollections':
 # Export networkX graph                                  
 g = ma.create_networkx_graph(graph_name = 'FraudDetection',  graph_attributes = attributes)
 
+# You can also provide valid Python-Arango AQL query options to the command above, like such:
+# g = ma.create_networkx_graph(graph_name = 'FraudDetection',  graph_attributes = attributes, ttl=1000, stream=True)
+
 # Use networkX
 nx.draw(g, with_labels=True)
 ```

diff --git a/adbnx_adapter/adbnx_adapter/arangoDB_networkx_adapter.py b/adbnx_adapter/adbnx_adapter/arangoDB_networkx_adapter.py
@@ -66,7 +66,7 @@ def is_valid_graph_attributes(self, graph_config):
 
         return valid_config
 
-    def create_networkx_graph(self, graph_name, graph_attributes):
+    def create_networkx_graph(self, graph_name, graph_attributes, **query_options):
 
         if self.is_valid_graph_attributes(graph_attributes):
             g = nx.DiGraph()
@@ -77,7 +77,7 @@ def create_networkx_graph(self, graph_name, graph_attributes):
                 csps = ','.join(cspl)
                 query = query + "RETURN { " + csps + "}"
 
-                cursor = self.db.aql.execute(query)
+                cursor = self.db.aql.execute(query, **query_options)
                 for doc in cursor:
                     g.add_node(doc['_id'], attr_dict=doc)
 
@@ -88,7 +88,7 @@ def create_networkx_graph(self, graph_name, graph_attributes):
                 csps = ','.join(cspl)
                 query = query + "RETURN { " + csps + "}"
 
-                cursor = self.db.aql.execute(query)
+                cursor = self.db.aql.execute(query, **query_options)
                 # breakpoint()
                 for doc in cursor:
                     g.add_edge(doc['_from'], doc['_to'])

diff --git a/adbnx_adapter/adbnx_adapter/dgl_arangoDB_networkx_adapter.py b/adbnx_adapter/adbnx_adapter/dgl_arangoDB_networkx_adapter.py
@@ -16,7 +16,7 @@
 
 class DGLArangoDB_Networkx_Adapter(ArangoDB_Networkx_Adapter):
 
-    def create_networkx_graph(self, graph_name, graph_attributes):
+    def create_networkx_graph(self, graph_name, graph_attributes, **query_options):
 
         if self.is_valid_graph_attributes(graph_attributes):
             edge_names = []
@@ -45,7 +45,7 @@ def create_networkx_graph(self, graph_name, graph_attributes):
                 ens = k.split('_', 1)
                 redge = ens[1] + '_' + ens[0]
                 rsgraph = rsgdata[redge]
-                cursor = self.db.aql.execute(query)
+                cursor = self.db.aql.execute(query, **query_options)
                 for doc in cursor:
                     nfrom = doc['_from']
                     nto = doc['_to']
@@ -67,7 +67,7 @@ def create_networkx_graph(self, graph_name, graph_attributes):
                 csps = ','.join(cspl)
                 query = query + "RETURN { " + csps + "}"
 
-                cursor = self.db.aql.execute(query)
+                cursor = self.db.aql.execute(query, **query_options)
                 for doc in cursor:
                     exclude_attr = ['_id', '_key', 'node_id']
                     if k == 'incident':

diff --git a/adbnx_adapter/adbnx_adapter/imdb_arangoDB_networkx_adapter.py b/adbnx_adapter/adbnx_adapter/imdb_arangoDB_networkx_adapter.py
@@ -13,7 +13,7 @@
 
 class IMDBArangoDB_Networkx_Adapter(ArangoDB_Networkx_Adapter):
 
-    def create_networkx_graph(self, graph_name, graph_attributes):
+    def create_networkx_graph(self, graph_name, graph_attributes, **query_options):
 
         if self.is_valid_graph_attributes(graph_attributes):
             g = nx.DiGraph()
@@ -24,7 +24,7 @@ def create_networkx_graph(self, graph_name, graph_attributes):
                 csps = ','.join(cspl)
                 query = query + "RETURN { " + csps + "}"
 
-                cursor = self.db.aql.execute(query)
+                cursor = self.db.aql.execute(query, **query_options)
                 for doc in cursor:
                     if k == "Users":
                         bip_key = 0
@@ -39,7 +39,7 @@ def create_networkx_graph(self, graph_name, graph_attributes):
                 csps = ','.join(cspl)
                 query = query + "RETURN { " + csps + "}"
 
-                cursor = self.db.aql.execute(query)
+                cursor = self.db.aql.execute(query, **query_options)
 
                 for doc in cursor:
                     g.add_edge(doc['_from'], doc['_to'])