ProtobufTable.py

protobuf_table.py

#
# ProtobufTable

import re
import os
import logging
import pickle
from threading import RLock


class ProtobufTable(object):
    """Simple database table-like structure for persisting protobuf objects to disk."""

    def __init__(self, clazz, db_dir, extension="object", persist=True):
        """Creates a new ProtobufTable.

        Args:
            clazz: the Class object for protobuf object type that this ProtobufTable will store
            db_dir: the directory that objects should be read from and written to
            extension: the filename extension that persisted objects should use.
                If multiple ProtobufTables share the same db_dir, they should have unique extensions.
                (default "object")
            persist: whether to write ProtobufTable objects to disk (default True)
        """

        self._lock = RLock()
        self._clazz = clazz
        self._db_dir = db_dir
        self._extension = extension
        self._info = None
        self._persist = persist

        # create the database directory if it doesn't exist
        if not os.path.exists(self._db_dir):
            os.makedirs(self._db_dir)

        # load our TableInfo file, if it exists
        if os.path.isfile(self._table_info_filename):
            with open(self._table_info_filename, "rb") as f:
                try:
                    self._info = pickle.load(f)
                except Exception, e:
                    logging.error("Error loading table data %s: %s" % (self._table_info_filename, str(e)))
                    self._info = None

        # load objects
        self._objects = {}
        highestId = 0
        for filename in self._find_files():
            try:
                with open(filename, "rb") as f:
                    obj = self._clazz()
                    obj.ParseFromString(f.read())
                    if not obj.IsInitialized():
                        raise RuntimeError("not fully initialized")

                    highestId = max(highestId, obj.id)
                    self._objects[obj.id] = obj
                    logging.info("Loaded %s" % filename)
            except Exception, e:
                logging.error("Error loading %s: %s" % (filename, str(e)))

        if self._info is None:
            self._info = TableInfo()
            self._info.nextId = highestId + 1
            self._save_table_info()

    @property
    def objects(self):
        """Returns all objects in the table."""
        with self._lock:
            return self._objects.values()

    def get_object(self, id):
        """Returns the object with the given id."""
        try:
            with self._lock:
                return self._objects[id]
        except Exception:
            return None

    def create_object(self):
        """Creates a new object in the table."""
        obj = self._clazz()
        with self._lock:
            obj.id = self._info.nextId
            self._info.nextId += 1
            self._objects[obj.id] = obj
        self._save_table_info()
        return obj

    def save_object(self, obj):
        """Persists the given object to disk.

        The object must be owned by this Table.
        (If the Table was created with persist=False, this method is a no-op).
        """
        if self.get_object(obj.id) is not obj:
            raise RuntimeError("Can't delete an object we don't own")

        if not self._persist:
            return

        data = obj.SerializeToString()
        filename = self._obj_filename(obj.id)
        try:
            with self._lock:
                with open(filename, "wb") as f:
                    f.write(data)
                    logging.info("Saved object %s" % filename)
        except Exception, e:
            logging.error("Error saving %s: %s" % (filename, str(e)))

    def delete_object(self, obj):
        """Deletes the given object from the Table. The object must be owned by this Table."""
        if self.get_object(obj.id) is not obj:
            raise RuntimeError("Can't delete an object we don't own")
        self.delete_object_with_id(obj.id)

    def delete_object_with_id(self, obj_id):
        """Deletes the object with the given id from the Table."""
        filename = self._obj_filename(obj_id)
        try:
            with self._lock:
                if obj_id in self._objects:
                    del self._objects[obj_id]
                    os.remove(filename)
                    logging.info("deleted object %s" % filename)
        except Exception, e:
            logging.error("Error deleting %s: %s" % (filename, str(e)))

    def delete_all_objects(self):
        """Clears all objects from the Table"""
        logging.info("Clearing table '%s'..." % self._extension)
        with self._lock:
            # remove all files, even corrupted ones that weren't loaded
            for fn in self._find_files():
                os.remove(fn)
            self._objects = {}
        logging.info("Cleared table '%s'" % self._extension)

    def _find_files(self):
        return _list_files(self._db_dir, re.compile(r'.+\.' + self._extension + '$'))

    def _obj_filename(self, obj_id):
        return os.path.join(self._db_dir, ("%d." % obj_id) + self._extension)

    def _save_table_info(self):
        try:
            with self._lock:
                with open(self._table_info_filename, "wb") as f:
                    pickle.dump(self._info, f)
        except Exception, e:
            logging.error("Error saving TableInfo: %s" % str(e))

    @property
    def _table_info_filename(self):
        return os.path.join(self._db_dir, self._extension + ".table_info")


def _list_files(dirname, regex):
    return [os.path.join(dirname, fn) for fn in os.listdir(dirname) if regex.match(fn)]


class TableInfo(object):
    def __init__(self):
        self.nextId = 1