Created
October 30, 2013 00:12
-
-
Save tconkling/7225045 to your computer and use it in GitHub Desktop.
ProtobufTable.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# | |
# 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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment