<commit_before>import os
import sys

from fabric.api import run, sudo, env, local, hide, settings
from fabric.contrib.files import append, sed, exists, contains
from fabric.context_managers import prefix
from fabric.operations import get, put
from fabric.context_managers import cd

from fabric.tasks import Task

from utils import random_password

class PostgresInstall(Task):
    """
    Install postgresql on server

    install postgresql package;
    enable postgres access from localhost without password;
    enable all other user access from other machines with password;
    setup a few parameters related with streaming replication;
    database server listen to all machines '*';
    create a user for database with password.
    """

    name = 'master_setup'
    db_version = '9.1'

    encrypt = 'md5'
    hba_txts = ('local   all    postgres                     ident\n'
                'host    replication replicator  0.0.0.0/0   md5\n'
                'local   all    all                          password\n'
                '# # IPv4 local connections:\n'
                'host    all    all         127.0.0.1/32     %(encrypt)s\n'
                '# # IPv6 local connections:\n'
                'host    all    all         ::1/128          %(encrypt)s\n'
                '# # IPv4 external\n'
                'host    all    all         0.0.0.0/0        %(encrypt)s\n')

    postgres_config = {
        'listen_addresses':  "'*'",
        'wal_level':         "hot_standby",
        'wal_keep_segments': "32",
        'max_wal_senders':   "5",
        'archive_mode':      "on" }

    def _get_data_dir(self, db_version):
        return os.path.join('/var', 'pgsql', 'data%s' %db_version)

    def _setup_parameter(self, file, **kwargs):
        for key, value in kwargs.items():
            origin = "#%s =" %key
            new = "%s = %s" %(key, value)
            sudo('sed -i "/%s/ c\%s" %s' %(origin, new, file))

    def _install_package(self, db_version=None):
        sudo("pkg_add postgresql%s-server" %db_version)
        sudo("pkg_add postgresql%s-replicationtools" %db_version)
        sudo("svcadm enable postgresql:pg%s" %db_version)

    def _setup_hba_config(self, data_dir=None, encrypt=None):
        """
        enable postgres access without password from localhost
        """
        hba_conf = os.path.join(data_dir, 'pg_hba.conf')
        kwargs = {'data_dir':data_dir, 'encrypt':encrypt}
        hba_txts = self.hba_txts % kwargs

        if exists(hba_conf, use_sudo=True):
            sudo("echo '%s' > %s" %(hba_txts, hba_conf))
        else:
            print ('Could not find file %s. Please make sure postgresql was '
                   'installed and data dir was created correctly.'%hba_conf)
            sys.exit()

    def _setup_postgres_config(self, data_dir=None, config=None):
        postgres_conf = os.path.join(data_dir, 'postgresql.conf')

        if exists(postgres_conf, use_sudo=True):
            self._setup_parameter(postgres_conf, **config)
        else:
            print ('Could not find file %s. Please make sure postgresql was '
                   'installed and data dir was created correctly.' %postgres_conf)
            sys.exit()

    def _setup_archive_dir(self, data_dir):
        archive_dir = os.path.join(data_dir, 'wal_archive')
        sudo("mkdir -p %s" %archive_dir)
        sudo("chown postgres:postgres %s" %archive_dir)

        return archive_dir

    def _setup_ssh_key(self):
        ssh_dir = '/var/pgsql/.ssh'

        sudo('mkdir -p %s' %ssh_dir)
        sudo('chown -R postgres:postgres %s' %ssh_dir)
        sudo('chmod -R og-rwx %s' %ssh_dir)
        rsa = os.path.join(ssh_dir, 'id_rsa')
        run('sudo su postgres -c "ssh-keygen -t rsa -f %s -N \'\'"' %rsa)

    def _restart_db_server(self, db_version):
        sudo('svcadm restart postgresql:pg%s' %db_version)

    def _create_user(self):
        username = raw_input("Now we are creating the database user, please "
                             "specify a username: ")
        # 'postgres' is postgresql superuser
        while username == 'postgres':
            username = raw_input("Sorry, you are not allowed to use postgres "
                                 "as username, please choose another one: ")
        db_out = run('echo "select usename from pg_shadow where usename=\'%s\'" |'
                     'sudo su postgres -c psql' %username)
        if username in db_out:
            print 'user %s already exists, skipping creating user.' %username
        else:
            run("sudo su postgres -c 'createuser -D -S -R -P %s'" %username)

        env.config_object.set('db-server', env.config_object.USERNAME, username)

        return username

    def _create_replicator(self, db_version):
        db_out = run("echo '\du replicator' | sudo su postgres -c 'psql'")
        if 'replicator' not in db_out:
            replicator_pass = <PASSWORD>(12)

            c1 = ('CREATE USER replicator REPLICATION LOGIN ENCRYPTED '
                  'PASSWORD \"\'%s\'\"' %replicator_pass)
            run("echo %s | sudo su postgres -c \'psql\'" %c1)
            history_file = os.path.join('/var', 'pgsql', '.psql_history')
            if exists(history_file):
                sudo('rm %s' %history_file)
            print "user replicator already exists, skipping creating user."

            env.config_object.set('db-server', env.config_object.REPLICATOR,
                                  'replicator')
            env.config_object.set('db-server', env.config_object.REPLICATOR_PASS,
                                  replicator_pass)

            return replicator_pass

    def run(self, db_version=None, encrypt=None, save_config=True, **kwargs):
        if not db_version:
            db_version = self.db_version
        db_version = ''.join(db_version.split('.')[:2])
        data_dir = self._get_data_dir(db_version)
        hba_conf= os.path.join(data_dir, 'pg_hba.conf')

        if not encrypt:
            encrypt = self.encrypt

        self._install_package(db_version=db_version)
        archive_dir = self._setup_archive_dir(data_dir)
        self.postgres_config['archive_command'] = ("'cp %s %s/wal_archive/%s'"
                                                   %('%p', data_dir, '%f'))

        self._setup_hba_config(data_dir, encrypt)
        self._setup_postgres_config(data_dir=data_dir,
                                    config=self.postgres_config)
        self._restart_db_server(db_version)
        self._setup_ssh_key()
        self._create_user()
        self._create_replicator(db_version)

        if save_config:
            print env.conf_filename
            env.config_object.save(env.conf_filename)


class SlaveSetup(PostgresInstall):
    """
    Set up master-slave streaming replication: slave node
    """

    name = 'slave_setup'

    postgres_config = {
        'listen_addresses': "'*'",
        'wal_level':      "hot_standby",
        'hot_standby':    "on"}

    def _get_master_db_version(self, master):
        command = ("ssh %s psql --version | head -1 | awk '{print $3}'" %master)
        version_string = local(command, capture=True)
        version = ''.join(version_string.split('.')[:2])

        return version

    def _get_replicator_pass(self, section='db-server'):
        password = env.config_object.get_list(section,
                                             env.config_object.REPLICATOR_PASS)
        return password[0]

    def _setup_recovery_conf(self, master_ip, password, data_dir):
        wal_dir = os.path.join(data_dir, 'wal_archive')<commit_msg>remove unnecessary print statement<commit_after>import os
import sys

from fabric.api import run, sudo, env, local, hide, settings
from fabric.contrib.files import append, sed, exists, contains
from fabric.context_managers import prefix
from fabric.operations import get, put
from fabric.context_managers import cd

from fabric.tasks import Task

from utils import random_password

class PostgresInstall(Task):
    """
    Install postgresql on server

    install postgresql package;
    enable postgres access from localhost without password;
    enable all other user access from other machines with password;
    setup a few parameters related with streaming replication;
    database server listen to all machines '*';
    create a user for database with password.
    """

    name = 'master_setup'
    db_version = '9.1'

    encrypt = 'md5'
    hba_txts = ('local   all    postgres                     ident\n'
                'host    replication replicator  0.0.0.0/0   md5\n'
                'local   all    all                          password\n'
                '# # IPv4 local connections:\n'
                'host    all    all         127.0.0.1/32     %(encrypt)s\n'
                '# # IPv6 local connections:\n'
                'host    all    all         ::1/128          %(encrypt)s\n'
                '# # IPv4 external\n'
                'host    all    all         0.0.0.0/0        %(encrypt)s\n')

    postgres_config = {
        'listen_addresses':  "'*'",
        'wal_level':         "hot_standby",
        'wal_keep_segments': "32",
        'max_wal_senders':   "5",
        'archive_mode':      "on" }

    def _get_data_dir(self, db_version):
        return os.path.join('/var', 'pgsql', 'data%s' %db_version)

    def _setup_parameter(self, file, **kwargs):
        for key, value in kwargs.items():
            origin = "#%s =" %key
            new = "%s = %s" %(key, value)
            sudo('sed -i "/%s/ c\%s" %s' %(origin, new, file))

    def _install_package(self, db_version=None):
        sudo("pkg_add postgresql%s-server" %db_version)
        sudo("pkg_add postgresql%s-replicationtools" %db_version)
        sudo("svcadm enable postgresql:pg%s" %db_version)

    def _setup_hba_config(self, data_dir=None, encrypt=None):
        """
        enable postgres access without password from localhost
        """
        hba_conf = os.path.join(data_dir, 'pg_hba.conf')
        kwargs = {'data_dir':data_dir, 'encrypt':encrypt}
        hba_txts = self.hba_txts % kwargs

        if exists(hba_conf, use_sudo=True):
            sudo("echo '%s' > %s" %(hba_txts, hba_conf))
        else:
            print ('Could not find file %s. Please make sure postgresql was '
                   'installed and data dir was created correctly.'%hba_conf)
            sys.exit()

    def _setup_postgres_config(self, data_dir=None, config=None):
        postgres_conf = os.path.join(data_dir, 'postgresql.conf')

        if exists(postgres_conf, use_sudo=True):
            self._setup_parameter(postgres_conf, **config)
        else:
            print ('Could not find file %s. Please make sure postgresql was '
                   'installed and data dir was created correctly.' %postgres_conf)
            sys.exit()

    def _setup_archive_dir(self, data_dir):
        archive_dir = os.path.join(data_dir, 'wal_archive')
        sudo("mkdir -p %s" %archive_dir)
        sudo("chown postgres:postgres %s" %archive_dir)

        return archive_dir

    def _setup_ssh_key(self):
        ssh_dir = '/var/pgsql/.ssh'

        sudo('mkdir -p %s' %ssh_dir)
        sudo('chown -R postgres:postgres %s' %ssh_dir)
        sudo('chmod -R og-rwx %s' %ssh_dir)
        rsa = os.path.join(ssh_dir, 'id_rsa')
        run('sudo su postgres -c "ssh-keygen -t rsa -f %s -N \'\'"' %rsa)

    def _restart_db_server(self, db_version):
        sudo('svcadm restart postgresql:pg%s' %db_version)

    def _create_user(self):
        username = raw_input("Now we are creating the database user, please "
                             "specify a username: ")
        # 'postgres' is postgresql superuser
        while username == 'postgres':
            username = raw_input("Sorry, you are not allowed to use postgres "
                                 "as username, please choose another one: ")
        db_out = run('echo "select usename from pg_shadow where usename=\'%s\'" |'
                     'sudo su postgres -c psql' %username)
        if username in db_out:
            print 'user %s already exists, skipping creating user.' %username
        else:
            run("sudo su postgres -c 'createuser -D -S -R -P %s'" %username)

        env.config_object.set('db-server', env.config_object.USERNAME, username)

        return username

    def _create_replicator(self, db_version):
        db_out = run("echo '\du replicator' | sudo su postgres -c 'psql'")
        if 'replicator' not in db_out:
            replicator_pass = <PASSWORD>)

            c1 = ('CREATE USER replicator REPLICATION LOGIN ENCRYPTED '
                  'PASSWORD \"\'%s\'\"' %replicator_pass)
            run("echo %s | sudo su postgres -c \'psql\'" %c1)
            history_file = os.path.join('/var', 'pgsql', '.psql_history')
            if exists(history_file):
                sudo('rm %s' %history_file)
            print "user replicator already exists, skipping creating user."

            env.config_object.set('db-server', env.config_object.REPLICATOR,
                                  'replicator')
            env.config_object.set('db-server', env.config_object.REPLICATOR_PASS,
                                  replicator_pass)

            return replicator_pass

    def run(self, db_version=None, encrypt=None, save_config=True, **kwargs):
        if not db_version:
            db_version = self.db_version
        db_version = ''.join(db_version.split('.')[:2])
        data_dir = self._get_data_dir(db_version)
        hba_conf= os.path.join(data_dir, 'pg_hba.conf')

        if not encrypt:
            encrypt = self.encrypt

        self._install_package(db_version=db_version)
        archive_dir = self._setup_archive_dir(data_dir)
        self.postgres_config['archive_command'] = ("'cp %s %s/wal_archive/%s'"
                                                   %('%p', data_dir, '%f'))

        self._setup_hba_config(data_dir, encrypt)
        self._setup_postgres_config(data_dir=data_dir,
                                    config=self.postgres_config)
        self._restart_db_server(db_version)
        self._setup_ssh_key()
        self._create_user()
        self._create_replicator(db_version)

        if save_config:
            env.config_object.save(env.conf_filename)


class SlaveSetup(PostgresInstall):
    """
    Set up master-slave streaming replication: slave node
    """

    name = 'slave_setup'

    postgres_config = {
        'listen_addresses': "'*'",
        'wal_level':      "hot_standby",
        'hot_standby':    "on"}

    def _get_master_db_version(self, master):
        command = ("ssh %s psql --version | head -1 | awk '{print $3}'" %master)
        version_string = local(command, capture=True)
        version = ''.join(version_string.split('.')[:2])

        return version

    def _get_replicator_pass(self, section='db-server'):
        password = env.config_object.get_list(section,
                                             env.config_object.REPLICATOR_PASS)
        return password[0]

    def _setup_recovery_conf(self, master_ip, password, data_dir):
        wal_dir = os.path.join(data_dir, 'wal_archive')
{'fraction_non_alphanumeric': 0.11136004217185029, 'fraction_numerical': 0.005271481286241434, 'mean_word_length': 2.1363918164910105, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 6, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2986999', 'n_tokens_mistral': 4791, 'n_tokens_neox': 4581, 'n_words': 1130}