Exporting Neo4J Aura graph into iGraph - python

neo4j_aura_to_igraph.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Jan 21 16:23:33 2023

@author: aas358
"""


from neo4j import GraphDatabase
from igraph import Graph
import pickle

class NEO4J_Connector:

    def __init__(self, uri, user, password):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))
        self.graph_nodes = list()
        self.graph_links = list()

    def close(self):
        # Don't forget to close the driver connection when you are finished with it
        self.driver.close()

            
    def retrieve_graph(self)->None:
        with self.driver.session(database="neo4j") as session:
            self.t_graph = session.execute_read(self._retrieve_graph)
            
            self.__process_nodes()
            self.__process_links()

    def __process_nodes(self)->None:
        for graph_node in self.t_graph.nodes:
            _node = dict()
            _node["label"] = list(set(graph_node.labels))[0]
            _node["id"] = graph_node.element_id
            for k, v in graph_node._properties.items():
                _node[k] = v 
            
            self.graph_nodes.append(_node)
            
    def __process_links(self)->None:
        for graph_link in self.t_graph.relationships:
            _link = dict()
            _link["type"] = graph_link.type
            _link["startNode"] = graph_link.nodes[0].element_id
            _link["endNode"] = graph_link.nodes[1].element_id
            for k, v in graph_link._properties.items():
                _link[k] = v 
            
            self.graph_links.append(_link)
        
    
    def get_nodes_and_links(self)->(list,list):
        return self.graph_nodes, self.graph_links
    
    @staticmethod
    def _retrieve_graph(tx):
        query = (
            "MATCH (head) OPTIONAL MATCH (head)-[relation]->(tail) RETURN head, relation, tail"
        )
        result = tx.run(query)
        return result.graph()
     

class IGRAPH:
    
    def __init__(self):
        self.g = Graph(directed=True)
        self.node_correspondeces = dict()
        
    def create_graph(self, graph_nodes_l:list, graph_links_l:list):
        ## ADD NODES
        for idx, node in enumerate(graph_nodes_l):
            self.g.add_vertices(1, node)
            self.node_correspondeces[node["id"]] = idx

        ## ADD LINKS
        for link in graph_links_l:
            self.g.add_edge(self.node_correspondeces[link["startNode"]], 
                            self.node_correspondeces[link["endNode"]], 
                            **link)
            
    def get_graph(self)->Graph:
        return self.g


if __name__ == "__main__":
    # Aura queries use an encrypted connection using the "neo4j+s" URI scheme    
    uri = "neo4j+s://XXXXXXXX.databases.neo4j.io"
    user = "neo4j"
    password = "XxXxXxXxXxX"
    n4jconn = NEO4J_Connector(uri, user, password)
    
    
    ######## RETRIEVE GRAPH FROM NEO4J AURA
    t_graph = n4jconn.retrieve_graph()
    
    graph_nodes, graph_links = n4jconn.get_nodes_and_links()
    
    n4jconn.close()
    
    
    ######## BACKUP OF NODES AND LINKS
    with open('simple_backup_of_graph.pkl', 'wb') as file:
        pickle.dump({"nodes": graph_nodes, "links":graph_links}, file)
    
    
    ######## BUILDING IGRAPH OBJECT
    igraph = IGRAPH()
    
    igraph.create_graph(graph_nodes, graph_links)
    
    g = igraph.get_graph()
    
    g.write_graphml("graph.graphml")