libs/openid/store/sqlstore.py

   1 """
   2 This module contains C{L{OpenIDStore}} implementations that use
   3 various SQL databases to back them.
   4
   5 Example of how to initialize a store database::
   6
   7     python -c 'from openid.store import sqlstore; import pysqlite2.dbapi2; sqlstore.SQLiteStore(pysqlite2.dbapi2.connect("cstore.db")).createTables()'
   8 """
   9 import re
  10 import time
  11
  12 from openid.association import Association
  13 from openid.store.interface import OpenIDStore
  14 from openid.store import nonce
  15
  16 def _inTxn(func):
  17     def wrapped(self, *args, **kwargs):
  18         return self._callInTransaction(func, self, *args, **kwargs)
  19
  20     if hasattr(func, '__name__'):
  21         try:
  22             wrapped.__name__ = func.__name__[4:]
  23         except TypeError:
  24             pass
  25
  26     if hasattr(func, '__doc__'):
  27         wrapped.__doc__ = func.__doc__
  28
  29     return wrapped
  30
  31 class SQLStore(OpenIDStore):
  32     """
  33     This is the parent class for the SQL stores, which contains the
  34     logic common to all of the SQL stores.
  35
  36     The table names used are determined by the class variables
  37     C{L{associations_table}} and
  38     C{L{nonces_table}}.  To change the name of the tables used, pass
  39     new table names into the constructor.
  40
  41     To create the tables with the proper schema, see the
  42     C{L{createTables}} method.
  43
  44     This class shouldn't be used directly.  Use one of its subclasses
  45     instead, as those contain the code necessary to use a specific
  46     database.
  47
  48     All methods other than C{L{__init__}} and C{L{createTables}}
  49     should be considered implementation details.
  50
  51
  52     @cvar associations_table: This is the default name of the table to
  53         keep associations in
  54
  55     @cvar nonces_table: This is the default name of the table to keep
  56         nonces in.
  57
  58
  59     @sort: __init__, createTables
  60     """
  61
  62     associations_table = 'oid_associations'
  63     nonces_table = 'oid_nonces'
  64
  65     def __init__(self, conn, associations_table=None, nonces_table=None):
  66         """
  67         This creates a new SQLStore instance.  It requires an
  68         established database connection be given to it, and it allows
  69         overriding the default table names.
  70
  71
  72         @param conn: This must be an established connection to a
  73             database of the correct type for the SQLStore subclass
  74             you're using.
  75
  76         @type conn: A python database API compatible connection
  77             object.
  78
  79
  80         @param associations_table: This is an optional parameter to
  81             specify the name of the table used for storing
  82             associations.  The default value is specified in
  83             C{L{SQLStore.associations_table}}.
  84
  85         @type associations_table: C{str}
  86
  87
  88         @param nonces_table: This is an optional parameter to specify
  89             the name of the table used for storing nonces.  The
  90             default value is specified in C{L{SQLStore.nonces_table}}.
  91
  92         @type nonces_table: C{str}
  93         """
  94         self.conn = conn
  95         self.cur = None
  96         self._statement_cache = {}
  97         self._table_names = {
  98             'associations': associations_table or self.associations_table,
  99             'nonces': nonces_table or self.nonces_table,
 100             }
 101         self.max_nonce_age = 6 * 60 * 60 # Six hours, in seconds
 102
 103         # DB API extension: search for "Connection Attributes .Error,
 104         # .ProgrammingError, etc." in
 105         # http://www.python.org/dev/peps/pep-0249/
 106         if (hasattr(self.conn, 'IntegrityError') and
 107             hasattr(self.conn, 'OperationalError')):
 108             self.exceptions = self.conn
 109
 110         if not (hasattr(self.exceptions, 'IntegrityError') and
 111                 hasattr(self.exceptions, 'OperationalError')):
 112             raise RuntimeError("Error using database connection module "
 113                                "(Maybe it can't be imported?)")
 114
 115     def blobDecode(self, blob):
 116         """Convert a blob as returned by the SQL engine into a str object.
 117
 118         str -> str"""
 119         return blob
 120
 121     def blobEncode(self, s):
 122         """Convert a str object into the necessary object for storing
 123         in the database as a blob."""
 124         return s
 125
 126     def _getSQL(self, sql_name):
 127         try:
 128             return self._statement_cache[sql_name]
 129         except KeyError:
 130             sql = getattr(self, sql_name)
 131             sql %= self._table_names
 132             self._statement_cache[sql_name] = sql
 133             return sql
 134
 135     def _execSQL(self, sql_name, *args):
 136         sql = self._getSQL(sql_name)
 137         # Kludge because we have reports of postgresql not quoting
 138         # arguments if they are passed in as unicode instead of str.
 139         # Currently the strings in our tables just have ascii in them,
 140         # so this ought to be safe.
 141         def unicode_to_str(arg):
 142             if isinstance(arg, unicode):
 143                 return str(arg)
 144             else:
 145                 return arg
 146         str_args = map(unicode_to_str, args)
 147         self.cur.execute(sql, str_args)
 148
 149     def __getattr__(self, attr):
 150         # if the attribute starts with db_, use a default
 151         # implementation that looks up the appropriate SQL statement
 152         # as an attribute of this object and executes it.
 153         if attr[:3] == 'db_':
 154             sql_name = attr[3:] + '_sql'
 155             def func(*args):
 156                 return self._execSQL(sql_name, *args)
 157             setattr(self, attr, func)
 158             return func
 159         else:
 160             raise AttributeError('Attribute %r not found' % (attr,))
 161
 162     def _callInTransaction(self, func, *args, **kwargs):
 163         """Execute the given function inside of a transaction, with an
 164         open cursor. If no exception is raised, the transaction is
 165         comitted, otherwise it is rolled back."""
 166         # No nesting of transactions
 167         self.conn.rollback()
 168
 169         try:
 170             self.cur = self.conn.cursor()
 171             try:
 172                 ret = func(*args, **kwargs)
 173             finally:
 174                 self.cur.close()
 175                 self.cur = None
 176         except:
 177             self.conn.rollback()
 178             raise
 179         else:
 180             self.conn.commit()
 181
 182         return ret
 183
 184     def txn_createTables(self):
 185         """
 186         This method creates the database tables necessary for this
 187         store to work.  It should not be called if the tables already
 188         exist.
 189         """
 190         self.db_create_nonce()
 191         self.db_create_assoc()
 192
 193     createTables = _inTxn(txn_createTables)
 194
 195     def txn_storeAssociation(self, server_url, association):
 196         """Set the association for the server URL.
 197
 198         Association -> NoneType
 199         """
 200         a = association
 201         self.db_set_assoc(
 202             server_url,
 203             a.handle,
 204             self.blobEncode(a.secret),
 205             a.issued,
 206             a.lifetime,
 207             a.assoc_type)
 208
 209     storeAssociation = _inTxn(txn_storeAssociation)
 210
 211     def txn_getAssociation(self, server_url, handle=None):
 212         """Get the most recent association that has been set for this
 213         server URL and handle.
 214
 215         str -> NoneType or Association
 216         """
 217         if handle is not None:
 218             self.db_get_assoc(server_url, handle)
 219         else:
 220             self.db_get_assocs(server_url)
 221
 222         rows = self.cur.fetchall()
 223         if len(rows) == 0:
 224             return None
 225         else:
 226             associations = []
 227             for values in rows:
 228                 assoc = Association(*values)
 229                 assoc.secret = self.blobDecode(assoc.secret)
 230                 if assoc.getExpiresIn() == 0:
 231                     self.txn_removeAssociation(server_url, assoc.handle)
 232                 else:
 233                     associations.append((assoc.issued, assoc))
 234
 235             if associations:
 236                 associations.sort()
 237                 return associations[-1][1]
 238             else:
 239                 return None
 240
 241     getAssociation = _inTxn(txn_getAssociation)
 242
 243     def txn_removeAssociation(self, server_url, handle):
 244         """Remove the association for the given server URL and handle,
 245         returning whether the association existed at all.
 246
 247         (str, str) -> bool
 248         """
 249         self.db_remove_assoc(server_url, handle)
 250         return self.cur.rowcount > 0 # -1 is undefined
 251
 252     removeAssociation = _inTxn(txn_removeAssociation)
 253
 254     def txn_useNonce(self, server_url, timestamp, salt):
 255         """Return whether this nonce is present, and if it is, then
 256         remove it from the set.
 257
 258         str -> bool"""
 259         if abs(timestamp - time.time()) > nonce.SKEW:
 260             return False
 261
 262         try:
 263             self.db_add_nonce(server_url, timestamp, salt)
 264         except self.exceptions.IntegrityError:
 265             # The key uniqueness check failed
 266             return False
 267         else:
 268             # The nonce was successfully added
 269             return True
 270
 271     useNonce = _inTxn(txn_useNonce)
 272
 273     def txn_cleanupNonces(self):
 274         self.db_clean_nonce(int(time.time()) - nonce.SKEW)
 275         return self.cur.rowcount
 276
 277     cleanupNonces = _inTxn(txn_cleanupNonces)
 278
 279     def txn_cleanupAssociations(self):
 280         self.db_clean_assoc(int(time.time()))
 281         return self.cur.rowcount
 282
 283     cleanupAssociations = _inTxn(txn_cleanupAssociations)
 284
 285
 286 class SQLiteStore(SQLStore):
 287     """
 288     This is an SQLite-based specialization of C{L{SQLStore}}.
 289
 290     To create an instance, see C{L{SQLStore.__init__}}.  To create the
 291     tables it will use, see C{L{SQLStore.createTables}}.
 292
 293     All other methods are implementation details.
 294     """
 295
 296     create_nonce_sql = """
 297     CREATE TABLE %(nonces)s (
 298         server_url VARCHAR,
 299         timestamp INTEGER,
 300         salt CHAR(40),
 301         UNIQUE(server_url, timestamp, salt)
 302     );
 303     """
 304
 305     create_assoc_sql = """
 306     CREATE TABLE %(associations)s
 307     (
 308         server_url VARCHAR(2047),
 309         handle VARCHAR(255),
 310         secret BLOB(128),
 311         issued INTEGER,
 312         lifetime INTEGER,
 313         assoc_type VARCHAR(64),
 314         PRIMARY KEY (server_url, handle)
 315     );
 316     """
 317
 318     set_assoc_sql = ('INSERT OR REPLACE INTO %(associations)s '
 319                      '(server_url, handle, secret, issued, '
 320                      'lifetime, assoc_type) '
 321                      'VALUES (?, ?, ?, ?, ?, ?);')
 322     get_assocs_sql = ('SELECT handle, secret, issued, lifetime, assoc_type '
 323                       'FROM %(associations)s WHERE server_url = ?;')
 324     get_assoc_sql = (
 325         'SELECT handle, secret, issued, lifetime, assoc_type '
 326         'FROM %(associations)s WHERE server_url = ? AND handle = ?;')
 327
 328     get_expired_sql = ('SELECT server_url '
 329                        'FROM %(associations)s WHERE issued + lifetime < ?;')
 330
 331     remove_assoc_sql = ('DELETE FROM %(associations)s '
 332                         'WHERE server_url = ? AND handle = ?;')
 333
 334     clean_assoc_sql = 'DELETE FROM %(associations)s WHERE issued + lifetime < ?;'
 335
 336     add_nonce_sql = 'INSERT INTO %(nonces)s VALUES (?, ?, ?);'
 337
 338     clean_nonce_sql = 'DELETE FROM %(nonces)s WHERE timestamp < ?;'
 339
 340     def blobDecode(self, buf):
 341         return str(buf)
 342
 343     def blobEncode(self, s):
 344         return buffer(s)
 345
 346     def useNonce(self, *args, **kwargs):
 347         # Older versions of the sqlite wrapper do not raise
 348         # IntegrityError as they should, so we have to detect the
 349         # message from the OperationalError.
 350         try:
 351             return super(SQLiteStore, self).useNonce(*args, **kwargs)
 352         except self.exceptions.OperationalError, why:
 353             if re.match('^columns .* are not unique$', why[0]):
 354                 return False
 355             else:
 356                 raise
 357
 358 class MySQLStore(SQLStore):
 359     """
 360     This is a MySQL-based specialization of C{L{SQLStore}}.
 361
 362     Uses InnoDB tables for transaction support.
 363
 364     To create an instance, see C{L{SQLStore.__init__}}.  To create the
 365     tables it will use, see C{L{SQLStore.createTables}}.
 366
 367     All other methods are implementation details.
 368     """
 369
 370     try:
 371         import MySQLdb as exceptions
 372     except ImportError:
 373         exceptions = None
 374
 375     create_nonce_sql = """
 376     CREATE TABLE %(nonces)s (
 377         server_url BLOB NOT NULL,
 378         timestamp INTEGER NOT NULL,
 379         salt CHAR(40) NOT NULL,
 380         PRIMARY KEY (server_url(255), timestamp, salt)
 381     )
 382     ENGINE=InnoDB;
 383     """
 384
 385     create_assoc_sql = """
 386     CREATE TABLE %(associations)s
 387     (
 388         server_url BLOB NOT NULL,
 389         handle VARCHAR(255) NOT NULL,
 390         secret BLOB NOT NULL,
 391         issued INTEGER NOT NULL,
 392         lifetime INTEGER NOT NULL,
 393         assoc_type VARCHAR(64) NOT NULL,
 394         PRIMARY KEY (server_url(255), handle)
 395     )
 396     ENGINE=InnoDB;
 397     """
 398
 399     set_assoc_sql = ('REPLACE INTO %(associations)s '
 400                      'VALUES (%%s, %%s, %%s, %%s, %%s, %%s);')
 401     get_assocs_sql = ('SELECT handle, secret, issued, lifetime, assoc_type'
 402                       ' FROM %(associations)s WHERE server_url = %%s;')
 403     get_expired_sql = ('SELECT server_url '
 404                        'FROM %(associations)s WHERE issued + lifetime < %%s;')
 405
 406     get_assoc_sql = (
 407         'SELECT handle, secret, issued, lifetime, assoc_type'
 408         ' FROM %(associations)s WHERE server_url = %%s AND handle = %%s;')
 409     remove_assoc_sql = ('DELETE FROM %(associations)s '
 410                         'WHERE server_url = %%s AND handle = %%s;')
 411
 412     clean_assoc_sql = 'DELETE FROM %(associations)s WHERE issued + lifetime < %%s;'
 413
 414     add_nonce_sql = 'INSERT INTO %(nonces)s VALUES (%%s, %%s, %%s);'
 415
 416     clean_nonce_sql = 'DELETE FROM %(nonces)s WHERE timestamp < %%s;'
 417
 418     def blobDecode(self, blob):
 419         if type(blob) is str:
 420             # Versions of MySQLdb >= 1.2.2
 421             return blob
 422         else:
 423             # Versions of MySQLdb prior to 1.2.2 (as far as we can tell)
 424             return blob.tostring()
 425
 426 class PostgreSQLStore(SQLStore):
 427     """
 428     This is a PostgreSQL-based specialization of C{L{SQLStore}}.
 429
 430     To create an instance, see C{L{SQLStore.__init__}}.  To create the
 431     tables it will use, see C{L{SQLStore.createTables}}.
 432
 433     All other methods are implementation details.
 434     """
 435
 436     try:
 437         import psycopg as exceptions
 438     except ImportError:
 439         # psycopg2 has the dbapi extension where the exception classes
 440         # are available on the connection object. A psycopg2
 441         # connection will use the correct exception classes because of
 442         # this, and a psycopg connection will fall through to use the
 443         # psycopg imported above.
 444         exceptions = None
 445
 446     create_nonce_sql = """
 447     CREATE TABLE %(nonces)s (
 448         server_url VARCHAR(2047) NOT NULL,
 449         timestamp INTEGER NOT NULL,
 450         salt CHAR(40) NOT NULL,
 451         PRIMARY KEY (server_url, timestamp, salt)
 452     );
 453     """
 454
 455     create_assoc_sql = """
 456     CREATE TABLE %(associations)s
 457     (
 458         server_url VARCHAR(2047) NOT NULL,
 459         handle VARCHAR(255) NOT NULL,
 460         secret BYTEA NOT NULL,
 461         issued INTEGER NOT NULL,
 462         lifetime INTEGER NOT NULL,
 463         assoc_type VARCHAR(64) NOT NULL,
 464         PRIMARY KEY (server_url, handle),
 465         CONSTRAINT secret_length_constraint CHECK (LENGTH(secret) <= 128)
 466     );
 467     """
 468
 469     def db_set_assoc(self, server_url, handle, secret, issued, lifetime, assoc_type):
 470         """
 471         Set an association.  This is implemented as a method because
 472         REPLACE INTO is not supported by PostgreSQL (and is not
 473         standard SQL).
 474         """
 475         result = self.db_get_assoc(server_url, handle)
 476         rows = self.cur.fetchall()
 477         if len(rows):
 478             # Update the table since this associations already exists.
 479             return self.db_update_assoc(secret, issued, lifetime, assoc_type,
 480                                         server_url, handle)
 481         else:
 482             # Insert a new record because this association wasn't
 483             # found.
 484             return self.db_new_assoc(server_url, handle, secret, issued,
 485                                      lifetime, assoc_type)
 486
 487     new_assoc_sql = ('INSERT INTO %(associations)s '
 488                      'VALUES (%%s, %%s, %%s, %%s, %%s, %%s);')
 489     update_assoc_sql = ('UPDATE %(associations)s SET '
 490                         'secret = %%s, issued = %%s, '
 491                         'lifetime = %%s, assoc_type = %%s '
 492                         'WHERE server_url = %%s AND handle = %%s;')
 493     get_assocs_sql = ('SELECT handle, secret, issued, lifetime, assoc_type'
 494                       ' FROM %(associations)s WHERE server_url = %%s;')
 495     get_expired_sql = ('SELECT server_url '
 496                        'FROM %(associations)s WHERE issued + lifetime < %%s;')
 497
 498     get_assoc_sql = (
 499         'SELECT handle, secret, issued, lifetime, assoc_type'
 500         ' FROM %(associations)s WHERE server_url = %%s AND handle = %%s;')
 501     remove_assoc_sql = ('DELETE FROM %(associations)s '
 502                         'WHERE server_url = %%s AND handle = %%s;')
 503
 504     clean_assoc_sql = 'DELETE FROM %(associations)s WHERE issued + lifetime < %%s;'
 505
 506     add_nonce_sql = 'INSERT INTO %(nonces)s VALUES (%%s, %%s, %%s);'
 507
 508     clean_nonce_sql = 'DELETE FROM %(nonces)s WHERE timestamp < %%s;'
 509
 510     def blobEncode(self, blob):
 511         try:
 512             from psycopg2 import Binary
 513         except ImportError:
 514             from psycopg import Binary
 515
 516         return Binary(blob)