From 6296b777ecf00baf5953f738a750a789a87f0427 Mon Sep 17 00:00:00 2001
From: john681611 <john681611@hotmail.com>
Date: Thu, 14 Sep 2023 11:07:43 +0100
Subject: [PATCH] Add CLI population command for DB

---
 application/cmd/cre_main.py | 9 +++++++++
 cre.py                      | 6 ++++++
 2 files changed, 15 insertions(+)

diff --git a/application/cmd/cre_main.py b/application/cmd/cre_main.py
index d0bca16d8..abae29e2c 100644
--- a/application/cmd/cre_main.py
+++ b/application/cmd/cre_main.py
@@ -426,6 +426,8 @@ def run(args: argparse.Namespace) -> None:  # pragma: no cover
         generate_embeddings(args.cache_file)
     if args.owasp_proj_meta:
         owasp_metadata_to_cre(args.owasp_proj_meta)
+    if args.populate_neo4j_db:
+        populate_neo4j_db(args.cache_file)
 
 
 def db_connect(path: str):
@@ -530,3 +532,10 @@ def owasp_metadata_to_cre(meta_file: str):
     },
     """
     raise NotImplementedError("someone needs to work on this")
+
+def populate_neo4j_db(cache: str):
+    logger.info(f"Populating neo4j DB: Connecting to SQL DB")
+    database = db_connect(path=cache)
+    logger.info(f"Populating neo4j DB: Populating")
+    database.neo_db.populate_DB(database.session)
+    logger.info(f"Populating neo4j DB: Complete")
\ No newline at end of file
diff --git a/cre.py b/cre.py
index 0feb76a62..c16c44da9 100644
--- a/cre.py
+++ b/cre.py
@@ -187,6 +187,12 @@ def main() -> None:
         action="store_true",
         help="for every node, download the text pointed to by the hyperlink and generate embeddings for the content of the specific node",
     )
+    parser.add_argument(
+        "--populate_neo4j_db",
+        action="store_true",
+        help="populate the neo4j db",
+    )
+
 
     args = parser.parse_args()