org.neo4j.driver.internal.async.InternalAsyncSession Maven / Gradle / Ivy

/*
 * Copyright (c) "Neo4j"
 * Neo4j Sweden AB [http://neo4j.com]
 *
 * This file is part of Neo4j.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.neo4j.driver.internal.async;

import static java.util.Collections.emptyMap;
import static org.neo4j.driver.internal.util.Futures.completedWithNull;
import static org.neo4j.driver.internal.util.Futures.failedFuture;

import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import org.neo4j.driver.AccessMode;
import org.neo4j.driver.Bookmark;
import org.neo4j.driver.Query;
import org.neo4j.driver.TransactionConfig;
import org.neo4j.driver.async.AsyncSession;
import org.neo4j.driver.async.AsyncTransaction;
import org.neo4j.driver.async.AsyncTransactionCallback;
import org.neo4j.driver.async.AsyncTransactionWork;
import org.neo4j.driver.async.ResultCursor;
import org.neo4j.driver.internal.InternalBookmark;
import org.neo4j.driver.internal.util.Futures;

public class InternalAsyncSession extends AsyncAbstractQueryRunner implements AsyncSession {
    private final NetworkSession session;

    public InternalAsyncSession(NetworkSession session) {
        this.session = session;
    }

    @Override
    public CompletionStage runAsync(Query query) {
        return runAsync(query, TransactionConfig.empty());
    }

    @Override
    public CompletionStage runAsync(String query, TransactionConfig config) {
        return runAsync(query, emptyMap(), config);
    }

    @Override
    public CompletionStage runAsync(
            String query, Map parameters, TransactionConfig config) {
        return runAsync(new Query(query, parameters), config);
    }

    @Override
    public CompletionStage runAsync(Query query, TransactionConfig config) {
        return session.runAsync(query, config);
    }

    @Override
    public CompletionStage closeAsync() {
        return session.closeAsync();
    }

    @Override
    public CompletionStage beginTransactionAsync() {
        return beginTransactionAsync(TransactionConfig.empty());
    }

    @Override
    public CompletionStage beginTransactionAsync(TransactionConfig config) {
        return session.beginTransactionAsync(config).thenApply(InternalAsyncTransaction::new);
    }

    @Override
    @Deprecated
    public  CompletionStage readTransactionAsync(AsyncTransactionWork> work) {
        return readTransactionAsync(work, TransactionConfig.empty());
    }

    @Override
    @Deprecated
    public  CompletionStage readTransactionAsync(
            AsyncTransactionWork> work, TransactionConfig config) {
        return transactionAsync(AccessMode.READ, work, config);
    }

    @Override
    public  CompletionStage executeReadAsync(
            AsyncTransactionCallback> callback, TransactionConfig config) {
        return readTransactionAsync(tx -> callback.execute(new DelegatingAsyncTransactionContext(tx)), config);
    }

    @Override
    @Deprecated
    public  CompletionStage writeTransactionAsync(AsyncTransactionWork> work) {
        return writeTransactionAsync(work, TransactionConfig.empty());
    }

    @Override
    @Deprecated
    public  CompletionStage writeTransactionAsync(
            AsyncTransactionWork> work, TransactionConfig config) {
        return transactionAsync(AccessMode.WRITE, work, config);
    }

    @Override
    public  CompletionStage executeWriteAsync(
            AsyncTransactionCallback> callback, TransactionConfig config) {
        return writeTransactionAsync(tx -> callback.execute(new DelegatingAsyncTransactionContext(tx)), config);
    }

    @Override
    @Deprecated
    public Bookmark lastBookmark() {
        return InternalBookmark.from(session.lastBookmarks());
    }

    @Override
    public Set lastBookmarks() {
        return new HashSet<>(session.lastBookmarks());
    }

    private  CompletionStage transactionAsync(
            AccessMode mode,
            @SuppressWarnings("deprecation") AsyncTransactionWork> work,
            TransactionConfig config) {
        return session.retryLogic().retryAsync(() -> {
            CompletableFuture resultFuture = new CompletableFuture<>();
            CompletionStage txFuture = session.beginTransactionAsync(mode, config);

            txFuture.whenComplete((tx, completionError) -> {
                Throwable error = Futures.completionExceptionCause(completionError);
                if (error != null) {
                    resultFuture.completeExceptionally(error);
                } else {
                    executeWork(resultFuture, tx, work);
                }
            });

            return resultFuture;
        });
    }

    private  void executeWork(
            CompletableFuture resultFuture,
            UnmanagedTransaction tx,
            @SuppressWarnings("deprecation") AsyncTransactionWork> work) {
        CompletionStage workFuture = safeExecuteWork(tx, work);
        workFuture.whenComplete((result, completionError) -> {
            Throwable error = Futures.completionExceptionCause(completionError);
            if (error != null) {
                closeTxAfterFailedTransactionWork(tx, resultFuture, error);
            } else {
                closeTxAfterSucceededTransactionWork(tx, resultFuture, result);
            }
        });
    }

    private  CompletionStage safeExecuteWork(
            UnmanagedTransaction tx, @SuppressWarnings("deprecation") AsyncTransactionWork> work) {
        // given work might fail in both async and sync way
        // async failure will result in a failed future being returned
        // sync failure will result in an exception being thrown
        try {
            CompletionStage result = work.execute(new InternalAsyncTransaction(tx));

            // protect from given transaction function returning null
            return result == null ? completedWithNull() : result;
        } catch (Throwable workError) {
            // work threw an exception, wrap it in a future and proceed
            return failedFuture(workError);
        }
    }

    private  void closeTxAfterFailedTransactionWork(
            UnmanagedTransaction tx, CompletableFuture resultFuture, Throwable error) {
        tx.closeAsync().whenComplete((ignored, rollbackError) -> {
            if (rollbackError != null) {
                error.addSuppressed(rollbackError);
            }
            resultFuture.completeExceptionally(error);
        });
    }

    private  void closeTxAfterSucceededTransactionWork(
            UnmanagedTransaction tx, CompletableFuture resultFuture, T result) {
        tx.closeAsync(true).whenComplete((ignored, completionError) -> {
            Throwable commitError = Futures.completionExceptionCause(completionError);
            if (commitError != null) {
                resultFuture.completeExceptionally(commitError);
            } else {
                resultFuture.complete(result);
            }
        });
    }
}