Merge pull request #9 from moevm/import-export

Import export: fix

gighunt/src/gighunt/modules/controller.py

@@ -226,6 +226,43 @@ def _get_dump(self) -> dict:
 
         return graph_data
 
+    def _count_imported_records(self, graph_data) -> int:
+        static = graph_data.get("static", [])
+        vertices = graph_data.get("vertices", [])
+        edges = graph_data.get("edges", [])
+
+        # Vertices
+        users = vertices["users"]
+        groups = vertices["groups"]
+        announcements = vertices["announcements"]
+        places = vertices["places"]
+
+        # Edges
+        user_groups = edges["users"]["groups"]
+        stars_to_users = edges["stars"]["users"]
+        stars_to_groups = edges["stars"]["groups"]
+        stars_to_announcements = edges["stars"]["announcements"]
+        comments_to_announcements = edges["comments"]["announcements"]
+        producer_announcements_from_users = edges["producer_announcements"]["users"]
+        producer_announcements_from_groups = edges["producer_announcements"][
+            "groups"
+        ]
+
+        amount = (
+            len(static)
+            + len(users)
+            + len(groups)
+            + len(announcements)
+            + len(places)
+            + len(user_groups)
+            + len(stars_to_users)
+            + len(stars_to_groups)
+            + len(stars_to_announcements)
+            + len(comments_to_announcements)
+            + len(producer_announcements_from_users)
+            + len(producer_announcements_from_groups)
+        )
+        return amount
 
     async def _import_data(self, file: UploadFile) -> Response:
         """
@@ -243,8 +280,11 @@ async def _import_data(self, file: UploadFile) -> Response:
             file_content = await file.read()
             data = json.loads(file_content)
             self._set_dump(data)
-
-            return {"message": "Graph imported successful!"}
+            number_of_records = self._count_imported_records(data)
+            return {
+                "message": "Данные успешно импортированы!",
+                "number_of_records": number_of_records
+            }
 
         except json.JSONDecodeError:
             raise HTTPException(status_code=400, detail="Invalid JSON format")

gighunt/src/gighunt/modules/use_cases/base_edge_use_cases.py

@@ -27,6 +27,17 @@ def create_new_entity(self, edge_data:Json, name:str)-> bool | Json:
     def get_entity(self, entity_id: str, name) -> Json | None:
         return self._graph.edge_collection(name).get(entity_id)
 
+    def insert_json_edge(self, name: str, json_edge: Json) -> None:
+        try:
+            self._graph.edge_collection(name).insert(json_edge)
+        except DocumentInsertError as err:
+            self._logger.error(f"Error inserting edge {json_edge}: {err}")
+
+    def clear(self) -> None:
+        if self.edge_collection_names:
+            for name in self.edge_collection_names:
+                self._graph.edge_collection(name.value).truncate()
+
     def delete_entity(self, entity_id: str, name)->Json|None:
         return self._graph.edge_collection(name).delete(entity_id)