aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/python/Twisted/py2/twisted/enterprise/adbapi.py
blob: e240a69fb591042007f68e2498c13ff17bef61de (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
# -*- 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>}.
"""

import sys

from twisted.internet import threads
from twisted.python import reflect, log, compat


class ConnectionLost(Exception):
    """
    This exception means that a db connection has been lost.  Client code may
    try again.
    """



class Connection(object):
    """
    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:
            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:
            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'})

        @param cp_min: the minimum number of connections in pool (default 3)

        @param cp_max: the maximum number of connections in pool (default 5)

        @param cp_noisy: generate informational log messages during operation
            (default C{False})

        @param 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.

        @param 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.

        @param cp_good_sql: an sql query which should always succeed and change
            no state (default C{'select 1'})

        @param 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 = 'cp_%s' % (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 threadpool
        from twisted.python import threadable

        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:
            excType, excValue, excTraceback = sys.exc_info()
            try:
                conn.rollback()
            except:
                log.err(None, "Rollback failed")
            compat.reraise(excValue, excTraceback)


    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('adbapi connecting: %s' % (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('adbapi closing: %s' % (self.dbapiName,))
        try:
            conn.close()
        except:
            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:
            excType, excValue, excTraceback = sys.exc_info()
            try:
                conn.rollback()
            except:
                log.err(None, "Rollback failed")
            compat.reraise(excValue, excTraceback)


    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']