diff --git a/README.md b/README.md index dab255d..9274f72 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ A high-performance dictionary database.

-The `flaxkv` provides an interface very similar to a dictionary for interacting with high-performance key-value databases. More importantly, as a persistent database, it offers performance close to that of native dictionaries (in-memory access). +The `flaxkv` provides an interface very similar to a dictionary for interacting with high-performance key-value databases. More importantly, as a persistent database, it offers performance close to that of native dictionaries (in-memory access). You can use it just like a Python dictionary without having to worry about blocking your user process when operating the database at any time. --- diff --git a/flaxkv/__init__.py b/flaxkv/__init__.py index 8be449d..0b22fd8 100644 --- a/flaxkv/__init__.py +++ b/flaxkv/__init__.py @@ -11,6 +11,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import os +import re from .core import LevelDBDict, LMDBDict from .serve.client import RemoteDictDB @@ -26,23 +28,26 @@ ] +url_pattern = re.compile(r'^(http://|https://|ftp://)') + + def dictdb( - path_or_url: str, + db_name: str, + root_path_or_url: str = ".", backend='lmdb', - remote=False, - db_name=None, rebuild=False, raw=False, **kwargs, ): - if remote: + if url_pattern.match(root_path_or_url): return RemoteDictDB( - path_or_url, db_name=db_name, rebuild=rebuild, backend=backend + root_path_or_url, db_name=db_name, rebuild=rebuild, backend=backend ) + db_path = os.path.join(root_path_or_url, db_name) if backend == 'lmdb': - return LMDBDict(path_or_url, rebuild=rebuild, raw=raw, **kwargs) + return LMDBDict(db_path, rebuild=rebuild, raw=raw, **kwargs) elif backend == 'leveldb': - return LevelDBDict(path_or_url, rebuild=rebuild, raw=raw, **kwargs) + return LevelDBDict(db_path, rebuild=rebuild, raw=raw, **kwargs) else: raise ValueError(f"Unsupported DB type {backend}.") diff --git a/flaxkv/serve/app.py b/flaxkv/serve/app.py index b2db039..bd86cd3 100644 --- a/flaxkv/serve/app.py +++ b/flaxkv/serve/app.py @@ -23,11 +23,15 @@ @post(path="/attach") async def attach(data: AttachRequest) -> dict: # todo switch `post` to `get` - db = db_manager.get(data.db_name) - if db is None or data.rebuild: - db_manager.set_db( - db_name=data.db_name, backend=data.backend, rebuild=data.rebuild - ) + try: + db = db_manager.get(data.db_name) + if db is None or data.rebuild: + db_manager.set_db( + db_name=data.db_name, backend=data.backend, rebuild=data.rebuild + ) + except Exception as e: + traceback.print_exc() + return {"success": False, "info": str(e)} return {"success": True} @@ -84,10 +88,10 @@ async def _get(db_name: str, request: Request) -> bytes: if db is None: raise ValueError("db not found") key = await request.body() - value = db.get(key) + value = db.get(key, None) if value is None: - return encode(None) - return db.get(key) + return b'iamnull123' + return value @post("/contains", media_type=MediaType.TEXT) diff --git a/flaxkv/serve/client.py b/flaxkv/serve/client.py index 7e591f5..65f9dd9 100644 --- a/flaxkv/serve/client.py +++ b/flaxkv/serve/client.py @@ -1,5 +1,3 @@ -import httpx - from ..pack import decode, decode_key, encode @@ -13,6 +11,8 @@ def __init__( timeout=6, **kwargs, ): + import httpx + self._url = url self._db_name = db_name self._client = kwargs.pop("client", httpx.Client(timeout=timeout)) @@ -33,10 +33,14 @@ def detach_db(self, db_name=None): response = self._client.post(url, json={"db_name": db_name}) return response.json() - def get(self, key): + def get(self, key, default=None): url = f"{self._url}/get_raw?db_name={self._db_name}" response = self._client.post(url, data=encode(key)) - return decode(response.read()) + raw_data = response.read() + if raw_data == b"iamnull123": + return default + value = decode(raw_data) + return value def set(self, key, value): url = f"{self._url}/set_raw?db_name={self._db_name}" diff --git a/flaxkv/serve/manager.py b/flaxkv/serve/manager.py index 59a5ce2..8f33613 100644 --- a/flaxkv/serve/manager.py +++ b/flaxkv/serve/manager.py @@ -17,7 +17,8 @@ def detach(self, db_name: str): def set_db(self, db_name: str, backend: str, rebuild: bool): db_path = self._root_path / db_name self._db_dict[db_name] = dictdb( - path_or_url=db_path.__str__(), + db_name=db_name, + root_path_or_url=self._root_path.__str__(), backend=backend, rebuild=rebuild, raw=self._raw_mode, diff --git a/pyproject.toml b/pyproject.toml index f3e0c39..bae0b16 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ authors = [ ] license-files = { paths = ["LICENSE"] } readme = "README.md" -keywords = ["database", "dict", "lmdb", "leveldb", "Machine Learning", "NLP"] +keywords = ["persistent-storage", "on-disk dict", "lmdb", "leveldb", "Machine Learning", "NLP"] classifiers = [ "Development Status :: 5 - Production/Stable", "Operating System :: OS Independent", @@ -56,6 +56,7 @@ database = [ server = [ "litestar", + "httpx", ]