tryton-upgrades/upgrade

#!/usr/bin/env python
import os
import sys
import yaml
import argparse
from itertools import chain
from urlparse import urlparse
import psycopg2
import subprocess
import ConfigParser

from blessings import Terminal

import trytond

trytond_version = '.'.join(trytond.__version__.split('.')[:2])


t = Terminal()

def get_url():
    config = ConfigParser.ConfigParser()
    config.read(config_file)
    url = urlparse(config.get('database', 'uri'))
    return url

def run(*args):
    print 'RUNING:', ' '.join(args)
    summary = set()
    process = subprocess.Popen(args, stdout=subprocess.PIPE,
        stderr=subprocess.PIPE, bufsize=1, shell=False)

    out, err = process.communicate()

    in_traceback = False
    for line in chain(out.split('\n'), err.split('\n')):
        line = line.strip()
        if 'ERROR' in line:
            s = line[line.index('ERROR'):]
            summary.add(t.yellow(s))
            line = t.red(line)
        elif 'WARNING' in line:
            s = line[line.index('WARNING'):]
            summary.add(t.yellow(s))
            line = t.yellow(line)
        elif 'Traceback' in line or in_traceback:
            #s = line[line.index('Traceback'):]
            in_traceback = True
            #summary.add(t.red(s))
            line = t.red(line)
            if line.startswith('Exception'):
                in_traceback = False
        print line

    process.stdout.close()
    process.stderr.close()
    process.wait()

    if summary:
        print t.bold('\nWARNING AND ERROR SUMMARY:')
        print '\n'.join(sorted(list(summary)))
    return process.returncode

def execute(query, *args, **kwargs):
    if not args:
        args = kwargs
    cursor.execute(query, args)

def run_trytond(to_install=None):
    to_run = ['trytond/bin/trytond-admin', '-v', '-c', config_file]
    if to_install:
        to_run += ['-u']
        to_run += to_install
    to_run.append('--all')
    to_run.append('-d')
    to_run.append(database_name)
    returncode = run(*to_run)
    if returncode:
        print t.red('Trytond update failed. Upgrade aborted.')
        sys.exit(1)
    return returncode

def table_exists(table):
    execute('SELECT count(*) FROM information_schema.tables '
        'WHERE table_name=%s', table)
    return bool(cursor.fetchone()[0])

def field_exists(field):
    table, field = field.split('.')
    execute('SELECT count(*) FROM information_schema.columns '
        'WHERE table_name=%s AND column_name=%s', table, field)
    return bool(cursor.fetchone()[0])

def uninstall_modules():
    to_uninstall = config.get('to_uninstall')
    module_table = None
    for table in ('ir_module_module', 'ir_module'):
        if table_exists(table):
            module_table = table
            break
    for module in to_uninstall:
        print 'Module:', module
        execute('DELETE FROM ' + module_table + '_dependency WHERE '
            'module IN (SELECT id FROM ' + module_table + ' WHERE name=%s)',
            module)
        execute('DELETE FROM ' + module_table + ' WHERE name=%s', module)

        execute('SELECT model, db_id FROM ir_model_data WHERE module=%s',
            module)
        for model, db_id in cursor.fetchall():
            print 'DELETING', model, db_id
            execute('DELETE FROM "' + model.replace('.', '_')
                + '" WHERE id=%s', db_id)

        execute('DELETE FROM ir_model_data WHERE module=%s', module)

    if table_exists('babi_report'):
        execute('DELETE from babi_filter_parameter where filter in'
            ' (SELECT id FROM babi_filter WHERE model IN (SELECT '
            'id FROM ir_model WHERE module NOT IN (SELECT name FROM %s)))' %
            module_table)
        execute('DELETE FROM babi_filter WHERE model IN (SELECT '
                'id FROM ir_model WHERE module NOT IN (SELECT name FROM %s))' %
                module_table)
        execute('DELETE from babi_order where report in'
            ' (SELECT id FROM babi_report WHERE model IN (SELECT '
            'id FROM ir_model WHERE module NOT IN (SELECT name FROM %s)))' %
            module_table)
        execute('DELETE from babi_measure where report in'
            ' (SELECT id FROM babi_report WHERE model IN (SELECT '
            'id FROM ir_model WHERE module NOT IN (SELECT name FROM %s)))' %
            module_table)
        execute('DELETE from babi_dimension where expression in'
            ' (SELECT id FROM babi_expression WHERE model IN (SELECT '
            'id FROM ir_model WHERE module NOT IN (SELECT name FROM %s)))' %
            module_table)
        execute('DELETE FROM babi_expression WHERE model IN (SELECT '
            'id FROM ir_model WHERE module NOT IN (SELECT name FROM %s))' %
            module_table)
        execute('DELETE FROM babi_report WHERE model IN (SELECT '
            'id FROM ir_model WHERE module NOT IN (SELECT name FROM %s))' %
            module_table)

    if table_exists('mass_editing'):
        execute('DELETE FROM mass_editing WHERE model IN (SELECT '
            'id FROM ir_model WHERE module NOT IN (SELECT name FROM %s))' %
            module_table)

    execute('DELETE FROM ir_trigger WHERE model IN (SELECT '
        'id FROM ir_model WHERE module NOT IN (SELECT name FROM %s))' %
        module_table)
    execute('DELETE FROM ir_action_act_window WHERE res_model IN (SELECT '
        'model FROM ir_model WHERE module NOT IN (SELECT name FROM %s))' %
        module_table)
    execute('DELETE FROM ir_action_wizard WHERE model in (SELECT model FROM '
        'ir_model WHERE module NOT IN (SELECT name FROM %s))' % module_table)
    execute('DELETE FROM ir_model WHERE module NOT IN (SELECT name FROM '
        '%s)' % module_table)
    execute('DELETE FROM ir_model_field WHERE module NOT IN (SELECT name FROM '
        '%s)' % module_table)
    execute('DELETE FROM ir_ui_view WHERE module NOT IN (SELECT name FROM '
        '%s)' % module_table)

def process_actions(actions):
    if not actions:
        return
    for action in actions:
        if isinstance(action, dict):
            tables = action.get('tables', '')
            fields = action.get('fields', '')
            version = action.get('version', '')
            query = action.get('query')
            script = action.get('script')

            # Check version
            if version <= trytond_version and version > from_version:
                continue

            # Check tables
            found = True
            tables = tables.split()
            for table in tables:
                if not table_exists(table):
                    print "TABLE '%s' NOT FOUND" % table
                    found = False
                    break
            if not found:
                continue

            # Check fields
            found = True
            fields = fields.split()
            for field in fields:
                if not field_exists(field):
                    print "FIELD '%s' NOT FOUND" % field
                    found = False
                    break
            if not found:
                continue
        else:
            query = action
            script = None

        if query:
            query = query.replace('%', '%%')
            print query
            execute(query)

        if script:
            if os.path.isfile(script):
                run(script, database_name, args.config)
            else:
                print t.red("Not found script: %s" % script)


parser = argparse.ArgumentParser(description='Upgrade a Tryton database to the '
    'version of the trytond library available.')
parser.add_argument('database', nargs=1, help='PostgreSQL database to upgrade')
parser.add_argument('from_version', nargs=1, help='Tryton version of the '
    'database to be migrated')
parser.add_argument('-c', '--config', default=None,
        help='path to the trytond configuration file')

args = parser.parse_args()

database_name, = args.database
from_version, = args.from_version
if not args.config:
    instance = os.path.basename(os.path.realpath(os.path.join(
                os.path.dirname(os.path.realpath(__file__)), '..')))
    paths = (
        '/etc/trytond/%s.conf' % instance,
        os.environ.get('TRYTOND_CONFIG'),
        )
    for config_file in paths:
        print 'Checking %s...' % config_file
        if os.path.exists(config_file):
            break
    print "Configuration file: %s" % config_file
else:
    config_file = args.config
    if not os.path.isfile(config_file):
        print 'Not found file %s' % config_file
        sys.exit(1)

url = get_url()

config = yaml.load(open('upgrades/config.yml', 'r').read())
config.setdefault('to_uninstall', [])
config.setdefault('to_install', [])
if os.path.exists('upgrade.yml'):
    override = yaml.load(open('upgrade.yml', 'r').read())
    config['to_install'] += override.get('to_install', [])
    config['before'] += override.get('before', [])
    config['after'] += override.get('after', [])

if url.username:
    connection = psycopg2.connect(dbname=database_name, host=url.hostname,
        port=url.port, user=url.username, password=url.password)
else:
    connection = psycopg2.connect(dbname=database_name)

cursor = connection.cursor()

print t.green('\nExecuting actions before update...')
process_actions(config.get('before'))

print t.green('\nUninstalling modules...')
uninstall_modules()

connection.commit()

print t.green('\nUpdating trytond...')
run_trytond(config.get('to_install'))

print t.green('\nExecuting actions after update...')
process_actions(config.get('after'))
connection.commit()

print t.green('\nUpdating trytond again...')
run_trytond()