apoc.export.graphml.ExportGraphML Maven / Gradle / Ivy

Go to download
package apoc.export.graphml;

import apoc.export.util.ExportConfig;
import apoc.export.util.NodesAndRelsSubGraph;
import apoc.export.util.ProgressReporter;
import apoc.result.ProgressInfo;
import apoc.util.FileUtils;
import apoc.util.Util;
import org.neo4j.cypher.export.CypherResultSubGraph;
import org.neo4j.cypher.export.DatabaseSubGraph;
import org.neo4j.cypher.export.SubGraph;
import org.neo4j.graphdb.GraphDatabaseService;
import org.neo4j.graphdb.Node;
import org.neo4j.graphdb.Relationship;
import org.neo4j.graphdb.Result;
import org.neo4j.helpers.collection.Iterables;
import org.neo4j.procedure.*;

import javax.xml.stream.XMLStreamException;
import java.io.PrintWriter;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

import static apoc.util.FileUtils.checkWriteAllowed;
import static apoc.util.FileUtils.getPrintWriter;

/**
 * @author mh
 * @since 22.05.16
 */
public class ExportGraphML {
    @Context
    public GraphDatabaseService db;

    @Procedure(name = "apoc.import.graphml",mode = Mode.WRITE)
    @Description("apoc.import.graphml(file,config) - imports graphml file")
    public Stream file(@Name("file") String fileName, @Name("config") Map config) throws Exception {
        ProgressInfo result =
        Util.inThread(() -> {
            ExportConfig exportConfig = new ExportConfig(config);
            ProgressReporter reporter = new ProgressReporter(null, null, new ProgressInfo(fileName, "file", "graphml"));
            XmlGraphMLReader graphMLReader = new XmlGraphMLReader(db).reporter(reporter)
                    .batchSize(exportConfig.getBatchSize())
                    .relType(exportConfig.defaultRelationshipType())
                    .nodeLabels(exportConfig.readLabels());

            if (exportConfig.storeNodeIds()) graphMLReader.storeNodeIds();


            graphMLReader.parseXML(FileUtils.readerFor(fileName));
            return reporter.getTotal();
        });
        return Stream.of(result);
    }
    @Procedure
    @Description("apoc.export.graphml.all(file,config) - exports whole database as graphml to the provided file")
    public Stream all(@Name("file") String fileName, @Name("config") Map config) throws Exception {

        String source = String.format("database: nodes(%d), rels(%d)", Util.nodeCount(db), Util.relCount(db));
        return exportGraphML(fileName, source, new DatabaseSubGraph(db), new ExportConfig(config));
    }

    @Procedure
    @Description("apoc.export.graphml.data(nodes,rels,file,config) - exports given nodes and relationships as graphml to the provided file")
    public Stream data(@Name("nodes") List nodes, @Name("rels") List rels, @Name("file") String fileName, @Name("config") Map config) throws Exception {

        String source = String.format("data: nodes(%d), rels(%d)", nodes.size(), rels.size());
        return exportGraphML(fileName, source, new NodesAndRelsSubGraph(db, nodes, rels), new ExportConfig(config));
    }
    @Procedure
    @Description("apoc.export.graphml.graph(graph,file,config) - exports given graph object as graphml to the provided file")
    public Stream graph(@Name("graph") Map graph, @Name("file") String fileName, @Name("config") Map config) throws Exception {

        Collection nodes = (Collection) graph.get("nodes");
        Collection rels = (Collection) graph.get("relationships");
        String source = String.format("graph: nodes(%d), rels(%d)", nodes.size(), rels.size());
        return exportGraphML(fileName, source, new NodesAndRelsSubGraph(db, nodes, rels), new ExportConfig(config));
    }

    @Procedure
    @Description("apoc.export.graphml.query(query,file,config) - exports nodes and relationships from the cypher statement as graphml to the provided file")
    public Stream query(@Name("query") String query, @Name("file") String fileName, @Name("config") Map config) throws Exception {
        ExportConfig c = new ExportConfig(config);
        Result result = db.execute(query);
        SubGraph graph = CypherResultSubGraph.from(result, db, c.getRelsInBetween());
        String source = String.format("statement: nodes(%d), rels(%d)",
                Iterables.count(graph.getNodes()), Iterables.count(graph.getRelationships()));
        return exportGraphML(fileName, source, graph, c);
    }

    private Stream exportGraphML(@Name("file") String fileName, String source, SubGraph graph, ExportConfig config) throws Exception, XMLStreamException {
        if (fileName != null) checkWriteAllowed();
        ProgressReporter reporter = new ProgressReporter(null, null, new ProgressInfo(fileName, source, "graphml"));
        PrintWriter printWriter = getPrintWriter(fileName, null);
        XmlGraphMLWriter exporter = new XmlGraphMLWriter();
        exporter.write(graph, printWriter, reporter, config);
        printWriter.flush();
        printWriter.close();
        return reporter.stream();
    }
}