tryton-upgrades/upgrade

#!/usr/bin/env python
import os
import sys
import yaml
import argparse
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(args.config)
    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)
    #for line in iter(process.stdout.readline, b''):
    for line in iter(process.stderr.readline, b''):
        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)
        print line
    process.stdout.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_run = ['trytond/bin/trytond-admin', '-v', '-c', config_file]
    to_install = config.get('to_install')
    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
            execute('DELETE FROM "' + model.replace('.', '_')
                + '" WHERE id=%s', db_id)
        execute('DELETE FROM ir_model_data WHERE module=%s', module)

    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:
            run('upgrades/after/' + script, database_name, args.config)


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='trytond.conf',
        help='path to the trytond configuration file')
parser.add_argument('-p', '--purge', action='store_true', dest='purge',
    help='remove all fields that drop information')

args = parser.parse_args()

database_name, = args.database
from_version, = args.from_version
config_file = args.config

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', [])


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()

if args.purge:
    process_actions(config.get('purge'))
    connection.commit()
    sys.exit(0)

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()

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

connection.commit()