Skip to content

Instantly share code, notes, and snippets.

@FxIII
Forked from ikwattro/schema-generator,py
Last active March 7, 2022 09:20
Show Gist options
  • Save FxIII/022cdcfffa757f13aefa8414af95c893 to your computer and use it in GitHub Desktop.
Save FxIII/022cdcfffa757f13aefa8414af95c893 to your computer and use it in GitHub Desktop.
Hume Schema Generator
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment