DocGen/docgen.py

import logging
from logging.handlers import RotatingFileHandler

class DocGen:
    def __init__(self,file):
        self.details = file
        self.logging = self.configure_logging()
        self.logging.info('Starting DocGen')

    def configure_logging(self):
        # Configure logger with rotating file handler and custom log levels
        logger = logging.getLogger(__name__)
        logger.setLevel(logging.DEBUG)

        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        formatter = logging.Formatter('[%(levelname)s] %(name)s.%(funcName)s: %(message)s')
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)

        # Rotate logs every 10 MB, keeping 5 backups
        file_handler = RotatingFileHandler('DocGen.log', maxBytes=1024*1024*10, backupCount=5)
        file_handler.setLevel(logging.DEBUG)
        formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(name)s.%(funcName)s: %(message)s')
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

        return logger

    def wiki_on_a_page(self) -> str:
        self.logging.info('Building Wiki on a page markdown file')
        master_file = ''
        self.logging.debug('Generating overview')
        overview = self.generate_overview()
        self.logging.debug('Generating md tables')
        md_tables = self.generate_md_tables()
        self.logging.debug('Generating entity relationship diagram')
        entity_relationship = self.generate_erd()
        self.logging.debug('Generating lineage report')
        lineage = self.generate_lineage()
        self.logging.debug('Generating markdown file')
        master_file += overview
        master_file += '\n## Table overview\n\n'
        master_file += md_tables
        master_file += '\n\n## Entity Relationship Diagram\n\n```mermaid\n'
        master_file += entity_relationship
        master_file += '\n```\n'
        master_file += '\n\n## Lineage Report\n\n```mermaid\n'
        master_file += lineage
        master_file += '\n```\n'
        self.logging.debug('Generating markdown file complete')
        return master_file

    def generate_erd(self) -> str:
        self.logging.info('starting erd generation')
        mermaid_erd = self._generate_erd_tables()
        self.logging.info('tables generated, generating relationships')
        relationships = self._generate_erd_relationships()
        self.logging.info('Relationships generated, combining and returning')
        mermaid_erd.extend(relationships)
        erd_code = "\n".join(mermaid_erd)
        return erd_code

    def generate_md_tables(self) -> str:
        self.logging.info('starting table generation')
        md_table = [
            "| Entity ID | Entity Name | Description | Entity Type | Primary Key(s) |",
            "|-----------|-------------|-------------|-------------|----------------|",
        ]
        for entity in self.details['entities']:
            if not entity.get('entity_columns',None):
                primary_key_string = ''
            else:
                pk_list = [column['column_name'] for column in entity['entity_columns'] if column['is_pk'] is True]
                primary_key_string = ", ".join(pk_list)
            md_table.append(
                f"| {entity['entity_id']} | {entity['entity_name']} | {entity.get('entity_description','')} | {entity['entity_type']} | {primary_key_string} |"
            )
        md_output = "\n".join(md_table)
        return md_output

    def generate_overview(self) -> str:
        self.logging.info('Generating Overview String')
        self.logging.debug(f'overview Product {self.details.get('product_name',"Missing Product Name")}')
        self.logging.debug(f'overview Team {self.details.get('responsible_team',"Missing Responsible Team")}')
        overview = f'''# {self.details.get('product_name',"Missing Product Name")}

## Maintained by: {self.details.get('responsible_team',"Missing Responsible Team")}
'''
        return overview

    def generate_lineage(self):
        mermaid_lines = [r'''%%{init: {"flowchart": {"defaultRenderer": "elk"}} }%%
flowchart LR
''']
        for entity in self.details['entities']:
            mermaid_lines.append(f"{entity['entity_id']}[{entity['entity_id']} {entity['entity_type']} {entity['entity_name']}]")

        mermaid_lines.append("\n")

        for entity in self.details['entities']:
            for dependant in entity.get('entity_dependencies',[]):
                mermaid_lines.append(f"{dependant} --> {entity['entity_id']}")

        mermaid = "\n".join(mermaid_lines)
        return mermaid

    def _generate_erd_tables(self) -> list:
        mermaid_erd = ["erDiagram"]
        self.logging.info('getting gold entity information')
        for entity in self.details['entities']:
            self.logging.debug(f'looking at {entity.get('entity_id')}, {entity.get('entity_name')}, {entity.get('entity_type')}')
            if entity.get('entity_type') != 'gold':
                continue
            self.logging.debug(f'entity {entity.get('entity_id')} is golden, getting column information')
            mermaid_erd.append(f"    {entity.get('entity_name')} {{")
            for column in entity.get('entity_columns'):
                self.logging.debug(column)
                col_name_with_pk = f'{column['column_name']} PK' if column['is_pk'] is True else column['column_name']
                mermaid_erd.append(f"        {
                    column['datatype']} {col_name_with_pk}")
            mermaid_erd.append("    }")
        return mermaid_erd
    
    def _generate_erd_relationships(self) -> set:
        relationships = set()
        for entity_a in self.details['entities']:
            self.logging.debug(f'looking at {entity_a.get('entity_id')}, {entity_a.get('entity_name')}, {entity_a.get('entity_type')}')
            if entity_a.get('entity_type') != 'gold':
                continue
            entity_a_columns = [column['column_name'] for column in entity_a['entity_columns']]
            for entity_b in self.details['entities']:
                self.logging.debug(f'looking at {entity_b.get('entity_id')}, {entity_b.get('entity_name')}, {entity_b.get('entity_type')}')
                if entity_b.get('entity_type') != 'gold':
                    continue
                if entity_a['entity_id'] == entity_b['entity_id']:
                    self.logging.debug('not handling self joins right now. 🫢')
                    continue
                entity_b_columns = [column['column_name'] for column in entity_b['entity_columns'] if column['is_pk'] is True]
                for column in entity_a_columns:
                    if column in entity_b_columns:
                        relationships.add(f'    {entity_a["entity_name"]} ||--|| {entity_b["entity_name"]} : "{column}"')
        return relationships