<commit_before>#!/usr/bin/python

""" Defines the backend connection class for PostgreSQL databases. """

# define authorship information
__authors__ = ['<NAME>']
__author__ = ','.join(__authors__)
__credits__ = []
__copyright__ = 'Copyright (c) 2011, Projex Software'
__license__ = 'LGPL'

# maintanence information
__maintainer__ = 'Projex Software'
__email__ = '<EMAIL>'

# define version information (major,minor,maintanence)
__depends__ = ['pyscopg2']
__version_info__ = (0, 0, 0)
__version__ = '%i.%i.%i' % __version_info__

# ------------------------------------------------------------------------------

import datetime
import logging
import os
import orb
import traceback

from orb import errors
from projex.text import nativestring as nstr

from ..abstractconnection import SQLConnection
from ..abstractsql import SQL

log = logging.getLogger(__name__)

try:
    import psycopg2 as pg
    from psycopg2.extras import DictCursor, register_hstore, register_json
    from psycopg2.extensions import QueryCanceledError

except ImportError:
    log.debug('For PostgreSQL backend, download the psycopg2 module')

    QueryCanceledError = errors.DatabaseError
    pg = None

# ----------------------------------------------------------------------

class PSQLConnection(SQLConnection):
    """ 
    Creates a PostgreSQL backend connection type for handling database
    connections to PostgreSQL databases.
    """

    # ----------------------------------------------------------------------
    # PROTECTED METHODS
    # ----------------------------------------------------------------------
    def _execute(self, command, data=None, autoCommit=True, autoClose=True,
                 returning=True, mapper=dict):
        """
        Executes the inputed command into the current \
        connection cursor.
        
        :param      command    | <str>
                    data       | <dict> || None
                    autoCommit | <bool> | commit database changes immediately
                    autoClose  | <bool> | closes connections immediately
        
        :return     [{<str> key: <variant>, ..}, ..], <int> count
        """
        rowcount = 0
        if data is None:
            data = {}

        # when in debug mode, simply log the command to the log
        elif self.database().commandsBlocked():
            log.info(command)
            return [], rowcount

        # create a new cursor for this transaction
        db = self.nativeConnection()
        if db is None:
            raise errors.ConnectionLost()

        cursor = db.cursor(cursor_factory=DictCursor)

        # register the hstore option
        try:
            register_hstore(cursor, unicode=True)
        except StandardError:
            pass

        # register the json option
        register_json(cursor)

        try:
            cursor.execute(command, data)
            rowcount = cursor.rowcount

        # look for a cancelled query
        except QueryCanceledError:
            try:
                db.rollback()
            except StandardError as err:
                log.error('Rollback error: {0}'.format(err))
            raise errors.Interruption()

        # look for a disconnection error
        except pg.InterfaceError:
            raise errors.ConnectionLost()

        # look for integrity errors
        except (pg.IntegrityError, pg.OperationalError), err:
            try:
                db.rollback()
            except StandardError:
                pass

            log.debug(traceback.print_exc())
            raise errors.QueryFailed(command, data, nstr(err))

        # connection has closed underneath the hood
        except pg.Error, err:
            try:
                db.rollback()
            except StandardError:
                pass

            log.error(traceback.print_exc())
            raise errors.QueryFailed(command, data, nstr(err))

        try:
            results = [mapper(record) for record in cursor.fetchall()]<commit_msg>added some more logging if a query times out<commit_after>#!/usr/bin/python

""" Defines the backend connection class for PostgreSQL databases. """

# define authorship information
__authors__ = ['<NAME>']
__author__ = ','.join(__authors__)
__credits__ = []
__copyright__ = 'Copyright (c) 2011, Projex Software'
__license__ = 'LGPL'

# maintanence information
__maintainer__ = 'Projex Software'
__email__ = '<EMAIL>'

# define version information (major,minor,maintanence)
__depends__ = ['pyscopg2']
__version_info__ = (0, 0, 0)
__version__ = '%i.%i.%i' % __version_info__

# ------------------------------------------------------------------------------

import datetime
import logging
import os
import orb
import traceback

from orb import errors
from projex.text import nativestring as nstr

from ..abstractconnection import SQLConnection
from ..abstractsql import SQL

log = logging.getLogger(__name__)

try:
    import psycopg2 as pg
    from psycopg2.extras import DictCursor, register_hstore, register_json
    from psycopg2.extensions import QueryCanceledError

except ImportError:
    log.debug('For PostgreSQL backend, download the psycopg2 module')

    QueryCanceledError = errors.DatabaseError
    pg = None

# ----------------------------------------------------------------------

class PSQLConnection(SQLConnection):
    """ 
    Creates a PostgreSQL backend connection type for handling database
    connections to PostgreSQL databases.
    """

    # ----------------------------------------------------------------------
    # PROTECTED METHODS
    # ----------------------------------------------------------------------
    def _execute(self, command, data=None, autoCommit=True, autoClose=True,
                 returning=True, mapper=dict):
        """
        Executes the inputed command into the current \
        connection cursor.
        
        :param      command    | <str>
                    data       | <dict> || None
                    autoCommit | <bool> | commit database changes immediately
                    autoClose  | <bool> | closes connections immediately
        
        :return     [{<str> key: <variant>, ..}, ..], <int> count
        """
        rowcount = 0
        if data is None:
            data = {}

        # when in debug mode, simply log the command to the log
        elif self.database().commandsBlocked():
            log.info(command)
            return [], rowcount

        # create a new cursor for this transaction
        db = self.nativeConnection()
        if db is None:
            raise errors.ConnectionLost()

        cursor = db.cursor(cursor_factory=DictCursor)

        # register the hstore option
        try:
            register_hstore(cursor, unicode=True)
        except StandardError:
            pass

        # register the json option
        register_json(cursor)

        try:
            cursor.execute(command, data)
            rowcount = cursor.rowcount

        # look for a cancelled query
        except QueryCanceledError:
            try:
                db.rollback()
            except StandardError as err:
                log.error('Rollback error: {0}'.format(err))
            log.critical(command)
            if data:
                log.critical(str(data))
            raise errors.Interruption()

        # look for a disconnection error
        except pg.InterfaceError:
            raise errors.ConnectionLost()

        # look for integrity errors
        except (pg.IntegrityError, pg.OperationalError), err:
            try:
                db.rollback()
            except StandardError:
                pass

            log.debug(traceback.print_exc())
            raise errors.QueryFailed(command, data, nstr(err))

        # connection has closed underneath the hood
        except pg.Error, err:
            try:
                db.rollback()
            except StandardError:
                pass

            log.error(traceback.print_exc())
            raise errors.QueryFailed(command, data, nstr(err))

        try:
            results = [mapper(record) for record in cursor.fetchall()]
{'fraction_non_alphanumeric': 0.1478786383818451, 'fraction_numerical': 0.003453379378391712, 'mean_word_length': 1.7147639772346837, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2758014', 'n_tokens_mistral': 2055, 'n_tokens_neox': 1968, 'n_words': 593}