Spaces:
Runtime error
Runtime error
File size: 4,050 Bytes
b115d50 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
"""A simple key-value store implemented atop Files and Tags."""
from typing import Any, Dict, List, Optional, Tuple
from steamship import Block, File, Steamship, Tag
KV_STORE_MARKER = "__init__"
class KeyValueStore:
"""A simple key value store implemented in Steamship.
Instances of the KeyValueStore are identified by its `namespace`.
This store_identifier corresponds to a File that will be created with a special tag identifying it.
Entries of the KeyValueStore are saved as `Tag` objects with:
* Kind = "KeyValueStore"
* Name = the key of the (kv) pair
* Value = a dict set to the value
Note that the value is always saved as a dict object. To save a string or int, wrap it in a dict.
WARNING:
This is essentially a clever hack atop Steamship's tag system to provide mutable key-value storage. It is in the
steamship.utils package because it's proven useful once or twice. But in general, if you find yourself heavily
relying upon it, consider reaching out to us at [email protected] to let us know, and we'll up-prioritize
adding a proper key-value API.
"""
client: Steamship
store_identifier: str
def __init__(self, client: Steamship, store_identifier: str = "KeyValueStore"):
"""Create a new KeyValueStore instance.
Args:
client (Steamship): The Steamship client.
store_identifier (str): The store_identifier which identifies this KeyValueStore instance. You can have multiple, separate KeyValueStore instances in a workspace using this implementation.
"""
self.client = client
self.store_identifier = f"kv-store-{store_identifier}"
def _get_file(self, or_create: bool = False) -> Optional[File]:
status_files = File.query(self.client, f'filetag and kind "{self.store_identifier}"').files
if len(status_files) == 0:
if not or_create:
return None
return File.create(
self.client,
blocks=[Block(text="")],
tags=[Tag(kind=self.store_identifier, name=KV_STORE_MARKER)],
)
else:
return status_files[0]
def get(self, key: str) -> Optional[Dict]:
"""Get the value represented by `key`."""
file = self._get_file()
if file is None:
return None
for tag in file.tags:
if tag.kind == self.store_identifier and tag.name == key:
return tag.value
def delete(self, key: str) -> bool:
"""Delete the entry represented by `key`"""
file = self._get_file()
if file is None:
return False
deleted = False
for tag in file.tags:
if tag.kind == self.store_identifier and tag.name == key:
tag.delete()
deleted = True
return deleted
def set(self, key: str, value: Dict[str, Any]):
"""Set the entry (key, value)."""
# First delete it if it exists to avoid duplicate tags.
self.delete(key)
# Now get/create the file
file = self._get_file(or_create=True)
req = Tag(file_id=file.id, kind=self.store_identifier, name=key, value=value)
return self.client.post("tag/create", req, expect=Tag)
def items(self, filter_keys: Optional[List[str]] = None) -> List[Tuple[str, Dict[str, Any]]]:
"""Return all key-value entries as a list of (key, value) tuples.
If `filter_keys` is provided, only returns keys within that list."""
file = self._get_file(or_create=True)
return [
(tag.name, tag.value)
for tag in file.tags
if (
tag.kind == self.store_identifier
and tag.name != KV_STORE_MARKER
and (filter_keys is None or tag.name in filter_keys)
)
]
def reset(self):
"""Delete all key-values."""
file = self._get_file()
if file is not None:
file.delete()
|