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
|
import logging
import os
import shelve
import shutil
from typing import Any
from errbot.storage.base import StorageBase, StoragePluginBase
log = logging.getLogger("errbot.storage.shelf")
class ShelfStorage(StorageBase):
def __init__(self, path):
log.debug("Open shelf storage %s", path)
self.shelf = shelve.DbfilenameShelf(path, protocol=2)
def get(self, key: str) -> Any:
return self.shelf[key]
def remove(self, key: str):
if key not in self.shelf:
raise KeyError(f"{key} doesn't exist.")
del self.shelf[key]
def set(self, key: str, value: Any) -> None:
self.shelf[key] = value
def len(self):
return len(self.shelf)
def keys(self):
return self.shelf.keys()
def close(self) -> None:
self.shelf.close()
self.shelf = None
class ShelfStoragePlugin(StoragePluginBase):
def __init__(self, bot_config):
super().__init__(bot_config)
if "basedir" not in self._storage_config:
self._storage_config["basedir"] = bot_config.BOT_DATA_DIR
def open(self, namespace: str) -> StorageBase:
config = self._storage_config
# Hack to port move old DBs to the new location.
new_spot = os.path.join(config["basedir"], namespace + ".db")
old_spot = os.path.join(config["basedir"], "plugins", namespace + ".db")
if os.path.isfile(old_spot):
if os.path.isfile(new_spot):
log.warning(
"You have an old v3 DB at %s and a duplicate new one at %s.",
old_spot,
new_spot,
)
log.warning(
"You need to either remove the old one or move it in place of the new one manually."
)
else:
log.info("Moving your old v3 DB from %s to %s.", old_spot, new_spot)
shutil.move(old_spot, new_spot)
return ShelfStorage(new_spot)
|