<commit_before>#! /usr/bin/env python
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

"""
Create, prepare and load schema for Socorro PostgreSQL database.
"""
import sys
import psycopg2
import psycopg2.extensions
from psycopg2 import ProgrammingError
import re
import logging

from socorro.app.generic_app import App, main
from configman import Namespace


class PostgreSQLManager(object):
    def __init__(self, dsn, logger):
        self.conn = psycopg2.connect(dsn)
        self.conn.set_isolation_level(
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
        self.logger = logger

    def execute(self, sql, allowable_errors=None):
        cur = self.conn.cursor()
        try:
            cur.execute(sql)
        except ProgrammingError, e:
            if not allowable_errors:
                raise
            dberr = e.pgerror.strip().split('ERROR:  ')[1]
            if allowable_errors:
                for err in allowable_errors:
                    if re.match(err, dberr):
                        self.logger.warning(dberr)
                    else:
                        raise

    def version(self):
        cur = self.conn.cursor()
        cur.execute("SELECT version()")
        version_info = cur.fetchall()[0][0].split()
        return version_info[1]

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.conn.close()


class SocorroDB(App):
    app_name = 'setupdb'
    app_version = '0.2'
    app_description = __doc__

    required_config = Namespace()

    required_config.add_option(
        name='database_name',
        default='',
        doc='Name of database to manage',
    )

    required_config.add_option(
        name='database_hostname',
        default='',
        doc='Hostname to connect to database',
    )

    required_config.add_option(
        name='database_username',
        default='',
        doc='Username to connect to database',
    )

    required_config.add_option(
        name='database_password',
        default='',
        doc='Password to connect to database',
    )

    required_config.add_option(
        name='database_port',
        default='',
        doc='Port to connect to database',
    )

    required_config.add_option(
        name='dropdb',
        default=False,
        doc='Whether or not to drop database_name',
        exclude_from_print_conf=True,
        exclude_from_dump_conf=True
    )

    required_config.add_option(
        name='no_schema',
        default=False,
        doc='Whether or not to load schema',
        exclude_from_print_conf=True,
        exclude_from_dump_conf=True
    )

    required_config.add_option(
        name='citext',
        default='/usr/share/postgresql/9.0/contrib/citext.sql',
        doc='Name of citext.sql file',
    )

    def main(self):

        self.database_name = self.config['database_name']
        if not self.database_name:
            print "Syntax error: --database_name required"
            return 1

        self.no_schema = self.config.get('no_schema')
        self.citext = self.config.get('citext')

        dsn_template = 'dbname=%s'

        self.database_username = self.config.get('database_username')
        if self.database_username:
            dsn_template += ' user=%s' % self.database_username
        self.database_password = self.config.get('database_password')
        if self.database_password:
            dsn_template += ' password=%s' % self.database_password
        self.database_hostname = self.config.get('database_hostname')
        if self.database_hostname:
            dsn_template += ' host=%s' % self.database_hostname
        self.database_port = self.config.get('database_port')
        if self.database_port:
            dsn_template += ' port=%s' % self.database_port

        dsn = dsn_template % 'template1'

        with PostgreSQLManager(dsn, self.config.logger) as db:
            db_version = db.version()
            if not re.match(r'9\.[01][.*]', db_version):
                print 'ERROR - unrecognized PostgreSQL vesion: %s' % db_version
                print 'Only 9.0.x and 9.1.x are supported at this time.'
                return 1
            if self.config.get('dropdb'):
                if 'test' not in self.database_name:
                    confirm = raw_input(
                        'drop database %s [y/N]: ' % self.database_name)
                    if not confirm == "y":
                        logging.warn('NOT dropping table')
                        return 2

                db.execute('DROP DATABASE %s' % self.database_name,
                    ['database "%s" does not exist' % self.database_name])
                db.execute('DROP SCHEMA pgx_diag', ['schema "pgx_diag" does not exist'])

            db.execute('CREATE DATABASE %s' % self.database_name,
                       ['database "%s" already exists' % self.database_name])

        dsn = dsn_template % self.database_name

        with PostgreSQLManager(dsn, self.config.logger) as db:
            db_version = db.version()
            with open('sql/roles.sql') as f:
                db.execute(f.read())

            if not self.no_schema:
                with open('sql/schema.sql') as f:
                    db.execute(f.read(), ['schema "pgx_diag" already exists'])

                db.execute('SELECT weekly_report_partitions()')

            else:
                db.execute(
                    'ALTER DATABASE %s OWNER TO breakpad_rw' %
                    self.database_name)
                if re.match(r'9\.0[.*]', db_version):
                    with open(self.citext) as f:
                        db.execute(f.read())
                elif re.match(r'9\.1[.*]', db_version):
                    db.execute('CREATE EXTENSION citext')
                

        return 0

if __name__ == '__main__':  # pragma: no cover
    sys.exit(main(SocorroDB))
<commit_msg>Add "from unpackaged" so that creating the `citext` extension works in PG 9.1<commit_after>#! /usr/bin/env python
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

"""
Create, prepare and load schema for Socorro PostgreSQL database.
"""
import sys
import psycopg2
import psycopg2.extensions
from psycopg2 import ProgrammingError
import re
import logging

from socorro.app.generic_app import App, main
from configman import Namespace


class PostgreSQLManager(object):
    def __init__(self, dsn, logger):
        self.conn = psycopg2.connect(dsn)
        self.conn.set_isolation_level(
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
        self.logger = logger

    def execute(self, sql, allowable_errors=None):
        cur = self.conn.cursor()
        try:
            cur.execute(sql)
        except ProgrammingError, e:
            if not allowable_errors:
                raise
            dberr = e.pgerror.strip().split('ERROR:  ')[1]
            if allowable_errors:
                for err in allowable_errors:
                    if re.match(err, dberr):
                        self.logger.warning(dberr)
                    else:
                        raise

    def version(self):
        cur = self.conn.cursor()
        cur.execute("SELECT version()")
        version_info = cur.fetchall()[0][0].split()
        return version_info[1]

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.conn.close()


class SocorroDB(App):
    app_name = 'setupdb'
    app_version = '0.2'
    app_description = __doc__

    required_config = Namespace()

    required_config.add_option(
        name='database_name',
        default='',
        doc='Name of database to manage',
    )

    required_config.add_option(
        name='database_hostname',
        default='',
        doc='Hostname to connect to database',
    )

    required_config.add_option(
        name='database_username',
        default='',
        doc='Username to connect to database',
    )

    required_config.add_option(
        name='database_password',
        default='',
        doc='Password to connect to database',
    )

    required_config.add_option(
        name='database_port',
        default='',
        doc='Port to connect to database',
    )

    required_config.add_option(
        name='dropdb',
        default=False,
        doc='Whether or not to drop database_name',
        exclude_from_print_conf=True,
        exclude_from_dump_conf=True
    )

    required_config.add_option(
        name='no_schema',
        default=False,
        doc='Whether or not to load schema',
        exclude_from_print_conf=True,
        exclude_from_dump_conf=True
    )

    required_config.add_option(
        name='citext',
        default='/usr/share/postgresql/9.0/contrib/citext.sql',
        doc='Name of citext.sql file',
    )

    def main(self):

        self.database_name = self.config['database_name']
        if not self.database_name:
            print "Syntax error: --database_name required"
            return 1

        self.no_schema = self.config.get('no_schema')
        self.citext = self.config.get('citext')

        dsn_template = 'dbname=%s'

        self.database_username = self.config.get('database_username')
        if self.database_username:
            dsn_template += ' user=%s' % self.database_username
        self.database_password = self.config.get('database_password')
        if self.database_password:
            dsn_template += ' password=%s' % self.database_password
        self.database_hostname = self.config.get('database_hostname')
        if self.database_hostname:
            dsn_template += ' host=%s' % self.database_hostname
        self.database_port = self.config.get('database_port')
        if self.database_port:
            dsn_template += ' port=%s' % self.database_port

        dsn = dsn_template % 'template1'

        with PostgreSQLManager(dsn, self.config.logger) as db:
            db_version = db.version()
            if not re.match(r'9\.[01][.*]', db_version):
                print 'ERROR - unrecognized PostgreSQL vesion: %s' % db_version
                print 'Only 9.0.x and 9.1.x are supported at this time.'
                return 1
            if self.config.get('dropdb'):
                if 'test' not in self.database_name:
                    confirm = raw_input(
                        'drop database %s [y/N]: ' % self.database_name)
                    if not confirm == "y":
                        logging.warn('NOT dropping table')
                        return 2

                db.execute('DROP DATABASE %s' % self.database_name,
                    ['database "%s" does not exist' % self.database_name])
                db.execute('DROP SCHEMA pgx_diag', ['schema "pgx_diag" does not exist'])

            db.execute('CREATE DATABASE %s' % self.database_name,
                       ['database "%s" already exists' % self.database_name])

        dsn = dsn_template % self.database_name

        with PostgreSQLManager(dsn, self.config.logger) as db:
            db_version = db.version()
            with open('sql/roles.sql') as f:
                db.execute(f.read())

            if not self.no_schema:
                with open('sql/schema.sql') as f:
                    db.execute(f.read(), ['schema "pgx_diag" already exists'])

                db.execute('SELECT weekly_report_partitions()')

            else:
                db.execute(
                    'ALTER DATABASE %s OWNER TO breakpad_rw' %
                    self.database_name)
                if re.match(r'9\.0[.*]', db_version):
                    with open(self.citext) as f:
                        db.execute(f.read())
                elif re.match(r'9\.1[.*]', db_version):
                    db.execute('CREATE EXTENSION citext from unpackaged')
                

        return 0

if __name__ == '__main__':  # pragma: no cover
    sys.exit(main(SocorroDB))

{'fraction_non_alphanumeric': 0.10480958634274458, 'fraction_numerical': 0.005581089954038083, 'mean_word_length': 1.9575242718446602, 'pattern_counts': {'":': 2, '<': 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': '6427999', 'n_tokens_mistral': 3617, 'n_tokens_neox': 3439, 'n_words': 814}