Skip to content

Instantly share code, notes, and snippets.

@ikasat ikasat/neo4j_dep.py
Last active Jul 20, 2018

Embed
What would you like to do?
#!/usr/bin/env python3
import re
import sys
from collections import namedtuple
from neo4j.v1 import GraphDatabase
Config = namedtuple('Config', ['url', 'username', 'password',
'node_label', 'rel_type', 'batch_size'])
VALID_NAME_RE = re.compile(r'^[A-Za-z][A-Za-z0-9_]*$')
def create_records(dep_dict, config):
db = GraphDatabase.driver(config.url, auth=(config.username, config.password))
with db.session() as session:
create_records_with_session(session, dep_dict, config)
def create_records_with_session(session, dep_dict, config):
if VALID_NAME_RE.match(config.node_label) is None:
raise ValueError('invalid Node label: ' + config.node_label)
if VALID_NAME_RE.match(config.rel_type) is None:
raise ValueError('invalid Relationship type: ' + config.rel_type)
session.run('CREATE INDEX ON :{node_label}({rel_type})'
.format(node_label=config.node_label, rel_type=config.rel_type))
tx = session.begin_transaction()
index = 0
mutation = '''
MERGE (src:{node_label} {{name: $src}})
MERGE (dest:{node_label} {{name: $dest}})
MERGE (src)-[dep:{rel_type}]->(dest)
'''.format(node_label=config.node_label, rel_type=config.rel_type)
for src, dests in dep_dict.items():
for dest in dests:
tx.run(mutation, src=src, dest=dest)
if index % config.batch_size == config.batch_size - 1:
print('{}: Committing...'.format(index + 1), file=sys.stderr)
tx.commit()
tx = session.begin_transaction()
index += 1
tx.commit()
def main():
import argparse
import json
import os
parser = argparse.ArgumentParser()
parser.add_argument('-U', '--url', default='bolt://localhost:7687')
parser.add_argument('-u', '--username')
parser.add_argument('-p', '--password')
parser.add_argument('-b', '--batchsize', type=int, default=10000)
parser.add_argument('-n', '--nodelabel', default='Package')
parser.add_argument('-r', '--reltype', default='PDEPENDS')
parser.add_argument('file', nargs='*')
args = parser.parse_args()
if args.username:
username = args.username
else:
username = os.getenv('NEO4J_USERNAME', 'neo4j')
if args.password:
password = args.password
else:
password = os.getenv('NEO4J_PASSWORD', 'neo4j')
config = Config(url=args.url, username=username, password=password,
node_label=args.nodelabel, rel_type=args.reltype,
batch_size=args.batchsize)
db = GraphDatabase.driver(config.url, auth=(config.username, config.password))
with db.session() as session:
if args.file:
for fpath in args.file:
with open(fpath) as f:
dep_dict = json.load(f)
create_records_with_session(session, dep_dict, config)
else:
dep_dict = json.load(sys.stdin)
create_records_with_session(session, dep_dict, config)
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.