refactor: optimize serve module (#11)

* Unified dbdict input parameters

* move httpx to __init__ method && update pypi keywords

README.md

@@ -38,7 +38,7 @@ A high-performance dictionary database.
 </p>
 
 
-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.
 
 ---

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}.")
 

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)

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}"

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,

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",
 ]