<commit_before>/*
 * Copyright (c) 2010-2020. Axon Framework
 *
 * 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.axonframework.queryhandling;

import org.axonframework.commandhandling.CommandResultMessage;
import org.axonframework.common.AxonConfigurationException;
import org.axonframework.common.Registration;
import org.axonframework.messaging.ResultMessage;
import org.axonframework.messaging.reactive.ReactiveMessageDispatchInterceptor;
import org.axonframework.messaging.reactive.ReactiveResultHandlerInterceptor;
import org.axonframework.messaging.responsetypes.ResponseType;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.function.Tuple2;

import java.util.List;
import java.util.Objects;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;

import static java.util.Arrays.asList;
import static org.axonframework.common.BuilderUtils.assertNonNull;
import static org.axonframework.messaging.GenericMessage.asMessage;

/**
 * Implementation of the {@link ReactiveQueryGateway} that uses Project Reactor to achieve reactiveness.
 *
 * @author <NAME>
 * @since 4.4
 */
public class ReactorQueryGateway implements ReactiveQueryGateway {

    private final List<ReactiveMessageDispatchInterceptor<QueryMessage<?, ?>>> dispatchInterceptors;
    private final List<ReactiveResultHandlerInterceptor<QueryMessage<?, ?>, ResultMessage<?>>> resultInterceptors;

    private final QueryBus queryBus;

    /**
     * Creates an instance of {@link ReactorQueryGateway} based on the fields contained in the {@link
     * Builder}.
     * <p>
     * Will assert that the {@link QueryBus} is not {@code null} and throws an {@link AxonConfigurationException} if
     * it is.
     * </p>
     *
     * @param builder the {@link Builder} used to instantiated a {@link ReactorQueryGateway} instance
     */
    protected ReactorQueryGateway(Builder builder) {
        builder.validate();
        this.queryBus = builder.queryBus;
        this.dispatchInterceptors = builder.dispatchInterceptors;
        this.resultInterceptors = builder.resultInterceptors;
    }

    /**
     * Instantiate a Builder to be able to create a {@link ReactorQueryGateway}.
     * <p>
     * The {@code dispatchInterceptors} are defaulted to an empty list.
     * The {@link QueryBus} is a <b>hard requirements</b> and as such should be provided.
     * </p>
     *
     * @return a Builder to be able to create a {@link ReactorQueryGateway}
     */
    public static Builder builder() {
        return new Builder();
    }

    @Override
    public Registration registerDispatchInterceptor(
            ReactiveMessageDispatchInterceptor<QueryMessage<?, ?>> interceptor) {
        dispatchInterceptors.add(interceptor);
        return () -> dispatchInterceptors.remove(interceptor);
    }

    @Override
    public Registration registerResultHandlerInterceptor(
            ReactiveResultHandlerInterceptor<QueryMessage<?, ?>, ResultMessage<?>> interceptor) {
        resultInterceptors.add(interceptor);
        return () -> resultInterceptors.remove(interceptor);
    }

    @Override
    public <R, Q> Mono<R> query(String queryName, Q query, ResponseType<R> responseType) {
        return Mono.<QueryMessage<?, ?>>fromCallable(() -> new GenericQueryMessage<>(asMessage(query),
                                                                                     queryName,
                                                                                     responseType))
                .transform(this::processDispatchInterceptors)
                .flatMap(this::dispatchQuery)
                .flatMapMany(this::processResultsInterceptors)
                .<R>transform(this::getPayload)
                .next();
    }

    @Override
    public <R, Q> Flux<R> scatterGather(String queryName, Q query, ResponseType<R> responseType, long timeout,
                                        TimeUnit timeUnit) {
        return Mono.<QueryMessage<?, ?>>fromCallable(() -> new GenericQueryMessage<>(asMessage(query),
                                                                                     queryName,
                                                                                     responseType))
                .transform(this::processDispatchInterceptors)
                .flatMap(q -> dispatchScatterGatherQuery(q, timeout, timeUnit))
                .flatMapMany(this::processResultsInterceptors)
                .transform(this::getPayload);
    }

    @Override
    public <Q, I, U> Mono<SubscriptionQueryResult<I, U>> subscriptionQuery(String queryName, Q query,
                                                                           ResponseType<I> initialResponseType,
                                                                           ResponseType<U> updateResponseType,
                                                                           SubscriptionQueryBackpressure backpressure,
                                                                           int updateBufferSize) {

        //noinspection unchecked
        return Mono.<QueryMessage<?, ?>>fromCallable(() -> new GenericSubscriptionQueryMessage<>(query,
                                                                                                 initialResponseType,
                                                                                                 updateResponseType))
                .transform(this::processDispatchInterceptors)
                .map(isq -> (SubscriptionQueryMessage<Q, U, I>) isq)
                .flatMap(isq -> dispatchSubscriptionQuery(isq, backpressure, updateBufferSize))
                .flatMap(processSubscriptionQueryResult());
    }


    private Mono<Tuple2<QueryMessage<?, ?>, Flux<ResultMessage<?>>>> dispatchQuery(QueryMessage<?, ?> queryMessage) {
        Flux<ResultMessage<?>> results = Flux
                .defer(() -> Mono.fromFuture(queryBus.query(queryMessage)));

        return Mono.<QueryMessage<?, ?>>just(queryMessage)
                .zipWith(Mono.just(results));
    }

    private Mono<Tuple2<QueryMessage<?, ?>, Flux<ResultMessage<?>>>> dispatchScatterGatherQuery(
            QueryMessage<?, ?> queryMessage, long timeout, TimeUnit timeUnit) {
        Flux<ResultMessage<?>> results = Flux
                .defer(() -> Flux.fromStream(queryBus.scatterGather(queryMessage,
                                                                    timeout,
                                                                    timeUnit)));

        return Mono.<QueryMessage<?, ?>>just(queryMessage)
                .zipWith(Mono.just(results));
    }

    private <Q, I, U> Mono<Tuple2<QueryMessage<Q, I>, Mono<SubscriptionQueryResult<QueryResponseMessage<I>, SubscriptionQueryUpdateMessage<U>>>>> dispatchSubscriptionQuery(
            SubscriptionQueryMessage<Q, I, U> queryMessage, SubscriptionQueryBackpressure backpressure,
            int updateBufferSize) {
        Mono<SubscriptionQueryResult<QueryResponseMessage<I>, SubscriptionQueryUpdateMessage<U>>> result = Mono
                .fromCallable(() -> {
                    System.out.println(Thread.currentThread() + " before dispatch");
                    return queryBus.subscriptionQuery(queryMessage, backpressure, updateBufferSize);
                });
        return Mono.<QueryMessage<Q, I>>just(queryMessage)
                .zipWith(Mono.just(result));
    }

    private Mono<QueryMessage<?, ?>> processDispatchInterceptors(Mono<QueryMessage<?, ?>> queryMessageMono) {
        return Flux.fromIterable(dispatchInterceptors)
                   .reduce(queryMessageMono, (queryMessage, interceptor) -> interceptor.intercept(queryMessage))
                   .flatMap(Mono::from);
    }

    private Flux<ResultMessage<?>> processResultsInterceptors(
            Tuple2<QueryMessage<?, ?>, Flux<ResultMessage<?>>> queryWithResponses) {
        QueryMessage<?, ?> queryMessage = queryWithResponses.getT1();
        Flux<ResultMessage<?>> queryResultMessage = queryWithResponses.getT2();

        return Flux.fromIterable(resultInterceptors)
                   .reduce(queryResultMessage,
                           (result, interceptor) -> interceptor.intercept(queryMessage, result))
                   .flatMapMany(it -> it);<commit_msg>Remove system.out<commit_after>/*
 * Copyright (c) 2010-2020. Axon Framework
 *
 * 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.axonframework.queryhandling;

import org.axonframework.commandhandling.CommandResultMessage;
import org.axonframework.common.AxonConfigurationException;
import org.axonframework.common.Registration;
import org.axonframework.messaging.ResultMessage;
import org.axonframework.messaging.reactive.ReactiveMessageDispatchInterceptor;
import org.axonframework.messaging.reactive.ReactiveResultHandlerInterceptor;
import org.axonframework.messaging.responsetypes.ResponseType;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.function.Tuple2;

import java.util.List;
import java.util.Objects;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;

import static java.util.Arrays.asList;
import static org.axonframework.common.BuilderUtils.assertNonNull;
import static org.axonframework.messaging.GenericMessage.asMessage;

/**
 * Implementation of the {@link ReactiveQueryGateway} that uses Project Reactor to achieve reactiveness.
 *
 * @author <NAME>
 * @since 4.4
 */
public class ReactorQueryGateway implements ReactiveQueryGateway {

    private final List<ReactiveMessageDispatchInterceptor<QueryMessage<?, ?>>> dispatchInterceptors;
    private final List<ReactiveResultHandlerInterceptor<QueryMessage<?, ?>, ResultMessage<?>>> resultInterceptors;

    private final QueryBus queryBus;

    /**
     * Creates an instance of {@link ReactorQueryGateway} based on the fields contained in the {@link
     * Builder}.
     * <p>
     * Will assert that the {@link QueryBus} is not {@code null} and throws an {@link AxonConfigurationException} if
     * it is.
     * </p>
     *
     * @param builder the {@link Builder} used to instantiated a {@link ReactorQueryGateway} instance
     */
    protected ReactorQueryGateway(Builder builder) {
        builder.validate();
        this.queryBus = builder.queryBus;
        this.dispatchInterceptors = builder.dispatchInterceptors;
        this.resultInterceptors = builder.resultInterceptors;
    }

    /**
     * Instantiate a Builder to be able to create a {@link ReactorQueryGateway}.
     * <p>
     * The {@code dispatchInterceptors} are defaulted to an empty list.
     * The {@link QueryBus} is a <b>hard requirements</b> and as such should be provided.
     * </p>
     *
     * @return a Builder to be able to create a {@link ReactorQueryGateway}
     */
    public static Builder builder() {
        return new Builder();
    }

    @Override
    public Registration registerDispatchInterceptor(
            ReactiveMessageDispatchInterceptor<QueryMessage<?, ?>> interceptor) {
        dispatchInterceptors.add(interceptor);
        return () -> dispatchInterceptors.remove(interceptor);
    }

    @Override
    public Registration registerResultHandlerInterceptor(
            ReactiveResultHandlerInterceptor<QueryMessage<?, ?>, ResultMessage<?>> interceptor) {
        resultInterceptors.add(interceptor);
        return () -> resultInterceptors.remove(interceptor);
    }

    @Override
    public <R, Q> Mono<R> query(String queryName, Q query, ResponseType<R> responseType) {
        return Mono.<QueryMessage<?, ?>>fromCallable(() -> new GenericQueryMessage<>(asMessage(query),
                                                                                     queryName,
                                                                                     responseType))
                .transform(this::processDispatchInterceptors)
                .flatMap(this::dispatchQuery)
                .flatMapMany(this::processResultsInterceptors)
                .<R>transform(this::getPayload)
                .next();
    }

    @Override
    public <R, Q> Flux<R> scatterGather(String queryName, Q query, ResponseType<R> responseType, long timeout,
                                        TimeUnit timeUnit) {
        return Mono.<QueryMessage<?, ?>>fromCallable(() -> new GenericQueryMessage<>(asMessage(query),
                                                                                     queryName,
                                                                                     responseType))
                .transform(this::processDispatchInterceptors)
                .flatMap(q -> dispatchScatterGatherQuery(q, timeout, timeUnit))
                .flatMapMany(this::processResultsInterceptors)
                .transform(this::getPayload);
    }

    @Override
    public <Q, I, U> Mono<SubscriptionQueryResult<I, U>> subscriptionQuery(String queryName, Q query,
                                                                           ResponseType<I> initialResponseType,
                                                                           ResponseType<U> updateResponseType,
                                                                           SubscriptionQueryBackpressure backpressure,
                                                                           int updateBufferSize) {

        //noinspection unchecked
        return Mono.<QueryMessage<?, ?>>fromCallable(() -> new GenericSubscriptionQueryMessage<>(query,
                                                                                                 initialResponseType,
                                                                                                 updateResponseType))
                .transform(this::processDispatchInterceptors)
                .map(isq -> (SubscriptionQueryMessage<Q, U, I>) isq)
                .flatMap(isq -> dispatchSubscriptionQuery(isq, backpressure, updateBufferSize))
                .flatMap(processSubscriptionQueryResult());
    }


    private Mono<Tuple2<QueryMessage<?, ?>, Flux<ResultMessage<?>>>> dispatchQuery(QueryMessage<?, ?> queryMessage) {
        Flux<ResultMessage<?>> results = Flux
                .defer(() -> Mono.fromFuture(queryBus.query(queryMessage)));

        return Mono.<QueryMessage<?, ?>>just(queryMessage)
                .zipWith(Mono.just(results));
    }

    private Mono<Tuple2<QueryMessage<?, ?>, Flux<ResultMessage<?>>>> dispatchScatterGatherQuery(
            QueryMessage<?, ?> queryMessage, long timeout, TimeUnit timeUnit) {
        Flux<ResultMessage<?>> results = Flux
                .defer(() -> Flux.fromStream(queryBus.scatterGather(queryMessage,
                                                                    timeout,
                                                                    timeUnit)));

        return Mono.<QueryMessage<?, ?>>just(queryMessage)
                .zipWith(Mono.just(results));
    }

    private <Q, I, U> Mono<Tuple2<QueryMessage<Q, I>, Mono<SubscriptionQueryResult<QueryResponseMessage<I>, SubscriptionQueryUpdateMessage<U>>>>> dispatchSubscriptionQuery(
            SubscriptionQueryMessage<Q, I, U> queryMessage, SubscriptionQueryBackpressure backpressure,
            int updateBufferSize) {
        Mono<SubscriptionQueryResult<QueryResponseMessage<I>, SubscriptionQueryUpdateMessage<U>>> result = Mono
                .fromCallable(() -> queryBus.subscriptionQuery(queryMessage, backpressure, updateBufferSize));
        return Mono.<QueryMessage<Q, I>>just(queryMessage)
                .zipWith(Mono.just(result));
    }

    private Mono<QueryMessage<?, ?>> processDispatchInterceptors(Mono<QueryMessage<?, ?>> queryMessageMono) {
        return Flux.fromIterable(dispatchInterceptors)
                   .reduce(queryMessageMono, (queryMessage, interceptor) -> interceptor.intercept(queryMessage))
                   .flatMap(Mono::from);
    }

    private Flux<ResultMessage<?>> processResultsInterceptors(
            Tuple2<QueryMessage<?, ?>, Flux<ResultMessage<?>>> queryWithResponses) {
        QueryMessage<?, ?> queryMessage = queryWithResponses.getT1();
        Flux<ResultMessage<?>> queryResultMessage = queryWithResponses.getT2();

        return Flux.fromIterable(resultInterceptors)
                   .reduce(queryResultMessage,
                           (result, interceptor) -> interceptor.intercept(queryMessage, result))
                   .flatMapMany(it -> it);
{'fraction_non_alphanumeric': 0.09454483657168754, 'fraction_numerical': 0.0023791990030023226, 'mean_word_length': 2.262010347376201, 'pattern_counts': {'":': 0, '<': 181, '<?xml version=': 0, '>': 209, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 74, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3754280', 'n_tokens_mistral': 4232, 'n_tokens_neox': 4054, 'n_words': 954}