neo4j importer

neo4j-importer.py

import os
import json

def dict_to_str(props, keys, DELIMITER, QUOTE):
    s = []
    
    for key in keys:
        value = props.get(key)
        
        if value:
            if type(value) is str:
                value = f'{QUOTE}{value}{QUOTE}'
                
            s.append(f'{value}')
        else:
            s.append('')
            
    return DELIMITER.join(s)

CSV_DIR = 'csv'
NODES_DIR = os.path.join(CSV_DIR, 'nodes')
ROLES_DIR = os.path.join(CSV_DIR, 'roles')
DELIMITER = ';'
ARRAY_DELIMITER = '|'
QUOTE = '"'

for d in (CSV_DIR, NODES_DIR, ROLES_DIR):
    os.makedirs(d, exist_ok=True)

with open('dump.json', 'r') as fh:
    graph = json.load(fh)
    
nodes_by_label = {}
properties_by_label = {}
import_command = ['neo4j-admin', 'import']
    
for node in graph['nodes']:
    node = node['data']
    label = node['label']
    
    nodes = nodes_by_label.setdefault(label, [])
    nodes.append(node)
    
    keys = node['props'].keys()
    properties = properties_by_label.get(label, [])
    
    if len(keys) > len(properties):
        properties_by_label[label] = list(keys)
    
for label, nodes in nodes_by_label.items():
    # create a csv file
    #with open(os.path.join(CSV_DIR, f'{label}.csv'), 'w') as fh:
    #print(label, properties_by_label[label])
    prop_keys = properties_by_label[label]
    path = os.path.join(NODES_DIR, f'{label}.csv')
    
    with open(path, 'w') as fh:
        fh.write(DELIMITER.join(['id:ID'] + prop_keys))
        fh.write('\n')
        
        for node in nodes:
            fh.write(f'{node["id"]}{DELIMITER}{dict_to_str(node["props"], prop_keys, DELIMITER, QUOTE)}')
            fh.write('\n')
    
    import_command.extend([f'--nodes:{label}', path])
    
roles_by_type = {}
properties_by_label = {}
        
for edge in graph['relationships']:
    edge = edge['data']
    label = edge['label']
    
    roles = roles_by_type.setdefault(label, [])
    roles.append(edge)
    
    keys = edge['props'].keys()
    properties = properties_by_label.get(label, [])
    
    if len(keys) > len(properties):
        properties_by_label[label] = list(keys)
        
for label, roles in roles_by_type.items():
    # create a csv file
    #with open(os.path.join(CSV_DIR, f'{label}.csv'), 'w') as fh:
    #print(label, properties_by_label[label])
    prop_keys = properties_by_label.get(label, [])
    path = os.path.join(ROLES_DIR, f'{label}.csv')
    with open(path, 'w') as fh:
        fh.write(DELIMITER.join([':START_ID'] + prop_keys + [':END_ID']))
        fh.write('\n')
        
        for role in roles:
            #p = f'{DELIMITER}{dict_to_str(role["props"], prop_keys, DELIMITER, QUOTE)}'
            #fh.write(f'{role["source"]}{DELIMITER}{role["target"]}')
            fh.write(f'{role["source"]}')
            fh.write(f'{DELIMITER}')
            
            if prop_keys:
                fh.write(dict_to_str(role["props"], prop_keys, DELIMITER, QUOTE))
                fh.write(f'{DELIMITER}')
                
            fh.write(f'{role["target"]}')
            fh.write('\n')
            
    import_command.extend([f'--relationships:{label}', path])
    
import_command.extend([
    '--delimiter', f'"{DELIMITER}"',
    '--array-delimiter', f'"{ARRAY_DELIMITER}"',
    '--quote', f'"\\{QUOTE}"'
])
    
print(' '.join(import_command))