-
Notifications
You must be signed in to change notification settings - Fork 3
/
import_supplement_edge.py
51 lines (34 loc) · 1.53 KB
/
import_supplement_edge.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
import json
import sys, os
from neo4j import GraphDatabase
import pandas as pd
#python import_node.py <hetionet_data_file> ip
class import_data:
def __init__(self, uri, user, password):
self.driver = GraphDatabase.driver(uri, auth=(user, password))
def close(self):
self.driver.close()
def add_nodes_connections(self, filename, nodetype_1, nodetype_2, relation_type):
with self.driver.session() as session:
with session.begin_transaction() as tx:
df = pd.read_csv(filename, sep=',')
for index, row in df.iterrows():
name_1 = row["from"].replace("'", "\\'")
name_2 = row["to"].replace("'", "\\'")
query = f"MATCH (s:`{nodetype_1}` " + "{" + f'name: "{name_1}"' + "}) "
query += f"MATCH (t:`{nodetype_2}` " + "{" + f'name: "{name_2}"' + "}) "
query += f"MERGE (s)-[r:`{relation_type}`]->(t);"
tx.run(query)
if index % 100 == 0:
tx.commit()
tx = session.begin_transaction()
tx.commit()
import_csv_file = sys.argv[1]
ip = sys.argv[2]
password = sys.argv[3]
nodetype_1 = sys.argv[4]
nodetype_2 = sys.argv[5]
relation_type = sys.argv[6]
connection = import_data(f"bolt://{ip}:7687", "neo4j", password)
connection.add_nodes_connections(import_csv_file, nodetype_1, nodetype_2, relation_type)
connection.close()