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.0'
  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.close()
  78
  79     def _register(self, class_):
  80         with self.lock:
  81             table, slots = self._schema(class_)
  82             cur = self.db.execute('PRAGMA table_info(%s)' % table)
  83             available = cur.fetchall()
  84
  85             if available:
  86                 available = [row[1] for row in available]
  87                 missing_slots = (s for s in slots if s not in available)
  88                 for slot in missing_slots:
  89                     self.db.execute('ALTER TABLE %s ADD COLUMN %s TEXT' % (table,
  90                         slot))
  91             else:
  92                 self.db.execute('CREATE TABLE %s (%s)' % (table,
  93                         ', '.join('%s TEXT'%s for s in slots)))
  94
  95     def convert(self, v):
  96         """Convert a value to its string representation"""
  97         if isinstance(v, str) or isinstance(v, unicode):
  98             return v
  99         else:
 100             return str(v)
 101
 102     def update(self, o, **kwargs):
 103         """Update fields of an object and store the changes
 104
 105         This will update named fields (specified by keyword
 106         arguments) inside the object and also store these
 107         changes in the database.
 108         """
 109         self.remove(o)
 110         for k, v in kwargs.items():
 111             setattr(o, k, v)
 112         self.save(o)
 113
 114     def save(self, o):
 115         """Save an object into the database
 116
 117         Save a newly-created object into the database. The
 118         object will always be newly created, never updated.
 119
 120         If you want to update an object inside the database,
 121         please use the "update" method instead.
 122         """
 123         if hasattr(o, '__iter__'):
 124             for child in o:
 125                 self.save(child)
 126             return
 127
 128         with self.lock:
 129             self._register(o.__class__)
 130             table, slots = self._schema(o.__class__)
 131
 132             # Only save values that have values set (non-None values)
 133             slots = [s for s in slots if getattr(o, s, None) is not None]
 134
 135             values = [self.convert(getattr(o, slot)) for slot in slots]
 136             self.db.execute('INSERT INTO %s (%s) VALUES (%s)' % (table,
 137                 ', '.join(slots), ', '.join('?'*len(slots))), values)
 138
 139     def delete(self, class_, **kwargs):
 140         """Delete objects from the database
 141
 142         Delete objects of type "class_" with the criteria
 143         specified in "kwargs". Please note that all objects
 144         that match the criteria will be deleted.
 145
 146         If you want to remove a specific object from the
 147         database, use "remove" instead.
 148         """
 149         with self.lock:
 150             self._register(class_)
 151             table, slots = self._schema(class_)
 152             sql = 'DELETE FROM %s' % (table,)
 153             if kwargs:
 154                 sql += ' WHERE %s' % (' AND '.join('%s=?' % k for k in kwargs))
 155             try:
 156                 self.db.execute(sql, kwargs.values())
 157                 return True
 158             except Exception, e:
 159                 return False
 160
 161     def remove(self, o):
 162         """Delete objects by template object
 163
 164         This will remove all objects from the database that
 165         compare to the given object (i.e. all attributes of
 166         "o" that are not None will match to the objects in
 167         the database).
 168
 169         This method should be used to remove specific object
 170         only. For bulk deletion based on some criteria, the
 171         "delete" method might be better suited.
 172         """
 173         if hasattr(o, '__iter__'):
 174             for child in o:
 175                 self.remove(child)
 176             return
 177
 178         with self.lock:
 179             self._register(o.__class__)
 180             table, slots = self._schema(o.__class__)
 181
 182             # Use "None" as wildcard selector in remove actions
 183             slots = [s for s in slots if getattr(o, s, None) is not None]
 184
 185             values = [self.convert(getattr(o, slot)) for slot in slots]
 186             self.db.execute('DELETE FROM %s WHERE %s' % (table,
 187                 ' AND '.join('%s=?'%s for s in slots)), values)
 188
 189     def load(self, class_, **kwargs):
 190         """Load objects of a given class
 191
 192         Return a list of objects from the database that are of
 193         type "class_". By default, all objects are returned,
 194         but a simple pre-selection can be made using keyword
 195         arguments.
 196         """
 197         with self.lock:
 198             self._register(class_)
 199             table, slots = self._schema(class_)
 200             sql = 'SELECT %s FROM %s' % (', '.join(slots), table)
 201             if kwargs:
 202                 sql += ' WHERE %s' % (' AND '.join('%s=?' % k for k in kwargs))
 203             try:
 204                 cur = self.db.execute(sql, kwargs.values())
 205             except Exception, e:
 206                 raise
 207             def apply(row):
 208                 o = class_.__new__(class_)
 209                 for attr, value in zip(slots, row):
 210                     try:
 211                         self._set(o, attr, value)
 212                     except ValueError, ve:
 213                         return None
 214                 return o
 215             return filter(lambda x: x is not None, [apply(row) for row in cur])
 216
 217     def get(self, class_, **kwargs):
 218         """Load one object of a given class
 219
 220         This is a convenience function that will load only a
 221         single object from the database, returning only that
 222         object or None when the object is not found.
 223
 224         This method only makes sense when using keyword
 225         arguments to select the object (i.e. using a
 226         unique set of attributes to retrieve it).
 227         """
 228         result = self.load(class_, **kwargs)
 229         if result:
 230             return result[0]
 231         else:
 232             return None
 233
 234 if __name__ == '__main__':
 235     class Person(object):
 236         __slots__ = {'username': str, 'id': int}
 237
 238         def __init__(self, username, id):
 239             self.username = username
 240             self.id = id
 241
 242         def __repr__(self):
 243             return '<Person "%s" (%d)>' % (self.username, self.id)
 244
 245     m = Store()
 246     m.save(Person('User %d' % x, x*20) for x in range(50))
 247
 248     p = m.get(Person, id=200)
 249     print p
 250     m.remove(p)
 251     p = m.get(Person, id=200)
 252
 253     # Remove some persons again (deletion by value!)
 254     m.remove(Person('User %d' % x, x*20) for x in range(40))
 255
 256     class Person(object):
 257         __slots__ = {'username': str, 'id': int, 'mail': str}
 258
 259         def __init__(self, username, id, mail):
 260             self.username = username
 261             self.id = id
 262             self.mail = mail
 263
 264         def __repr__(self):
 265             return '<Person "%s" (%s)>' % (self.username, self.mail)
 266
 267     # A schema update takes place here
 268     m.save(Person('User %d' % x, x*20, 'user@home.com') for x in range(50))
 269     print m.load(Person)
 270