genmod / application_neo4j /load_data.py
Quentin Lhoest
add app
f0806e2
from neo4j import GraphDatabase, basic_auth
from graphdatascience import GraphDataScience
import utils as u
import argparse
# --- Configuration ---
SOURCE_CSV = 'models_with_all_information.csv'
NEO4J_URI = "bolt://localhost:7687"
# --- Configuration des arguments du script ---
parser = argparse.ArgumentParser(description="Script pour peupler une base de données Neo4j avec des données sur des modèles.")
parser.add_argument(
'reset', # Le nom de l'argument (positionnel, car sans tirets)
type=str, # On attend une chaîne de caractères
choices=['true', 'false'], # L'utilisateur DOIT choisir entre 'true' et 'false'
help="Spécifiez 'true' pour réinitialiser la base de données, 'false' pour continuer."
)
parser.add_argument(
'neo4j_user', # Le nom de l'argument (positionnel, car sans tirets)
type=str, # On attend une chaîne de caractères
help="Nom de votre instance Neo4J"
)
parser.add_argument(
'neo4j_password', # Le nom de l'argument (positionnel, car sans tirets)
type=str, # On attend une chaîne de caractères
help="Mot de passe de votre instance Neo4J"
)
args = parser.parse_args()
# Conversion de l'argument string 'true'/'false' en un vrai booléen Python True/False
RESET = args.reset.lower() == 'true'
NEO4J_PASSWORD = args.neo4j_password
NEO4J_USER = args.neo4j_user
# --- Initialisation et chargement de la base Neo4j ---
# Connexion au serveur Neo4j via le driver natif
# C'est une bonne pratique de garder le driver et de ne pas créer la session tout de suite
driver = GraphDatabase.driver(NEO4J_URI, auth=basic_auth(NEO4J_USER, NEO4J_PASSWORD))
# Réinitialisation de la base si activé
if RESET:
u.reset_database(driver) # Cette fonction supprime maintenant la DB + le fichier
processed_ids = set() # On part d'un set vide
print("Reset done")
else:
# Si on ne réinitialise pas, ALORS on charge les IDs existants pour continuer
processed_ids = u.load_processed_ids()
# Création des index pour améliorer les performances des requêtes
with driver.session() as session:
session.run("CREATE INDEX IF NOT EXISTS FOR (m:Model) ON (m.name);")
session.run("CREATE INDEX IF NOT EXISTS FOR (a:Author) ON (a.name);")
session.run("CREATE INDEX IF NOT EXISTS FOR (d:Dataset) ON (d.name);")
print("Indexes created.")
# Insertion des données dans la base Neo4j depuis le fichier CSV en parallèle
u.insert_parallel('models_with_all_information.csv', driver, processed_ids)
# --- Finalisation ---
# Fermer la connexion du driver à la toute fin du script
driver.close()
print("\nScript finished. Neo4j driver connection closed.")