From 1e5b9afd0f2fdd36bee255c75be1c45c8847e83f Mon Sep 17 00:00:00 2001 From: liyaxing Date: Thu, 7 Sep 2023 20:19:18 +0800 Subject: [PATCH] Add chunk relate impl and test case --- src/knowledge/__init__.py | 5 ++- src/knowledge/data/__init__.py | 2 + src/knowledge/data/chunk.py | 43 +++++++++++++++++++ src/knowledge/data/tracker.py | 77 ++++++++++++++++++++++++++++++++++ test/test_chunk.py | 36 ++++++++++++++++ 5 files changed, 161 insertions(+), 2 deletions(-) create mode 100644 src/knowledge/data/__init__.py create mode 100644 src/knowledge/data/chunk.py create mode 100644 src/knowledge/data/tracker.py create mode 100644 test/test_chunk.py diff --git a/src/knowledge/__init__.py b/src/knowledge/__init__.py index c879518..6edb83e 100644 --- a/src/knowledge/__init__.py +++ b/src/knowledge/__init__.py @@ -1,2 +1,3 @@ -from .object import EmailObject, ImageObject, TextChunkObject, ObjectID -from .vector import * \ No newline at end of file +from .object import * +from .vector import * +from .data import * \ No newline at end of file diff --git a/src/knowledge/data/__init__.py b/src/knowledge/data/__init__.py new file mode 100644 index 0000000..81c672a --- /dev/null +++ b/src/knowledge/data/__init__.py @@ -0,0 +1,2 @@ +from .chunk import ChunkID, PositionType, PositionFileRange +from .tracker import ChunkTracker \ No newline at end of file diff --git a/src/knowledge/data/chunk.py b/src/knowledge/data/chunk.py new file mode 100644 index 0000000..ff355bd --- /dev/null +++ b/src/knowledge/data/chunk.py @@ -0,0 +1,43 @@ +from enum import Enum +from ..object import ObjectID + +ChunkID = ObjectID + +class PositionType(Enum): + Unknown = 1 + Device = 2 + File = 3 + FileRange = 4 + ChunkManager = 5 + + +class PositionFileRange: + def __init__(self, path: str, range_begin: int, range_end: int): + self.path = path + self.range_begin = range_begin + self.range_end = range_end + + def encode(self): + return f"{self.range_begin}:{self.range_end}:{self.path}" + + @staticmethod + def decode(value: str): + parts = value.split(":") + if len(parts) < 3: + raise ValueError("Invalid input string") + + try: + range_begin = int(parts[0]) + range_end = int(parts[1]) + except ValueError as e: + raise ValueError("Invalid range_begin or range_end string") from e + + path = ":".join(parts[2:]) + return PositionFileRange(path, range_begin, range_end) + + def __str__(self): + return self.encode() + + @staticmethod + def from_string(value: str): + return PositionFileRange.decode(value) diff --git a/src/knowledge/data/tracker.py b/src/knowledge/data/tracker.py new file mode 100644 index 0000000..efe4ede --- /dev/null +++ b/src/knowledge/data/tracker.py @@ -0,0 +1,77 @@ +import sqlite3 +import time +import logging +import os +from .chunk import ChunkID, PositionType, PositionFileRange + + +class ChunkTracker: + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + directory = os.path.join( + os.path.dirname(__file__), "../../../rootfs/data/chunk/" + ) + + if not os.path.exists(directory): + os.mkdir(directory) + file = os.path.join(directory, "chunk_tracker.db") + logging.info(f"will init chunk tracker, db={file}") + cls._instance.__singleton_init__(file) + + return cls._instance + + def __singleton_init__(self, db_path: str): + logging.info(f"will init chunk tracker, db={db_path}") + + self.conn = sqlite3.connect(db_path) + self.cursor = self.conn.cursor() + self.cursor.execute( + """ + CREATE TABLE IF NOT EXISTS chunks ( + id TEXT NOT NULL, + pos TEXT NOT NULL, + pos_type TINYINT NOT NULL, + insert_time UNSIGNED BIG INT NOT NULL, + update_time UNSIGNED BIG INT NOT NULL, + flags INTEGER DEFAULT 0, + PRIMARY KEY(id, pos, pos_type) + ) + """ + ) + self.conn.commit() + + def add_position(self, chunk_id: ChunkID, position: str, position_type: PositionType): + logging.debug(f"add chunk position: {chunk_id}, {position}, {position_type}") + + insert_time = update_time = int(time.time()) + self.cursor.execute( + """ + INSERT OR REPLACE INTO chunks (id, pos, pos_type, insert_time, update_time) + VALUES (?, ?, ?, ?, ?) + """, + (str(chunk_id), position, position_type.value, insert_time, update_time,), + ) + self.conn.commit() + + def remove_position(self, chunk_id: ChunkID): + logging.info(f"remove chunk position: {chunk_id}") + + self.cursor.execute( + """ + DELETE FROM chunks WHERE id = ? + """, + (str(chunk_id),), + ) + self.conn.commit() + + def get_position(self, chunk_id: ChunkID) -> (str, PositionType): + self.cursor.execute( + """ + SELECT pos, pos_type FROM chunks WHERE id = ? + """, + (str(chunk_id),), + ) + return self.cursor.fetchone() diff --git a/test/test_chunk.py b/test/test_chunk.py new file mode 100644 index 0000000..40a7aeb --- /dev/null +++ b/test/test_chunk.py @@ -0,0 +1,36 @@ +import sys +import os + +dir_path = os.path.dirname(os.path.realpath(__file__)) +print(dir_path) + +sys.path.append("{}/../src/".format(dir_path)) +print(sys.path) + +from knowledge import ChunkTracker, ChunkID, HashValue, PositionType + + +import asyncio +import unittest + + +class TestChunk(unittest.TestCase): + def test_chunk_tracker(self): + tracker = ChunkTracker() + + hash = HashValue.hash_data("1234567890".encode("utf-8")); + cid = ChunkID.new_chunk_id(hash) + print(cid) + + tracker.add_position(cid, "/tmp/1", PositionType.File) + ret = tracker.get_position(cid) + print(ret[0]) + + tracker.remove_position(cid) + ret = tracker.get_position(cid) + self.assertEqual(ret, None) + + + +if __name__ == "__main__": + unittest.main()