Python psycopg2.IntegrityError() Examples

def simple_db_mutate_returning_item(result_cls, context, mutation_query, *,
                                          item_query, item_cls):
    async with context['dbpool'].acquire() as conn, conn.begin():
        try:
            result = await conn.execute(mutation_query)
            if result.rowcount > 0:
                result = await conn.execute(item_query)
                item = await result.first()
                return result_cls(True, 'success', item_cls.from_row(item))
            else:
                return result_cls(False, 'no matching record', None)
        except (pg.IntegrityError, sa.exc.IntegrityError) as e:
            return result_cls(False, f'integrity error: {e}', None)
        except (asyncio.CancelledError, asyncio.TimeoutError):
            raise
        except Exception as e:
            return result_cls(False, f'unexpected error: {e}', None) 

def mutate(cls, root, info, gid):
        async with info.context['dbpool'].acquire() as conn, conn.begin():
            try:
                # query = groups.delete().where(groups.c.id == gid)
                query = groups.update().values(is_active=False,
                                               integration_id=None).where(groups.c.id == gid)
                result = await conn.execute(query)
                if result.rowcount > 0:
                    return cls(ok=True, msg='success')
                else:
                    return cls(ok=False, msg='no such group')
            except (pg.IntegrityError, sa.exc.IntegrityError) as e:
                return cls(ok=False, msg=f'integrity error: {e}')
            except (asyncio.CancelledError, asyncio.TimeoutError):
                raise
            except Exception as e:
                return cls(ok=False, msg=f'unexpected error: {e}') 

def lock(account, state = None):
	if state == None:
		cur = database().cursor()
		cur.execute("SELECT * FROM locked WHERE account = %s", (account,))
		return not not cur.rowcount
	elif state == True:
		db = database()
		cur = db.cursor()
		try:
			cur.execute("INSERT INTO locked VALUES (%s)", (account,))
			db.commit()
		except psycopg2.IntegrityError as e:
			pass
	elif state == False:
		db = database()
		cur = db.cursor()
		cur.execute("DELETE FROM locked WHERE account = %s", (account,))
		db.commit() 

def withdraw(token, account, address, amount): 
	db = database()
	cur = db.cursor()
	try:
		cur.execute("UPDATE accounts SET balance = balance - %s WHERE account = %s", (amount + 1, account))
	except psycopg2.IntegrityError as e:
		raise NotEnoughMoney()
	if not cur.rowcount:
		raise NotEnoughMoney()
	try:
		tx = daemon().sendtoaddress(address, amount, comment = "sent with Doger")
	except InsufficientFunds:
		raise
	except:
		Logger.irclog("Emergency lock on account '%s'" % (account))
		lock(account, True)
		raise
	db.commit()
	txlog(cur, token, amount + 1, tx = tx.encode("ascii"), address = address, src = account)
	db.commit()
	return tx.encode("ascii") 

def tip_multiple(token, source, dict):
	db = database()
	cur = db.cursor()
	cur.execute("SELECT * FROM accounts WHERE account = ANY(%s) FOR UPDATE", (sorted(dict.keys() + [source]),))
	spent = 0
	for target in dict:
		amount = dict[target]
		try:
			cur.execute("UPDATE accounts SET balance = balance - %s WHERE account = %s", (amount, source))
		except psycopg2.IntegrityError as e:
			raise NotEnoughMoney()
		if not cur.rowcount:
			raise NotEnoughMoney()
		spent += amount
		cur.execute("UPDATE accounts SET balance = balance + %s WHERE account = %s", (amount, target)) 
		if not cur.rowcount:
			cur.execute("INSERT INTO accounts VALUES (%s, %s)", (target, amount))
	for target in dict:
		txlog(cur, token, dict[target], src = source, dest = target)
	db.commit() 

def aiopg_exception_handling(exception):
    err_msg = str(exception)
    body = {"err_msg": err_msg}
    if isinstance(exception, psycopg2.IntegrityError):
        if "duplicate key" in err_msg:
            return DBResponse(response_code=409, body=json.dumps(body))
        elif "foreign key" in err_msg:
            return DBResponse(response_code=404, body=json.dumps(body))
        else:
            return DBResponse(response_code=500, body=json.dumps(body))
    elif isinstance(exception, psycopg2.errors.UniqueViolation):
        return DBResponse(response_code=409, body=json.dumps(body))
    elif isinstance(exception, IndexError):
        return DBResponse(response_code=404, body="{}")
    else:
        return DBResponse(response_code=500, body=json.dumps(body)) 

def test_error(self):
        cur = self.conn.cursor()
        cur.execute("insert into table1 values (%s)", (1, ))
        self.wait(cur)
        cur.execute("insert into table1 values (%s)", (1, ))
        # this should fail
        self.assertRaises(psycopg2.IntegrityError, self.wait, cur)
        cur.execute("insert into table1 values (%s); "
                    "insert into table1 values (%s)", (2, 2))
        # this should fail as well
        self.assertRaises(psycopg2.IntegrityError, self.wait, cur)
        # but this should work
        cur.execute("insert into table1 values (%s)", (2, ))
        self.wait(cur)
        # and the cursor should be usable afterwards
        cur.execute("insert into table1 values (%s)", (3, ))
        self.wait(cur)
        cur.execute("select * from table1 order by id")
        self.wait(cur)
        self.assertEquals(cur.fetchall(), [(1, ), (2, ), (3, )])
        cur.execute("delete from table1")
        self.wait(cur) 

def _create_database(engine: "Engine", db: Text):
        """Create database `db` on `engine` if it does not exist."""

        import psycopg2

        conn = engine.connect()

        cursor = conn.connection.cursor()
        cursor.execute("COMMIT")
        cursor.execute(f"SELECT 1 FROM pg_catalog.pg_database WHERE datname = '{db}'")
        exists = cursor.fetchone()
        if not exists:
            try:
                cursor.execute(f"CREATE DATABASE {db}")
            except psycopg2.IntegrityError as e:
                logger.error(f"Could not create database '{db}': {e}")

        cursor.close()
        conn.close() 

def insert_phash_link_many(self, link_tuples):
		try:
			with self.transaction() as cur:
				for item_1, item_2, distance in link_tuples:
					item_1, item_2 = min(item_1, item_2), max(item_1, item_2)
					cur.execute(''' INSERT INTO
									{table}_plink (item_1_link, item_2_link, distance)
								VALUES (%s, %s, %s)
								ON CONFLICT DO NOTHING
							'''.format(table=self.tableName), (item_1, item_2, distance))
		except (psycopg2.OperationalError, psycopg2.IntegrityError):
			for item_1, item_2, distance in link_tuples:
				item_1, item_2 = min(item_1, item_2), max(item_1, item_2)
				try:
					with self.transaction() as cur:
						cur.execute(''' INSERT INTO
										{table}_plink (item_1_link, item_2_link, distance)
									VALUES (%s, %s, %s)
									ON CONFLICT DO NOTHING
								'''.format(table=self.tableName), (item_1, item_2, distance))
				except (psycopg2.OperationalError, psycopg2.IntegrityError):
					self.log.error("Failure inserting link between %s and %s with distance %s. Skipping", distance, item_1, item_2) 

def tip(token, source, target, amount): 
	db = database()
	cur = db.cursor()
	cur.execute("SELECT * FROM accounts WHERE account = ANY(%s) FOR UPDATE", (sorted([target, source]),))
	try:
		cur.execute("UPDATE accounts SET balance = balance - %s WHERE account = %s", (amount, source))
	except psycopg2.IntegrityError as e:
		raise NotEnoughMoney()
	if not cur.rowcount:
		raise NotEnoughMoney()
	cur.execute("UPDATE accounts SET balance = balance + %s WHERE account = %s", (amount, target)) 
	if not cur.rowcount:
		cur.execute("INSERT INTO accounts VALUES (%s, %s)", (target, amount))
	txlog(cur, token, amount, src = source, dest = target)
	db.commit() 

def update_user(self, user: User) -> User:
        async with AsgardDBSession() as s:
            user_db, userTable = await user.to_alchemy_obj()
            update = (
                userTable.__table__.update()
                .where(userTable.id == user.id)
                .values(tx_name=user.name, tx_email=user.email)
            )
            try:
                await s.execute(update)
            except psycopg2.IntegrityError as e:
                raise DuplicateEntity(e.pgerror)
            return user 

def execute(self, query, args=None):
        try:
            return self.cursor.execute(query, args)
        except Database.IntegrityError as e:
            six.reraise(utils.IntegrityError, utils.IntegrityError(*tuple(e.args)), sys.exc_info()[2])
        except Database.DatabaseError as e:
            six.reraise(utils.DatabaseError, utils.DatabaseError(*tuple(e.args)), sys.exc_info()[2]) 

def executemany(self, query, args):
        try:
            return self.cursor.executemany(query, args)
        except Database.IntegrityError as e:
            six.reraise(utils.IntegrityError, utils.IntegrityError(*tuple(e.args)), sys.exc_info()[2])
        except Database.DatabaseError as e:
            six.reraise(utils.DatabaseError, utils.DatabaseError(*tuple(e.args)), sys.exc_info()[2]) 

def _commit(self):
        if self.connection is not None:
            try:
                return self.connection.commit()
            except Database.IntegrityError as e:
                six.reraise(utils.IntegrityError, utils.IntegrityError(*tuple(e.args)), sys.exc_info()[2]) 

def upsert(self, *args, **kwargs):
		try:
			self.insertIntoDb(*args, **kwargs)
		except psycopg2.IntegrityError:
			self.updateDbEntry(*args, **kwargs)

	# Insert new item into DB.
	# MASSIVELY faster if you set commit=False (it doesn't flush the write to disk), but that can open a transaction which locks the DB.
	# Only pass commit=False if the calling code can gaurantee it'll call commit() itself within a reasonable timeframe. 

def test_failed_commit(self):
        # Test that we can recover from a failed commit.
        # We use a deferred constraint to cause a failure on commit.
        curs = self.conn.cursor()
        curs.execute('SET CONSTRAINTS table2__table1_id__fk DEFERRED')
        curs.execute('INSERT INTO table2 VALUES (2, 42)')
        # The commit should fail, and move the cursor back to READY state
        self.assertEqual(self.conn.status, STATUS_BEGIN)
        self.assertRaises(psycopg2.IntegrityError, self.conn.commit)
        self.assertEqual(self.conn.status, STATUS_READY)
        # The connection should be ready to use for the next transaction:
        curs.execute('SELECT 1')
        self.assertEqual(curs.fetchone()[0], 1) 

def create_user(self, user: User) -> User:
        user_db, userTable = await user.to_alchemy_obj()
        try:
            async with AsgardDBSession() as s:
                returned_values = await s.execute(
                    userTable.__table__.insert()
                    .values(tx_name=user.name, tx_email=user.email)
                    .return_defaults(userTable.id)
                )
                created_id = list(returned_values)[0].id
                return User(id=created_id, name=user.name, email=user.email)
        except psycopg2.IntegrityError as e:
            raise DuplicateEntity(e.pgerror) 

def make_unique_violation():
    """Make a serialization exception.

    Artificially construct an exception that resembles what Django's ORM would
    raise when PostgreSQL fails a transaction because of a unique violation.

    :returns: an instance of :py:class:`IntegrityError` that will pass the
        `is_unique_violation` predicate.
    """
    exception = IntegrityError()
    exception.__cause__ = UniqueViolation()
    assert is_unique_violation(exception)
    return exception 

def make_foreign_key_violation():
    """Make a serialization exception.

    Artificially construct an exception that resembles what Django's ORM would
    raise when PostgreSQL fails a transaction because of a foreign key
    violation.

    :returns: an instance of :py:class:`IntegrityError` that will pass the
        `is_foreign_key_violation` predicate.
    """
    exception = IntegrityError()
    exception.__cause__ = ForeignKeyViolation()
    assert is_foreign_key_violation(exception)
    return exception 

def allow_pgcodes(cr, *codes):
    """Context manager that will omit specified error codes.

    E.g., suppose you expect a migration to produce unique constraint
    violations and you want to ignore them. Then you could just do::

        with allow_pgcodes(cr, psycopg2.errorcodes.UNIQUE_VIOLATION):
            cr.execute("INSERT INTO me (name) SELECT name FROM you")

    .. warning::
        **All** sentences inside this context will be rolled back if **a single
        error** is raised, so the above example would insert **nothing** if a
        single row violates a unique constraint.

        This would ignore duplicate files but insert the others::

            cr.execute("SELECT name FROM you")
            for row in cr.fetchall():
                with allow_pgcodes(cr, psycopg2.errorcodes.UNIQUE_VIOLATION):
                    cr.execute("INSERT INTO me (name) VALUES (%s)", row[0])

    :param *str codes:
        Undefined amount of error codes found in :mod:`psycopg2.errorcodes`
        that are allowed. Codes can have either 2 characters (indicating an
        error class) or 5 (indicating a concrete error). Any other errors
        will be raised.
    """
    try:
        with cr.savepoint():
            with core.tools.mute_logger('odoo.sql_db'):
                yield
    except (ProgrammingError, IntegrityError) as error:
        msg = "Code: {code}. Class: {class_}. Error: {error}.".format(
            code=error.pgcode,
            class_=errorcodes.lookup(error.pgcode[:2]),
            error=errorcodes.lookup(error.pgcode))
        if error.pgcode in codes or error.pgcode[:2] in codes:
            logger.info(msg)
        else:
            logger.exception(msg)
            raise 

def logged_query(cr, query, args=None, skip_no_result=False):
    """
    Logs query and affected rows at level DEBUG.

    :param query: a query string suitable to pass to cursor.execute()
    :param args: a list, tuple or dictionary passed as substitution values
      to cursor.execute().
    :param skip_no_result: If True, then logging details are only shown
      if there are affected records.
    """
    if args is None:
        args = ()
    args = tuple(args) if type(args) == list else args
    log_level = _logging_module.DEBUG
    log_msg = False
    start = datetime.now()
    try:
        cr.execute(query, args)
    except (ProgrammingError, IntegrityError):
        log_level = _logging_module.ERROR
        log_msg = "Error after %(duration)s running %(full_query)s"
        raise
    else:
        if not skip_no_result or cr.rowcount:
            log_msg = ('%(rowcount)d rows affected after '
                       '%(duration)s running %(full_query)s')
    finally:
        duration = datetime.now() - start
        if log_msg:
            try:
                full_query = tools.ustr(cr._obj.query)
            except AttributeError:
                full_query = tools.ustr(cr.mogrify(query, args))
            logger.log(log_level, log_msg, {
                "full_query": full_query,
                "rowcount": cr.rowcount,
                "duration": duration,
            })
    return cr.rowcount 

def simple_db_mutate(result_cls, context, mutation_query):
    async with context['dbpool'].acquire() as conn, conn.begin():
        try:
            result = await conn.execute(mutation_query)
            if result.rowcount > 0:
                return result_cls(True, 'success')
            else:
                return result_cls(False, 'no matching record')
        except (pg.IntegrityError, sa.exc.IntegrityError) as e:
            return result_cls(False, f'integrity error: {e}')
        except (asyncio.CancelledError, asyncio.TimeoutError):
            raise
        except Exception as e:
            return result_cls(False, f'unexpected error: {e}') 

def update_questionnaire(self, github_id, q1, q2, q3, q4, q5):
		try:
			query = 'INSERT INTO questionnaire (github_id, q1, q2, q3, q4, q5) VALUES(%s, %s, %s, %s, %s, %s);'
			yield self.execute(query, github_id, q1, q2, q3, q4, q5)
		except psycopg2.IntegrityError:
			query = '''
				UPDATE questionnaire
				SET q1 = %s, q2 = %s, q3 = %s, q4 = %s, q5 = %s
				WHERE github_id = %s;
			'''
			yield self.execute(query, q1, q2, q3, q4, q5, github_id) 

def set_contact_info(self, github_id, info_type, info):
		try:
			query = 'INSERT INTO contact_info (github_id, type, info) VALUES(%s, %s, %s);'
			yield self.execute(query, github_id, info_type, info)
		except psycopg2.IntegrityError:
			query = 'UPDATE contact_info SET info = %s WHERE github_id = %s AND type = %s;'
			yield self.execute(query, info, github_id, info_type) 

def populate_fixture(db_connection, fixture_data, *,
                     ignore_unique_violation: bool = False):

    def insert(table, row):
        # convert enumtype to native values
        for col in table.columns:
            if isinstance(col.type, EnumType):
                row[col.name] = col.type._enum_cls[row[col.name]]
            elif isinstance(col.type, EnumValueType):
                row[col.name] = col.type._enum_cls(row[col.name])
        db_connection.execute(table.insert(), [row])

    for table_name, rows in fixture_data.items():
        table = getattr(models, table_name)
        pk_cols = table.primary_key.columns
        for row in rows:
            if len(pk_cols) == 0:
                # some tables may not have primary keys.
                # (e.g., m2m relationship)
                try:
                    insert(table, row)
                except sa.exc.IntegrityError as e:
                    if ignore_unique_violation and isinstance(e.orig, pg.errors.UniqueViolation):
                        continue
                    raise
                continue
            # compose pk match where clause
            pk_match = functools.reduce(lambda x, y: x & y, [
                (col == row[col.name])
                for col in pk_cols
            ])
            ret = db_connection.execute(
                sa.select(pk_cols).select_from(table).where(pk_match))
            if ret.rowcount == 0:
                insert(table, row)
            else:
                pk_tuple = tuple(row[col.name] for col in pk_cols)
                log.info('skipped inserting {} to {} as the row already exists.',
                         f"[{','.join(pk_tuple)}]", table_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 mutate(cls, root, info, name, props):
        async with info.context['dbpool'].acquire() as conn, conn.begin():
            assert _rx_slug.search(name) is not None, 'invalid name format. slug format required.'
            data = {
                'name': name,
                'description': props.description,
                'is_active': props.is_active,
                'domain_name': props.domain_name,
                'total_resource_slots': ResourceSlot.from_user_input(
                    props.total_resource_slots, None),
                'allowed_vfolder_hosts': props.allowed_vfolder_hosts,
                'integration_id': props.integration_id,
            }
            query = (groups.insert().values(data))
            try:
                result = await conn.execute(query)
                if result.rowcount > 0:
                    checkq = groups.select().where((groups.c.name == name) &
                                                   (groups.c.domain_name == props.domain_name))
                    result = await conn.execute(checkq)
                    o = Group.from_row(info.context, await result.first())
                    return cls(ok=True, msg='success', group=o)
                else:
                    return cls(ok=False, msg='failed to create group', group=None)
            except (pg.IntegrityError, sa.exc.IntegrityError) as e:
                return cls(ok=False, msg=f'integrity error: {e}', group=None)
            except (asyncio.CancelledError, asyncio.TimeoutError):
                raise
            except Exception as e:
                return cls(ok=False, msg=f'unexpected error: {e}', group=None) 

def save_to_db(key, content):
    """Returns succcess and error code if operation was not succesfull"""
    try:
        with psycopg2.connect(DB) as conn:
            with conn.cursor() as curs:
                curs.execute(
                    "INSERT INTO bot_memory (key, content) VALUES (%s, %s);",
                    (key, content),
                )
                logger.info(f"Content successfully added to db with key '{key}'")
                conn.commit()
                success = (True, None)

    except psycopg2.IntegrityError as e:
        if e.pgcode == UNIQUE_VIOLATION:
            logger.info("Snippet key already exists on db")
            success = (False, DUPLICATE_KEY_MESSAGE)
        else:
            logger.exception("Snippet could not be saved!")
            success = (False, DEFAULT_ERROR_MESSAGE)

    except Exception:
        logger.exception("Error writing to db")
        success = (False, repr(Exception))

    return success 

def delete_invitation(request: web.Request, params: Any) -> web.Response:
    dbpool = request.app['dbpool']
    access_key = request['keypair']['access_key']
    request_email = request['user']['email']
    inv_id = params['inv_id']
    log.info('VFOLDER.DELETE_INVITATION (ak:{}, inv:{})', access_key, inv_id)
    try:
        async with dbpool.acquire() as conn:
            query = (sa.select([vfolder_invitations.c.inviter,
                                vfolder_invitations.c.invitee])
                       .select_from(vfolder_invitations)
                       .where((vfolder_invitations.c.id == inv_id) &
                              (vfolder_invitations.c.state == VFolderInvitationState.PENDING)))
            result = await conn.execute(query)
            row = await result.first()
            if row is None:
                raise GenericNotFound('No such vfolder invitation')
            if request_email == row.inviter:
                state = VFolderInvitationState.CANCELED
            elif request_email == row.invitee:
                state = VFolderInvitationState.REJECTED
            else:
                raise GenericForbidden('Cannot change other user\'s invitaiton')
            query = (vfolder_invitations
                     .update()
                     .where(vfolder_invitations.c.id == inv_id)
                     .values(state=state))
            await conn.execute(query)
    except (psycopg2.IntegrityError, sa.exc.IntegrityError) as e:
        raise InternalServerError(f'integrity error: {e}')
    except (asyncio.CancelledError, asyncio.TimeoutError):
        raise
    except Exception as e:
        raise InternalServerError(f'unexpected error: {e}')
    return web.json_response({}) 

def mutate(cls, root, info, email):
        async with info.context['dbpool'].acquire() as conn, conn.begin():
            try:
                # Make all user keypairs inactive.
                from ai.backend.manager.models import keypairs
                query = (
                    keypairs.update()
                    .values(is_active=False)
                    .where(keypairs.c.user_id == email)
                )
                await conn.execute(query)
                # Mark user as deleted.
                query = (
                    users.update()
                    .values(status=UserStatus.DELETED)
                    .where(users.c.email == email)
                )
                result = await conn.execute(query)
                if result.rowcount > 0:
                    return cls(ok=True, msg='success')
                else:
                    return cls(ok=False, msg='no such user')
            except (pg.IntegrityError, sa.exc.IntegrityError) as e:
                return cls(ok=False, msg=f'integrity error: {e}')
            except (asyncio.CancelledError, asyncio.TimeoutError):
                raise
            except Exception as e:
                return cls(ok=False, msg=f'unexpected error: {e}')