synapse->neo4j

synapse2neo4j.py

from py2neo import Graph, Node, Relationship

from synapse.events import FrozenEvent

import json
import sqlite3
import sys


def main(db_name, room_id):
    conn = sqlite3.connect(db_name)

    sql = (
        "SELECT json FROM event_json as j"
        " INNER JOIN events as e ON e.event_id = j.event_id"
        " WHERE j.room_id = ?"
        " ORDER BY topological_ordering ASC"
    )

    c = conn.cursor()

    print "Loading from DB"

    c.execute(sql, (room_id,))

    events = c.fetchall()

    print "Loaded from DB"

    events[:] = [
        FrozenEvent(json.loads(e))
        for e, in events
    ]

    print "Parsed to FrozenEvents"

    graph = Graph()

    i = 0

    tx = graph.cypher.begin()

    create_cypher = "CREATE (n:Event {event_id:{id}, room_id:{room}, sender:{sender}, type:{type}}) return n"
    state_cypher = "CREATE (n:Event {event_id:{id}, room_id:{room}, sender:{sender}, type:{type}, state_key:{state}}) return n"

    rel_cypher = (
        "MATCH (a:Event), (b:Event)"
        " WHERE a.event_id = {id1} AND b.event_id = {id2}"
        " CREATE (a)-[r:%s]->(b)"
        " RETURN r"
    )

    nodes = set()
    for e in events:
        if i % 100 == 0:
            print "Sending..."
            tx.process()
            print i, len(events)
        i += 1

        if e.is_state():
            tx.append(create_cypher, {
                "id": e.event_id,
                "room": e.room_id,
                "sender": e.sender,
                "type": e.type,
                "state": e.state_key,
            })
        else:
            tx.append(create_cypher, {
                "id": e.event_id,
                "room": e.room_id,
                "sender": e.sender,
                "type": e.type,
            })

        for attr, label in (("auth_events", "AUTH"), ("prev_events", "REFERENCES"),):
            for e_id, e_hash in getattr(e, attr):
                if e_id in nodes:
                    tx.append(rel_cypher % (label,), {
                        "id1": e.event_id,
                        "id2": e_id,
                    })

        nodes.add(e.event_id)
    tx.commit()


if __name__ == "__main__":
    db_name, room_id, = sys.argv[1:]

    main(db_name, room_id)