Clear dirty flag in each migration

clear-dirty-flag.py

#!/usr/bin/env python3

import os
from colored import fg, bg, stylize, attr


def wrap_in_transaction(path, content):
    lines = content.splitlines()

    if 'COMMIT;' in lines:
        return {'result': 'unmodified'}

    if all("CREATE INDEX CONCURRENTLY" in line for line in lines):
        print(
            stylize(
                f"Skipping {path} (all it does is CREATE INDEX CONCURRENTLY)",
                attr('dim')))
        return {'result': 'unmodified'}

    if all("Nothing" in line or "empty" in line for line in lines):
        print(stylize(f"Skipping {path} (it's empty)", attr('dim')))
        return {'result': 'unmodified'}

    if any("ALTER SYSTEM" in line for line in lines):
        print(
            stylize(
                f"Skipping {path} (ALTER SYSTEM cannot run in a transaction)",
                attr('dim')))
        return {'result': 'unmodified'}

    print(f"{stylize(path, fg('yellow'))} is missing BEGIN and COMMIT:")
    prefix = stylize("  | ", fg('yellow'))
    print(prefix + stylize("BEGIN:", bg('dark_green')))
    print(prefix)
    print('\n'.join([prefix + line for line in content.splitlines()]))
    print(prefix)
    print(prefix + stylize("COMMIT:", bg('dark_green')))

    s = ""
    s += "BEGIN;\n"
    s += "\n"
    s += content
    s += "\n"
    s += "COMMIT;\n"

    return {'result': 'modified', 'content': s}


dir_to_table = {
    'frontend': 'schema_migrations',
    'codeintel': 'codeintel_schema_migrations',
    'codeinsights': 'codeinsights_schema_migrations',
}


def clear_dirty_flag(path, content):
    lines = content.splitlines()

    if "COMMIT;" not in lines:
        return {'result': 'unmodified'}

    dir = os.path.basename(os.path.dirname(path))
    update_statement = f"UPDATE {dir_to_table[dir]} SET dirty = 'f'"
    if update_statement in lines:
        return {'result': 'unmodified'}

    print(f"{stylize(path, fg('yellow'))} ")

    comment = "-- Clear the dirty flag in case the operator timed out and isn't around to clear it."

    return {
        'result':
        'modified',
        'content':
        content.replace("COMMIT;",
                        comment + "\n" + update_statement + "\nCOMMIT;")
    }


def usage():
    print("Usage: clear-dirty-flag <dry-run | write>")
    exit(1)


def transform(path, f, write):
    with open(path, 'r') as file:
        content = file.read()

    result = f(path, content)

    if result['result'] == 'modified':
        if write:
            print("  Modifying the file...", end='')
            with open(path, 'w') as file:
                file.write(result['content'])
            print('done')
        return 1
    else:
        return 0


def main():
    if len(os.sys.argv) != 2:
        usage()
    if os.sys.argv[1] == "dry-run":
        write = False
    elif os.sys.argv[1] == "write":
        write = True
    else:
        usage()

    modified_count = 0

    for root, dirs, files in os.walk('migrations'):
        for file in files:
            if file.endswith('.sql'):
                path = os.path.join(root, file)
                modified_count += transform(path, wrap_in_transaction, write)
                modified_count += transform(path, clear_dirty_flag, write)

    if write:
        print(f"{modified_count} files modified.")
    else:
        print(f"Would modify {modified_count} files.")


if __name__ == "__main__":
    main()