memgraph/tools/neo_to_memgraph.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

'''
A script for transfering all data from a Neo4j database
into a Memgraph database.
'''

import logging
import json
import os
from time import time
from datetime import datetime
from argparse import ArgumentParser

from neo4j.v1 import GraphDatabase, basic_auth

log = logging.getLogger(__name__)


TEMP_ID = "__memgraph_temp_id_314235423"
TEMP_LABEL = "__memgraph_temp_label_1414213562"


def parse_args():
    argp = ArgumentParser(description=__doc__)
    argp.add_argument("--neo-url", default = "127.0.0.1:7687",
                      help = "Neo4j url, default 127.0.0.1:7687.")
    argp.add_argument("--neo-user", default = "neo4j",
                      help = "Neo4j username, default neo4j.")
    argp.add_argument("--neo-password", default = "1234",
                      help = "Neo4j password, default 1234.")
    argp.add_argument("--neo-ssl", default = False, choices = [True, False],
                      help = "Encryption for neo4j auth data, default False")
    argp.add_argument("--memgraph-url", default = "127.0.0.1:7688",
                      help = "Memgraph url, default 127.0.0.1:7688.")
    argp.add_argument("--logging", default = "DEBUG", choices = ["INFO", "DEBUG"],
                      help = "Logging level, default debug.") 
    argp.add_argument("--json-storage", default = "json",
                      help = "Storage for JSON files.")
    return argp.parse_args()


def create_vertex_cypher(vertex):
    """
    Helper function that generates a cypher query for creting
    a vertex based on the given Bolt vertex.
    """
    labels = ""
    if vertex.labels:
        labels += ":" + ":".join(vertex.labels)
    vertex.properties[TEMP_ID] = vertex.id
    properties = ", ".join('%s: %r' % kv for kv
                           in vertex.properties.items())
    return "CREATE (%s {%s})" % (labels, properties)


def create_edge_cypher(edge, edge_num):
    """
    Helper function that generates a cypher query for creting
    a edge based on the given Bolt edge.
    """
    properties = ", ".join('%s: %r' % kv for kv
                           in edge.properties.items())
    return ["(from%s {%s: %r}), (to%s {%s: %r})" \
                % (edge_num, TEMP_ID, edge.start, edge_num, TEMP_ID, edge.end),
            "CREATE (from%s)-[:%s {%s}]->(to%s)" \
                % (edge_num, edge.type, properties, edge_num)]

def vertex_to_dict(vertex):
    """
    Returns dictionary which represents one vertex with labels and
    properties.

    :param vertex: Graph vertex
    """
    return {"labels": list(vertex.labels), "properties": vertex.properties}

def edge_to_dict(edge):
    """
    Returns dictionary which represents one edge with type, start vertex,
    end vertex and properties.

    :param edge: Graph edge
    """
    return {"type": edge.type, "properties": edge.properties, \
            "start": edge.start, "end": edge.end}

def create_json_file(storage, timestamp, element, batch_index, content):
    """
    Creates json file with given content and path
    storage/timestamp/element and file name is batch_count.json. Creates
    directories where files are stored if directories don't exist.

    :param storage: str, path where all json files are stored
    :param timestamp: str, timestamp of the current transfer
    :param element: str, expected vertex or edge, which elements are
                    stored in json
    :param batch_index: int, index of the current batch, used in file name
    :param content: list, contet which will be dumped in file
    """
    json_file = os.path.join(storage, timestamp, element,
                        str(batch_index) + ".json")
    os.makedirs(os.path.dirname(json_file), exist_ok = True)
    print(content)
    with open(json_file, 'w') as f:
        json.dump(content, f, indent = 2)


def transfer(storage, neo_driver, memgraph_driver):
    """ Copies all the data from Neo4j to Memgraph. """

    # TODO add error handling
    neo_session = neo_driver.session()
    memgraph_session = memgraph_driver.session()

    # Creating index
    log.debug("Creating memgraph index on TEMP_LABEL and TEMP_ID.")
    memgraph_session.run("CREATE INDEX ON :%s(%s)" % (TEMP_LABEL, TEMP_ID))
    neo_session.run("MATCH(n) SET n :%s, n.%s = ID(n)" % (TEMP_LABEL, TEMP_ID))
    neo_session.run("CREATE INDEX ON :%s(%s)" % (TEMP_LABEL, TEMP_ID))

    read_vertex_batch = 2
    write_vertex_batch = 3
    read_edge_batch = 2
    write_edge_batch = 3
    vertex_count = 0
    edge_count = 0

    cypher_query = ""
    batch_count = 0
    timestamp = datetime.fromtimestamp(time()).strftime("%Y_%m_%d__%H_%M_%S")

    def write_vertices(vertices):
        nonlocal batch_count
        cypher_query = ""
        vertices_list = []
        for vertex in vertices:
            cypher_query += create_vertex_cypher(vertex)
            vertices_list.append(vertex_to_dict(vertex))
        create_json_file(storage, timestamp, "vertices", batch_count, vertices_list)
        log.debug("Vertex create on cypher: %s" % (cypher_query))
        memgraph_session.run(cypher_query).consume()
        batch_count += 1
        vertices[:] = []

    def write_edges(edges):
        nonlocal batch_count
        cypher_query = ""
        edge_num = 0
        edges_list = []
        for edge in edges:
            edges_list.append(edge_to_dict(edge))
            edge_queries = create_edge_cypher(edge, edge_num)
            edge_num += 1
            if cypher_query:
                cypher_query = edge_queries[0] + ", " + cypher_query + \
                        " " + edge_queries[1]
            else:
                cypher_query = ' '.join(edge_queries)
        create_json_file(storage, timestamp, "edges", batch_count, edges_list)
        cypher_query = "MATCH " + cypher_query
        log.debug("Edge create on cypher: %s" % (cypher_query))
        memgraph_session.run(cypher_query).consume()
        batch_count += 1
        edges[:] = []

    # Vertex transfer
    start_id = 0
    vertices_batch = []
    while True:
        read_vertices_in_batch = 0
        vertices = neo_session.run("MATCH(n) WHERE n.%s>=%s RETURN n " \
                      "ORDER BY ID(n) LIMIT %s" % (TEMP_ID, start_id, read_vertex_batch))
        for vertex in vertices:
            vertex = vertex['n']
            vertices_batch.append(vertex)
            if len(vertices_batch) >= write_vertex_batch:
                write_vertices(vertices_batch)
            start_id = vertex.id
            read_vertices_in_batch += 1
            vertex_count += 1
        start_id += 1
        if read_vertices_in_batch != read_vertex_batch:
            break
    if len(vertices_batch) > 0:
        write_vertices(vertices_batch)

    max_id = neo_session.run("MATCH(n) RETURN MAX(ID(n)) AS id").peek()['id']

    start_id = 0
    batch_count = 0
    edges_batch = []
    while start_id <= max_id:
        edges = neo_session.run("MATCH (n)-[r]->() WHERE n.%s>=%s AND " \
              "n.%s<%s RETURN r" % (TEMP_ID, start_id, TEMP_ID, start_id + read_edge_batch))
        start_id = start_id + read_edge_batch
        for edge in edges:
            edge_count += 1
            edge = edge['r']
            edges_batch.append(edge)
            if len(edges_batch) >= write_edge_batch:
                write_edges(edges_batch)
    if len(edges_batch) > 0:
        write_edges(edges_batch)

    # TODO Drop index in memgraph when it will be supported
    log.debug("Removing TEMP_LABEL and TEMP_ID")
    memgraph_session.run("MATCH (n) REMOVE n:%s, n.%s" % (TEMP_LABEL, TEMP_ID))
    neo_session.run("MATCH (n) REMOVE n:%s, n.%s" % (TEMP_LABEL, TEMP_ID))
    neo_session.run("DROP INDEX ON :%s(%s)" % (TEMP_LABEL, TEMP_ID))
    log.info("Created %d vertiecs and %d edges", vertex_count, edge_count)


def main():
    args = parse_args()
    if args.logging:
        logging.basicConfig(level=args.logging)
        logging.getLogger("neo4j").setLevel(logging.WARNING)

    log.info("Memgraph from Neo4j data import tool")

    neo_driver = GraphDatabase.driver(
        "bolt://" + args.neo_url,
        auth=basic_auth(args.neo_user, args.neo_password),
        encrypted=args.neo_ssl)
    memgraph_driver = GraphDatabase.driver(
        "bolt://" + args.memgraph_url,
        auth=basic_auth("", ""),
        encrypted=False)

    start_time = time()
    transfer(args.json_storage, neo_driver, memgraph_driver)
    log.info("Import complete in %.2f seconds", time() - start_time)

    pass

if __name__ == '__main__':
    main()