generate_db_schema_documentation.py

#!/usr/bin/env python3

import io
import md_toc
from os import listdir
import os
import re
import subprocess
import sys

PATH_TO_TEMP_FOLDER = "/tmp/tbls/doc/schema/"
PATH_TO_GITHUB_WORKSPACE = "/github/workspace/"
DATABASE_TABLE_PREFIX = "public"
FILE_EXTENSION = ".md"
PATH_TO_DB_SCHEMA_FILE = os.environ["INPUT_PATH_TO_DB_SCHEMA_FILE"]
PATH_TO_GENERATED_DB_SCHEMA_FILE = os.environ["INPUT_PATH_TO_GENERATED_DB_SCHEMA_FILE"]
INPUT_DATABASE_HOST = os.environ["INPUT_DATABASE_HOST"]
INPUT_DATABASE_NAME = os.environ["INPUT_DATABASE_NAME"]
INPUT_DATABASE_PASSWORD = os.environ["INPUT_DATABASE_PASSWORD"]
INPUT_DATABASE_PORT = os.environ["INPUT_DATABASE_PORT"]
INPUT_DATABASE_USER_NAME = os.environ["INPUT_DATABASE_USER_NAME"]
TBLS_YML = f"""
dsn: postgres://{INPUT_DATABASE_USER_NAME}:{INPUT_DATABASE_PASSWORD}@{INPUT_DATABASE_HOST}:{INPUT_DATABASE_PORT}/{INPUT_DATABASE_NAME}?sslmode=disable
er:
  skip: true
exclude:
  - public.flyway_schema_history

docPath: /tmp/tbls/doc/schema
"""

def write_tbls_yml():
    with open(PATH_TO_GITHUB_WORKSPACE + "/tbls.yml", 'w') as fout:
        fout.write(TBLS_YML)

def delete_tbls_yml():
    os.remove(PATH_TO_GITHUB_WORKSPACE + "/tbls.yml")

def run_proc(cmd):
    return subprocess.run(cmd, shell=True, check=True, text=True).stdout

def list_of_md_files(directory, extension):
    return (f for f in listdir(directory) if f.endswith(extension) and not f.startswith("README"))

def generate_table_of_contents():
    toc = md_toc.build_toc(PATH_TO_GITHUB_WORKSPACE + PATH_TO_GENERATED_DB_SCHEMA_FILE, keep_header_levels=6, parser="github")
    with open(PATH_TO_GITHUB_WORKSPACE + PATH_TO_GENERATED_DB_SCHEMA_FILE,"r") as fin:
        filedata = fin.read()
    filedata = filedata.replace("<!-- TOC -->", toc)
    with open(PATH_TO_GITHUB_WORKSPACE + PATH_TO_GENERATED_DB_SCHEMA_FILE,"w") as fout:
        fout.write(filedata)
    
def list_tables():
    list_of_all_tables = []
    table_files = list_of_md_files(PATH_TO_TEMP_FOLDER,FILE_EXTENSION)
    for file_name in table_files:
        table_name = file_name.split(".md")[0]
        if DATABASE_TABLE_PREFIX in table_name:
            list_of_all_tables.append(table_name.split(f"{DATABASE_TABLE_PREFIX}.")[1])
    return list_of_all_tables

def filter_table(table):
    """Remove link to other table under children and parents columns which is generated by default """
    remove_left_bracket = re.sub(r"\["," ",table)
    remove_right_bracket = re.sub(r"\]"," ",remove_left_bracket)
    remove_link_name = re.sub(r"\(" + DATABASE_TABLE_PREFIX + r"\..*\.md\)", " ",remove_right_bracket)
    return remove_link_name.replace("##","###")

def get_table(filename):
    """Extract table, constraints and indexes from the generated file and returns the table as string"""
    file_path = PATH_TO_TEMP_FOLDER + f"{DATABASE_TABLE_PREFIX}." + filename + FILE_EXTENSION
    with open(file_path,"r") as table_file:
        table_file_read = table_file.read()
        table_file_read_split_columns = table_file_read.split("## Columns")
        table_file_read_tables = table_file_read_split_columns[1]
        tables = table_file_read_tables.split("---\n")
        table = tables[0]
        return filter_table(table)

def get_replace_line(key):
    return get_table(key)

def main():
    try: 
        write_tbls_yml()
        run_proc(f"tbls doc")
        undocumented_tables = list_tables()
        with open(PATH_TO_GITHUB_WORKSPACE + PATH_TO_DB_SCHEMA_FILE,"r") as db_schema_file:
            db_schema_contents = db_schema_file.readlines()
        with open (PATH_TO_GITHUB_WORKSPACE + PATH_TO_GENERATED_DB_SCHEMA_FILE,"w") as generated_db_schema:
            generated_db_schema.write("This file is auto-generated. Please do not modify this file.\n<!-- TOC -->\n")
            try:
                for line in db_schema_contents:
                    if(line.startswith("<<<-") and line.endswith("->>>\n")):
                        start = line.index("-")
                        end = line.index("-",start+1)
                        replace_line = get_replace_line(line[start+1:end])
                        if line[start+1:end] in undocumented_tables:
                            undocumented_tables.remove(line[start+1:end])
                        generated_db_schema.writelines(replace_line)
                    else:
                        generated_db_schema.writelines(line)
                if(len(undocumented_tables) != 0):
                    undocumented_tables.sort()
                    generated_db_schema.write("\n## Undocumented Tables\nThe following tables have not been documented. Please document them if needed.\n")
                    for table in undocumented_tables:
                        generated_db_schema.writelines("- ")
                        generated_db_schema.writelines(table)
                        generated_db_schema.writelines("\n")

            except Exception as e:
                print("Database schema generation failed due to error:",e)
                sys.exit(1)
        generate_table_of_contents()
    finally:
        delete_tbls_yml()

if __name__ == "__main__":
    main()