Source code for fictive.sqlalchemy.ext.errors.psycopg2

# pylint: disable=missing-docstring
# pylint: disable=wrong-import-order

import re

import psycopg2

from fictive.sqlalchemy.ext.errors import IntegrityErrorData


class Psycopg2IntegrityErrorData(IntegrityErrorData):

    ERROR_CLASS = psycopg2.errors.IntegrityError

    @property
    def code(self):
        return type(self.integrity_error.orig).__name__

    @property
    def table(self):
        return self.integrity_error.orig.diag.table_name

    @property
    def message(self):
        return self.integrity_error.orig.diag.message_primary


class Psycopg2NotNullViolationData(Psycopg2IntegrityErrorData):

    ERROR_CLASS = psycopg2.errors.NotNullViolation

    @property
    def fields(self):
        return {self.integrity_error.orig.diag.column_name: None}


class Psycopg2UniqueViolationData(Psycopg2IntegrityErrorData):

    ERROR_CLASS = psycopg2.errors.UniqueViolation

    FIELDS_RE = re.compile(r'\((?P<keys>.*?)\)=\((?P<values>.*?)\)')

    @property
    def fields(self):
        match = self.FIELDS_RE.search(
            self.integrity_error.orig.diag.message_detail)
        keys = map(str.strip, match.group('keys').split(','))
        values = map(str.strip, match.group('values').split(','))
        return dict(zip(keys, values))


class Psycopg2CheckViolationData(Psycopg2IntegrityErrorData):

    ERROR_CLASS = psycopg2.errors.CheckViolation


class Psycopg2ForeignKeyViolationData(Psycopg2IntegrityErrorData):

    ERROR_CLASS = psycopg2.errors.ForeignKeyViolation

    FIELDS_RE = re.compile(r'\((?P<keys>.*?)\)=\((?P<values>.*?)\)')

    @property
    def fields(self):
        match = self.FIELDS_RE.search(
            self.integrity_error.orig.diag.message_detail)
        keys = map(str.strip, match.group('keys').split(','))
        values = map(str.strip, match.group('values').split(','))
        return dict(zip(keys, values))