diff --git a/src/knowledge/object/relation.py b/src/knowledge/object/relation.py index f6eca2e..34c0911 100644 --- a/src/knowledge/object/relation.py +++ b/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() \ No newline at end of file + self.conn.commit() diff --git a/test/test_object.py b/test/test_object.py index 1e1bc56..0f5620e 100644 --- a/test/test_object.py +++ b/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()