org.neo4j.fabric.stream.StatementResults Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of neo4j-fabric Show documentation
Neo4j - Fabric
There is a newer version: 5.25.1
/*
 * Copyright (c) "Neo4j"
 * Neo4j Sweden AB [https://neo4j.com]
 *
 * This file is part of Neo4j.
 *
 * Neo4j is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see .
 */
package org.neo4j.fabric.stream;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;
import java.util.function.UnaryOperator;
import org.neo4j.fabric.stream.summary.Summary;
import org.neo4j.graphdb.QueryExecutionType;
import org.neo4j.kernel.impl.query.QueryExecution;
import org.neo4j.kernel.impl.query.QuerySubscriber;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

public final class StatementResults {
    private StatementResults() {}

    public static StatementResult map(StatementResult statementResult, UnaryOperator> func) {
        return new BasicStatementResult(
                statementResult.columns(),
                func.apply(statementResult.records()),
                statementResult.summary(),
                statementResult.executionType());
    }

    public static StatementResult initial() {
        return new BasicStatementResult(
                Collections.emptyList(), Flux.just(Records.empty()), Mono.empty(), Mono.empty());
    }

    public static StatementResult connectVia(SubscribableExecution execution, QuerySubject subject) {
        QueryExecution queryExecution = execution.subscribe(subject);
        subject.setQueryExecution(queryExecution);
        return create(
                Arrays.asList(queryExecution.fieldNames()),
                Flux.from(subject),
                subject.getSummary(),
                Mono.just(queryExecution.executionType()));
    }

    public static StatementResult create(
            List columns, Flux records, Mono summary, Mono executionType) {
        return new BasicStatementResult(columns, records, summary, executionType);
    }

    public static  StatementResult withErrorMapping(
            StatementResult statementResult, Class type, Function mapper) {
        var records = statementResult.records().onErrorMap(type, mapper);
        var summary = statementResult.summary().onErrorMap(type, mapper);
        var executionType = statementResult.executionType().onErrorMap(type, mapper);

        return create(statementResult.columns(), records, summary, executionType);
    }

    public static StatementResult error(Throwable err) {
        return new BasicStatementResult(Collections.emptyList(), Flux.error(err), Mono.error(err), Mono.error(err));
    }

    public static StatementResult trace(StatementResult input) {
        return new BasicStatementResult(
                input.columns(),
                input.records().doOnEach(signal -> {
                    if (signal.hasValue()) {
                        System.out.println(String.join(", ", signal.getType().toString(), Records.show(signal.get())));
                    } else if (signal.hasError()) {
                        System.out.println(String.join(
                                ", ",
                                signal.getType().toString(),
                                signal.getThrowable().toString()));
                    } else {
                        System.out.println(String.join(", ", signal.getType().toString()));
                    }
                }),
                input.summary(),
                input.executionType());
    }

    private static class BasicStatementResult implements StatementResult {
        private final List columns;
        private final Flux records;
        private final Mono summary;
        private final Mono executionType;

        BasicStatementResult(
                List columns,
                Flux records,
                Mono summary,
                Mono executionType) {
            this.columns = columns;
            this.records = records;
            this.summary = summary;
            this.executionType = executionType;
        }

        @Override
        public List columns() {
            return columns;
        }

        @Override
        public Flux records() {
            return records;
        }

        @Override
        public Mono summary() {
            return summary;
        }

        @Override
        public Mono executionType() {
            return executionType;
        }
    }

    @FunctionalInterface
    public interface SubscribableExecution {
        QueryExecution subscribe(QuerySubscriber subscriber);
    }
}