<commit_before>import logging
import os
import psycopg2
import sys
import time

is_py3 = sys.hexversion >= 0x03000000

if is_py3:
    from urllib.parse import urlparse
else:
    from urlparse import urlparse


logger = logging.getLogger(__name__)


def parseurl(url):
    r = urlparse(url)
    return {
        'host': r.hostname,
        'port': r.port or 5432,
        'user': r.username,
        'password': <PASSWORD>,
        'database': r.path[1:],
        'fallback_application_name': 'Governor',
        'connect_timeout': 5,
    }


class Postgresql:

    def __init__(self, config):
        self.config = config
        self.name = config['name']
        self.listen_addresses, self.port = config['listen'].split(':')
        self.data_dir = config['data_dir']
        self.replication = config['replication']
        self.recovery_conf = os.path.join(self.data_dir, 'recovery.conf')
        self.pid_path = os.path.join(self.data_dir, 'postmaster.pid')
        self._pg_ctl = 'pg_ctl -w -D ' + self.data_dir

        self.local_address = self.get_local_address()
        connect_address = config.get('connect_address', None) or self.local_address
        self.connection_string = 'postgres://{username}:{password}@{connect_address}/postgres'.format(
            connect_address=connect_address, **self.replication)

        self._connection = None
        self._cursor_holder = None
        self.members = []  # list of already existing replication slots

    def get_local_address(self):
        # TODO: try to get unix_socket_directory from postmaster.pid
        return self.listen_addresses.split(',')[0].strip() + ':' + self.port

    def connection(self):
        if not self._connection or self._connection.closed != 0:
            self._connection = psycopg2.connect('postgres://{}/postgres'.format(self.local_address))
            self._connection.autocommit = True
        return self._connection

    def _cursor(self):
        if not self._cursor_holder or self._cursor_holder.closed:
            self._cursor_holder = self.connection().cursor()
        return self._cursor_holder

    def disconnect(self):
        self._connection and self._connection.close()
        self._connection = self._cursor_holder = None

    def query(self, sql, *params):
        max_attempts = 0
        while True:
            ex = None
            try:
                cursor = self._cursor()
                cursor.execute(sql, params)
                return cursor
            except psycopg2.InterfaceError as e:
                ex = e<commit_msg>Use connect_timeout=3s and statement_timeout=2s everywhere (even for local connections)<commit_after>import logging
import os
import psycopg2
import sys
import time

is_py3 = sys.hexversion >= 0x03000000

if is_py3:
    from urllib.parse import urlparse
else:
    from urlparse import urlparse


logger = logging.getLogger(__name__)


def parseurl(url):
    r = urlparse(url)
    ret = {
        'host': r.hostname,
        'port': r.port or 5432,
        'database': r.path[1:],
        'fallback_application_name': 'Governor',
        'connect_timeout': 3,
        'options': '-c statement_timeout=2000',
    }
    if r.username:
        ret['user'] = r.username
    if r.password:
        ret['password'] = <PASSWORD>
    return ret


class Postgresql:

    def __init__(self, config):
        self.config = config
        self.name = config['name']
        self.listen_addresses, self.port = config['listen'].split(':')
        self.data_dir = config['data_dir']
        self.replication = config['replication']
        self.recovery_conf = os.path.join(self.data_dir, 'recovery.conf')
        self.pid_path = os.path.join(self.data_dir, 'postmaster.pid')
        self._pg_ctl = 'pg_ctl -w -D ' + self.data_dir

        self.local_address = self.get_local_address()
        connect_address = config.get('connect_address', None) or self.local_address
        self.connection_string = 'postgres://{username}:{password}@{connect_address}/postgres'.format(
            connect_address=connect_address, **self.replication)

        self._connection = None
        self._cursor_holder = None
        self.members = []  # list of already existing replication slots

    def get_local_address(self):
        # TODO: try to get unix_socket_directory from postmaster.pid
        return self.listen_addresses.split(',')[0].strip() + ':' + self.port

    def connection(self):
        if not self._connection or self._connection.closed != 0:
            r = parseurl('postgres://{}/postgres'.format(self.local_address))
            self._connection = psycopg2.connect(**r)
            self._connection.autocommit = True
        return self._connection

    def _cursor(self):
        if not self._cursor_holder or self._cursor_holder.closed:
            self._cursor_holder = self.connection().cursor()
        return self._cursor_holder

    def disconnect(self):
        self._connection and self._connection.close()
        self._connection = self._cursor_holder = None

    def query(self, sql, *params):
        max_attempts = 0
        while True:
            ex = None
            try:
                cursor = self._cursor()
                cursor.execute(sql, params)
                return cursor
            except psycopg2.InterfaceError as e:
                ex = e
{'fraction_non_alphanumeric': 0.1092752534735261, 'fraction_numerical': 0.00976342470897484, 'mean_word_length': 2.4977019041365724, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6775579', 'n_tokens_mistral': 1591, 'n_tokens_neox': 1525, 'n_words': 346}