import functools import heapq import logging import threading import time from collections import namedtuple from peewee import MySQLDatabase from peewee import PostgresqlDatabase from peewee import SqliteDatabase logger = logging.getLogger('peewee.pool') def make_int(val): if val is not None and not isinstance(val, (int, float)): return int(val) return val class MaxConnectionsExceeded(ValueError): pass PoolConnection = namedtuple('PoolConnection', ('timestamp', 'connection', 'checked_out')) class _sentinel(object): def __lt__(self, other): return True def locked(fn): @functools.wraps(fn) def inner(self, *args, **kwargs): with self._pool_lock: return fn(self, *args, **kwargs) return inner class PooledDatabase(object): def __init__(self, database, max_connections=20, stale_timeout=None, timeout=None, **kwargs): self._max_connections = make_int(max_connections) self._stale_timeout = make_int(stale_timeout) self._wait_timeout = make_int(timeout) if self._wait_timeout == 0: self._wait_timeout = float('inf') # Lock for pool operations and condition for notifying when connection # is released back to pool. self._pool_lock = threading.RLock() self._pool_available = threading.Condition(self._pool_lock) # Available / idle connections stored in a heap, sorted oldest first. self._connections = [] # Counter used for tie-breaker in heap (so we don't try comparing # connection against connection). self._heap_counter = 0 # Mapping of connection id to PoolConnection. Ordinarily we would want # to use something like a WeakKeyDictionary, but Python typically won't # allow us to create weak references to connection objects. self._in_use = {} # Use the memory address of the connection as the key in the event the # connection object is not hashable. Connections will not get # garbage-collected, however, because a reference to them will persist # in "_in_use" as long as the conn has not been closed. self.conn_key = id super(PooledDatabase, self).__init__(database, **kwargs) def init(self, database, max_connections=None, stale_timeout=None, timeout=None, **connect_kwargs): super(PooledDatabase, self).init(database, **connect_kwargs) if max_connections is not None: self._max_connections = make_int(max_connections) if stale_timeout is not None: self._stale_timeout = make_int(stale_timeout) if timeout is not None: self._wait_timeout = make_int(timeout) if self._wait_timeout == 0: self._wait_timeout = float('inf') def connect(self, reuse_if_open=False): if not self._wait_timeout: return super(PooledDatabase, self).connect(reuse_if_open) deadline = time.monotonic() + self._wait_timeout while True: try: return super(PooledDatabase, self).connect(reuse_if_open) except MaxConnectionsExceeded: remaining = deadline - time.monotonic() if remaining <= 0: raise MaxConnectionsExceeded( 'Max connections exceeded, timed out attempting to ' 'connect.') with self._pool_available: self._pool_available.wait(timeout=min(remaining, 1.0)) @locked def _connect(self): while self._connections: try: # Remove the oldest connection from the heap. ts, _counter, conn = heapq.heappop(self._connections) except IndexError: break key = self.conn_key(conn) if self._is_closed(conn): # Connection closed either by user or by driver - discard. logger.debug('Connection %s was closed, discarding.', key) continue if self._stale_timeout and self._is_stale(ts): logger.debug('Connection %s was stale, closing.', key) self._close_raw(conn) continue # Connection OK to use. self._in_use[key] = PoolConnection(ts, conn, time.time()) return conn if self._max_connections and ( len(self._in_use) >= self._max_connections): raise MaxConnectionsExceeded('Exceeded maximum connections.') conn = super(PooledDatabase, self)._connect() ts = time.time() key = self.conn_key(conn) logger.debug('Created new connection %s.', key) self._in_use[key] = PoolConnection(ts, conn, time.time()) return conn def _is_stale(self, timestamp): # Called on check-out and check-in to ensure the connection has # not outlived the stale timeout. return (time.time() - timestamp) > self._stale_timeout def _is_closed(self, conn): return False def _can_reuse(self, conn): # Called on check-in to make sure the connection can be re-used. return True def _close_raw(self, conn): try: super(PooledDatabase, self)._close(conn) except Exception: logger.debug('Error closing connection %s.', self.conn_key(conn), exc_info=True) @locked def _close(self, conn, close_conn=False): # if close_conn == True, close underlying driver connection and remove # from _in_use tracking. Do not return to available conns. key = self.conn_key(conn) if close_conn: self._in_use.pop(key, None) self._close_raw(conn) return if key not in self._in_use: logger.debug('Connection %s not in use, ignoring close.', key) return pool_conn = self._in_use.pop(key) if self._stale_timeout and self._is_stale(pool_conn.timestamp): logger.debug('Closing stale connection %s on check-in.', key) self._close_raw(conn) elif not self._can_reuse(conn): logger.debug('Connection %s not reusable, closing.', key) self._close_raw(conn) else: logger.debug('Returning %s to pool.', key) self._heap_counter += 1 heapq.heappush(self._connections, (pool_conn.timestamp, self._heap_counter, conn)) # Wake up thread that may be waiting on connection. self._pool_available.notify() def manual_close(self): """ Close the underlying connection without returning it to the pool. """ if self.is_closed(): return False # Obtain reference to the connection in-use by the calling thread. conn = self.connection() key = self.conn_key(conn) # Remove from _in_use so that subsequent self.close() won't try to # restore it to the pool. pool lock must be released before calling # self.close(), since close acquires the database lock. with self._pool_lock: self._in_use.pop(key, None) self.close() self._close_raw(conn) @locked def close_idle(self): # Close any open connections that are not currently in-use. idle = self._connections self._connections = [] for _, _, conn in idle: self._close_raw(conn) @locked def close_stale(self, age=600): # Close any connections that are in-use but were checked out quite some # time ago and can be considered stale. May close connections in use by # running threads. cutoff = time.time() - age n = 0 for key, pool_conn in list(self._in_use.items()): if pool_conn.checked_out < cutoff: self._close_raw(pool_conn.connection) del self._in_use[key] n += 1 self._pool_available.notify_all() return n def close_all(self): # Close all connections -- available and in-use. Warning: may break any # active connections used by other threads. # self.close() acquires the database lock, calling it while holding # the pool lock would invert the lock order used by Database.connect() # (db lock, then pool lock via _connect) and deadlock. self.close() with self._pool_lock: self.close_idle() in_use, self._in_use = self._in_use, {} for pool_conn in in_use.values(): self._close_raw(pool_conn.connection) self._pool_available.notify_all() class _PooledMySQLDatabase(PooledDatabase): def _is_closed(self, conn): if self.server_version[0] == 8: args = () else: args = (False,) try: conn.ping(*args) except: return True return False class PooledMySQLDatabase(_PooledMySQLDatabase, MySQLDatabase): pass class _PooledPostgresqlDatabase(PooledDatabase): def _is_closed(self, conn): if conn.closed: return True return self._adapter.is_connection_closed(conn) def _can_reuse(self, conn): return self._adapter.is_connection_reusable(conn) class PooledPostgresqlDatabase(_PooledPostgresqlDatabase, PostgresqlDatabase): pass class _PooledSqliteDatabase(PooledDatabase): def _is_closed(self, conn): try: conn.total_changes except: return True return False class PooledSqliteDatabase(_PooledSqliteDatabase, SqliteDatabase): pass