### # Copyright (c) 2002-2004, Jeremiah Fincher # All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: # # * Redistributions of source code must retain the above copyright notice, # this list of conditions, and the following disclaimer. # * Redistributions in binary form must reproduce the above copyright notice, # this list of conditions, and the following disclaimer in the # documentation and/or other materials provided with the distribution. # * Neither the name of the author of this software nor the name of # contributors to this software may be used to endorse or promote products # derived from this software without specific prior written consent. # # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE # POSSIBILITY OF SUCH DAMAGE. ### """ Module for some slight database-independence for simple databases. """ import supybot.fix as fix import csv import math import sets import random from itertools import ilen import supybot.cdb as cdb import supybot.utils as utils class Error(Exception): """General error for this module.""" class NoRecordError(KeyError): pass class InvalidDBError(Exception): pass class MappingInterface(object): """This is a class to represent the underlying representation of a map from integer keys to strings.""" def __init__(self, filename, **kwargs): """Feel free to ignore the filename.""" raise NotImplementedError def get(id): """Gets the record matching id. Raises NoRecordError otherwise.""" raise NotImplementedError def set(id, s): """Sets the record matching id to s.""" raise NotImplementedError def add(self, s): """Adds a new record, returning a new id for it.""" raise NotImplementedError def remove(self, id): "Returns and removes the record with the given id from the database." raise NotImplementedError def __iter__(self): "Return an iterator over (id, s) pairs. Not required to be ordered." raise NotImplementedError def flush(self): """Flushes current state to disk.""" raise NotImplementedError def close(self): """Flushes current state to disk and invalidates the Mapping.""" raise NotImplementedError def vacuum(self): "Cleans up in the database, if possible. Not required to do anything." pass class DirMapping(MappingInterface): def __init__(self, filename, **kwargs): self.dirname = filename if not os.path.exists(self.dirname): os.mkdir(self.dirname) if not os.path.exists(os.path.join(self.dirname, 'max')): self._setMax(1) def _setMax(self, id): fd = file(os.path.join(self.dirname, 'max'), 'w') try: fd.write(str(id)) finally: fd.close() def _getMax(self): fd = file(os.path.join(self.dirname, 'max')) try: i = int(fd.read()) return i finally: fd.close() def _makeFilename(self, id): return os.path.join(self.dirname, str(id)) def get(id): try: fd = file(self._makeFilename(id)) return fd.read() except EnvironmentError, e: exn = NoRecordError(id) exn.realException = e raise exn def set(id, s): fd = file(self._makeFilename(id), 'w') fd.write(s) fd.close() def add(self, s): id = self._getMax() fd = file(self._makeFilename(id), 'w') try: fd.write(s) return id finally: fd.close() def remove(self, id): try: os.remove(self._makeFilename(id)) except EnvironmentError, e: raise NoRecordError, id class FlatfileMapping(MappingInterface): def __init__(self, filename, maxSize=10**6): self.filename = filename try: fd = file(self.filename) strId = fd.readline().rstrip() self.maxSize = len(strId) try: self.currentId = int(strId) except ValueError: raise Error, 'Invalid file for FlatfileMapping: %s' % filename except EnvironmentError, e: # File couldn't be opened. self.maxSize = int(math.log10(maxSize)) self.currentId = 0 self._incrementCurrentId() def _canonicalId(self, id): if id is not None: return str(id).zfill(self.maxSize) else: return '-'*self.maxSize def _incrementCurrentId(self, fd=None): fdWasNone = fd is None if fdWasNone: fd = file(self.filename, 'a') fd.seek(0) self.currentId += 1 fd.write(self._canonicalId(self.currentId)) fd.write('\n') if fdWasNone: fd.close() def _splitLine(self, line): line = line.rstrip('\r\n') (id, s) = line.split(':', 1) return (id, s) def _joinLine(self, id, s): return '%s:%s\n' % (self._canonicalId(id), s) def add(self, s): line = self._joinLine(self.currentId, s) fd = file(self.filename, 'r+') try: fd.seek(0, 2) # End. fd.write(line) return self.currentId finally: self._incrementCurrentId(fd) fd.close() def get(self, id): strId = self._canonicalId(id) try: fd = file(self.filename) fd.readline() # First line, nextId. for line in fd: (lineId, s) = self._splitLine(line) if lineId == strId: return s raise NoRecordError, id finally: fd.close() # XXX This assumes it's not been given out. We should make sure that our # maximum id remains accurate if this is some value we've never given # out -- i.e., self.maxid = max(self.maxid, id) or something. def set(self, id, s): strLine = self._joinLine(id, s) try: fd = file(self.filename, 'r+') self.remove(id, fd) fd.seek(0, 2) # End. fd.write(strLine) finally: fd.close() def remove(self, id, fd=None): fdWasNone = fd is None strId = self._canonicalId(id) try: if fdWasNone: fd = file(self.filename, 'r+') fd.seek(0) fd.readline() # First line, nextId pos = fd.tell() line = fd.readline() while line: (lineId, _) = self._splitLine(line) if lineId == strId: fd.seek(pos) fd.write(self._canonicalId(None)) fd.seek(pos) fd.readline() # Same line we just rewrote the id for. pos = fd.tell() line = fd.readline() # We should be at the end. finally: if fdWasNone: fd.close() def __iter__(self): fd = file(self.filename) fd.readline() # First line, nextId. for line in fd: (id, s) = self._splitLine(line) if not id.startswith('-'): yield (int(id), s) fd.close() def vacuum(self): infd = file(self.filename) outfd = utils.transactionalFile(self.filename, makeBackupIfSmaller=False) outfd.write(infd.readline()) # First line, nextId. for line in infd: if not line.startswith('-'): outfd.write(line) infd.close() outfd.close() def flush(self): pass # No-op, we maintain no open files. def close(self): self.vacuum() # Should we do this? It should be fine. class CdbMapping(MappingInterface): def __init__(self, filename, **kwargs): self.filename = filename self._openCdb() # So it can be overridden later. if 'nextId' not in self.db: self.db['nextId'] = '1' def _openCdb(self, *args, **kwargs): self.db = cdb.open(filename, 'c', **kwargs) def _getNextId(self): i = int(self.db['nextId']) self.db['nextId'] = str(i+1) return i def get(self, id): try: return self.db[str(id)] except KeyError: raise NoRecordError, id # XXX Same as above. def set(self, id, s): self.db[str(id)] = s def add(self, s): id = self._getNextId() self.set(id, s) return id def remove(self, id): del self.db[str(id)] def __iter__(self): for (id, s) in self.db.iteritems(): if id != 'nextId': yield (int(id), s) def flush(self): self.db.flush() def close(self): self.db.close() class DB(object): Mapping = 'flat' # This is a good, sane default. Record = None def __init__(self, filename, Mapping=None, Record=None): if Record is not None: self.Record = Record if Mapping is not None: self.Mapping = Mapping if isinstance(self.Mapping, basestring): self.Mapping = Mappings[self.Mapping] self.map = self.Mapping(filename) def _newRecord(self, id, s): record = self.Record(id=id) record.deserialize(s) return record def get(self, id): s = self.map.get(id) return self._newRecord(id, s) def set(self, id, record): s = record.serialize() self.map.set(id, s) def add(self, record): s = record.serialize() id = self.map.add(s) record.id = id return id def remove(self, id): self.map.remove(id) def __iter__(self): for (id, s) in self.map: # We don't need to yield the id because it's in the record. yield self._newRecord(id, s) def select(self, p): for record in self: if p(record): yield record def random(self): try: return self._newRecord(*random.choice(self.map)) except IndexError: return None def size(self): return ilen(self.map) def flush(self): self.map.flush() def vacuum(self): self.map.vacuum() def close(self): self.map.close() Mappings = { 'cdb': CdbMapping, 'flat': FlatfileMapping, } class Record(object): def __init__(self, id=None, **kwargs): if id is not None: assert isinstance(id, int), 'id must be an integer.' self.id = id self.fields = [] self.defaults = {} self.converters = {} for name in self.__fields__: if isinstance(name, tuple): (name, spec) = name else: spec = utils.safeEval assert name != 'id' self.fields.append(name) if isinstance(spec, tuple): (converter, default) = spec else: converter = spec default = None self.defaults[name] = default self.converters[name] = converter seen = sets.Set() for (name, value) in kwargs.iteritems(): assert name in self.fields, 'name must be a record value.' seen.add(name) setattr(self, name, value) for name in self.fields: if name not in seen: default = self.defaults[name] if callable(default): default = default() setattr(self, name, default) def serialize(self): return csv.join([repr(getattr(self, name)) for name in self.fields]) def deserialize(self, s): unseenRecords = sets.Set(self.fields) for (name, strValue) in zip(self.fields, csv.split(s)): setattr(self, name, self.converters[name](strValue)) unseenRecords.remove(name) for name in unseenRecords: setattr(self, name, self.defaults[name]) # vim:set shiftwidth=4 tabstop=8 expandtab textwidth=78: