from __future__ import generators True, False = 1==1, 0==1 import threading import re import warnings import atexit import os import new import sqlbuilder from cache import CacheSet import col from joins import sorter from converters import sqlrepr import urllib warnings.filterwarnings("ignore", "DB-API extension cursor.lastrowid used") _connections = {} class DBConnection: def __init__(self, name=None, debug=False, debugOutput=False, cache=True, style=None, autoCommit=True, debugThreading=False): self.name = name self.debug = debug self.debugOutput = debugOutput self.debugThreading = debugThreading self.cache = CacheSet(cache=cache) self.doCache = cache self.style = style self._connectionNumbers = {} self._connectionCount = 1 self.autoCommit = autoCommit registerConnectionInstance(self) def uri(self): auth = self.user or '' if auth: if self.password: auth = auth + '@' + self.password auth = auth + ':' else: assert not password, 'URIs cannot express passwords without usernames' uri = '%s://%s' % (self.dbName, auth) if self.host: uri += self.host + '/' if path.startswith('/'): path = path[1:] return uri + path def isSupported(cls): raise NotImplemented isSupported = classmethod(isSupported) def connectionFromURI(cls, uri): raise NotImplemented connectionFromURI = classmethod(connectionFromURI) def _parseURI(uri): schema, rest = uri.split(':', 1) assert rest.startswith('/'), "URIs must start with scheme:/ -- you did not include a / (in %r)" % rest if rest.startswith('/') and not rest.startswith('//'): host = None rest = rest[1:] elif rest.startswith('///'): host = None rest = rest[3:] else: rest = rest[2:] if rest.find('/') == -1: host = rest rest = '' else: host, rest = rest.split('/', 1) if host and host.find('@') != -1: user, host = host.split('@', 1) if user.find(':') != -1: user, password = user.split(':', 1) else: password = None else: user = password = None path = '/' + rest args = {} if path.find('?') != -1: path, arglist = path.split('?', 1) arglist = arglist.split('&') for single in arglist: argname, argvalue = single.split('=', 1) argvalue = urllib.unquote(argvalue) args[argname] = argvalue return user, password, host, path, args _parseURI = staticmethod(_parseURI) class DBAPI(DBConnection): """ Subclass must define a `makeConnection()` method, which returns a newly-created connection object. ``queryInsertID`` must also be defined. """ dbName = None def __init__(self, **kw): self._pool = [] self._poolLock = threading.Lock() DBConnection.__init__(self, **kw) def _runWithConnection(self, meth, *args): conn = self.getConnection() try: val = meth(conn, *args) finally: self.releaseConnection(conn) return val def getConnection(self): self._poolLock.acquire() try: if not self._pool: newConn = self.makeConnection() self._pool.append(newConn) self._connectionNumbers[id(newConn)] = self._connectionCount self._connectionCount += 1 val = self._pool.pop() if self.debug: s = 'ACQUIRE' if self._pool is not None: s += ' pool=[%s]' % ', '.join([str(self._connectionNumbers[id(v)]) for v in self._pool]) self.printDebug(val, s, 'Pool') return val finally: self._poolLock.release() def releaseConnection(self, conn, explicit=False): if self.debug: if explicit: s = 'RELEASE (explicit)' else: s = 'RELEASE (implicit, autocommit=%s)' % self.autoCommit if self._pool is None: s += ' no pooling' else: s += ' pool=[%s]' % ', '.join([str(self._connectionNumbers[id(v)]) for v in self._pool]) self.printDebug(conn, s, 'Pool') if self.supportTransactions and not explicit: if self.autoCommit == 'exception': if self.debug: self.printDebug(conn, 'auto/exception', 'ROLLBACK') conn.rollback() raise Exception, 'Object used outside of a transaction; implicit COMMIT or ROLLBACK not allowed' elif self.autoCommit: if self.debug: self.printDebug(conn, 'auto', 'COMMIT') conn.commit() else: if self.debug: self.printDebug(conn, 'auto', 'ROLLBACK') conn.rollback() if self._pool is not None: if conn not in self._pool: # @@: We can get duplicate releasing of connections with # the __del__ in Iteration (unfortunately, not sure why # it happens) self._pool.append(conn) def printDebug(self, conn, s, name, type='query'): if type == 'query': sep = ': ' else: sep = '->' s = repr(s) n = self._connectionNumbers[id(conn)] spaces = ' '*(8-len(name)) if self.debugThreading: threadName = threading.currentThread().getName() threadName = (':' + threadName + ' '*(8-len(threadName))) else: threadName = '' print '%(n)2i%(threadName)s/%(name)s%(spaces)s%(sep)s %(s)s' % locals() def _executeRetry(self, conn, cursor, query): return cursor.execute(query) def _query(self, conn, s): if self.debug: self.printDebug(conn, s, 'Query') self._executeRetry(conn, conn.cursor(), s) def query(self, s): return self._runWithConnection(self._query, s) def _queryAll(self, conn, s): if self.debug: self.printDebug(conn, s, 'QueryAll') c = conn.cursor() self._executeRetry(conn, c, s) value = c.fetchall() if self.debugOutput: self.printDebug(conn, value, 'QueryAll', 'result') return value def queryAll(self, s): return self._runWithConnection(self._queryAll, s) def _queryOne(self, conn, s): if self.debug: self.printDebug(conn, s, 'QueryOne') c = conn.cursor() self._executeRetry(conn, c, s) value = c.fetchone() if self.debugOutput: self.printDebug(conn, value, 'QueryOne', 'result') return value def queryOne(self, s): return self._runWithConnection(self._queryOne, s) def _insertSQL(self, table, names, values): return ("INSERT INTO %s (%s) VALUES (%s)" % (table, ', '.join(names), ', '.join([self.sqlrepr(v) for v in values]))) def transaction(self): return Transaction(self) def queryInsertID(self, soInstance, id, names, values): return self._runWithConnection(self._queryInsertID, soInstance, id, names, values) def iterSelect(self, select): return Iteration(self, self.getConnection(), select, keepConnection=False) def accumulateSelect(self, select, expression): """ Apply an accumulate function (like SUM, COUNT, ..) to the select object. Return the value resulting from the SQL accumulate function as an integer. """ q = "SELECT %s" % expression q += " FROM %s WHERE" % ", ".join(select.tables) q = self._addWhereClause(select, q, limit=0, order=0) val = int(self.queryOne(q)[0]) return val def queryForSelect(self, select): ops = select.ops cls = select.sourceClass if ops.get('lazyColumns', 0): q = "SELECT %s.%s FROM %s WHERE " % \ (cls._table, cls._idName, ", ".join(select.tables)) else: q = "SELECT %s.%s, %s FROM %s WHERE " % \ (cls._table, cls._idName, ", ".join(["%s.%s" % (cls._table, col.dbName) for col in cls._SO_columns]), ", ".join(select.tables)) return self._addWhereClause(select, q) def _addWhereClause(self, select, startSelect, limit=1, order=1): q = select.clause if type(q) not in [type(''), type(u'')]: q = self.sqlrepr(q) ops = select.ops def clauseList(lst, desc=False): if type(lst) not in (type([]), type(())): lst = [lst] lst = [clauseQuote(i) for i in lst] if desc: lst = [sqlbuilder.DESC(i) for i in lst] return ', '.join([self.sqlrepr(i) for i in lst]) def clauseQuote(s): if type(s) is type(""): if s.startswith('-'): desc = True s = s[1:] else: desc = False assert sqlbuilder.sqlIdentifier(s), "Strings in clauses are expected to be column identifiers. I got: %r" % s if select.sourceClass._SO_columnDict.has_key(s): s = select.sourceClass._SO_columnDict[s].dbName if desc: return sqlbuilder.DESC(sqlbuilder.SQLConstant(s)) else: return sqlbuilder.SQLConstant(s) else: return s if order and ops.get('dbOrderBy'): q = "%s ORDER BY %s" % (q, clauseList(ops['dbOrderBy'], ops.get('reversed', False))) start = ops.get('start', 0) end = ops.get('end', None) q = startSelect + ' ' + q if limit and (start or end): # @@: Raising an error might be an annoyance, but some warning is # in order. #assert ops.get('orderBy'), "Getting a slice of an unordered set is unpredictable!" q = self._queryAddLimitOffset(q, start, end) return q def _SO_createJoinTable(self, join): self.query(self._SO_createJoinTableSQL(join)) def _SO_createJoinTableSQL(self, join): return ('CREATE TABLE %s (\n%s %s,\n%s %s\n)' % (join.intermediateTable, join.joinColumn, self.joinSQLType(join), join.otherColumn, self.joinSQLType(join))) def _SO_dropJoinTable(self, join): self.query("DROP TABLE %s" % join.intermediateTable) def createTable(self, soClass): self.query(self.createTableSQL(soClass)) def createTableSQL(self, soClass): return ('CREATE TABLE %s (\n%s\n)' % (soClass._table, self.createColumns(soClass))) def createColumns(self, soClass): columnDefs = [self.createIDColumn(soClass)] \ + [self.createColumn(soClass, col) for col in soClass._SO_columns] return ",\n".join([" %s" % c for c in columnDefs]) def createColumn(self, soClass, col): assert 0, "Implement in subclasses" def dropTable(self, tableName, cascade=False): self.query("DROP TABLE %s" % tableName) def clearTable(self, tableName): # 3-03 @@: Should this have a WHERE 1 = 1 or similar # clause? In some configurations without the WHERE clause # the query won't go through, but maybe we shouldn't override # that. self.query("DELETE FROM %s" % tableName) # The _SO_* series of methods are sorts of "friend" methods # with SQLObject. They grab values from the SQLObject instances # or classes freely, but keep the SQLObject class from accessing # the database directly. This way no SQL is actually created # in the SQLObject class. def _SO_update(self, so, values): self.query("UPDATE %s SET %s WHERE %s = %s" % (so._table, ", ".join(["%s = %s" % (dbName, self.sqlrepr(value)) for dbName, value in values]), so._idName, self.sqlrepr(so.id))) def _SO_selectOne(self, so, columnNames): return self.queryOne("SELECT %s FROM %s WHERE %s = %s" % (", ".join(columnNames), so._table, so._idName, self.sqlrepr(so.id))) def _SO_selectOneAlt(self, cls, columnNames, column, value): return self.queryOne("SELECT %s FROM %s WHERE %s = %s" % (", ".join(columnNames), cls._table, column, self.sqlrepr(value))) def _SO_delete(self, so): self.query("DELETE FROM %s WHERE %s = %s" % (so._table, so._idName, self.sqlrepr(so.id))) def _SO_selectJoin(self, soClass, column, value): return self.queryAll("SELECT %s FROM %s WHERE %s = %s" % (soClass._idName, soClass._table, column, self.sqlrepr(value))) def _SO_intermediateJoin(self, table, getColumn, joinColumn, value): return self.queryAll("SELECT %s FROM %s WHERE %s = %s" % (getColumn, table, joinColumn, self.sqlrepr(value))) def _SO_intermediateDelete(self, table, firstColumn, firstValue, secondColumn, secondValue): self.query("DELETE FROM %s WHERE %s = %s AND %s = %s" % (table, firstColumn, self.sqlrepr(firstValue), secondColumn, self.sqlrepr(secondValue))) def _SO_intermediateInsert(self, table, firstColumn, firstValue, secondColumn, secondValue): self.query("INSERT INTO %s (%s, %s) VALUES (%s, %s)" % (table, firstColumn, secondColumn, self.sqlrepr(firstValue), self.sqlrepr(secondValue))) def _SO_columnClause(self, soClass, kw): return ' AND '.join(['%s = %s' % (soClass._SO_columnDict[key].dbName, self.sqlrepr(value)) for key, value in kw.items()]) def sqlrepr(self, v): return sqlrepr(v, self.dbName) class Iteration(object): def __init__(self, dbconn, rawconn, select, keepConnection=False): self.dbconn = dbconn self.rawconn = rawconn self.select = select self.keepConnection = keepConnection self.cursor = rawconn.cursor() self.query = self.dbconn.queryForSelect(select) if dbconn.debug: dbconn.printDebug(rawconn, self.query, 'Select') self.dbconn._executeRetry(self.rawconn, self.cursor, self.query) def __iter__(self): return self def next(self): result = self.cursor.fetchone() if result is None: self._cleanup() raise StopIteration if self.select.ops.get('lazyColumns', 0): obj = self.select.sourceClass.get(result[0], connection=self.dbconn) return obj else: obj = self.select.sourceClass.get(result[0], selectResults=result[1:], connection=self.dbconn) return obj def _cleanup(self): if getattr(self, 'query', None) is None: # already cleaned up return self.query = None if not self.keepConnection: self.dbconn.releaseConnection(self.rawconn) self.dbconn = self.rawconn = self.select = self.cursor = None def __del__(self): self._cleanup() class Transaction(object): def __init__(self, dbConnection): self._dbConnection = dbConnection self._connection = dbConnection.getConnection() self._dbConnection._setAutoCommit(self._connection, 0) self.cache = CacheSet(cache=dbConnection.doCache) self._obsolete = False def assertActive(self): assert not self._obsolete, "This transaction has already gone through COMMIT/ROLLBACK; create another transaction" def query(self, s): self.assertActive() return self._dbConnection._query(self._connection, s) def queryAll(self, s): self.assertActive() return self._dbConnection._queryAll(self._connection, s) def queryOne(self, s): self.assertActive() return self._dbConnection._queryOne(self._connection, s) def queryInsertID(self, soInstance, id, names, values): self.assertActive() return self._dbConnection._queryInsertID( self._connection, soInstance, id, names, values) def iterSelect(self, select): self.assertActive() # We can't keep the cursor open with results in a transaction, # because we might want to use the connection while we're # still iterating through the results. # @@: But would it be okay for psycopg, with threadsafety # level 2? return iter(list(Iteration(self, self._connection, select, keepConnection=True))) def commit(self): if self._obsolete: # @@: is it okay to get extraneous commits? return if self._dbConnection.debug: self._dbConnection.printDebug(self._connection, '', 'COMMIT') self._connection.commit() def rollback(self): if self._obsolete: # @@: is it okay to get extraneous rollbacks? return if self._dbConnection.debug: self._dbConnection.printDebug(self._connection, '', 'ROLLBACK') subCaches = [(sub, sub.allIDs()) for sub in self.cache.allSubCaches()] self._connection.rollback() for subCache, ids in subCaches: for id in ids: inst = subCache.tryGet(id) if inst is not None: inst.expire() self._makeObsolete() def __getattr__(self, attr): """ If nothing else works, let the parent connection handle it. Except with this transaction as 'self'. Poor man's acquisition? Bad programming? Okay, maybe. """ self.assertActive() attr = getattr(self._dbConnection, attr) try: func = attr.im_func except AttributeError: return attr else: meth = new.instancemethod(func, self, self.__class__) return meth def _makeObsolete(self): self._obsolete = True self._dbConnection.releaseConnection(self._connection, explicit=True) self._connection = None def begin(self): # @@: Should we do this, or should begin() be a no-op when we're # not already obsolete? assert self._obsolete, "You cannot begin a new transaction session without committing or rolling back the transaction" self._obsolete = False self._connection = self._dbConnection.getConnection() def __del__(self): if self._obsolete: return self.rollback() class ConnectionURIOpener(object): def __init__(self): self.schemeBuilders = {} self.schemeSupported = {} self.instanceNames = {} self.cachedURIs = {} def registerConnection(self, schemes, builder, isSupported): for uriScheme in schemes: assert not self.schemeBuilders.has_key(uriScheme) \ or self.schemeBuilders[uriScheme] is builder, \ "A driver has already been registered for the URI scheme %s" % uriScheme self.schemeBuilders[uriScheme] = builder self.schemeSupported = isSupported def registerConnectionInstance(self, inst): if inst.name: assert not self.instanceNames.has_key(inst.name) \ or self.instanceNames[inst.name] is cls, \ "A instance has already been registered with the name %s" % inst.name assert inst.name.find(':') == -1, "You cannot include ':' in your class names (%r)" % cls.name self.instanceNames[inst.name] = inst def connectionForURI(self, uri): if self.cachedURIs.has_key(uri): return self.cachedURIs[uri] if uri.find(':') != -1: scheme, rest = uri.split(':', 1) assert self.schemeBuilders.has_key(scheme), \ "No SQLObject driver exists for %s" % scheme conn = self.schemeBuilders[scheme]().connectionFromURI(uri) else: # We just have a name, not a URI assert self.instanceNames.has_key(uri), \ "No SQLObject driver exists under the name %s" % uri conn = self.instanceNames[uri] # @@: Do we care if we clobber another connection? self.cachedURIs[uri] = conn return conn TheURIOpener = ConnectionURIOpener() registerConnection = TheURIOpener.registerConnection registerConnectionInstance = TheURIOpener.registerConnectionInstance connectionForURI = TheURIOpener.connectionForURI