Changeset 461 for trunk/module/pg.py


Ignore:
Timestamp:
Nov 1, 2012, 7:40:55 AM (7 years ago)
Author:
cito
Message:

Add transaction handling methods and context managers to pg.

File:
1 edited

Legend:

Unmodified
Added
Removed
  • trunk/module/pg.py

    r460 r461  
    2929except ImportError:  # Python < 2.4
    3030    pass
     31try:
     32    from contextlib import contextmanager
     33except ImportError:  # Python < 2.5
     34    def contextmanager(f):
     35        raise NotImplementedError
    3136try:
    3237    from collections import namedtuple
     
    166171        self._pkeys = {}
    167172        self._privileges = {}
     173        self._transaction = False
    168174        self._args = args, kw
    169175        self.debug = None  # For debugging scripts, this can be set
    170176            # * to a string format specification (e.g. in CGI set to "%s<BR>"),
    171             # * to a file object to write debug statements,
    172             # * to a callable object which takes a string argument or
    173             # * to any other true value to just print debug statements.
     177            # * to a file object to write debug statements or
     178            # * to a callable object which takes a string argument
     179            # * to any other true value to just print debug statements
    174180
    175181    def __getattr__(self, name):
     
    179185        else:
    180186            raise _int_error('Connection is not valid')
     187
     188    # Context manager methods
     189
     190    def __enter__(self):
     191        if self._transaction:
     192            self.begin()
     193        return self
     194
     195    def __exit__(self, typ, val, tb):
     196        if self._transaction:
     197            self._transaction = None
     198            if tb is None:
     199                self.commit()
     200            else:
     201                self.rollback()
     202        else:
     203            self.close()
    181204
    182205    # Auxiliary methods
     
    335358                self.db.close()
    336359            self.db = db
     360
     361    def begin(self, mode=None):
     362        """Begin a transaction."""
     363        qstr = 'BEGIN'
     364        if mode:
     365            qstr += ' ' + mode
     366        return self.query(qstr)
     367
     368    start = begin
     369
     370    def commit(self):
     371        """Commit the current transaction."""
     372        return self.query('COMMIT')
     373
     374    end = commit
     375
     376    def rollback(self, name=None):
     377        """Rollback the current transaction."""
     378        qstr = 'ROLLBACK'
     379        if name:
     380            qstr += ' TO ' + name
     381        return self.query('ROLLBACK')
     382
     383    def savepoint(self, name=None):
     384        """Define a new savepoint within the current transaction."""
     385        qstr = 'SAVEPOINT'
     386        if name:
     387            qstr += ' ' + name
     388        return self.query(qstr)
     389
     390    def release(self, name):
     391        """Destroy a previously defined savepoint."""
     392        return self.query('RELEASE ' + name)
     393
     394    @property
     395    def transaction(self):
     396        """Return a context manager for running a transaction."""
     397        self._transaction = True
     398        return self
    337399
    338400    def query(self, qstr, *args):
Note: See TracChangeset for help on using the changeset viewer.