Merge branch 'main' of https://github.com/seung-lab/KVDbClient into main

Author: akhileshh

MANIFEST.in

@@ -0,0 +1 @@
+include requirements.txt

kvdbclient/__init__.py

@@ -0,0 +1,28 @@
+from collections import namedtuple
+
+from .bigtable.client import Client as BigTableClient
+from .bigtable.attributes import Attribute
+from .serializers import Serializer
+from .serializers import UInt64String
+
+
+_backend_clientinfo_fields = ("TYPE", "CONFIG")
+_backend_clientinfo_defaults = (None, None)
+BackendClientInfo = namedtuple(
+    "BackendClientInfo",
+    _backend_clientinfo_fields,
+    defaults=_backend_clientinfo_defaults,
+)
+
+
+def get_default_client_info():
+    """
+    Load client from env variables.
+    """
+
+    # TODO make dynamic after multiple platform support is added
+    from .bigtable import get_client_config as get_bigtable_client_config
+
+    return BackendClientInfo(
+        CONFIG=get_bigtable_client_config(admin=True, read_only=False)
+    )

kvdbclient/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

kvdbclient/base.py

@@ -0,0 +1,95 @@
+from abc import ABC
+from typing import Any
+from typing import Dict
+from typing import Tuple
+from typing import Callable
+from typing import Iterable
+from typing import Optional
+from datetime import datetime
+
+from .serializers import Serializer
+from .serializers import UInt64String
+
+
+class SimpleClient(ABC):
+    from abc import abstractmethod
+
+    """
+    Abstract class for interacting with backend data storage system.
+    Eg., BigTableClient for using big table as storage.
+    """
+
+    @abstractmethod
+    def create_table(self) -> None:
+        """Initialize the table and store associated meta."""
+
+    @abstractmethod
+    def write_metadata(self, metadata):
+        """Update stored metadata."""
+
+    @abstractmethod
+    def read_metadata(self):
+        """Read stored metadata."""
+
+    @abstractmethod
+    def read_entries(
+        self,
+        entry_ids,
+        properties=None,
+        start_time=None,
+        end_time=None,
+        end_time_inclusive=False,
+    ):
+        """Read entries and their properties."""
+
+    @abstractmethod
+    def read_entry(
+        self,
+        entry_id,
+        properties=None,
+        start_time=None,
+        end_time=None,
+        end_time_inclusive=False,
+    ):
+        """Read a single entry and it's properties."""
+
+    @abstractmethod
+    def write_entries(self, entries):
+        """Writes/updates entries (IDs along with properties)."""
+
+
+class EntryKey:
+    def __init__(self, key: Any, serializer: Optional[Serializer] = UInt64String()):
+        self._key = key
+        self._serializer = serializer.serialize
+
+    def serialize(self) -> Any:
+        return self._serializer(self._key)
+
+
+class Entry:
+    """
+    Represents a single entry/record/entity in the database.
+    """
+
+    def __init__(
+        self,
+        key: EntryKey,
+        val_dict: Dict[Any, Any],
+        timestamp: Optional[datetime] = None,
+    ):
+        self._key = key
+        self._val_dict = val_dict
+        self._timestamp = timestamp
+
+    @property
+    def key(self) -> Any:
+        return self._key.serialize()
+
+    @property
+    def values(self) -> Dict:
+        return self._val_dict
+
+    @property
+    def timestamp(self):
+        return self._timestamp

kvdbclient/bigtable/__init__.py

@@ -0,0 +1,53 @@
+from collections import namedtuple
+from typing import Any
+from typing import Iterable
+from typing import Tuple
+
+from .attributes import Attribute
+
+DEFAULT_PROJECT = "neuromancer-seung-import"
+DEFAULT_INSTANCE = "pychunkedgraph"
+
+_bigtableconfig_fields = (
+    "PROJECT",
+    "INSTANCE",
+    "ADMIN",
+    "READ_ONLY",
+    "CREDENTIALS",
+)
+_bigtableconfig_defaults = (
+    DEFAULT_PROJECT,
+    DEFAULT_INSTANCE,
+    True,
+    False,
+    None,
+)
+BigTableConfig = namedtuple(
+    "BigTableConfig", _bigtableconfig_fields, defaults=_bigtableconfig_defaults
+)
+
+
+def get_client_config(
+    project: str = None,
+    instance: str = None,
+    admin: bool = False,
+    read_only: bool = True,
+):
+    """Helper function to load config from env."""
+    from os import environ
+
+    _project = environ.get("BIGTABLE_PROJECT", DEFAULT_PROJECT)
+    if project:
+        _project = project
+
+    _instance = environ.get("BIGTABLE_INSTANCE", DEFAULT_INSTANCE)
+    if instance:
+        _instance = instance
+
+    kwargs = {
+        "PROJECT": _project,
+        "INSTANCE": _instance,
+        "ADMIN": admin,
+        "READ_ONLY": read_only,
+    }
+    return BigTableConfig(**kwargs)

kvdbclient/bigtable/attributes.py

@@ -0,0 +1,76 @@
+from typing import NamedTuple
+
+from .. import serializers
+
+
+class _AttributeType(NamedTuple):
+    key: bytes
+    family_id: str
+    serializer: serializers.Serializer
+
+
+class Attribute(_AttributeType):
+    __slots__ = ()
+    _attributes = {}
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        Attribute._attributes[(kwargs["family_id"], kwargs["key"])] = self
+
+    def serialize(self, obj):
+        return self.serializer.serialize(obj)
+
+    def deserialize(self, stream):
+        return self.serializer.deserialize(stream)
+
+    @property
+    def basetype(self):
+        return self.serializer.basetype
+
+    @property
+    def index(self):
+        return int(self.key.decode("utf-8").split("_")[-1])
+
+
+class AttributeArray:
+    _attributearrays = {}
+
+    def __init__(self, pattern, family_id, serializer):
+        self._pattern = pattern
+        self._family_id = family_id
+        self._serializer = serializer
+        AttributeArray._attributearrays[(family_id, pattern)] = self
+
+        # TODO: Add missing check in `fromkey(family_id, key)` and remove this
+        #       loop (pre-creates `Attributes`, so that the inverse lookup works)
+        for i in range(20):
+            self[i]  # pylint: disable=W0104
+
+    def __getitem__(self, item):
+        return Attribute(
+            key=self.pattern % item,
+            family_id=self._family_id,
+            serializer=self._serializer,
+        )
+
+    @property
+    def pattern(self):
+        return self._pattern
+
+    @property
+    def serialize(self):
+        return self._serializer.serialize
+
+    @property
+    def deserialize(self):
+        return self._serializer.deserialize
+
+    @property
+    def basetype(self):
+        return self._serializer.basetype
+
+
+class TableMeta:
+    key = b"meta"
+    family_id = "0"
+    data = Attribute(key=key, family_id=family_id, serializer=serializers.Pickle())