diff --git a/norduniclient/models.py b/norduniclient/models.py index f50c654..4765e37 100644 --- a/norduniclient/models.py +++ b/norduniclient/models.py @@ -9,6 +9,8 @@ # Python 3 from norduniclient import core +import six + __author__ = 'lundberg' @@ -189,6 +191,29 @@ def delete(self): def reload(self, node=None): return core.get_node_model(self.manager, self.handle_id, node=node) + def add_property(self, property, value): + if isinstance(value, six.string_types): + value = "'{}'".format(value) + + q = """ + MATCH (n:Node {{handle_id: {{handle_id}}}}) + SET n.{property} = {value} + RETURN n + """.format(property=property, value=value) + with self.manager.session as s: + node = s.run(q, {'handle_id': self.handle_id}).single()['n'] + return self.reload(node=node) + + def remove_property(self, property): + q = """ + MATCH (n:Node {{handle_id: {{handle_id}}}}) + REMOVE n.{property} + RETURN n + """.format(property=property) + with self.manager.session as s: + node = s.run(q, {'handle_id': self.handle_id}).single()['n'] + return self.reload(node=node) + class CommonQueries(BaseNodeModel): @@ -216,7 +241,7 @@ def get_child_form_data(self, node_type): def get_relations(self): q = """ - MATCH (n:Node {handle_id: {handle_id}})<-[r:Owns|Uses|Provides|Responsible_for]-(node) + MATCH (n:Node {handle_id: {handle_id}})<-[r:Owns|Uses|Provides|Responsible_for|Works_for|Parent_of|Member_of|Uses_a]-(node) RETURN r, node """ return self._basic_read_query_to_dict(q) @@ -277,7 +302,6 @@ def get_ports(self): """ return core.query_to_list(self.manager, q, handle_id=self.handle_id) - class LogicalModel(CommonQueries): def get_part_of(self): @@ -892,6 +916,314 @@ class FPCModel(SubEquipmentModel): class CustomerModel(RelationModel): pass +class OrganizationModel(RelationModel): + def set_parent(self, org_handle_id): + q = """ + MATCH (n:Node:Organization {handle_id: {handle_id}}), (m:Node:Organization {handle_id: {org_handle_id}}) + MERGE (m)-[r:Parent_of]->(n) + RETURN m as created, r, n as node + """ + return self._basic_write_query_to_dict(q, org_handle_id=org_handle_id) + + def set_child(self, org_handle_id): + q = """ + MATCH (n:Node:Organization {handle_id: {handle_id}}), (m:Node:Organization {handle_id: {org_handle_id}}) + MERGE (n)-[r:Parent_of]->(m) + RETURN m as created, r, n as node + """ + return self._basic_write_query_to_dict(q, org_handle_id=org_handle_id) + + def add_procedure(self, proc_handle_id): + q = """ + MATCH (n:Node:Organization {handle_id: {handle_id}}), (m:Node:Procedure {handle_id: {proc_handle_id}}) + MERGE (n)-[r:Uses_a]->(m) + RETURN m as created, r, n as node + """ + return self._basic_write_query_to_dict(q, proc_handle_id=proc_handle_id) + + def get_outgoing_relations(self): + q = """ + MATCH (n:Node:Organization {handle_id: {handle_id}})-[r:Parent|Uses_a]->(node) + RETURN r, node + """ + return self._basic_read_query_to_dict(q) + + def get_contacts(self): + q = """ + MATCH (c:Node:Contact)-[:Works_for]->(o:Node:Organization) + WHERE o.handle_id = {handle_id} + RETURN DISTINCT c.handle_id as handle_id, c.name as name + """ + return core.query_to_list(self.manager, q, handle_id=self.handle_id) + class ProviderModel(RelationModel): pass + + +class ContactModel(RelationModel): + def add_group(self, group_handle_id): + q = """ + MATCH (n:Node:Contact {handle_id: {handle_id}}), (m:Node:Group {handle_id: {group_handle_id}}) + MERGE (n)-[r:Member_of]->(m) + RETURN m as created, r, n as node + """ + return self._basic_write_query_to_dict(q, group_handle_id=group_handle_id) + + def get_outgoing_relations(self): + q = """ + MATCH (n:Node:Contact {handle_id: {handle_id}})-[r:Works_for|Member_of]->(node) + RETURN r, node + """ + return self._basic_read_query_to_dict(q) + +class GroupModel(LogicalModel): + def add_member(self, contact_handle_id): + q = """ + MATCH (n:Node:Contact {handle_id: {contact_handle_id}}), (m:Node:Group {handle_id: {handle_id}}) + MERGE (n)-[r:Member_of]->(m) + RETURN m as created, r, n as node + """ + return self._basic_write_query_to_dict(q, contact_handle_id=contact_handle_id) + + +class RoleRelationship(BaseRelationshipModel): + RELATION_NAME = 'Works_for' + DEFAULT_ROLE_NAME = 'Employee' + + def __init__(self, manager): + super(RoleRelationship, self).__init__(manager) + self.type = RoleRelationship.RELATION_NAME + self.name = None + self.handle_id = None + + def load(self, relationship_bundle): + super(RoleRelationship, self).load(relationship_bundle) + self.type = RoleRelationship.RELATION_NAME + self.name = self.data.get('name', None) + self.handle_id = self.data.get('handle_id', None) + + return self + + @classmethod + def get_manager(cls, manager): + if not manager: + manager = core.GraphDB.get_instance().manager + + return manager + + @classmethod + def link_contact_organization(cls, contact_id, organization_id, role_name, manager=None): + if isinstance(contact_id, six.string_types): + contact_id = "'{}'".format(contact_id) + + if isinstance(organization_id, six.string_types): + organization_id = "'{}'".format(organization_id) + + if not role_name: + role_name = cls.DEFAULT_ROLE_NAME + + # create relation + manager = cls.get_manager(manager) + + q = """ + MATCH (c:Contact), (o:Organization) + WHERE c.handle_id = {contact_id} AND o.handle_id = {organization_id} + MERGE (c)-[r:Works_for {{ name: '{role_name}'}}]->(o) + RETURN ID(r) as relation_id + """.format(contact_id=contact_id, organization_id=organization_id, role_name=role_name) + ret = core.query_to_dict(manager, q) + + # load and return + if ret: + relation_id = ret['relation_id'] + relation = cls.get_relationship_model(manager, relationship_id=relation_id) + + return relation + + @classmethod + def get_role_relation_from_organization(cls, organization_id, role_name, manager=None): + if isinstance(organization_id, six.string_types): + organization_id = "'{}'".format(organization_id) + + manager = cls.get_manager(manager) + + q = """ + MATCH (c:Node:Contact)-[r:Works_for]->(o:Node:Organization) + WHERE r.name = "{role_name}" AND o.handle_id = {organization_id} + RETURN ID(r) as relation_id + """.format(role_name=role_name, organization_id=organization_id) + ret = core.query_to_dict(manager, q) + + if ret: + relation_id = ret['relation_id'] + relation = cls.get_relationship_model(manager, relationship_id=relation_id) + + return relation + + @classmethod + def get_contact_with_role_in_organization(cls, organization_id, role_name, manager=None): + if isinstance(organization_id, six.string_types): + organization_id = "'{}'".format(organization_id) + + manager = cls.get_manager(manager) + + q = """ + MATCH (c:Node:Contact)-[r:Works_for]->(o:Node:Organization) + WHERE r.name = "{role_name}" AND o.handle_id = {organization_id} + RETURN c.handle_id as contact_handle_id + """.format(role_name=role_name, organization_id=organization_id) + ret = core.query_to_dict(manager, q) + + if ret: + return ret['contact_handle_id'] + + @classmethod + def get_role_relation_from_contact_organization(cls, organization_id, role_name, contact_id, manager=None): + if isinstance(contact_id, six.string_types): + contact_id = "'{}'".format(contact_id) + + if isinstance(organization_id, six.string_types): + organization_id = "'{}'".format(organization_id) + + manager = cls.get_manager(manager) + + q = """ + MATCH (c:Node:Contact)-[r:Works_for]->(o:Node:Organization) + WHERE c.handle_id = {contact_id} + AND r.name = "{role_name}" + AND o.handle_id = {organization_id} + RETURN ID(r) as relation_id + """.format(contact_id=contact_id, role_name=role_name, + organization_id=organization_id) + ret = core.query_to_dict(manager, q) + + if ret: + relation_id = ret['relation_id'] + relation = cls.get_relationship_model(manager, relationship_id=relation_id) + + return relation + + @classmethod + def unlink_contact_with_role_organization(cls, contact_id, organization_id, role_name, manager=None): + if isinstance(contact_id, six.string_types): + contact_id = "'{}'".format(contact_id) + + if isinstance(organization_id, six.string_types): + organization_id = "'{}'".format(organization_id) + + manager = cls.get_manager(manager) + + q = ''' + MATCH (c:Contact)-[r:Works_for]->(o:Organization) + WHERE c.handle_id = {contact_id} + AND r.name = "{role_name}" + AND o.handle_id = {organization_id} + DELETE r RETURN c + '''.format(contact_id=contact_id, role_name=role_name, \ + organization_id=organization_id) + ret = core.query_to_dict(manager, q) + + @classmethod + def update_roles_withname(cls, role_name, new_name, manager=None): + manager = cls.get_manager(manager) + + q = """ + MATCH (c:Contact)-[r:Works_for]->(o:Organization) + WHERE r.name = "{role_name}" + SET r.name = "{new_name}" + RETURN r + """.format(role_name=role_name, new_name=new_name) + + core.query_to_dict(manager, q) + + @classmethod + def delete_roles_withname(cls, role_name, manager=None): + manager = cls.get_manager(manager) + + q = """ + MATCH (c:Contact)-[r:Works_for]->(o:Organization) + WHERE r.name = "{role_name}" + DELETE r + """.format(role_name=role_name) + + core.query_to_dict(manager, q) + + def load_from_nodes(self, contact_id, organization_id): + if isinstance(contact_id, six.string_types): + contact_id = "'{}'".format(contact_id) + + if isinstance(organization_id, six.string_types): + organization_id = "'{}'".format(organization_id) + + q = """ + MATCH (c:Contact)-[r:Works_for]->(o:Organization) + WHERE c.handle_id = {contact_id} AND o.handle_id = {organization_id} + RETURN ID(r) as relation_id + """.format(contact_id=contact_id, organization_id=organization_id) + + ret = core.query_to_dict(self.manager, q) + + bundle = core.get_relationship_bundle(self.manager, ret['relation_id']) + self.load(bundle) + + @classmethod + def get_relationship_model(cls, manager, relationship_id): + """ + :param manager: Context manager to handle transactions + :type manager: Neo4jDBSessionManager + :param relationship_id: Internal Neo4j relationship id + :type relationship_id: int + :return: Relationship model + :rtype: models.BaseRelationshipModel + """ + manager = cls.get_manager(manager) + + bundle = core.get_relationship_bundle(manager, relationship_id) + return cls(manager).load(bundle) + + @classmethod + def get_all_role_names(cls, manager=None): + manager = cls.get_manager(manager) + + q = """MATCH (n:Contact)-[r:Works_for]->(m:Organization) + WHERE r.name IS NOT NULL + RETURN DISTINCT r.name as role_name""" + + result = core.query_to_list(manager, q) + endresult = [] + for r in result: + endresult.append(r['role_name']) + + return endresult + + @classmethod + def get_contacts_with_role_name(cls, role_name, manager=None): + manager = cls.get_manager(manager) + + q = """ + MATCH (c:Contact)-[r:Works_for]->(o:Organization) + WHERE r.name = "{role_name}" + RETURN c, o + """.format(role_name=role_name) + + result = core.query_to_list(manager, q) + contact_list = [] + + for node in result: + contact = ContactModel(manager) + contact.data = {} + contact.data['handle_id'] = node['c'].properties['handle_id'] + contact.reload(node['c']) + + organization = OrganizationModel(manager) + organization.data = {} + organization.data['handle_id'] = node['o'].properties['handle_id'] + organization.reload(node['o']) + + contact_list.append((contact, organization)) + + return contact_list + +class ProcedureModel(LogicalModel): + pass diff --git a/norduniclient/tests/test_models.py b/norduniclient/tests/test_models.py index 4f3f88e..bce0378 100644 --- a/norduniclient/tests/test_models.py +++ b/norduniclient/tests/test_models.py @@ -144,6 +144,26 @@ def setUp(self): (physical2)<-[:Connected_to]-(physical4)-[:Connected_to]->(physical3) """ + self.role_name_1 = u'IT-Manager' + self.role_name_2 = u'Abuse Management' + + q3 = """ + // Create organization and contact nodes + CREATE (organization1:Node:Relation:Organization{name:'Organization1', handle_id:'113'}), + (organization2:Node:Relation:Organization{name:'Organization2', handle_id:'114'}), + (contact1:Node:Relation:Contact{name:'Contact1', handle_id:'115'}), + (contact2:Node:Relation:Contact{name:'Contact2', handle_id:'116'}), + (procedure1:Node:Logical:Procedure{name:'Procedure1', handle_id:'119'}), + (procedure2:Node:Logical:Procedure{name:'Procedure2', handle_id:'120'}), + (group1:Node:Logical:Group{name:'Group1', handle_id:'121'}), + + + // Create relationships + (contact1)-[:Works_for {name: 'IT-Manager' }]->(organization1), + (contact2)-[:Works_for {name: 'Abuse Management' }]->(organization2), + (organization1)-[:Uses_a]->(procedure1) + """ + # Insert mocked network with self.neo4jdb.session as s: s.run(q1) @@ -152,6 +172,10 @@ def setUp(self): with self.neo4jdb.session as s: s.run(q2) + # Insert organizations and contacts + with self.neo4jdb.session as s: + s.run(q3) + def test_base_node_model(self): node_model_1 = core.get_node_model(self.neo4jdb, handle_id='101') node_model_2 = core.get_node_model(self.neo4jdb, handle_id='102') @@ -191,6 +215,20 @@ def test_remove_label(self): new_labels = node_model_1.labels self.assertEqual(sorted(new_labels), sorted(initial_labels)) + def test_add_remove_property(self): + node_model_1 = core.get_node_model(self.neo4jdb, handle_id='115') + initial_data = node_model_1.data + first_name = 'Smith' + node_model_1 = node_model_1.add_property('first_name', first_name) + new_data = node_model_1.data + new_property = {'first_name': first_name} + expected_data = initial_data.copy() + expected_data.update(new_property) + self.assertEqual(sorted(new_data), sorted(expected_data)) + node_model_1 = node_model_1.remove_property('first_name') + new_data = node_model_1.data + self.assertEqual(sorted(new_data), sorted(initial_data)) + def test_change_meta_type(self): node_model_1 = core.get_node_model(self.neo4jdb, handle_id='101') self.assertEqual(node_model_1.meta_type, 'Physical') @@ -290,6 +328,10 @@ def test_get_relations(self): relations = location2.get_relations() self.assertIsInstance(relations['Responsible_for'][0]['node'], models.RelationModel) + organization1 = core.get_node_model(self.neo4jdb, handle_id='113') + relations = organization1.get_relations() + self.assertIsInstance(relations['Works_for'][0]['node'], models.RelationModel) + def test_get_dependencies(self): logical3 = core.get_node_model(self.neo4jdb, handle_id='107') dependencies = logical3.get_dependencies() @@ -552,6 +594,103 @@ def test_set_responsible_for_location_model(self): relations = rack_4.get_relations() self.assertEqual(len(relations['Responsible_for']), 1) + def test_set_parent(self): + organization1 = core.get_node_model(self.neo4jdb, handle_id='113') + organization2 = core.get_node_model(self.neo4jdb, handle_id='114') + organization2.set_parent(organization1.handle_id) + relations = organization2.get_relations() + self.assertIsInstance(relations['Parent_of'][0]['node'], models.OrganizationModel) + + def test_get_outgoing_relations(self): + contact1 = core.get_node_model(self.neo4jdb, handle_id='115') + relations = contact1.get_outgoing_relations() + self.assertIsInstance(relations['Works_for'][0]['node'], models.OrganizationModel) + + expected_value = self.role_name_1 + self.assertEquals(relations['Works_for'][0]['relationship']['name'], expected_value) + + def test_contact_role_org(self): + contact1 = core.get_node_model(self.neo4jdb, handle_id='115') + organization1 = core.get_node_model(self.neo4jdb, handle_id='113') + + # unlink + models.RoleRelationship.unlink_contact_with_role_organization(contact1.handle_id, + organization1.handle_id, self.role_name_1, self.neo4jdb) + + relations = contact1.get_outgoing_relations() + self.assertEquals(relations, {}) + + # relink + models.RoleRelationship.link_contact_organization(contact1.handle_id, + organization1.handle_id, self.role_name_1, self.neo4jdb) + relations = contact1.get_outgoing_relations() + expected_value = self.role_name_1 + self.assertEquals(relations['Works_for'][0]['relationship']['name'], expected_value) + + # get contact which holds this role in this organization + contact_handle_id = models.RoleRelationship.get_contact_with_role_in_organization( + organization1.handle_id, self.role_name_1, self.neo4jdb) + self.assertEqual(contact_handle_id, contact1.handle_id) + + # get the relation of a organization with a specific role + relation = models.RoleRelationship.get_role_relation_from_organization( + organization1.handle_id, self.role_name_1, self.neo4jdb) + self.assertEqual(relations['Works_for'][0]['relationship_id'], relation.id) + + # get the relation between contact and organization with a specific role + relation = models.RoleRelationship.get_role_relation_from_contact_organization( + organization1.handle_id, self.role_name_1, contact1.handle_id, self.neo4jdb) + self.assertEqual(relations['Works_for'][0]['relationship_id'], relation.id) + + # check role list + role_list = models.RoleRelationship.get_all_role_names(self.neo4jdb) + self.assertEquals(role_list, [self.role_name_1, self.role_name_2]) + + # role name change + new_role_name = u"Abuse Manager" + models.RoleRelationship.update_roles_withname(self.role_name_2, + new_role_name, self.neo4jdb) + role_list = models.RoleRelationship.get_all_role_names(self.neo4jdb) + self.assertEquals(role_list, [self.role_name_1, new_role_name]) + + # delete role + models.RoleRelationship.delete_roles_withname(new_role_name, + self.neo4jdb) + role_list = models.RoleRelationship.get_all_role_names(self.neo4jdb) + self.assertEquals(role_list, [self.role_name_1]) + + def test_uses_a_procedure(self): + organization1 = core.get_node_model(self.neo4jdb, handle_id='113') + organization2 = core.get_node_model(self.neo4jdb, handle_id='114') + procedure1 = core.get_node_model(self.neo4jdb, handle_id='119') + procedure2 = core.get_node_model(self.neo4jdb, handle_id='120') + + relations1 = procedure1.get_relations() + self.assertIsInstance(relations1['Uses_a'][0]['node'], models.OrganizationModel) + + organization2.add_procedure(procedure2.handle_id) + relations2 = procedure2.get_relations() + self.assertIsInstance(relations2['Uses_a'][0]['node'], models.OrganizationModel) + + def test_organization_outgoingrel(self): + organization1 = core.get_node_model(self.neo4jdb, handle_id='113') + relations1 = organization1.get_outgoing_relations() + self.assertIsInstance(relations1['Uses_a'][0]['node'], models.ProcedureModel) + + def test_organization_contacts(self): + organization1 = core.get_node_model(self.neo4jdb, handle_id='113') + contacts = organization1.get_contacts() + contact1 = core.get_node_model(self.neo4jdb, handle_id='115') + self.assertEqual(contact1.handle_id, contacts[0]['handle_id']) + self.assertEqual(contact1.data['name'], contacts[0]['name']) + + def test_groups(self): + group1 = core.get_node_model(self.neo4jdb, handle_id='121') + contact1 = core.get_node_model(self.neo4jdb, handle_id='115') + group1.add_member(contact1.handle_id) + relations = contact1.get_outgoing_relations() + self.assertIsInstance(relations['Member_of'][0]['node'], models.GroupModel) + # TODO: EquipmentModel get_ports should probably work as CommonQueries get_ports def test_get_ports_equipment_model(self): odf1 = core.get_node_model(self.neo4jdb, handle_id='23')