Hume Schema Generator

schema-generator.py

import os

from neo4j import GraphDatabase
import json
import random
import logging
import uuid

# default parameters
URI = "bolt://localhost:7687"
AUTH = ("neo4j", "neo4j")
DATABASE = "neo4j"


class Colors:
    _cur = 0
    colors = ['ef5350', 'd50000', 'f44336', 'e53935', 'c62828', 'ec407a', 'c51162', 'e91e63', 'd81b60', 'ad1457',
              '7e57c2', '6200ea', '673ab7', '5e35b1', '4527a0', '5c6bc0', '304ffe', '3f51b5', '3949ab', '283593',
              '42a5f5', '2962ff', '2196f3', '1e88e5', '1565c0', '26c6da', '00b8d4', '00bcd4', '00acc1', '00838f',
              '26a69a', '00bfa5', '009688', '00897b', '00695c', '66bb6a', '00c853', '4caf50', '43a047', '2e7d32',
              '9ccc65', '64dd17', '8bc34a', '7cb342', '558b2f', 'ffca28', 'ffab00', 'ffc107', 'ffb300', 'ff8f00',
              'ff7043', 'dd2c00', 'ff5722', 'f4511e', 'd84315']
    random.shuffle(colors)

    @classmethod
    def next(cls):
        c = cls.colors[cls._cur]
        cls._cur = (cls._cur + 1) % len(cls.colors)
        return "#" + c


class Converter(object):
    """Convert database schema to Hume Schema import file format """

    # cypher to Hume data types mapping
    TYPES_MAPPING = {'String': 'STRING', 'Long': 'NUMBER', 'Double': "DOUBLE", "StringArray": "STRING",
                     'Date': 'DATE', 'Point': 'STRING', 'Boolean': 'BOOLEAN',
                     'FloatArray': 'NUMBER', 'DoubleArray': 'NUMBER'}

    def __init__(self, uri, auth, database="neo4j"):
        logging.basicConfig(level=os.environ.get("DEUG_LEVEL", "INFO"))
        self.logger = logging.getLogger()

        # connect to database
        self.logger.info(f"connecting to {uri}/{database}")
        self.driver = GraphDatabase.driver(uri, auth=auth, database=database)
        self.session = self.driver.session()

        # class id to internal uuid mapping
        self._class_uuid = {}

        # fetch classes and relationships
        classes, relationships = self.populate_schema()
        self.logger.info(f"got {len(classes)} classes and {len(relationships)} relationships")

        # add attriute to classes
        self.collect_attributes(classes)

        self.schema = {'classes': list(classes.values()),
                       'relationships': relationships}

    def populate_schema(self):
        """use db.schema.visualization() to grab information about classes and relationship amongst classess
            :returns found classes, relationships
        """
        query = "CALL db.schema.visualization()"
        schema = list(self.session.run(query))[0]
        classes = {node["name"]: self._make_class(node) for node in schema['nodes']}
        relationships = [self._make_rel(rel) for rel in schema['relationships']
                         if not (rel.type == "HAS_ASSOCIATION" and
                                 "Association" not in [list(n.labels)[0] for n in rel.nodes])]
        return classes, relationships

    def collect_attributes(self, classes):
        """use db.schema.nodeTypeProperties() to attach attribute information to the classes
        """
        query = "CALL db.schema.nodeTypeProperties()"
        node_type_properties = list(self.session.run(query))

        attributes_map = {}

        for node in node_type_properties:
            self.logger.info(f"parsing {node}")
            if node['propertyTypes'] is None:
                self.logger.warning(f"Skipping node {node}")
                continue
            labels = node['nodeLabels']
            property_name = node['propertyName']

            # apply the type to all labels
            if len(labels) > 1:
                self.logger.warning(f"apply {property_name} to labels {labels}")

            for label in labels:
                humeType = self._to_hume_type(node)
                attributes_map.setdefault(label, []).append({'label': property_name,
                                                             'type': humeType})

        for (key, attributes) in attributes_map.items():
            classes[key]['attributes'] = list(attributes)

    def _make_class(self, node):
        """create a class object out of the node. It updates the class.id to uuid relation"""
        class_uuid = str(uuid.uuid4())
        self._class_uuid[node.id] = class_uuid
        return {'label': node['name'],
                'canvasPosition': self._random_canvas_position(),
                'icon': 'mdi-circle-outline',
                'color': Colors.next(),  # '#aaa',
                'uuid': class_uuid}

    def _make_rel(self, relation):
        """create a relation object out of the relation entry."""
        node_from = relation.nodes[0]
        node_to = relation.nodes[1]
        return {'uuid': str(uuid.uuid4()),
                'start': self._class_uuid[node_from.id],
                'startLabel': node_from['name'],
                'endLabel': node_to['name'],
                'endId': self._class_uuid[node_to.id],
                'label': relation.type}

    def _to_hume_type(self, node):
        """extract the HUME type for the node"""
        nodeType = node['nodeType']
        property_types = node['propertyTypes']
        property_name = node['propertyName']
        property_type = property_types[0]
        if len(property_types) > 1:
            self.logger.warning(f"multiple property types for {nodeType} ({property_types})")
        if property_type not in self.TYPES_MAPPING:
            self.logger.warning(
                f"unsupported type {property_type} for {nodeType}.{property_name} falling back to STRING")
        return self.TYPES_MAPPING.get(property_type, "STRING")

    @staticmethod
    def _random_canvas_position():
        """return a random canvas position"""
        return {'x': random.randint(100, 1200), 'y': random.randint(50, 800)}


if __name__ == "__main__":
    converter = Converter(URI, AUTH, DATABASE)
    with open('schema-generated.json', 'w') as fp:
        json.dump(converter.schema, fp)