Last active
September 18, 2021 08:18
-
-
Save chrismwendt/80c8f68ce12bc9c9e6638bd22971d581 to your computer and use it in GitHub Desktop.
Clear dirty flag in each migration
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment