Add object relation store impl and test case

src/knowledge/object/relation.py

@@ -1,64 +1,104 @@
 # define a relation store class
-import sqlite3
-import os
-import logging
 from .object_id import ObjectID
+import sqlite3
+from typing import List, Tuple, Optional
+import logging
+import os
+from enum import Enum
+
+
+class ObjectRelationType(Enum):
+    Parent = 1
 
 
 class ObjectRelationStore:
     def __init__(self, root_dir: str):
         if not os.path.exists(root_dir):
             os.makedirs(root_dir)
-        file = os.path.join(root_dir, "relationships.db")
-        logging.info(f"will init chunk tracker, db={file}")
+        file = os.path.join(root_dir, "relation.db")
+        logging.info(f"will init object relation store, db={file}")
 
         self.conn = sqlite3.connect(file)
         self.cursor = self.conn.cursor()
         self.cursor.execute(
             """
-            CREATE TABLE IF NOT EXISTS relationships (
-                id TEXT NOT NULL,
-                parent TEXT NOT NULL,
-                PRIMARY KEY(id, parent)
+            CREATE TABLE IF NOT EXISTS relations (
+                object_id TEXT,
+                assoc_id TEXT,
+                relation_type TEXT,
+                PRIMARY KEY (object_id, assoc_id, relation_type)
             )
         """
         )
-        self.conn.commit()
 
-    def add_relation(self, object: ObjectID, parent: ObjectID):
+    def add_relation(
+        self,
+        object_id: ObjectID,
+        assoc_id: ObjectID,
+        relation_type: ObjectRelationType = ObjectRelationType.Parent,
+    ):
+        if relation_type == None:
+            relation_type = ObjectRelationType.Parent
+
         self.cursor.execute(
             """
-            INSERT INTO relationships (id, parent)
-            VALUES (?, ?, ?, ?, ?)
+            INSERT OR IGNORE INTO relations (object_id, assoc_id, relation_type)
+            VALUES (?, ?, ?)
         """,
-            (
-                str(object),
-                str(parent),
-            ),
+            (str(object_id), str(assoc_id), relation_type.value),
         )
         self.conn.commit()
 
-    def get_related_objects(self, object: ObjectID) -> [ObjectID]:
-        parents = []
-        cur = object
-        while True:
+    def get_related_objects(
+        self, object_id: ObjectID, relation_type: Optional[ObjectRelationType] = None
+    ) -> List[ObjectID]:
+        if relation_type:
             self.cursor.execute(
                 """
-                SELECT id, parent FROM relationships WHERE id = ?
+                SELECT assoc_id FROM relations WHERE object_id = ? AND relation_type = ?
             """,
-                (str(cur),),
+                (str(object_id), relation_type.value),
             )
-            parent = self.cursor.fetchone()
-            if parent is None:
-                break
-            parents.append(parent)
-            cur = parent
+        else:
+            self.cursor.execute(
+                """
+                SELECT assoc_id FROM relations WHERE object_id = ?
+            """,
+                (str(object_id),),
+            )
+        return [ObjectID.from_base58(row[0]) for row in self.cursor.fetchall()]
 
-    def delete_relation(self, object_id: ObjectID) -> [ObjectID]:
+    def get_related_root_objects(
+        self, object_id: ObjectID, relation_type: Optional[ObjectRelationType] = None
+    ) -> List[ObjectID]:
+        root_objects = []
+        related_objects = self.get_related_objects(object_id, relation_type)
+        history = []
+        history.append(object_id)
+
+        while related_objects:
+            for obj in related_objects:
+                next_related_objects = self.get_related_objects(obj, relation_type)
+                if not next_related_objects:
+                    if obj not in root_objects:
+                        root_objects.append(obj)
+                else:
+                    for related_object in next_related_objects:
+                        if obj not in history:
+                            related_objects.append(related_object)
+                        else:
+                            logging.warning(
+                                f"loop detected: {obj} <-> {related_object}"
+                            )
+            related_objects = next_related_objects
+
+        return root_objects
+
+    def delete_relation(self, object_id: ObjectID):
         self.cursor.execute(
             """
-            DELETE FROM relationships WHERE id = ?
+            DELETE FROM relations WHERE object_id = ?
         """,
             (str(object_id),),
         )
-        self.conn.commit()
+        self.conn.commit()