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 [http://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 reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

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.graphdb.QueryStatistics;
import org.neo4j.kernel.api.exceptions.Status;
import org.neo4j.kernel.impl.query.QueryExecution;
import org.neo4j.kernel.impl.query.QuerySubscriber;

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( Flux.empty(), 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(
                Flux.fromArray( queryExecution.fieldNames() ),
                Flux.from( subject ),
                subject.getSummary(),
                Mono.just( queryExecution.executionType() )
        );
    }

    public static StatementResult create( Flux 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 columns = statementResult.columns().onErrorMap( type, mapper );
        var records = statementResult.records().onErrorMap( type, mapper );
        var summary = statementResult.summary().onErrorMap( type, mapper );
        var executionType = statementResult.executionType().onErrorMap( type, mapper );

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

    public static StatementResult error( Throwable err )
    {
        return new BasicStatementResult( Flux.error( err ), 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 Flux columns;
        private final Flux records;
        private final Mono summary;
        private final Mono executionType;

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

        @Override
        public Flux 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 );
    }
}