In [32]:
import sys
from pprint import pprint as pp

from arango import ArangoClient
from arango_orm.collections import Collection, Relation
from arango_orm.database import Database
from arango_orm.graph import Graph, GraphConnection
from marshmallow import Schema
from marshmallow.fields import String, Integer, Boolean


class Student(Collection):
    __collection__ = 'students'

    class _Schema(Schema):
        _key = String(required=True)  # registration number
        name = String(required=True, allow_none=False)
        age = Integer()

    def __str__(self):
        return "<Student({},{})>".format(self._key, self.name)


class Teacher(Collection):
    __collection__ = 'teachers'

    class _Schema(Schema):
        _key = String(required=True)  # employee id
        name = String(required=True)

    def __str__(self):
        return "<Teacher({})>".format(self.name)


class Subject(Collection):
    __collection__ = 'subjects'

    class _Schema(Schema):
        _key = String(required=True)  # subject code
        name = String(required=True)
        credit_hours = Integer()
        has_labs = Boolean(missing=True)

    def __str__(self):
        return "<Subject({})>".format(self.name)


class Area(Collection):
    __collection__ = 'areas'

    class _Schema(Schema):
        _key = String(required=True)  # area name


class SpecializesIn(Relation):
    __collection__ = 'specializes_in'

    class _Schema(Schema):
        expertise_level = String(required=True, options=["expert", "medium", "basic"])
#         _key = String(required=True)

    def __str__(self):
        return "<SpecializesIn(_key={}, expertise_level={}, _from={}, _to={})>".format(
            self._key, self.expertise_level, self._from, self._to)


class UniversityGraph(Graph):
    __graph__ = 'university_graph'

    graph_connections = [
        # Using general Relation class for relationship
        GraphConnection(Student, Relation("studies"), Subject),
        GraphConnection(Teacher, Relation("teaches"), Subject),
        GraphConnection(Teacher, Relation("teacher"), Student),

        # Using specific classes for vertex and edges
        GraphConnection(Teacher, SpecializesIn, Subject),
        GraphConnection([Teacher, Student], Relation("resides_in"), Area)

    ]
In [33]:
db_name = 'db_test'
username = 'root'
password = ''

#initalize Arango Client
client = ArangoClient(username=username, password=password, verify=True)
In [49]:
# delete Test Database
try:
    client.delete_database(name=db_name)
except: pass
In [50]:
# Creat database or using exist
try:
    test_db = client.create_database(name=db_name)
except:
    test_db = client.db(db_name)

# db object
db = Database(test_db)
In [51]:
# Create a graph connection object
uni_graph = UniversityGraph(connection=db)
In [52]:
# Create graph and collection
db.create_graph(uni_graph)
In [53]:
# Add Test data in base

students_data = [
    Student(_key='S1001', name='John Wayne', age=30),
    Student(_key='S1002', name='Lilly Parker', age=22),
    Student(_key='S1003', name='Cassandra Nix', age=25),
    Student(_key='S1004', name='Peter Parker', age=20)
]

teachers_data = [
    Teacher(_key='T001', name='Bruce Wayne'),
    Teacher(_key='T002', name='Barry Allen'),
    Teacher(_key='T003', name='Amanda Waller')
]

subjects_data = [
    Subject(_key='ITP101', name='Introduction to Programming', credit_hours=4, has_labs=True),
    Subject(_key='CS102', name='Computer History', credit_hours=3, has_labs=False),
    Subject(_key='CSOOP02', name='Object Oriented Programming', credit_hours=3, has_labs=True),
]

specializations_data = [
    SpecializesIn(_from="teachers/T001", _to="subjects", expertise_level="medium")
]

areas_data = [
    Area(_key="Gotham"),
    Area(_key="Metropolis"),
    Area(_key="StarCity")
]

for L in [students_data,teachers_data,subjects_data,areas_data]:
    for v in L:
        db.add(v)
In [54]:
# Get the document object added in the previous step

gotham = db.query(Area).by_key("Gotham")
metropolis = db.query(Area).by_key("Metropolis")
star_city = db.query(Area).by_key("StarCity")

john_wayne = db.query(Student).by_key("S1001")
lilly_parker = db.query(Student).by_key("S1002")
cassandra_nix = db.query(Student).by_key("S1003")
peter_parker = db.query(Student).by_key("S1004")

intro_to_prog = db.query(Subject).by_key("ITP101")
comp_history = db.query(Subject).by_key("CS102")
oop = db.query(Subject).by_key("CSOOP02")

barry_allen = db.query(Teacher).by_key("T002")
bruce_wayne = db.query(Teacher).by_key("T001")
amanda_waller = db.query(Teacher).by_key("T003")
In [62]:
# Create EDGE. of the structure in dictionary. or the same step by step in the commented out lines below
REL = {
    Relation:{
        "studies":{
            peter_parker:[
                oop,
                intro_to_prog,
            ],
            john_wayne:[
                oop,
                comp_history,
            ],
            lilly_parker:[
                intro_to_prog,
                comp_history,
            ],
            cassandra_nix:[
                oop,
                intro_to_prog,
            ],
        },
        "teacher":{
            bruce_wayne:[
                peter_parker,
                john_wayne,
                lilly_parker,
            ],
        },
        "teaches":{
            bruce_wayne:[oop],
            barry_allen:[intro_to_prog],
            amanda_waller:[comp_history],
        },
        "resides_in":{
            bruce_wayne:[gotham],
            barry_allen:[star_city],
            amanda_waller:[metropolis],
            john_wayne:[gotham],
            lilly_parker:[metropolis],
            cassandra_nix:[star_city],
            peter_parker:[metropolis],
        }
    },
    SpecializesIn:{
        "expert":{
            barry_allen:[
                oop,
                intro_to_prog,
            ],
            bruce_wayne:[
                comp_history,            
            ],
        },
        "medium":{
            bruce_wayne:[
                oop,
            ]
        },
        "basic":{
            amanda_waller:[
                intro_to_prog,
            ],
        },
    },
}

def traverse_dict(D):
    for RelationType in D:
        for RelationName in D[RelationType]:
            if RelationType == SpecializesIn:
                val = {'expertise_level':RelationName}
            else:
                val = {'collection_name':RelationName}
            for _from in D[RelationType][RelationName]:
                for _to in D[RelationType][RelationName][_from]:
                    db.add(uni_graph.relation(_from,RelationType(**val),_to))


traverse_dict(REL)
                    
# db.add(uni_graph.relation(peter_parker, Relation("studies"), oop))
# db.add(uni_graph.relation(peter_parker, Relation("studies"), intro_to_prog))
# db.add(uni_graph.relation(john_wayne, Relation("studies"), oop))
# db.add(uni_graph.relation(john_wayne, Relation("studies"), comp_history))
# db.add(uni_graph.relation(lilly_parker, Relation("studies"), intro_to_prog))
# db.add(uni_graph.relation(lilly_parker, Relation("studies"), comp_history))
# db.add(uni_graph.relation(cassandra_nix, Relation("studies"), oop))
# db.add(uni_graph.relation(cassandra_nix, Relation("studies"), intro_to_prog))

# db.add(uni_graph.relation(barry_allen, SpecializesIn(expertise_level="expert"), oop))
# db.add(uni_graph.relation(barry_allen, SpecializesIn(expertise_level="expert"), intro_to_prog))
# db.add(uni_graph.relation(bruce_wayne, SpecializesIn(expertise_level="medium"), oop))
# db.add(uni_graph.relation(bruce_wayne, SpecializesIn(expertise_level="expert"), comp_history))
# db.add(uni_graph.relation(amanda_waller, SpecializesIn(expertise_level="basic"), intro_to_prog))
# db.add(uni_graph.relation(amanda_waller, SpecializesIn(expertise_level="medium"), comp_history))

# db.add(uni_graph.relation(bruce_wayne, Relation("teacher"), peter_parker))
# db.add(uni_graph.relation(bruce_wayne, Relation("teacher"), john_wayne))
# db.add(uni_graph.relation(bruce_wayne, Relation("teacher"), lilly_parker))

# db.add(uni_graph.relation(bruce_wayne, Relation("teaches"), oop))
# db.add(uni_graph.relation(barry_allen, Relation("teaches"), intro_to_prog))
# db.add(uni_graph.relation(amanda_waller, Relation("teaches"), comp_history))

# db.add(uni_graph.relation(bruce_wayne, Relation("resides_in"), gotham))
# db.add(uni_graph.relation(barry_allen, Relation("resides_in"), star_city))
# db.add(uni_graph.relation(amanda_waller, Relation("resides_in"), metropolis))
# db.add(uni_graph.relation(john_wayne, Relation("resides_in"), gotham))
# db.add(uni_graph.relation(lilly_parker, Relation("resides_in"), metropolis))
# db.add(uni_graph.relation(cassandra_nix, Relation("resides_in"), star_city))
# db.add(uni_graph.relation(peter_parker, Relation("resides_in"), metropolis))
In [64]:
bruce = db.query(Teacher).by_key("T001")
uni_graph.expand(bruce, depth=1, direction='any')

# bruce relations
pp(bruce._relations)
{'resides_in': [<Relation(_key=5690800, _from=teachers/T001, _to=areas/Gotham)>,
                <Relation(_key=5690912, _from=teachers/T001, _to=areas/Gotham)>],
 'specializes_in': [<SpecializesIn(_key=5690854, expertise_level=medium, _from=teachers/T001, _to=subjects/CSOOP02)>,
                    <SpecializesIn(_key=5690973, expertise_level=expert, _from=teachers/T001, _to=subjects/CS102)>,
                    <SpecializesIn(_key=5690966, expertise_level=medium, _from=teachers/T001, _to=subjects/CSOOP02)>,
                    <SpecializesIn(_key=5690861, expertise_level=expert, _from=teachers/T001, _to=subjects/CS102)>],
 'teacher': [<Relation(_key=5690819, _from=teachers/T001, _to=students/S1004)>,
             <Relation(_key=5690938, _from=teachers/T001, _to=students/S1002)>,
             <Relation(_key=5690935, _from=teachers/T001, _to=students/S1001)>,
             <Relation(_key=5690931, _from=teachers/T001, _to=students/S1004)>,
             <Relation(_key=5690826, _from=teachers/T001, _to=students/S1002)>,
             <Relation(_key=5690823, _from=teachers/T001, _to=students/S1001)>],
 'teaches': [<Relation(_key=5690809, _from=teachers/T001, _to=subjects/CSOOP02)>,
             <Relation(_key=5690921, _from=teachers/T001, _to=subjects/CSOOP02)>]}
In [65]:
# All Students of Bruce
students_of_bruce = [r._next for r in bruce._relations['teacher']]
for s in students_of_bruce:
    pp(s)
<Student(S1004,Peter Parker)>
<Student(S1002,Lilly Parker)>
<Student(S1001,John Wayne)>
<Student(S1004,Peter Parker)>
<Student(S1002,Lilly Parker)>
<Student(S1001,John Wayne)>
In [21]:
# delete graph and colections
db.drop_graph(uni_graph)
In [ ]:
# delete Test Database
try:
    client.delete_database(name=db_name)
except: pass