minidb.py

   1 #!/usr/bin/python
   2 # -*- coding: utf-8 -*-
   3 #
   4 # minidb - A simple SQLite3 store for Python objects
   5 # (based on "ORM wie eine Kirchenmaus" by thp, 2009-11-29)
   6 #
   7 # Copyright 2009-2010 Thomas Perl <thp.io>. All rights reserved.
   8 #
   9 # Permission to use, copy, modify, and/or distribute this software for any
  10 # purpose with or without fee is hereby granted, provided that the above
  11 # copyright notice and this permission notice appear in all copies.
  12 #
  13 # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
  14 # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
  15 # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
  16 # ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
  17 # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
  18 # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
  19 # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
  20 #
  21
  22 """A simple SQLite3 store for Python objects"""
  23
  24 # For Python 2.5, we need to request the "with" statement
  25 from __future__ import with_statement
  26
  27 __author__ = 'Thomas Perl <m@thp.io>'
  28 __version__ = '1.1'
  29 __website__ = 'http://thp.io/2010/minidb/'
  30 __license__ = 'ISC'
  31
  32 try:
  33     import sqlite3.dbapi2 as sqlite
  34 except ImportError:
  35     try:
  36         from pysqlite2 import dbapi2 as sqlite
  37     except ImportError:
  38         raise Exception('Please install SQLite3 support.')
  39
  40
  41 import threading
  42
  43 class Store(object):
  44     def __init__(self, filename=':memory:'):
  45         """Create (or load) a new minidb storage
  46
  47         Without arguments, this will create an in-memory
  48         database that will be deleted when closed. If you
  49         pass an argument, it should be the filename of the
  50         database file (which will be created if it does
  51         not yet exist).
  52         """
  53         self.db = sqlite.connect(filename, check_same_thread=False)
  54         self.lock = threading.RLock()
  55
  56     def _schema(self, class_):
  57         return class_.__name__, list(sorted(class_.__slots__))
  58
  59     def _set(self, o, slot, value):
  60         # Set a slot on the given object to value, doing a cast if
  61         # necessary. The value None is special-cased and never cast.
  62         cls = o.__class__.__slots__[slot]
  63         if value is not None:
  64             if isinstance(value, unicode):
  65                 value = value.decode('utf-8')
  66             value = cls(value)
  67         setattr(o, slot, value)
  68
  69     def commit(self):
  70         """Commit changes into the database"""
  71         with self.lock:
  72             self.db.commit()
  73
  74     def close(self):
  75         """Close the underlying database file"""
  76         with self.lock:
  77             self.db.execute('VACUUM')
  78             self.db.close()
  79
  80     def _register(self, class_):
  81         with self.lock:
  82             table, slots = self._schema(class_)
  83             cur = self.db.execute('PRAGMA table_info(%s)' % table)
  84             available = cur.fetchall()
  85
  86             if available:
  87                 available = [row[1] for row in available]
  88                 missing_slots = (s for s in slots if s not in available)
  89                 for slot in missing_slots:
  90                     self.db.execute('ALTER TABLE %s ADD COLUMN %s TEXT' % (table,
  91                         slot))
  92             else:
  93                 self.db.execute('CREATE TABLE %s (%s)' % (table,
  94                         ', '.join('%s TEXT'%s for s in slots)))
  95
  96     def convert(self, v):
  97         """Convert a value to its string representation"""
  98         if isinstance(v, unicode):
  99             return v
 100         elif isinstance(v, str):
 101             return v.decode('utf-8')
 102         else:
 103             return str(v)
 104
 105     def update(self, o, **kwargs):
 106         """Update fields of an object and store the changes
 107
 108         This will update named fields (specified by keyword
 109         arguments) inside the object and also store these
 110         changes in the database.
 111         """
 112         self.remove(o)
 113         for k, v in kwargs.items():
 114             setattr(o, k, v)
 115         self.save(o)
 116
 117     def save(self, o):
 118         """Save an object into the database
 119
 120         Save a newly-created object into the database. The
 121         object will always be newly created, never updated.
 122
 123         If you want to update an object inside the database,
 124         please use the "update" method instead.
 125         """
 126         if hasattr(o, '__iter__'):
 127             for child in o:
 128                 self.save(child)
 129             return
 130
 131         with self.lock:
 132             self._register(o.__class__)
 133             table, slots = self._schema(o.__class__)
 134
 135             # Only save values that have values set (non-None values)
 136             slots = [s for s in slots if getattr(o, s, None) is not None]
 137
 138             values = [self.convert(getattr(o, slot)) for slot in slots]
 139             self.db.execute('INSERT INTO %s (%s) VALUES (%s)' % (table,
 140                 ', '.join(slots), ', '.join('?'*len(slots))), values)
 141
 142     def delete(self, class_, **kwargs):
 143         """Delete objects from the database
 144
 145         Delete objects of type "class_" with the criteria
 146         specified in "kwargs". Please note that all objects
 147         that match the criteria will be deleted.
 148
 149         If you want to remove a specific object from the
 150         database, use "remove" instead.
 151         """
 152         with self.lock:
 153             self._register(class_)
 154             table, slots = self._schema(class_)
 155             sql = 'DELETE FROM %s' % (table,)
 156             if kwargs:
 157                 sql += ' WHERE %s' % (' AND '.join('%s=?' % k for k in kwargs))
 158             try:
 159                 self.db.execute(sql, kwargs.values())
 160                 return True
 161             except Exception, e:
 162                 return False
 163
 164     def remove(self, o):
 165         """Delete objects by template object
 166
 167         This will remove all objects from the database that
 168         compare to the given object (i.e. all attributes of
 169         "o" that are not None will match to the objects in
 170         the database).
 171
 172         This method should be used to remove specific object
 173         only. For bulk deletion based on some criteria, the
 174         "delete" method might be better suited.
 175         """
 176         if hasattr(o, '__iter__'):
 177             for child in o:
 178                 self.remove(child)
 179             return
 180
 181         with self.lock:
 182             self._register(o.__class__)
 183             table, slots = self._schema(o.__class__)
 184
 185             # Use "None" as wildcard selector in remove actions
 186             slots = [s for s in slots if getattr(o, s, None) is not None]
 187
 188             values = [self.convert(getattr(o, slot)) for slot in slots]
 189             self.db.execute('DELETE FROM %s WHERE %s' % (table,
 190                 ' AND '.join('%s=?'%s for s in slots)), values)
 191
 192     def load(self, class_, **kwargs):
 193         """Load objects of a given class
 194
 195         Return a list of objects from the database that are of
 196         type "class_". By default, all objects are returned,
 197         but a simple pre-selection can be made using keyword
 198         arguments.
 199         """
 200         with self.lock:
 201             self._register(class_)
 202             table, slots = self._schema(class_)
 203             sql = 'SELECT %s FROM %s' % (', '.join(slots), table)
 204             if kwargs:
 205                 sql += ' WHERE %s' % (' AND '.join('%s=?' % k for k in kwargs))
 206             try:
 207                 cur = self.db.execute(sql, kwargs.values())
 208             except Exception, e:
 209                 raise
 210             def apply(row):
 211                 o = class_.__new__(class_)
 212                 for attr, value in zip(slots, row):
 213                     try:
 214                         self._set(o, attr, value)
 215                     except ValueError, ve:
 216                         return None
 217                 return o
 218             return filter(lambda x: x is not None, [apply(row) for row in cur])
 219
 220     def get(self, class_, **kwargs):
 221         """Load one object of a given class
 222
 223         This is a convenience function that will load only a
 224         single object from the database, returning only that
 225         object or None when the object is not found.
 226
 227         This method only makes sense when using keyword
 228         arguments to select the object (i.e. using a
 229         unique set of attributes to retrieve it).
 230         """
 231         result = self.load(class_, **kwargs)
 232         if result:
 233             return result[0]
 234         else:
 235             return None
 236
 237 if __name__ == '__main__':
 238     class Person(object):
 239         __slots__ = {'username': str, 'id': int}
 240
 241         def __init__(self, username, id):
 242             self.username = username
 243             self.id = id
 244
 245         def __repr__(self):
 246             return '<Person "%s" (%d)>' % (self.username, self.id)
 247
 248     m = Store()
 249     m.save(Person('User %d' % x, x*20) for x in range(50))
 250
 251     p = m.get(Person, id=200)
 252     print p
 253     m.remove(p)
 254     p = m.get(Person, id=200)
 255
 256     # Remove some persons again (deletion by value!)
 257     m.remove(Person('User %d' % x, x*20) for x in range(40))
 258
 259     class Person(object):
 260         __slots__ = {'username': str, 'id': int, 'mail': str}
 261
 262         def __init__(self, username, id, mail):
 263             self.username = username
 264             self.id = id
 265             self.mail = mail
 266
 267         def __repr__(self):
 268             return '<Person "%s" (%s)>' % (self.username, self.mail)
 269
 270     # A schema update takes place here
 271     m.save(Person('User %d' % x, x*20, 'user@home.com') for x in range(50))
 272     print m.load(Person)
 273