colorize a DBeaver ERD file

colorize-dbeaver-erd-diagram.py

#!/usr/bin/env python
import xml.etree.ElementTree as ET
import sys
import os
import re

# python colorize-diagram.py ./samplehub_gitlab.erd

# Check if the input file is provided
if len(sys.argv) != 2:
    print("Usage: python colorize-diagram.py <input_file.erd>")
    sys.exit(1)

input_file = sys.argv[1]

# Generate the output file name
base, ext = os.path.splitext(input_file)
output_file = f"{base}-colored{ext}"

# Load the XML file
tree = ET.parse(input_file)
root = tree.getroot()

prefix_colors = {
    "project" : "153,193,241",  # blue
    "archive" : "246,245,244",  # grey
    "sample"  : "143,240,164",  # green
    "core"    : "143,240,164",  # green
    "analysis": "255,230,245",  # pink
    "geology" : "220,138,221",  # violet
    "storage" : "255,165,0",    # orange
    "curation": "255,165,0",    # orange
    "contact" : "249,240,107"   # yellow
}

prefix_regex = re.compile(r'^(.*?)_')

# Iterate over the entities and set the color-bg attribute
for entity in root.findall('.//entity'):
    entity_name = entity.get('name')
    match = prefix_regex.match(entity_name)
    if match:
        prefix = match.group(1)
        if prefix in prefix_colors:
            entity.set('color-bg', prefix_colors[prefix])

tree.write(output_file)

print(f"Colored diagram saved to {output_file}")

inject-xml-fragment.py

#!/usr/bin/env python3
# Write a <notes> fragment to be injected between the <entities> and <relations> tags
import xml.etree.ElementTree as ET
import sys
import os
import datetime

# Check if the input file is provided
if len(sys.argv) < 2 or len(sys.argv) > 4:
    print("Usage: python inject-fragment.py <input_file.erd> <pipeline-id> [annotation]")
    sys.exit(1)


input_file = sys.argv[1]
pipeline_id = sys.argv[2] if len(sys.argv) >= 3 else "(pipeline-id)"

annotation = sys.argv[3] if len(sys.argv) == 4 else datetime.datetime.now().strftime("%b %Y")

# Generate the output file name
base, ext = os.path.splitext(input_file)
output_file = f"{base}-with-notes{ext}"

# Load the XML file



def inject_fragment(input_file, fragment):
    tree = ET.parse(input_file)
    root = tree.getroot()
    entities = root.find('entities')
    relations = root.find('relations')
    
    if entities is not None and relations is not None:
        
        fragment_element = ET.fromstring(fragment)
        children = list(root)
        index = children.index(relations)
        root.insert(index, fragment_element)
    
    return ET.tostring(root, encoding='unicode')

# inject this fragment between the <entities> and the <relations> tags:
xml_note = f'''
<notes>
<note id="123" order="122" color-bg="255,255,255" color-fg="0,0,0" border-width="0" font="Sans:20:1" x="-2443" y="58" w="394" h="114">  mDIS-Samplehub {pipeline_id}
  {annotation}
  Simplified Datamodel</note>
</notes>
'''

result = inject_fragment(input_file, xml_note)

# Write the modified XML back to the output file
with open(output_file, 'w') as f:
    f.write(result)

print(f"Diagram with note saved to {output_file}")

make-colorized-diagram.sh

#!/usr/bin/env bash
# colorize ER--diagram, flip all items, add a note and save as PNG

# Usage: ./make-colorized-diagram.sh [CIPIPE_ID] [ANNOTATION]
# Example: ./make-colorized-diagram.sh 77777 "SampleHub ERD"

# ??? maybe obsolete:
# You should also export the CIPIPE_ID variable in your shell:
#  export CIPIPE_ID=77777
# because the CIPIPE_ID will be used for searching the file name

CIPIPE_ID=${1:-77777}
ANNOTATION=${2:-$(date +"%b %Y")}
# check if the file exists
if [ ! -f "wb31-mdis-samplehub-${CIPIPE_ID}.erd" ]; then
  echo "Error: Input file 'wb31-mdis-samplehub-${CIPIPE_ID}.erd' does not exist."
  exit 1
fi

./scripts/colorize-dbeaver-erd-diagram.py wb31-mdis-samplehub-${CIPIPE_ID}.erd

./scripts/translate-dbeaver-erd-diagram.py wb31-mdis-samplehub-${CIPIPE_ID}-colored.erd 

./scripts/inject-xml-fragment.py wb31-mdis-samplehub-${CIPIPE_ID}-colored-translated.erd ${CIPIPE_ID} "$ANNOTATION"

translate-items-DBeaver-ERD-file.py

#!/usr/bin/env python
import xml.etree.ElementTree as ET
import sys
import os

# Check if the input file is provided
if len(sys.argv) != 2:
    print("Usage: python translate-diagram.py <input_file.erd>")
    sys.exit(1)

input_file = sys.argv[1]

# Generate the output file name
base, ext = os.path.splitext(input_file)
output_file = f"{base}-translated{ext}"

# Load the XML file
tree = ET.parse(input_file)
root = tree.getroot()

# Functions to translate the x and y attribute (origin is top-left corner)
def translate_y(element, offset):
    y = element.get('y')
    if y is not None:
        new_y = int(y) + offset
        element.set('y', str(new_y))

def translate_x(element, offset):
    x = element.get('x')
    if x is not None:
        new_x = int(x) + offset
        element.set('x', str(new_x))


def find_rightmost_x(root):
    max_x = float('-inf')
    for item in root.findall(".//bend") + root.findall(".//entity") + root.findall(".//note"):
        # if none, set to 0
        x = int(item.get('x'))
        if x is not None:
            if x > max_x:
                max_x = x
    return max_x

def find_lowest_y(root):
    lowest_y = 0
    for item in root.findall(".//bend") + root.findall(".//entity") + root.findall(".//note"):
        y = item.get('y')
        if y is not None:
            y = int(y)
            if y > lowest_y:
                lowest_y = y
    return lowest_y

def flip_horizontally(root, offset=100):
    rightmost_x = find_rightmost_x(root)
    mirror_axis = rightmost_x + offset

    for item in root.findall(".//bend") + root.findall(".//entity") + root.findall(".//note"):
        x = item.get('x')
        if x is not None:
            x = int(x)
            mirrored_x = 2 * mirror_axis - x
            item.set('x', str(mirrored_x))

def flip_vertically(root, offset=100):
    lowest_y = find_lowest_y(root)
    mirror_axis = lowest_y + offset

    for item in root.findall(".//bend") + root.findall(".//entity") + root.findall(".//note"):
        y = item.get('y')
        if y is not None:
            y = int(y)
            mirrored_y = 2 * mirror_axis - y
            item.set('y', str(mirrored_y))

tree = ET.parse(input_file)
root = tree.getroot()

# flip horizontally, and move back to the left near the origin
flip_horizontally(root)
for element in root.findall('.//entity') + root.findall('.//note') + root.findall('.//bend'):
    translate_x(element, -5000)

# flip vertically, and move back to the top near the origin
flip_vertically(root)
for element in root.findall('.//entity') + root.findall('.//note') + root.findall('.//bend'):
    translate_y(element, -2000)


# Save the modified XML file
tree.write(output_file)

print(f"Translated diagram saved to {output_file}")