davebshow · October 16, 2017 19:00 · Oct 16, 2017
diff --git a/janus_schema.py b/janus_schema.py
@@ -0,0 +1,94 @@
+import collections
+import datetime
+import logging
+
+from goblin import properties
+
+logger = logging.getLogger(__name__)
+
+
+CARD_MAPPING = {'Cardinality.single': 'Cardinality.SINGLE',
+                'Cardinality.list_':'Cardinality.LIST',
+                'Cardinality.set_': 'Cardinality.SET'}
+
+
+DATA_TYPE_MAPPING = {properties.Integer: 'Integer.class',
+                     properties.Float: 'Float.class',
+                     properties.String: 'String.class',
+                     properties.Boolean: 'Boolean.class'}
+
+
+prop_keys = {}
+
+
+PropertyKey = collections.namedtuple('PropertyKey', ['name', 'data_type', 'card'])
+
+
+async def create_schema(app, indices, cluster):
+    client = await cluster.connect()
+    schema_definition = get_schema(app, indices)
+    start_time = datetime.datetime.now()
+    logger.info("Processing schema....")
+    resp = await client.submit(schema_definition)
+    await resp.all()
+    logger.info("Processed schema in {}".format(datetime.datetime.now() - start_time))
+
+
+def get_schema(app, indices=None):
+    if not indices:
+        indices = []
+    schema_definition = """graph.tx().rollback()
+                           mgmt = graph.openManagement()\n"""
+    for label, vertex in app.vertices.items():
+        schema_definition += get_vertex_schema(label, vertex)
+    schema_definition += "// Edge schema\n"
+    for label, edge in app.edges.items():
+        schema_definition += get_edge_schema(label, edge)
+    # Need to register vertex props with app TODO Fix in Goblin
+
+    schema_definition += get_indices_schema(indices)
+    schema_definition += "mgmt.commit()"
+    return schema_definition
+
+
+def get_vertex_schema(label, vertex):
+    vertex_schema = "// Schema for vertex label: {}\n".format(label)
+    vertex_schema += "{} = mgmt.makeVertexLabel('{}').make()\n".format(label, label)
+    mapping = vertex.__mapping__
+    properties = vertex.__properties__
+    for db_name, (ogm_name, _) in mapping.db_properties.items():
+        prop = properties[ogm_name]
+
+        # Get cardinality
+        if hasattr(prop, 'cardinality'):
+            card = str(prop.cardinality)
+        else:
+            card = 'Cardinality.single'
+        mapped_card = CARD_MAPPING[card]
+
+        # Get data type
+        data_type = prop.data_type
+        mapped_data_type = DATA_TYPE_MAPPING[data_type.__class__]
+        prop_key = PropertyKey(db_name, mapped_data_type, mapped_card)
+        if db_name in prop_keys:
+            assert prop_key == prop_keys[db_name]
+        else:
+            prop_keys[db_name] = prop_key
+        prop_key_string = "{} = mgmt.makePropertyKey('{}').dataType({}).cardinality({}).make()\n".format(
+            prop_key.name, prop_key.name, prop_key.data_type, prop_key.card)
+        vertex_schema += prop_key_string
+    vertex_schema += "\n"
+    return vertex_schema
+
+
+def get_indices_schema(indices):
+    indices_schema = "// Indices ...\n"
+    for index in indices:
+        indices_schema += "mgmt.buildIndex('by_{}', Vertex.class).addKey({}).buildCompositeIndex()\n".format(index, index)
+    return indices_schema
+
+
+def get_edge_schema(label, edge):
+    edge_schema = "{} = mgmt.makeEdgeLabel('{}').multiplicity(SIMPLE).make()\n".format(label, label)
+    #TODO edge prop keys
+    return edge_schema