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()

test/test_object.py

@@ -17,25 +17,56 @@ handler.setFormatter(formatter)
 root.addHandler(handler)
 
 
-from knowledge import ObjectID, HashValue, EmailObjectBuilder
+from knowledge import (
+    ObjectID,
+    HashValue,
+    EmailObjectBuilder,
+    ObjectRelationStore,
+    KnowledgeStore,
+)
 import asyncio
 import unittest
 
 
 class TestVectorSTorage(unittest.TestCase):
     def test_object(self):
-        data = HashValue.hash_data("1233".encode("utf-8"));
+        data = HashValue.hash_data("1233".encode("utf-8"))
         print(data.to_base58())
         print(data.to_base36())
-        
+
         data2 = HashValue.from_base58(data.to_base58())
         self.assertEqual(data.to_base36(), data2.to_base36())
-        
+
         data2 = HashValue.from_base36(data.to_base36())
         self.assertEqual(data.to_base58(), data2.to_base58())
 
         email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d"
         email_object = EmailObjectBuilder({}, email_folder).build()
-        
+
+    def test_relation(self):
+        obj1 = ObjectID.hash_data("12345".encode("utf-8"))
+        obj2 = ObjectID.hash_data("67890".encode("utf-8"))
+        obj3 = ObjectID.hash_data("abcde".encode("utf-8"))
+        obj4 = ObjectID.hash_data("fghij".encode("utf-8"))
+        print(obj1.to_base58(), obj2.to_base58(), obj3.to_base58())
+        relation_store = KnowledgeStore().get_relation_store()
+        relation_store.add_relation(obj1, obj2)
+        relation_store.add_relation(obj1, obj2)
+        relation_store.add_relation(obj2, obj3)
+
+        relation_store.add_relation(obj1, obj3)
+        relation_store.add_relation(obj1, obj4)
+
+        objs = relation_store.get_related_objects(obj2)
+        self.assertEqual(len(objs), 1)
+        self.assertEqual(objs[0], obj3)
+
+        objs = relation_store.get_related_root_objects(obj1)
+        self.assertEqual(len(objs), 2)
+        self.assertEqual(obj3 in objs, True)
+        self.assertEqual(obj4 in objs, True)
+        # self.assertCountEqual(objs, [obj3, obj4])
+
+
 if __name__ == "__main__":
     unittest.main()