Source code for postgres.context_managers

from psycopg2 import InterfaceError


[docs]class CursorContextManager: """Instantiated once per :func:`~postgres.Postgres.get_cursor` call. :param pool: see :mod:`psycopg2_pool` :param bool autocommit: see :attr:`psycopg2:connection.autocommit` :param bool readonly: see :attr:`psycopg2:connection.readonly` :param cursor_kwargs: passed to :meth:`psycopg2:connection.cursor` During construction, a connection is checked out of the connection pool and its :attr:`autocommit` and :attr:`readonly` attributes are set, then a :class:`psycopg2:cursor` is created from that connection. Upon exit of the ``with`` block, the connection is rolled back if an exception was raised, or committed otherwise. There are two exceptions to this: 1. if :attr:`autocommit` is :obj:`True`, then the connection is neither rolled back nor committed; 2. if :attr:`readonly` is :obj:`True`, then the connection is always rolled back, never committed. In all cases the cursor is closed and the connection is put back in the pool. """ __slots__ = ('pool', 'conn', 'cursor') def __init__(self, pool, autocommit=False, readonly=False, **cursor_kwargs): self.pool = pool conn = self.pool.getconn() conn.autocommit = autocommit conn.readonly = readonly self.cursor = conn.cursor(**cursor_kwargs) self.conn = conn def __enter__(self): return self.cursor def __exit__(self, exc_type, exc_val, exc_tb): """Put our connection back in the pool. """ self.cursor.close() self.conn.__exit__(exc_type, exc_val, exc_tb) self.pool.putconn(self.conn)
[docs]class ConnectionCursorContextManager: """Creates a cursor from the given connection, then wraps it in a context manager that automatically commits or rolls back the changes on exit. :param conn: a :class:`psycopg2:connection` :param bool autocommit: see :attr:`psycopg2:connection.autocommit` :param bool readonly: see :attr:`psycopg2:connection.readonly` :param cursor_kwargs: passed to :meth:`psycopg2:connection.cursor` During construction, the connection's :attr:`autocommit` and :attr:`readonly` attributes are set, then :meth:`psycopg2:connection.cursor` is called with `cursor_kwargs`. Upon exit of the ``with`` block, the connection is rolled back if an exception was raised, or committed otherwise. There are two exceptions to this: 1. if :attr:`autocommit` is :obj:`True`, then the connection is neither rolled back nor committed; 2. if :attr:`readonly` is :obj:`True`, then the connection is always rolled back, never committed. In all cases the cursor is closed. """ __slots__ = ('conn', 'cursor') def __init__(self, conn, autocommit=False, readonly=False, **cursor_kwargs): conn.autocommit = autocommit conn.readonly = readonly self.conn = conn self.cursor = conn.cursor(**cursor_kwargs) def __enter__(self): return self.cursor def __exit__(self, exc_type, exc_val, exc_tb): self.cursor.close() self.conn.__exit__(exc_type, exc_val, exc_tb)
[docs]class CursorSubcontextManager: """Wraps a cursor so that it can be used for a subtransaction. See :meth:`~postgres.Postgres.get_cursor` for an explanation of subtransactions. :param cursor: the :class:`psycopg2:cursor` to wrap :param back_as: temporarily overwrites the cursor's :attr:`~postgres.cursors.SimpleCursorBase.back_as` attribute """ __slots__ = ('cursor', 'back_as', 'outer_back_as') PRESERVE = object() def __init__(self, cursor, back_as=PRESERVE): self.cursor = cursor self.back_as = back_as def __enter__(self): if self.back_as is not self.PRESERVE: self.outer_back_as = self.cursor.back_as self.cursor.back_as = self.back_as return self.cursor def __exit__(self, exc_type, exc_val, exc_tb): if self.back_as is not self.PRESERVE: self.cursor.back_as = self.outer_back_as
[docs]class ConnectionContextManager: """Instantiated once per :func:`~postgres.Postgres.get_connection` call. :param pool: see :mod:`psycopg2_pool` :param bool autocommit: see :attr:`psycopg2:connection.autocommit` :param bool readonly: see :attr:`psycopg2:connection.readonly` This context manager checks out a connection out of the specified pool, sets its :attr:`autocommit` and :attr:`readonly` attributes. The :meth:`__enter__` method returns the :class:`~postgres.Connection`. The :meth:`__exit__` method rolls back the connection and puts it back in the pool. """ __slots__ = ('pool', 'conn') def __init__(self, pool, autocommit=False, readonly=False): self.pool = pool conn = self.pool.getconn() conn.autocommit = autocommit conn.readonly = readonly self.conn = conn def __enter__(self): return self.conn def __exit__(self, *exc_info): """Put our connection back in the pool. """ try: self.conn.rollback() except InterfaceError: pass self.pool.putconn(self.conn)