123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478 |
- # -*- test-case-name: twisted.test.test_adbapi -*-
- # Copyright (c) Twisted Matrix Laboratories.
- # See LICENSE for details.
-
- """
- An asynchronous mapping to U{DB-API
- 2.0<http://www.python.org/topics/database/DatabaseAPI-2.0.html>}.
- """
-
-
- from twisted.internet import threads
- from twisted.python import log, reflect
-
-
- class ConnectionLost(Exception):
- """
- This exception means that a db connection has been lost. Client code may
- try again.
- """
-
-
- class Connection:
- """
- A wrapper for a DB-API connection instance.
-
- The wrapper passes almost everything to the wrapped connection and so has
- the same API. However, the L{Connection} knows about its pool and also
- handle reconnecting should when the real connection dies.
- """
-
- def __init__(self, pool):
- self._pool = pool
- self._connection = None
- self.reconnect()
-
- def close(self):
- # The way adbapi works right now means that closing a connection is
- # a really bad thing as it leaves a dead connection associated with
- # a thread in the thread pool.
- # Really, I think closing a pooled connection should return it to the
- # pool but that's handled by the runWithConnection method already so,
- # rather than upsetting anyone by raising an exception, let's ignore
- # the request
- pass
-
- def rollback(self):
- if not self._pool.reconnect:
- self._connection.rollback()
- return
-
- try:
- self._connection.rollback()
- curs = self._connection.cursor()
- curs.execute(self._pool.good_sql)
- curs.close()
- self._connection.commit()
- return
- except BaseException:
- log.err(None, "Rollback failed")
-
- self._pool.disconnect(self._connection)
-
- if self._pool.noisy:
- log.msg("Connection lost.")
-
- raise ConnectionLost()
-
- def reconnect(self):
- if self._connection is not None:
- self._pool.disconnect(self._connection)
- self._connection = self._pool.connect()
-
- def __getattr__(self, name):
- return getattr(self._connection, name)
-
-
- class Transaction:
- """
- A lightweight wrapper for a DB-API 'cursor' object.
-
- Relays attribute access to the DB cursor. That is, you can call
- C{execute()}, C{fetchall()}, etc., and they will be called on the
- underlying DB-API cursor object. Attributes will also be retrieved from
- there.
- """
-
- _cursor = None
-
- def __init__(self, pool, connection):
- self._pool = pool
- self._connection = connection
- self.reopen()
-
- def close(self):
- _cursor = self._cursor
- self._cursor = None
- _cursor.close()
-
- def reopen(self):
- if self._cursor is not None:
- self.close()
-
- try:
- self._cursor = self._connection.cursor()
- return
- except BaseException:
- if not self._pool.reconnect:
- raise
- else:
- log.err(None, "Cursor creation failed")
-
- if self._pool.noisy:
- log.msg("Connection lost, reconnecting")
-
- self.reconnect()
- self._cursor = self._connection.cursor()
-
- def reconnect(self):
- self._connection.reconnect()
- self._cursor = None
-
- def __getattr__(self, name):
- return getattr(self._cursor, name)
-
-
- class ConnectionPool:
- """
- Represent a pool of connections to a DB-API 2.0 compliant database.
-
- @ivar connectionFactory: factory for connections, default to L{Connection}.
- @type connectionFactory: any callable.
-
- @ivar transactionFactory: factory for transactions, default to
- L{Transaction}.
- @type transactionFactory: any callable
-
- @ivar shutdownID: L{None} or a handle on the shutdown event trigger which
- will be used to stop the connection pool workers when the reactor
- stops.
-
- @ivar _reactor: The reactor which will be used to schedule startup and
- shutdown events.
- @type _reactor: L{IReactorCore} provider
- """
-
- CP_ARGS = "min max name noisy openfun reconnect good_sql".split()
-
- noisy = False # If true, generate informational log messages
- min = 3 # Minimum number of connections in pool
- max = 5 # Maximum number of connections in pool
- name = None # Name to assign to thread pool for debugging
- openfun = None # A function to call on new connections
- reconnect = False # Reconnect when connections fail
- good_sql = "select 1" # A query which should always succeed
-
- running = False # True when the pool is operating
- connectionFactory = Connection
- transactionFactory = Transaction
-
- # Initialize this to None so it's available in close() even if start()
- # never runs.
- shutdownID = None
-
- def __init__(self, dbapiName, *connargs, **connkw):
- """
- Create a new L{ConnectionPool}.
-
- Any positional or keyword arguments other than those documented here
- are passed to the DB-API object when connecting. Use these arguments to
- pass database names, usernames, passwords, etc.
-
- @param dbapiName: an import string to use to obtain a DB-API compatible
- module (e.g. C{'pyPgSQL.PgSQL'})
-
- @keyword cp_min: the minimum number of connections in pool (default 3)
-
- @keyword cp_max: the maximum number of connections in pool (default 5)
-
- @keyword cp_noisy: generate informational log messages during operation
- (default C{False})
-
- @keyword cp_openfun: a callback invoked after every C{connect()} on the
- underlying DB-API object. The callback is passed a new DB-API
- connection object. This callback can setup per-connection state
- such as charset, timezone, etc.
-
- @keyword cp_reconnect: detect connections which have failed and reconnect
- (default C{False}). Failed connections may result in
- L{ConnectionLost} exceptions, which indicate the query may need to
- be re-sent.
-
- @keyword cp_good_sql: an sql query which should always succeed and change
- no state (default C{'select 1'})
-
- @keyword cp_reactor: use this reactor instead of the global reactor
- (added in Twisted 10.2).
- @type cp_reactor: L{IReactorCore} provider
- """
- self.dbapiName = dbapiName
- self.dbapi = reflect.namedModule(dbapiName)
-
- if getattr(self.dbapi, "apilevel", None) != "2.0":
- log.msg("DB API module not DB API 2.0 compliant.")
-
- if getattr(self.dbapi, "threadsafety", 0) < 1:
- log.msg("DB API module not sufficiently thread-safe.")
-
- reactor = connkw.pop("cp_reactor", None)
- if reactor is None:
- from twisted.internet import reactor
- self._reactor = reactor
-
- self.connargs = connargs
- self.connkw = connkw
-
- for arg in self.CP_ARGS:
- cpArg = f"cp_{arg}"
- if cpArg in connkw:
- setattr(self, arg, connkw[cpArg])
- del connkw[cpArg]
-
- self.min = min(self.min, self.max)
- self.max = max(self.min, self.max)
-
- # All connections, hashed on thread id
- self.connections = {}
-
- # These are optional so import them here
- from twisted.python import threadable, threadpool
-
- self.threadID = threadable.getThreadID
- self.threadpool = threadpool.ThreadPool(self.min, self.max)
- self.startID = self._reactor.callWhenRunning(self._start)
-
- def _start(self):
- self.startID = None
- return self.start()
-
- def start(self):
- """
- Start the connection pool.
-
- If you are using the reactor normally, this function does *not*
- need to be called.
- """
- if not self.running:
- self.threadpool.start()
- self.shutdownID = self._reactor.addSystemEventTrigger(
- "during", "shutdown", self.finalClose
- )
- self.running = True
-
- def runWithConnection(self, func, *args, **kw):
- """
- Execute a function with a database connection and return the result.
-
- @param func: A callable object of one argument which will be executed
- in a thread with a connection from the pool. It will be passed as
- its first argument a L{Connection} instance (whose interface is
- mostly identical to that of a connection object for your DB-API
- module of choice), and its results will be returned as a
- L{Deferred}. If the method raises an exception the transaction will
- be rolled back. Otherwise, the transaction will be committed.
- B{Note} that this function is B{not} run in the main thread: it
- must be threadsafe.
-
- @param args: positional arguments to be passed to func
-
- @param kw: keyword arguments to be passed to func
-
- @return: a L{Deferred} which will fire the return value of
- C{func(Transaction(...), *args, **kw)}, or a
- L{twisted.python.failure.Failure}.
- """
- return threads.deferToThreadPool(
- self._reactor, self.threadpool, self._runWithConnection, func, *args, **kw
- )
-
- def _runWithConnection(self, func, *args, **kw):
- conn = self.connectionFactory(self)
- try:
- result = func(conn, *args, **kw)
- conn.commit()
- return result
- except BaseException:
- try:
- conn.rollback()
- except BaseException:
- log.err(None, "Rollback failed")
- raise
-
- def runInteraction(self, interaction, *args, **kw):
- """
- Interact with the database and return the result.
-
- The 'interaction' is a callable object which will be executed in a
- thread using a pooled connection. It will be passed an L{Transaction}
- object as an argument (whose interface is identical to that of the
- database cursor for your DB-API module of choice), and its results will
- be returned as a L{Deferred}. If running the method raises an
- exception, the transaction will be rolled back. If the method returns a
- value, the transaction will be committed.
-
- NOTE that the function you pass is *not* run in the main thread: you
- may have to worry about thread-safety in the function you pass to this
- if it tries to use non-local objects.
-
- @param interaction: a callable object whose first argument is an
- L{adbapi.Transaction}.
-
- @param args: additional positional arguments to be passed to
- interaction
-
- @param kw: keyword arguments to be passed to interaction
-
- @return: a Deferred which will fire the return value of
- C{interaction(Transaction(...), *args, **kw)}, or a
- L{twisted.python.failure.Failure}.
- """
- return threads.deferToThreadPool(
- self._reactor,
- self.threadpool,
- self._runInteraction,
- interaction,
- *args,
- **kw,
- )
-
- def runQuery(self, *args, **kw):
- """
- Execute an SQL query and return the result.
-
- A DB-API cursor which will be invoked with C{cursor.execute(*args,
- **kw)}. The exact nature of the arguments will depend on the specific
- flavor of DB-API being used, but the first argument in C{*args} be an
- SQL statement. The result of a subsequent C{cursor.fetchall()} will be
- fired to the L{Deferred} which is returned. If either the 'execute' or
- 'fetchall' methods raise an exception, the transaction will be rolled
- back and a L{twisted.python.failure.Failure} returned.
-
- The C{*args} and C{**kw} arguments will be passed to the DB-API
- cursor's 'execute' method.
-
- @return: a L{Deferred} which will fire the return value of a DB-API
- cursor's 'fetchall' method, or a L{twisted.python.failure.Failure}.
- """
- return self.runInteraction(self._runQuery, *args, **kw)
-
- def runOperation(self, *args, **kw):
- """
- Execute an SQL query and return L{None}.
-
- A DB-API cursor which will be invoked with C{cursor.execute(*args,
- **kw)}. The exact nature of the arguments will depend on the specific
- flavor of DB-API being used, but the first argument in C{*args} will be
- an SQL statement. This method will not attempt to fetch any results
- from the query and is thus suitable for C{INSERT}, C{DELETE}, and other
- SQL statements which do not return values. If the 'execute' method
- raises an exception, the transaction will be rolled back and a
- L{Failure} returned.
-
- The C{*args} and C{*kw} arguments will be passed to the DB-API cursor's
- 'execute' method.
-
- @return: a L{Deferred} which will fire with L{None} or a
- L{twisted.python.failure.Failure}.
- """
- return self.runInteraction(self._runOperation, *args, **kw)
-
- def close(self):
- """
- Close all pool connections and shutdown the pool.
- """
- if self.shutdownID:
- self._reactor.removeSystemEventTrigger(self.shutdownID)
- self.shutdownID = None
- if self.startID:
- self._reactor.removeSystemEventTrigger(self.startID)
- self.startID = None
- self.finalClose()
-
- def finalClose(self):
- """
- This should only be called by the shutdown trigger.
- """
- self.shutdownID = None
- self.threadpool.stop()
- self.running = False
- for conn in self.connections.values():
- self._close(conn)
- self.connections.clear()
-
- def connect(self):
- """
- Return a database connection when one becomes available.
-
- This method blocks and should be run in a thread from the internal
- threadpool. Don't call this method directly from non-threaded code.
- Using this method outside the external threadpool may exceed the
- maximum number of connections in the pool.
-
- @return: a database connection from the pool.
- """
-
- tid = self.threadID()
- conn = self.connections.get(tid)
- if conn is None:
- if self.noisy:
- log.msg(f"adbapi connecting: {self.dbapiName}")
- conn = self.dbapi.connect(*self.connargs, **self.connkw)
- if self.openfun is not None:
- self.openfun(conn)
- self.connections[tid] = conn
- return conn
-
- def disconnect(self, conn):
- """
- Disconnect a database connection associated with this pool.
-
- Note: This function should only be used by the same thread which called
- L{ConnectionPool.connect}. As with C{connect}, this function is not
- used in normal non-threaded Twisted code.
- """
- tid = self.threadID()
- if conn is not self.connections.get(tid):
- raise Exception("wrong connection for thread")
- if conn is not None:
- self._close(conn)
- del self.connections[tid]
-
- def _close(self, conn):
- if self.noisy:
- log.msg(f"adbapi closing: {self.dbapiName}")
- try:
- conn.close()
- except BaseException:
- log.err(None, "Connection close failed")
-
- def _runInteraction(self, interaction, *args, **kw):
- conn = self.connectionFactory(self)
- trans = self.transactionFactory(self, conn)
- try:
- result = interaction(trans, *args, **kw)
- trans.close()
- conn.commit()
- return result
- except BaseException:
- try:
- conn.rollback()
- except BaseException:
- log.err(None, "Rollback failed")
- raise
-
- def _runQuery(self, trans, *args, **kw):
- trans.execute(*args, **kw)
- return trans.fetchall()
-
- def _runOperation(self, trans, *args, **kw):
- trans.execute(*args, **kw)
-
- def __getstate__(self):
- return {
- "dbapiName": self.dbapiName,
- "min": self.min,
- "max": self.max,
- "noisy": self.noisy,
- "reconnect": self.reconnect,
- "good_sql": self.good_sql,
- "connargs": self.connargs,
- "connkw": self.connkw,
- }
-
- def __setstate__(self, state):
- self.__dict__ = state
- self.__init__(self.dbapiName, *self.connargs, **self.connkw)
-
-
- __all__ = ["Transaction", "ConnectionPool"]
|