source: AnnotationManagerN4J/src/main/java/de/mpiwg/itgroup/annotations/neo4j/AnnotationStore.java @ 24:e208a7b1a37a

/**
 * 
 */
package de.mpiwg.itgroup.annotations.neo4j;

import java.util.ArrayList;
import java.util.Calendar;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.apache.log4j.Logger;
import org.neo4j.graphdb.Direction;
import org.neo4j.graphdb.GraphDatabaseService;
import org.neo4j.graphdb.Node;
import org.neo4j.graphdb.Relationship;
import org.neo4j.graphdb.RelationshipType;
import org.neo4j.graphdb.Transaction;
import org.neo4j.graphdb.index.Index;
import org.neo4j.graphdb.index.IndexHits;

import de.mpiwg.itgroup.annotations.Actor;
import de.mpiwg.itgroup.annotations.Annotation;
import de.mpiwg.itgroup.annotations.Annotation.FragmentTypes;
import de.mpiwg.itgroup.annotations.Group;
import de.mpiwg.itgroup.annotations.Person;

/**
 * @author casties
 * 
 */
public class AnnotationStore {

    protected static Logger logger = Logger.getLogger(AnnotationStore.class);

    protected GraphDatabaseService graphDb;

    public static enum NodeTypes {
        ANNOTATION, PERSON, TARGET, GROUP, TAG
    }

    protected List<Index<Node>> nodeIndexes;

    public static enum RelationTypes implements RelationshipType {
        ANNOTATES, CREATED, PERMITS_ADMIN, PERMITS_DELETE, PERMITS_UPDATE, PERMITS_READ, MEMBER_OF, HAS_TAG
    }

    public static String ANNOTATION_URI_BASE = "http://entities.mpiwg-berlin.mpg.de/annotations/";

    public AnnotationStore(GraphDatabaseService graphDb) {
        super();
        this.graphDb = graphDb;
        nodeIndexes = new ArrayList<Index<Node>>(5);
        // List.set(enum.ordinal(), val) seems not to work.
        nodeIndexes.add(NodeTypes.ANNOTATION.ordinal(), graphDb.index().forNodes("annotations"));
        nodeIndexes.add(NodeTypes.PERSON.ordinal(), graphDb.index().forNodes("persons"));
        nodeIndexes.add(NodeTypes.TARGET.ordinal(), graphDb.index().forNodes("targets"));
        nodeIndexes.add(NodeTypes.GROUP.ordinal(), graphDb.index().forNodes("groups"));
        nodeIndexes.add(NodeTypes.TAG.ordinal(), graphDb.index().forNodes("tags"));
    }

    protected Index<Node> getNodeIndex(NodeTypes type) {
        return nodeIndexes.get(type.ordinal());
    }

    /**
     * @param userUri
     * @return
     */
    public Node getPersonNodeByUri(String userUri) {
        if (userUri == null) return null;
        Node person = getNodeIndex(NodeTypes.PERSON).get("uri", userUri).getSingle();
        return person;
    }

    public List<Actor> getActors(String key, String query, NodeTypes type) {
        ArrayList<Actor> actors = new ArrayList<Actor>();
        Index<Node> idx = getNodeIndex(type);
        if (key == null) {
            key = "uri";
            query = "*";
        }
        IndexHits<Node> actorNodes = idx.query(key, query);
        for (Node actorNode : actorNodes) {
            Actor actor = createActorFromNode(actorNode);
            actors.add(actor);
        }
        return actors;
    }

    /**
     * Returns List of Groups.
     * Key has to be indexed.
     * 
     * @param key
     * @param query
     * @return
     */
    public List<Group> getGroups(String key, String query) {
        ArrayList<Group> groups = new ArrayList<Group>();
        Index<Node> idx = getNodeIndex(NodeTypes.GROUP);
        if (key == null) {
            key = "uri";
            query = "*";
        }
        IndexHits<Node> groupNodes = idx.query(key, query);
        for (Node groupNode : groupNodes) {
            Actor group = createActorFromNode(groupNode);
            groups.add((Group) group);
        }
        return groups;
    }

    /**
     * Returns List of Groups the person is member of.
     * 
     * @param person
     * @return
     */
    public List<Group> getGroupsForPersonNode(Node person) {
        ArrayList<Group> groups = new ArrayList<Group>();
        Iterable<Relationship> rels = person.getRelationships(RelationTypes.MEMBER_OF);
        for (Relationship rel : rels) {
            Node groupNode = rel.getEndNode();
            Actor group = createActorFromNode(groupNode);
            // make sure we're getting a group
            if (!(group instanceof Group)) {
                logger.error("target of MEMBER_OF is not GROUP! rel=" + rel);
                continue;
            }
            groups.add((Group) group);
        }
        return groups;
    }

    /**
     * Returns if person with uri is in Group group.
     * 
     * @param person
     * @param group
     * @return
     */
    public boolean isPersonInGroup(Person person, Group group) {
        Node pn = getPersonNodeByUri(person.getUriString());
        if (pn == null) return false;
        // optimized version of getGroupsForPersonNode
        Iterable<Relationship> rels = pn.getRelationships(RelationTypes.MEMBER_OF);
        for (Relationship rel : rels) {
            Node gn = rel.getEndNode();
            if (gn.getProperty("uri", "").equals(group.getUriString()) || gn.getProperty("id", "").equals(group.getId())) {
                return true;
            }
        }
        return false;
    }

    /**
     * Returns the members of the group.
     * 
     * @param group
     * @return
     */
    public List<Person> getMembersOfGroup(Group group) {
        ArrayList<Person> members = new ArrayList<Person>();
        Node gn = getActorNode(group);
        Iterable<Relationship> rels = gn.getRelationships(RelationTypes.MEMBER_OF);
        for (Relationship rel : rels) {
            Node memberNode = rel.getStartNode();
            Actor member = createActorFromNode(memberNode);
            // make sure we're getting a group
            if (!(member instanceof Person)) {
                logger.error("source of MEMBER_OF is not PERSON! rel=" + rel);
                continue;
            }
            members.add((Person) member);
        }
        return members;
    }
    
    /**
     * Returns the stored Actor matching the given one.
     * 
     * @param actor
     * @return
     */
    public Actor getActor(Actor actor) {
        Node actorNode = getActorNode(actor);
        Actor storedActor = createActorFromNode(actorNode);
        return storedActor;
    }
    
    /**
     * Stores an Actor (Person or Group). Creates a new actor Node or returns an existing one.
     * 
     * @param actor
     * @return
     */
    public Actor storeActor(Actor actor) {
        Node actorNode = getOrCreateActorNode(actor);
        Actor storedActor = createActorFromNode(actorNode);
        return storedActor;
    }
    
    /**
     * Returns the Annotation with the given id.
     * 
     * @param id
     * @return
     */
    public Annotation getAnnotationById(String id) {
        Node annotNode = getNodeIndex(NodeTypes.ANNOTATION).get("id", id).getSingle();
        Annotation annot = createAnnotationFromNode(annotNode);
        return annot;
    }

    /**
     * Returns an Annotation object from an annotation-Node.
     * 
     * @param annotNode
     * @return
     */
    public Annotation createAnnotationFromNode(Node annotNode) {
        Annotation annot = new Annotation();
        annot.setUri((String) annotNode.getProperty("id", null));
        annot.setBodyText((String) annotNode.getProperty("bodyText", null));
        annot.setBodyUri((String) annotNode.getProperty("bodyUri", null));
        /*
         * get annotation target from relation
         */
        Relationship targetRel = getRelation(annotNode, RelationTypes.ANNOTATES, null);
        if (targetRel != null) {
            Node target = targetRel.getEndNode();
            annot.setTargetBaseUri((String) target.getProperty("uri", null));
        } else {
            logger.error("annotation " + annotNode + " has no target node!");
        }
        annot.setTargetFragment((String) annotNode.getProperty("targetFragment", null));
        String ft = (String) annotNode.getProperty("fragmentType", null);
        if (ft != null) {
            annot.setFragmentType(FragmentTypes.valueOf(ft));
        }
        /*
         * get creator from relation
         */
        Relationship creatorRel = getRelation(annotNode, RelationTypes.CREATED, null);
        if (creatorRel != null) {
            Node creatorNode = creatorRel.getStartNode();
            Actor creator = createActorFromNode(creatorNode);
            annot.setCreator(creator);
        } else {
            logger.error("annotation " + annotNode + " has no creator node!");
        }
        /*
         * get creation date
         */
        annot.setCreated((String) annotNode.getProperty("created", null));
        /*
         * get permissions
         */
        Relationship adminRel = getRelation(annotNode, RelationTypes.PERMITS_ADMIN, null);
        if (adminRel != null) {
            Node adminNode = adminRel.getEndNode();
            Actor admin = createActorFromNode(adminNode);
            annot.setAdminPermission(admin);
        }
        Relationship deleteRel = getRelation(annotNode, RelationTypes.PERMITS_DELETE, null);
        if (deleteRel != null) {
            Node deleteNode = deleteRel.getEndNode();
            Actor delete = createActorFromNode(deleteNode);
            annot.setDeletePermission(delete);
        }
        Relationship updateRel = getRelation(annotNode, RelationTypes.PERMITS_UPDATE, null);
        if (updateRel != null) {
            Node updateNode = updateRel.getEndNode();
            Actor update = createActorFromNode(updateNode);
            annot.setUpdatePermission(update);
        }
        Relationship readRel = getRelation(annotNode, RelationTypes.PERMITS_READ, null);
        if (readRel != null) {
            Node readNode = readRel.getEndNode();
            Actor read = createActorFromNode(readNode);
            annot.setReadPermission(read);
        }
        /*
         * get tags
         */
        Set<String> tags = new HashSet<String>();
        for (Relationship rel : annotNode.getRelationships(RelationTypes.HAS_TAG)) {
            String tag = (String) rel.getEndNode().getProperty("name", null);
            if (tag != null) {
                tags.add(tag);
            }
        }
        annot.setTags(tags);

        return annot;
    }

    /**
     * Returns an Actor object from a node.
     * 
     * @param actorNode
     * @return
     */
    protected Actor createActorFromNode(Node actorNode) {
        String id = (String) actorNode.getProperty("id", null);
        String uri = (String) actorNode.getProperty("uri", null);
        String name = (String) actorNode.getProperty("name", null);
        String type = (String) actorNode.getProperty("TYPE", null);
        if (type != null && type.equals("PERSON")) {
            return new Person(id, uri, name);
        } else if (type != null && type.equals("GROUP")) {
            return new Group(id, uri, name);
        }
        return null;
    }

    /**
     * Store a new annotation in the store or update an existing one. Returns
     * the stored annotation.
     * 
     * @param annot
     * @return
     */
    public Annotation storeAnnotation(Annotation annot) {
        Node annotNode = null;
        Transaction tx = graphDb.beginTx();
        try {
            /*
             * create or get the annotation
             */
            String id = annot.getUri();
            if (id == null) {
                id = createRessourceURI("annot:");
            }
            annotNode = getOrCreateAnnotationNode(id);

            /*
             * the annotation body
             */
            String bodyText = annot.getBodyText();
            if (bodyText != null) {
                annotNode.setProperty("bodyText", bodyText);
            }
            String bodyUri = annot.getBodyUri();
            if (bodyUri != null) {
                annotNode.setProperty("bodyUri", bodyUri);
            }

            /*
             * the annotation target
             */
            String targetBaseUri = annot.getTargetBaseUri();
            if (targetBaseUri != null) {
                Node target = getOrCreateTargetNode(targetBaseUri);
                getOrCreateRelation(annotNode, RelationTypes.ANNOTATES, target);
            }

            /*
             * The fragment part of the annotation target.
             */
            String targetFragment = annot.getTargetFragment();
            FragmentTypes fragmentType = annot.getFragmentType();
            if (targetFragment != null) {
                annotNode.setProperty("targetFragment", targetFragment);
                annotNode.setProperty("fragmentType", fragmentType.name());
            }

            /*
             * The creator of this annotation.
             */
            Actor creator = annot.getCreator();
            if (creator != null) {
                Node creatorNode = getOrCreateActorNode(creator);
                getOrCreateRelation(creatorNode, RelationTypes.CREATED, annotNode);
            }

            /*
             * The creation date of this annotation.
             */
            String created = annot.getCreated();
            if (created != null) {
                annotNode.setProperty("created", created);
            }

            /*
             * Permissions for this annotation.
             */
            setPermissionRelation(annotNode, RelationTypes.PERMITS_ADMIN, annot.getAdminPermission());
            setPermissionRelation(annotNode, RelationTypes.PERMITS_DELETE, annot.getDeletePermission());
            setPermissionRelation(annotNode, RelationTypes.PERMITS_UPDATE, annot.getUpdatePermission());
            setPermissionRelation(annotNode, RelationTypes.PERMITS_READ, annot.getReadPermission());

            /*
             * Tags on this annotation.
             */
            Set<String> newTags = annot.getTags();
            // we ignore existing tags if tags == null
            if (newTags != null) {
                List<Relationship> oldHasTags = new ArrayList<Relationship>();
                for (Relationship rel : annotNode.getRelationships(RelationTypes.HAS_TAG)) {
                    oldHasTags.add(rel);
                }
                // adjust to new tags
                if (newTags.isEmpty()) {
                    // remove old tags
                    if (!oldHasTags.isEmpty()) {
                        for (Relationship rel : oldHasTags) {
                            rel.delete();
                            // TODO: should we delete orphan nodes too?
                        }
                    }
                } else {
                    if (!oldHasTags.isEmpty()) {
                        // adjust old tags
                        for (Relationship rel : oldHasTags) {
                            String oldTag = (String) rel.getEndNode().getProperty("name", null);
                            if (newTags.contains(oldTag)) {
                                // tag exists
                                newTags.remove(oldTag);
                            } else {
                                // tag exists no longer
                                rel.delete();
                                // TODO: should we delete orphan nodes too?
                            }
                        }
                    }
                    if (!newTags.isEmpty()) {
                        // still tags to add
                        for (String tag : newTags) {
                            // create new tag
                            Node tagNode = getOrCreateTagNode(tag);
                            getOrCreateRelation(annotNode, RelationTypes.HAS_TAG, tagNode);
                        }
                    }

                }
            }
            tx.success();
        } finally {
            tx.finish();
        }

        // re-read and return annotation
        Annotation storedAnnot = createAnnotationFromNode(annotNode);
        return storedAnnot;
    }

    /**
     * Deletes the annotation with the given id.
     * 
     * @param id
     */
    public void deleteById(String id) {
        Node annotNode = getNodeIndex(NodeTypes.ANNOTATION).get("id", id).getSingle();
        if (annotNode != null) {
            // delete related objects
            Transaction tx = graphDb.beginTx();
            try {
                for (Relationship rel : annotNode.getRelationships()) {
                    // delete relation and the related node if it has no other
                    // relations
                    Node other = rel.getOtherNode(annotNode);
                    rel.delete();
                    if (!other.hasRelationship()) {
                        deleteNode(other);
                    }
                }
                if (!annotNode.hasRelationship()) {
                    deleteNode(annotNode);
                } else {
                    logger.error("deleteById: unable to delete: Node still has relations.");
                }
                tx.success();
            } finally {
                tx.finish();
            }
        }
    }

    /**
     * Returns all annotations with the given uri and/or user.
     * 
     * @param uri
     * @param userUri
     * @param limit
     * @param offset
     * @return
     */
    public List<Annotation> searchByUriUser(String targetUri, String userUri, String limit, String offset) {
        List<Annotation> annotations = new ArrayList<Annotation>();
        if (targetUri != null) {
            // there should be only one
            Node target = getNodeIndex(NodeTypes.TARGET).get("uri", targetUri).getSingle();
            if (target != null) {
                Iterable<Relationship> relations = target.getRelationships(RelationTypes.ANNOTATES);
                for (Relationship relation : relations) {
                    Node ann = relation.getStartNode();
                    if (ann.getProperty("TYPE", "").equals("ANNOTATION")) {
                        Annotation annot = createAnnotationFromNode(ann);
                        annotations.add(annot);
                    } else {
                        logger.error("ANNOTATES relation does not start with ANNOTATION: " + ann);
                    }
                }
            }
        }
        if (userUri != null) {
            // there should be only one
            Node person = getPersonNodeByUri(userUri);
            if (person != null) {
                Iterable<Relationship> relations = person.getRelationships(RelationTypes.CREATED);
                for (Relationship relation : relations) {
                    Node ann = relation.getEndNode();
                    if (ann.getProperty("TYPE", "").equals("ANNOTATION")) {
                        Annotation annot = createAnnotationFromNode(ann);
                        annotations.add(annot);
                    } else {
                        logger.error("CREATED relation does not end with ANNOTATION: " + ann);
                    }
                }
            }
        }
        // TODO: if both uri and user are given we should intersect
        return annotations;
    }

    /**
     * Returns Relationship of type from Node start to Node end. Creates one if
     * it doesn't exist.
     * 
     * @param start
     * @param type
     * @param end
     * @return
     */
    protected Relationship getOrCreateRelation(Node start, RelationshipType type, Node end) {
        if (start.hasRelationship()) {
            // there are relations
            Iterable<Relationship> rels = start.getRelationships(type, Direction.OUTGOING);
            for (Relationship rel : rels) {
                if (rel.getEndNode().equals(end)) {
                    // relation exists
                    return rel;
                }
            }
        }
        // create new one
        Relationship rel;
        Transaction tx = graphDb.beginTx();
        try {
            rel = start.createRelationshipTo(end, type);
            tx.success();
        } finally {
            tx.finish();
        }
        return rel;
    }

    protected Node getOrCreateAnnotationNode(String id) {
        Index<Node> idx = getNodeIndex(NodeTypes.ANNOTATION);
        IndexHits<Node> annotations = idx.get("id", id);
        Node annotation = annotations.getSingle();
        if (annotation == null) {
            // does not exist yet
            Transaction tx = graphDb.beginTx();
            try {
                annotation = graphDb.createNode();
                annotation.setProperty("TYPE", NodeTypes.ANNOTATION.name());
                annotation.setProperty("id", id);
                idx.add(annotation, "id", id);
                tx.success();
            } finally {
                tx.finish();
            }
        }
        return annotation;
    }

    protected Node getOrCreateTargetNode(String uri) {
        Index<Node> idx = getNodeIndex(NodeTypes.TARGET);
        IndexHits<Node> targets = idx.get("uri", uri);
        Node target = targets.getSingle();
        if (target == null) {
            // does not exist yet
            Transaction tx = graphDb.beginTx();
            try {
                target = graphDb.createNode();
                target.setProperty("TYPE", NodeTypes.TARGET.name());
                target.setProperty("uri", uri);
                idx.add(target, "uri", uri);
                tx.success();
            } finally {
                tx.finish();
            }
        }
        return target;
    }

    protected Node getActorNode(Actor actor) {
        // Person/Group is identified by URI or id
        String uri = actor.getUriString();
        Index<Node> idx;
        if (actor.isGroup()) {
            idx = getNodeIndex(NodeTypes.GROUP);
        } else {
            idx = getNodeIndex(NodeTypes.PERSON);
        }
        IndexHits<Node> persons = idx.get("uri", uri);
        Node person = persons.getSingle();
        return person;
    }
    
    protected Node getOrCreateActorNode(Actor actor) {
        // Person/Group is identified by URI or id
        String uri = actor.getUriString();
        String name = actor.getName();
        String id = actor.getId();
        Index<Node> idx;
        if (actor.isGroup()) {
            idx = getNodeIndex(NodeTypes.GROUP);
        } else {
            idx = getNodeIndex(NodeTypes.PERSON);
        }
        IndexHits<Node> persons = idx.get("uri", uri);
        Node person = persons.getSingle();
        if (person == null) {
            // does not exist yet
            Transaction tx = graphDb.beginTx();
            try {
                person = graphDb.createNode();
                if (actor.isGroup()) {
                    person.setProperty("TYPE", NodeTypes.GROUP.name());
                } else {
                    person.setProperty("TYPE", NodeTypes.PERSON.name());
                }
                person.setProperty("uri", uri);
                idx.add(person, "uri", uri);
                if (name != null) {
                    person.setProperty("name", name);
                }
                if (id != null) {
                    person.setProperty("id", id);
                }
                tx.success();
            } finally {
                tx.finish();
            }
        }
        return person;
    }

    protected Node getOrCreateTagNode(String tagname) {
        Index<Node> idx = getNodeIndex(NodeTypes.TAG);
        IndexHits<Node> tags = idx.get("name", tagname);
        Node tag = tags.getSingle();
        if (tag == null) {
            // does not exist yet
            Transaction tx = graphDb.beginTx();
            try {
                tag = graphDb.createNode();
                tag.setProperty("TYPE", NodeTypes.TAG.name());
                tag.setProperty("name", tagname);
                idx.add(tag, "name", tagname);
                tx.success();
            } finally {
                tx.finish();
            }
        }
        return tag;
    }

    /**
     * Create or update permissions relations for an annotation.
     * 
     * @param annotNode
     * @param type
     * @param annot
     */
    protected void setPermissionRelation(Node annotNode, RelationTypes type, Actor actor) {
        Node newActorNode = null;
        if (actor != null) {
            newActorNode = getOrCreateActorNode(actor);
        }
        Relationship rel = getRelation(annotNode, type, null);
        if (rel != null) {
            // relation exists
            Node oldActorNode = rel.getEndNode();
            if (!oldActorNode.equals(newActorNode)) {
                // new admin is different
                rel.delete();
                if (newActorNode != null) {
                    rel = getOrCreateRelation(annotNode, type, newActorNode);
                }
            }
        } else {
            // no relation yet
            if (newActorNode != null) {
                rel = getOrCreateRelation(annotNode, type, newActorNode);
            }
        }
    }

    /**
     * Unindexes and deletes given Node if it has no relations.
     * 
     * @param node
     */
    protected void deleteNode(Node node) {
        Transaction tx = graphDb.beginTx();
        try {
            if (node.hasRelationship()) {
                logger.error("deleteNode: unable to delete: Node still has relations.");
            } else {
                String ts = (String) node.getProperty("TYPE", null);
                try {
                    NodeTypes type = NodeTypes.valueOf(ts);
                    getNodeIndex(type).remove(node);
                } catch (Exception e) {
                    logger.error("deleteNode: unable to get TYPE of node: " + node);
                }
                node.delete();
            }
            tx.success();
        } finally {
            tx.finish();
        }
    }

    /**
     * returns the (first) Relationship of RelationTypes type from Node start.
     * 
     * @param start
     * @param type
     * @param direction
     * @return
     */
    protected Relationship getRelation(Node start, RelationTypes type, Direction direction) {
        Iterable<Relationship> rels;
        if (direction == null) {
            // ignore direction
            rels = start.getRelationships(type);
        } else {
            rels = start.getRelationships(type, direction);
        }
        for (Relationship rel : rels) {
            // just the first one
            return rel;
        }
        return null;
    }

    /**
     * Erzeuge eine urn aus der aktuellen Zeit in millis
     * 
     * @return
     */
    private String createRessourceURI(String prefix) {

        Calendar cal = Calendar.getInstance();

        long time = cal.getTimeInMillis();

        return String.format("%s%s%s", ANNOTATION_URI_BASE, prefix, time);

    }

}