eu.drus.jpa.unit.neo4j.dataset.DatabaseReader Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of jpa-unit-neo4j Show documentation
JUnit extension for simple testing of JPA entities and components
The newest version!
package eu.drus.jpa.unit.neo4j.dataset;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.jgrapht.Graph;
import org.jgrapht.graph.ClassBasedEdgeFactory;
import org.jgrapht.graph.DefaultDirectedGraph;

public class DatabaseReader {

    private GraphElementFactory factory;

    public DatabaseReader(final GraphElementFactory factory) {
        this.factory = factory;
    }

    public Graph readGraph(final Connection connection) throws SQLException {
        final List nodes = new ArrayList<>();
        final List edges = new ArrayList<>();

        readGraphElements(connection, edges, nodes);

        final DefaultDirectedGraph graph = new DefaultDirectedGraph<>(new ClassBasedEdgeFactory<>(Edge.class));
        nodes.forEach(graph::addVertex);
        edges.forEach(e -> graph.addEdge(e.getSourceNode(), e.getTargetNode(), e));

        return graph;
    }

    @SuppressWarnings("unchecked")
    private void readGraphElements(final Connection connection, final List edgeList, final List nodeList) throws SQLException {
        final Map nodes = new HashMap<>();

        try (PreparedStatement ps = connection
                .prepareStatement("MATCH (n) RETURN { id: id(n), labels: labels(n), attributes: properties(n) } as node")) {
            try (final ResultSet rs = ps.executeQuery()) {
                while (rs.next()) {
                    try {
                        final Map n = (Map) rs.getObject("node");
                        final Node node = toNode(n);
                        nodes.put(n.get("id"), node);
                        nodeList.add(node);
                    } catch (final NoSuchClassException e1) {
                        // happens, when the db contains nodes which are not part of the given
                        // domain model. E.g. like special objects from the persistence provider.
                        // Such nodes have to be ignored
                    }
                }
            }
        }

        try (PreparedStatement ps = connection.prepareStatement(
                "MATCH ()-[r]->() RETURN {id: id(r), label: type(r), attributes: properties(r), from: id(startNode(r)), to: id(endNode(r))} as relation")) {
            try (final ResultSet rs = ps.executeQuery()) {
                while (rs.next()) {
                    final Map r = (Map) rs.getObject("relation");

                    final Edge edge = toEdge(nodes, r);
                    if (edge != null) {
                        edgeList.add(edge);
                    }
                }
            }
        }
    }

    @SuppressWarnings("unchecked")
    private Node toNode(final Map node) throws NoSuchClassException {
        final String id = node.get("id").toString();
        final List labels = (List) node.get("labels");
        final Map attributes = (Map) node.get("attributes");

        return factory.createNode(id, labels, attributes);
    }

    @SuppressWarnings("unchecked")
    private Edge toEdge(final Map nodes, final Map edge) {
        final String id = edge.get("id").toString();
        final List labels = Arrays.asList(edge.get("label").toString());
        final Map attributes = (Map) edge.get("attributes");

        final Node from = nodes.get(edge.get("from"));
        final Node to = nodes.get(edge.get("to"));

        if (to == null || from == null) {
            // happens, when the edge is between nodes, which are not part of the given domain model
            return null;
        }

        return factory.createEdge(from, to, id, labels, attributes);
    }
}