<commit_before>"""
PSQL command.

Runs the 'psql' tool with the appropriate parameters.
"""

import contextlib
import os
import shlex
import signal
import subprocess
import sys
import tempfile
import termios

from . import subcommand
from . import util


class PSQL(subcommand.SubCommand):
    def __init__(self, config, **args):
        self.config = config
        self.action = args['action']

    @staticmethod
    def argparse_register(subparser):
        subparser.add_argument(
            'action',
            choices=['rebuild', 'load-funcs', 'attach'],
            nargs='?'
        )

    async def run(self):
        """
        CLI entry point
        """
        project_folder = os.path.dirname(os.path.dirname(__file__))
        db_folder = os.path.join(project_folder, 'db')

        with contextlib.ExitStack() as exitstack:
            env = dict(os.environ)
            env['PGDATABASE'] = self.config['database']['dbname']

            if self.config['database']['user'] is None:
                if self.config['database']['host'] is not None:
                    raise ValueError("database user is None, but host is set")
                if self.config['database']['password'] is not None:
                    raise ValueError("database user is None, "
                                     "but password is set")
            else:
                def escape_colon(str):
                    return str.replace('\\', '\\\\').replace(':', '\\:')

                passfile = exitstack.enter_context(tempfile.NamedTemporaryFile('w'))
                os.chmod(passfile.name, 0o600)

                passfile.write(':'.join([
                    escape_colon(self.config['database']['host']),
                    '*',
                    escape_colon(self.config['database']['dbname']),
                    escape_colon(self.config['database']['user']),
                    escape_colon(self.config['database']['password'])
                ]))
                passfile.write('\n')
                passfile.flush()

                env['PGHOST'] = self.config['database']['host']
                env['PGUSER'] = self.config['database']['user']
                env['PGPASSFILE'] = passfile.name

            command = ['psql', '--variable', 'ON_ERROR_STOP=1']
            if self.action == 'rebuild':
                command.extend(['--file', 'rebuild.sql'])
            elif self.action == 'load-funcs':
                command.extend(['--file', 'funcs.sql'])
            elif self.action == 'attach':
                pass
            else:
                raise Exception(f'unknown action {self.action}')

            ret = util.run_as_fg_process(command, env=env, cwd=db_folder)

            if ret != 0:
                print(util.format_error('psql failed'))
            exit(ret)
<commit_msg>Add pgcli to psql attach<commit_after>"""
PSQL command.

Runs the 'psql' tool with the appropriate parameters.
"""

import contextlib
import os
import shlex
import shutil
import signal
import subprocess
import sys
import tempfile
import termios

from . import subcommand
from . import util


class PSQL(subcommand.SubCommand):
    def __init__(self, config, **args):
        self.config = config
        self.action = args['action']

    @staticmethod
    def argparse_register(subparser):
        subparser.add_argument(
            'action',
            choices=['rebuild', 'load-funcs', 'attach'],
            nargs='?'
        )

    async def run(self):
        """
        CLI entry point
        """
        project_folder = os.path.dirname(os.path.dirname(__file__))
        db_folder = os.path.join(project_folder, 'db')

        with contextlib.ExitStack() as exitstack:
            env = dict(os.environ)
            env['PGDATABASE'] = self.config['database']['dbname']

            if self.config['database']['user'] is None:
                if self.config['database']['host'] is not None:
                    raise ValueError("database user is None, but host is set")
                if self.config['database']['password'] is not None:
                    raise ValueError("database user is None, "
                                     "but password is set")
            else:
                def escape_colon(str):
                    return str.replace('\\', '\\\\').replace(':', '\\:')

                passfile = exitstack.enter_context(tempfile.NamedTemporaryFile('w'))
                os.chmod(passfile.name, 0o600)

                passfile.write(':'.join([
                    escape_colon(self.config['database']['host']),
                    '*',
                    escape_colon(self.config['database']['dbname']),
                    escape_colon(self.config['database']['user']),
                    escape_colon(self.config['database']['password'])
                ]))
                passfile.write('\n')
                passfile.flush()

                env['PGHOST'] = self.config['database']['host']
                env['PGUSER'] = self.config['database']['user']
                env['PGPASSFILE'] = passfile.name

            command = ['psql', '--variable', 'ON_ERROR_STOP=1']
            if self.action == 'rebuild':
                command.extend(['--file', 'rebuild.sql'])
            elif self.action == 'load-funcs':
                command.extend(['--file', 'funcs.sql'])
            elif self.action == 'attach':
                if shutil.which('pgcli') is not None:
                    # if pgcli is installed, use that instead!
                    command = ['pgcli']
            else:
                raise Exception(f'unknown action {self.action}')

            ret = util.run_as_fg_process(command, env=env, cwd=db_folder)

            if ret != 0:
                print(util.format_error('psql failed'))
            exit(ret)

{'fraction_non_alphanumeric': 0.1362143474503025, 'fraction_numerical': 0.002074330164217805, 'mean_word_length': 1.7803940413262855, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2083737', 'n_tokens_mistral': 1582, 'n_tokens_neox': 1503, 'n_words': 337}