docker-postgres-backup/migrate.py

#! /usr/bin/env python3

import argparse
import sys
import os
import os.path
import re
import distutils.version
import subprocess
import datetime


def file_runner_pgsql(folder, file):
    subprocess.call(['psql', '-f', os.path.join(folder, file)])

def file_runner_pgsql_gz(folder, file):
    p1 = subprocess.Popen(['gunzip', '-c', os.path.join(folder, file)], stdout=subprocess.PIPE)
    p2 = subprocess.Popen(["psql"], stdin=p1.stdout)
    p1.stdout.close()
    p2.communicate()

def file_runner_exec(folder, file):
    subprocess.call([os.path.join(folder, file)])


def list_files(folder):
    files = [f for f in os.listdir(folder) if os.path.isfile(os.path.join(folder, f))]
    files.sort()
    return files

def sort_versions(versions):
    versions.sort(key=distutils.version.StrictVersion)

def list_versions(files):
    versions = []
    for file in files:
        version = file.split('_')[0]
        if re.match('^[0-9]+\.[0-9]+\.[0-9]+$', version):
            if version not in versions:
                versions.append(version)
    sort_versions(versions)
    return versions

def run_file(folder, file, no_run, file_runners):
    can_run = False
    for file_runner_ext in file_runners:
        if file.endswith(file_runner_ext):
            print('Running file %s' % (file))
            can_run = True
            if not no_run:
                try:
                    sys.stdout.flush()
                    file_runners[file_runner_ext](folder, file)
                except Exception as e:
                    print('Failed to run file: %s' % (e))
    if not can_run:
        print('Ignoring file %s' % (file))

def run_migration(folder, files, version, no_run, file_runners):
    for file in files:
        if file.startswith('%s_' % (version)):
            run_file(folder, file, no_run, file_runners)

def run_migrations(folder, version_from, version_to, no_run, file_runners, log_file):
    files = list_files(folder)
    versions = list_versions(files)

    if not versions:
        print('No migration available. Exiting.')
        return 0

    if version_to is None:
        version_to = versions[-1]
        print('Defaulting VERSION_TO to %s' % (version_to))
    if version_to not in versions:
        print('Could not find VERSION_TO %s' % (version_to))
        return 1
    if version_from is not None and version_from not in versions:
        print('Could not find VERSION_FROM %s' % (version_from))
        return 1

    if no_run:
        print('Not running because --no-run was specified')

    if version_from is None:
        msg = 'Initialising from %s (inclusive) to %s (inclusive)' % (versions[0], version_to)
    else:
        msg = 'Migrating from %s (exclusive) to %s (inclusive)' % (version_from, version_to)

    print(msg)

    if not no_run:
        try:
            with open(log_file, "a") as f:
                f.write('%s: %s in %s\n' % (datetime.datetime.now(), msg, folder))
        except Exception as e:
            print('Failed to write logs in %: %s' % (log_file, e))

    version_from_obj = distutils.version.StrictVersion('0.0.0' if version_from is None else version_from)
    version_to_obj = distutils.version.StrictVersion(version_to)

    last_version = version_from
    for version in versions:
        version_obj = distutils.version.StrictVersion(version)
        if version_from_obj < version_obj and version_obj <= version_to_obj:
            if last_version is None:
                print('Initialising %s' % (version))
            else:
                print('Migrating from %s to %s' % (last_version, version))
            run_migration(folder, files, version, no_run, file_runners)
            last_version = version

    return 0

def main():
    parser = argparse.ArgumentParser(description='Migrate container configuration/data')
    parser.add_argument('--folder', dest='folder', default='/docker-entrypoint-initdb.d/', help='Folder to use for migration')
    parser.add_argument('--init', dest='is_init', default=False, action='store_true', help='Run all migrations')
    parser.add_argument('--migrate', dest='is_migrate', default=False, action='store_true', help='Run all migrations between VERSION_FROM (exclusive) and VERSION_TO (inclusive)')
    parser.add_argument('--no-run', dest='no_run', default=False, action='store_true', help='Do not run migration, just print')
    parser.add_argument('--version-from', dest='version_from', default=None, help='Current version, required if using --migrate, can not be used --init')
    parser.add_argument('--version-to', dest='version_to', default=None, help='Final version, default to last available version is not given')
    parser.add_argument('--log-file', dest='log_file', default='/var/lib/postgresql/data/migrate.py.log', help='File to log migrations')

    args = parser.parse_args()

    file_runners = {
            'sql': file_runner_pgsql,
            'sql.gz': file_runner_pgsql_gz,
            'sh': file_runner_exec
            }

    if args.is_init and args.is_migrate:
        print('--init and --migrate can not be used together')
        return 64
    if args.is_init:
        if args.version_from is not None:
            print('--version-from can not be used with --init')
            return 64
        return run_migrations(args.folder, None, args.version_to, args.no_run, file_runners, args.log_file)
    elif args.is_migrate:
        if args.version_from is None:
            print('--version-from is required. Use --init to run all migrations')
            return 64
        return run_migrations(args.folder, args.version_from, args.version_to, args.no_run, file_runners, args.log_file)
    else:
        print('Missing --init or --migrate')
        return 64

if __name__ == '__main__':
    sys.exit(main())