embed_schemas.py

#!/usr/bin/env python3
 
import io
import sys
import os
 
heading_text = '''
// This is file was autogenerated, do not modify it manually
#include <string>
#include <vector>
#include <map>
'''
 
 
def escape_string(sql):
    return sql.translate(str.maketrans({'"':  r'\"', '\n':  '\\n'}))
 
 
def apend_migration(migration_path, header_file):
    migration_file = open(migration_path, 'r')
    header_file.write("\"")
    migration_content = migration_file.read()
    header_file.write(escape_string(migration_content))
 
 
if __name__ == '__main__':
    if len(sys.argv) != 4:
        print("\nIncorrect arguments")
        print("Usage:\n {} {} {} {}\n".format(sys.argv[0], "folder_with_schemas",
                                              "generated_file_path", "name_prefix"))
        sys.exit(-1)
 
    sql_folder = sys.argv[1]
    schemas_header = sys.argv[2]
    prefix = sys.argv[3]
    migration_folder = os.path.join(sql_folder, 'migration')
    rollback_folder = os.path.join(sql_folder, 'rollback')
    migration_list = sorted(os.listdir(migration_folder))
    rollback_migrations_list = sorted(os.listdir(rollback_folder))
 
    with io.BytesIO() as outf:
        out_text = io.TextIOWrapper(outf, write_through=True)
        out_text.write(heading_text)
        out_text.write("extern const std::vector<std::string> {}_schema_migrations = {{".format(prefix))
        for migration in migration_list[:-1]:
            apend_migration(os.path.join(migration_folder, migration), out_text)
            out_text.write("\",\n")
        apend_migration(os.path.join(migration_folder, migration_list[-1]), out_text)
        out_text.write("\"\n};\n")
 
        out_text.write("extern const std::vector<std::string> {}_schema_rollback_migrations = {{".format(prefix))
        if len(rollback_migrations_list) > 0:
            ver = int(rollback_migrations_list[0].split(".")[1])
            for i in range(ver):
                out_text.write("\"\",\n")
            for migration in rollback_migrations_list[:-1]:
                apend_migration(os.path.join(rollback_folder, migration), out_text)
                out_text.write("\",\n")
            version = int(rollback_migrations_list[-1].split(".")[1])
            apend_migration(os.path.join(rollback_folder, rollback_migrations_list[-1]), out_text)
            out_text.write("\"")
        out_text.write("\n};\n")
 
        current_schema = open(os.path.join(sql_folder, "schema.sql"), 'r').read()
        current_schema_escaped = escape_string(current_schema)
        out_text.write('extern const std::string %s_current_schema = "%s";' % (prefix, current_schema_escaped));
        out_text.write('extern const int %s_current_schema_version = %u;' % (prefix ,(len(migration_list)-1)));
 
        # compare with current and only overwrite if it was changed
        new_content = outf.getvalue()
        try:
            with open(schemas_header, 'r') as hf:
                current_content = hf.read().encode()
        except Exception:
            current_content = b''
        if new_content != current_content:
            with open(schemas_header, 'wb') as hf:
                hf.write(new_content)