<commit_before>import multiprocessing
import os
import warnings

from datetime import datetime

import six

from fabric import api as fab
from fabric.contrib import files

import fabricio

from fabricio import docker, utils
from fabricio.docker.container import Attribute
from fabricio.utils import Options


class PostgresqlBackupMixin(docker.BaseService):
    """
    Your Docker image must have pg_dump and pg_restore installed in order
    to run backup and restore respectively
    (usually this requires `postgresql-client-common` package for Ubuntu/Debian)
    """

    db_name = Attribute()

    db_user = Attribute(default='postgres')

    db_host = Attribute()

    db_port = Attribute()

    @Attribute
    def db_backup_folder(self):
        warnings.warn(
            'db_backup_folder is deprecated and will be removed in ver. 0.4, '
            'use db_backup_dir instead', DeprecationWarning,
        )
        return None

    @Attribute
    def db_backup_dir(self):
        warnings.warn(
            'db_backup_folder is deprecated and will be removed in ver. 0.4, '
            'use db_backup_dir instead', RuntimeWarning,
        )
        return self.db_backup_folder

    db_backup_format = Attribute(default='c')

    db_backup_compress_level = Attribute()  # 0-9 (0 - no compression, 9 - max)

    db_backup_workers = Attribute(default=1)

    @Attribute
    def db_backup_name(self):
        warnings.warn(
            'db_backup_name is deprecated and will be removed in ver. 0.4, '
            'use db_backup_filename instead', DeprecationWarning,
        )
        return '{datetime:%Y-%m-%dT%H:%M:%S.%f}.dump'

    db_restore_workers = Attribute(default=4)

    @Attribute
    def db_backup_filename(self):
        warnings.warn(
            'db_backup_name is deprecated and will be removed in ver. 0.4, '
            'use db_backup_filename instead', RuntimeWarning,
        )
        return self.db_backup_name

    @property
    def db_connection_options(self):
        return Options([
            ('username', self.db_user),
            ('host', self.db_host),
            ('port', self.db_port),
        ])

    @property
    def db_backup_options(self):
        return Options([
            ('if-exists', True),
            ('create', True),
            ('clean', True),
        ])

    def make_backup_command(self):
        options = Options(self.db_connection_options)
        options.update(self.db_backup_options)
        options.update([
            ('format', self.db_backup_format),
            ('dbname', self.db_name),
            ('compress', self.db_backup_compress_level),
            ('jobs', self.db_backup_workers),
            ('file', os.path.join(
                self.db_backup_dir,
                self.db_backup_filename.format(datetime=datetime.utcnow())
            )),
        ])
        return 'pg_dump {options}'.format(options=options)

    @utils.once_per_command
    def backup(self):
        if self.db_backup_dir is None:
            fab.abort('db_backup_dir not set, can\'t continue with backup')
        command = self.make_backup_command()
        self.image.run(
            command=command,
            quiet=False,
            options=self.safe_options,
        )

    @property
    def db_restore_options(self):
        return self.db_backup_options

    def make_restore_command(self, backup_filename):
        options = Options(self.db_connection_options)
        options.update(self.db_restore_options)
        options.update([
            ('dbname', 'template1'),  # use any existing DB
            ('jobs', self.db_restore_workers),
            ('file', os.path.join(self.db_backup_dir, backup_filename)),
        ])
        return 'pg_restore {options}'.format(options=options)

    @utils.once_per_command
    def restore(self, backup_name=None):
        """
        Before run this method you have somehow to disable incoming connections,
        e.g. by stopping all database client containers:

            client_container.stop()
            pg_container.restore()
            client_container.start()
        """
        if self.db_backup_dir is None:
            fab.abort('db_backup_dir not set, can\'t continue with restore')

        if backup_name is None:
            raise ValueError('backup_filename not provided')

        command = self.make_restore_command(backup_name)
        self.image.run(
            command=command,
            quiet=False,
            options=self.safe_options,
        )


class PostgresqlContainer(docker.Container):

    @Attribute
    def postgresql_conf(self):
        warnings.warn(
            'postgresql_conf is deprecated and will be removed in ver. 0.4, '
            'use pg_conf instead', DeprecationWarning,
        )
        return 'postgresql.conf'

    @Attribute
    def pg_conf(self):
        warnings.warn(
            'postgresql_conf is deprecated and will be removed in ver. 0.4, '
            'use pg_conf instead', RuntimeWarning,
        )
        return self.postgresql_conf

    @Attribute
    def pg_hba_conf(self):
        warnings.warn(
            'pg_hba_conf is deprecated and will be removed in ver. 0.4, '
            'use pg_hba instead', DeprecationWarning,
        )
        return 'pg_hba.conf'

    @Attribute
    def pg_hba(self):
        warnings.warn(
            'pg_hba_conf is deprecated and will be removed in ver. 0.4, '
            'use pg_hba instead', RuntimeWarning,
        )
        return self.pg_hba_conf

    @Attribute
    def data(self):
        warnings.warn(
            'data is deprecated and will be removed in ver. 0.4, '
            'use pg_data instead', DeprecationWarning,
        )
        return NotImplemented

    @Attribute
    def pg_data(self):
        warnings.warn(
            'data is deprecated and will be removed in ver. 0.4, '
            'use pg_data instead', RuntimeWarning,
        )
        return self.data

    stop_signal = 'INT'

    stop_timeout = 30

    def __init__(self, *args, **kwargs):
        super(PostgresqlContainer, self).__init__(*args, **kwargs)
        assert self.volumes, 'provide volume for your data'

    @staticmethod
    def update_config(content, path):
        old_file = six.BytesIO()
        if files.exists(path, use_sudo=True):
            fab.get(remote_path=path, local_path=old_file, use_sudo=True)
        old_content = old_file.getvalue()
        need_update = content != old_content
        if need_update:
            fabricio.move_file(
                path_from=path,
                path_to=path + '.backup',
                sudo=True,
                ignore_errors=True,
            )
            fab.put(six.BytesIO(content), path, use_sudo=True, mode='0644')
            fabricio.log('{path} updated'.format(path=path))
        else:
            fabricio.log('{path} not changed'.format(path=path))
        return need_update

    def db_exists(self):
        return files.exists(
            os.path.join(self.pg_data, 'PG_VERSION'),
            use_sudo=True,
        )

    def create_db(self, tag=None, registry=None, account=None):
        fabricio.log('PostgreSQL database not found, creating new...')
        self.image[registry:tag:account].run(
            # official PostgreSQL image executes 'postgres initdb' before
            # any 'postgres' command (see /docker-entrypoint.sh),
            # therefore if you use image other then official, you may need
            # implement your own `create_db()`
            'postgres --version',
            options=self.safe_options,
            quiet=False,
        )

    def update(self, tag=None, registry=None, account=None, force=False):
        if not self.db_exists():
            self.create_db(tag=tag, registry=registry, account=account)

        main_conf = os.path.join(self.pg_data, 'postgresql.conf')
        hba_conf = os.path.join(self.pg_data, 'pg_hba.conf')

        main_config_updated = self.update_config(
            content=open(self.pg_conf).read(),
            path=main_conf,
        )
        hba_config_updated = self.update_config(
            content=open(self.pg_hba).read(),<commit_msg>fixed DB data 'volume' check<commit_after>import multiprocessing
import os
import warnings

from datetime import datetime

import six

from fabric import api as fab
from fabric.contrib import files

import fabricio

from fabricio import docker, utils
from fabricio.docker.container import Attribute
from fabricio.utils import Options


class PostgresqlBackupMixin(docker.BaseService):
    """
    Your Docker image must have pg_dump and pg_restore installed in order
    to run backup and restore respectively
    (usually this requires `postgresql-client-common` package for Ubuntu/Debian)
    """

    db_name = Attribute()

    db_user = Attribute(default='postgres')

    db_host = Attribute()

    db_port = Attribute()

    @Attribute
    def db_backup_folder(self):
        warnings.warn(
            'db_backup_folder is deprecated and will be removed in ver. 0.4, '
            'use db_backup_dir instead', DeprecationWarning,
        )
        return None

    @Attribute
    def db_backup_dir(self):
        warnings.warn(
            'db_backup_folder is deprecated and will be removed in ver. 0.4, '
            'use db_backup_dir instead', RuntimeWarning,
        )
        return self.db_backup_folder

    db_backup_format = Attribute(default='c')

    db_backup_compress_level = Attribute()  # 0-9 (0 - no compression, 9 - max)

    db_backup_workers = Attribute(default=1)

    @Attribute
    def db_backup_name(self):
        warnings.warn(
            'db_backup_name is deprecated and will be removed in ver. 0.4, '
            'use db_backup_filename instead', DeprecationWarning,
        )
        return '{datetime:%Y-%m-%dT%H:%M:%S.%f}.dump'

    db_restore_workers = Attribute(default=4)

    @Attribute
    def db_backup_filename(self):
        warnings.warn(
            'db_backup_name is deprecated and will be removed in ver. 0.4, '
            'use db_backup_filename instead', RuntimeWarning,
        )
        return self.db_backup_name

    @property
    def db_connection_options(self):
        return Options([
            ('username', self.db_user),
            ('host', self.db_host),
            ('port', self.db_port),
        ])

    @property
    def db_backup_options(self):
        return Options([
            ('if-exists', True),
            ('create', True),
            ('clean', True),
        ])

    def make_backup_command(self):
        options = Options(self.db_connection_options)
        options.update(self.db_backup_options)
        options.update([
            ('format', self.db_backup_format),
            ('dbname', self.db_name),
            ('compress', self.db_backup_compress_level),
            ('jobs', self.db_backup_workers),
            ('file', os.path.join(
                self.db_backup_dir,
                self.db_backup_filename.format(datetime=datetime.utcnow())
            )),
        ])
        return 'pg_dump {options}'.format(options=options)

    @utils.once_per_command
    def backup(self):
        if self.db_backup_dir is None:
            fab.abort('db_backup_dir not set, can\'t continue with backup')
        command = self.make_backup_command()
        self.image.run(
            command=command,
            quiet=False,
            options=self.safe_options,
        )

    @property
    def db_restore_options(self):
        return self.db_backup_options

    def make_restore_command(self, backup_filename):
        options = Options(self.db_connection_options)
        options.update(self.db_restore_options)
        options.update([
            ('dbname', 'template1'),  # use any existing DB
            ('jobs', self.db_restore_workers),
            ('file', os.path.join(self.db_backup_dir, backup_filename)),
        ])
        return 'pg_restore {options}'.format(options=options)

    @utils.once_per_command
    def restore(self, backup_name=None):
        """
        Before run this method you have somehow to disable incoming connections,
        e.g. by stopping all database client containers:

            client_container.stop()
            pg_container.restore()
            client_container.start()
        """
        if self.db_backup_dir is None:
            fab.abort('db_backup_dir not set, can\'t continue with restore')

        if backup_name is None:
            raise ValueError('backup_filename not provided')

        command = self.make_restore_command(backup_name)
        self.image.run(
            command=command,
            quiet=False,
            options=self.safe_options,
        )


class PostgresqlContainer(docker.Container):

    @Attribute
    def postgresql_conf(self):
        warnings.warn(
            'postgresql_conf is deprecated and will be removed in ver. 0.4, '
            'use pg_conf instead', DeprecationWarning,
        )
        return 'postgresql.conf'

    @Attribute
    def pg_conf(self):
        warnings.warn(
            'postgresql_conf is deprecated and will be removed in ver. 0.4, '
            'use pg_conf instead', RuntimeWarning,
        )
        return self.postgresql_conf

    @Attribute
    def pg_hba_conf(self):
        warnings.warn(
            'pg_hba_conf is deprecated and will be removed in ver. 0.4, '
            'use pg_hba instead', DeprecationWarning,
        )
        return 'pg_hba.conf'

    @Attribute
    def pg_hba(self):
        warnings.warn(
            'pg_hba_conf is deprecated and will be removed in ver. 0.4, '
            'use pg_hba instead', RuntimeWarning,
        )
        return self.pg_hba_conf

    @Attribute
    def data(self):
        warnings.warn(
            'data is deprecated and will be removed in ver. 0.4, '
            'use pg_data instead', DeprecationWarning,
        )
        return NotImplemented

    @Attribute
    def pg_data(self):
        warnings.warn(
            'data is deprecated and will be removed in ver. 0.4, '
            'use pg_data instead', RuntimeWarning,
        )
        return self.data

    stop_signal = 'INT'

    stop_timeout = 30

    @staticmethod
    def update_config(content, path):
        old_file = six.BytesIO()
        if files.exists(path, use_sudo=True):
            fab.get(remote_path=path, local_path=old_file, use_sudo=True)
        old_content = old_file.getvalue()
        need_update = content != old_content
        if need_update:
            fabricio.move_file(
                path_from=path,
                path_to=path + '.backup',
                sudo=True,
                ignore_errors=True,
            )
            fab.put(six.BytesIO(content), path, use_sudo=True, mode='0644')
            fabricio.log('{path} updated'.format(path=path))
        else:
            fabricio.log('{path} not changed'.format(path=path))
        return need_update

    def db_exists(self):
        return files.exists(
            os.path.join(self.pg_data, 'PG_VERSION'),
            use_sudo=True,
        )

    def create_db(self, tag=None, registry=None, account=None):
        fabricio.log('PostgreSQL database not found, creating new...')
        self.image[registry:tag:account].run(
            # official PostgreSQL image executes 'postgres initdb' before
            # any 'postgres' command (see /docker-entrypoint.sh),
            # therefore if you use image other then official, you may need
            # implement your own `create_db()`
            'postgres --version',
            options=self.safe_options,
            quiet=False,
        )

    def update(self, tag=None, registry=None, account=None, force=False):
        if 'volume' not in self.options:
            raise ValueError('Make sure you provide volume for DB data')
        if not self.db_exists():
            self.create_db(tag=tag, registry=registry, account=account)

        main_conf = os.path.join(self.pg_data, 'postgresql.conf')
        hba_conf = os.path.join(self.pg_data, 'pg_hba.conf')

        main_config_updated = self.update_config(
            content=open(self.pg_conf).read(),
            path=main_conf,
        )
        hba_config_updated = self.update_config(
            content=open(self.pg_hba).read(),
{'fraction_non_alphanumeric': 0.09500924214417746, 'fraction_numerical': 0.004066543438077634, 'mean_word_length': 2.190682130921958, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6173285', 'n_tokens_mistral': 4712, 'n_tokens_neox': 4576, 'n_words': 1115}