Python psycopg2.errorcodes() Examples

def test_delete_slot(slot):
    with patch.object(psycopg2.ProgrammingError, 'pgcode',
                      new_callable=PropertyMock,
                      return_value=psycopg2.errorcodes.UNDEFINED_OBJECT):
        pe = psycopg2.ProgrammingError()
        slot._repl_cursor.drop_replication_slot = Mock(side_effect=pe)
        slot.delete_slot()
    slot._repl_cursor.drop_replication_slot.assert_called_with('pg2kinesis')

    with patch.object(psycopg2.ProgrammingError, 'pgcode',
                      new_callable=PropertyMock,
                      return_value=-1):
        pe = psycopg2.ProgrammingError()
        slot._repl_cursor.create_replication_slot = Mock(side_effect=pe)
        with pytest.raises(psycopg2.ProgrammingError) as e_info:
            slot.delete_slot()
            slot._repl_cursor.drop_replication_slot.assert_called_with('pg2kinesis')

            assert e_info.value.pgcode == -1

    slot._repl_cursor.create_replication_slot = Mock(side_effect=Exception)
    with pytest.raises(Exception):
        slot.delete_slot()
        slot._repl_cursor.drop_replication_slot.assert_called_with('pg2kinesis') 

def _txn_retry_loop(conn, callback, max_retries, max_backoff):
    """Inner transaction retry loop.

    ``conn`` may be either a Connection or a Session, but they both
    have compatible ``begin()`` and ``begin_nested()`` methods.
    """
    retry_count = 0
    with conn.begin():
        while True:
            try:
                with _NestedTransaction(conn):
                    ret = callback(conn)
                    return ret
            except sqlalchemy.exc.DatabaseError as e:
                if max_retries is not None and retry_count >= max_retries:
                    raise
                retry_count += 1
                if isinstance(e.orig, psycopg2.OperationalError):
                    if e.orig.pgcode == psycopg2.errorcodes.SERIALIZATION_FAILURE:
                        if max_backoff > 0:
                            retry_exponential_backoff(retry_count, max_backoff)
                        continue
                raise 

def test_create_slot(slot):

    with patch.object(psycopg2.ProgrammingError, 'pgcode',
                      new_callable=PropertyMock,
                      return_value=psycopg2.errorcodes.DUPLICATE_OBJECT):
        pe = psycopg2.ProgrammingError()


        slot._repl_cursor.create_replication_slot = Mock(side_effect=pe)
        slot.create_slot()
        slot._repl_cursor.create_replication_slot.assert_called_with('pg2kinesis',
                                                                     slot_type=psycopg2.extras.REPLICATION_LOGICAL,
                                                                     output_plugin=u'test_decoding')
    with patch.object(psycopg2.ProgrammingError, 'pgcode',
                          new_callable=PropertyMock,
                          return_value=-1):
        pe = psycopg2.ProgrammingError()
        slot._repl_cursor.create_replication_slot = Mock(side_effect=pe)

        with pytest.raises(psycopg2.ProgrammingError) as e_info:
            slot.create_slot()
            slot._repl_cursor.create_replication_slot.assert_called_with('pg2kinesis',
                                                                         slot_type=psycopg2.extras.REPLICATION_LOGICAL,
                                                                         output_plugin=u'test_decoding')
        assert e_info.value.pgcode == -1

        slot._repl_cursor.create_replication_slot = Mock(side_effect=Exception)
    with pytest.raises(Exception):
        slot.create_slot()
        slot._repl_cursor.create_replication_slot.assert_called_with('pg2kinesis',
                                                                         slot_type=psycopg2.extras.REPLICATION_LOGICAL,
                                                                         output_plugin=u'test_decoding') 

def create_slot(self):
        logger.info('Creating slot %s' % self.slot_name)
        try:
            self._repl_cursor.create_replication_slot(self.slot_name,
                                                      slot_type=psycopg2.extras.REPLICATION_LOGICAL,
                                                      output_plugin=self.output_plugin)
        except psycopg2.ProgrammingError as p:
            # Will be raised if slot exists already.
            if p.pgcode != psycopg2.errorcodes.DUPLICATE_OBJECT:
                logger.error(p)
                raise
            else:
                logger.info('Slot %s is already present.' % self.slot_name) 

def delete_slot(self):
        logger.info('Deleting slot %s' % self.slot_name)
        try:
            self._repl_cursor.drop_replication_slot(self.slot_name)
        except psycopg2.ProgrammingError as p:
            # Will be raised if slot exists already.
            if p.pgcode != psycopg2.errorcodes.UNDEFINED_OBJECT:
                logger.error(p)
                raise
            else:
                logger.info('Slot %s was not found.' % self.slot_name) 

def _store_nobatch(self, conn, items):
        """Store the items one by one, without batching.

        This is to be used in the anomalous cases where inserting the whole
        batch throws an IntegrityError - we fall back to inserting the items
        one by one, so that only the errorneous message is dropped.
        """
        for message, exchange, ack in items:
            item = self._get_db_item(conn, message, exchange)
            if item is None:
                continue
            insert = (self._insert_events if exchange == EVENTS_EXCHANGE_NAME
                      else self._insert_logs)
            try:
                with conn.cursor() as cur:
                    insert(cur, [item])
                conn.commit()
            except psycopg2.OperationalError as e:
                self.on_db_connection_error(e)
            except (psycopg2.IntegrityError, ValueError):
                logger.debug('Error storing %s: %s', exchange, item)
                conn.rollback()
            except psycopg2.ProgrammingError as e:
                if e.pgcode == psycopg2.errorcodes.UNDEFINED_COLUMN:
                    logger.debug('Error storing %s: %s (undefined column)',
                                 exchange, item)
                else:
                    logger.exception('Error storing %s: %s (ProgrammingError)',
                                     exchange, item)
                conn.rollback()
            except Exception:
                logger.exception('Unexpected error while storing %s: %s',
                                 exchange, item)
                conn.rollback() 

def test_cleanup_on_badconn_close(self):
        # ticket #148
        conn = self.conn
        cur = conn.cursor()
        try:
            cur.execute("select pg_terminate_backend(pg_backend_pid())")
        except psycopg2.OperationalError, e:
            if e.pgcode != psycopg2.errorcodes.ADMIN_SHUTDOWN:
                raise 

def skip_if_no_superuser(f):
    """Skip a test if the database user running the test is not a superuser"""
    @wraps(f)
    def skip_if_no_superuser_(self):
        from psycopg2 import ProgrammingError
        try:
            return f(self)
        except ProgrammingError, e:
            import psycopg2.errorcodes
            if e.pgcode == psycopg2.errorcodes.INSUFFICIENT_PRIVILEGE:
                self.skipTest("skipped because not superuser")
            else:
                raise