<commit_before>/*
 * Copyright (C) 2017-2017 DataStax Inc.
 *
 * 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 com.datastax.oss.driver.internal.core.adminrequest;

import com.datastax.oss.driver.api.core.DriverException;
import com.datastax.oss.driver.api.core.ProtocolVersion;
import com.datastax.oss.driver.internal.core.adminrequest.codec.TypeCodecs;
import com.datastax.oss.driver.internal.core.channel.DriverChannel;
import com.datastax.oss.driver.internal.core.channel.ResponseCallback;
import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions;
import com.datastax.oss.protocol.internal.Frame;
import com.datastax.oss.protocol.internal.Message;
import com.datastax.oss.protocol.internal.ProtocolConstants;
import com.datastax.oss.protocol.internal.request.Prepare;
import com.datastax.oss.protocol.internal.request.Query;
import com.datastax.oss.protocol.internal.request.query.QueryOptions;
import com.datastax.oss.protocol.internal.response.result.Prepared;
import com.datastax.oss.protocol.internal.response.result.Rows;
import com.google.common.collect.Maps;
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.ScheduledFuture;
import java.net.InetAddress;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Handles the lifecyle of an admin request. */
public class AdminRequestHandler implements ResponseCallback {
  private static final Logger LOG = LoggerFactory.getLogger(AdminRequestHandler.class);

  public static AdminRequestHandler query(
      DriverChannel channel,
      String query,
      Map<String, Object> parameters,
      Duration timeout,
      int pageSize) {
    Query message =
        new Query(
            query,
            buildQueryOptions(pageSize, serialize(parameters, channel.protocolVersion()), null));
    return new AdminRequestHandler(channel, message, timeout);
  }

  public static AdminRequestHandler query(
      DriverChannel channel, String query, Duration timeout, int pageSize) {
    return query(channel, query, Collections.emptyMap(), timeout, pageSize);
  }

  public static AdminRequestHandler prepare(DriverChannel channel, String query, Duration timeout) {
    return new AdminRequestHandler(channel, new Prepare(query), timeout);
  }

  private final DriverChannel channel;
  private final Message message;
  private final Duration timeout;
  private final CompletableFuture<AdminResult> result = new CompletableFuture<>();

  // This is only ever accessed on the channel's event loop, so it doesn't need to be volatile
  private ScheduledFuture<?> timeoutFuture;

  private AdminRequestHandler(DriverChannel channel, Message message, Duration timeout) {
    this.channel = channel;
    this.message = message;
    this.timeout = timeout;
  }

  public CompletionStage<AdminResult> start() {
    if (LOG.isDebugEnabled()) {
      if (message instanceof Query) {
        Query query = (Query) this.message;
        LOG.debug("Executing query '{}' with values {}", query.query, query.options.namedValues);
      } else if (message instanceof Prepare) {
        LOG.debug("Preparing {}", ((Prepare) message).cqlQuery);
      }
    }
    channel.write(message, false, Frame.NO_PAYLOAD, this).addListener(this::onWriteComplete);
    return result;
  }

  private void onWriteComplete(Future<? super Void> future) {
    if (future.isSuccess()) {
      timeoutFuture =
          channel.eventLoop().schedule(this::fireTimeout, timeout.toNanos(), TimeUnit.NANOSECONDS);
      timeoutFuture.addListener(UncaughtExceptions::log);
    } else {
      result.completeExceptionally(future.cause());
    }
  }

  private void fireTimeout() {
    result.completeExceptionally(new DriverException("Query timed out"));
  }

  @Override
  public void onFailure(Throwable error) {
    timeoutFuture.cancel(true);
    result.completeExceptionally(error);
  }

  @Override
  public void onResponse(Frame responseFrame) {
    timeoutFuture.cancel(true);
    Message message = responseFrame.message;
    if (message instanceof Rows) {
      Rows rows = (Rows) message;
      ByteBuffer pagingState = rows.metadata.pagingState;
      AdminRequestHandler nextHandler = (pagingState == null) ? null : this.copy(pagingState);
      result.complete(new AdminResult(rows, nextHandler, channel.protocolVersion()));
    } else if (message instanceof Prepared) {
      // Internal prepares are only "reprepare on up" types of queries, where we only care about
      // success, not the actual result, so this is good enough:<commit_msg>Check that timeout future is not null before cancelling it<commit_after>/*
 * Copyright (C) 2017-2017 DataStax Inc.
 *
 * 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 com.datastax.oss.driver.internal.core.adminrequest;

import com.datastax.oss.driver.api.core.DriverException;
import com.datastax.oss.driver.api.core.ProtocolVersion;
import com.datastax.oss.driver.internal.core.adminrequest.codec.TypeCodecs;
import com.datastax.oss.driver.internal.core.channel.DriverChannel;
import com.datastax.oss.driver.internal.core.channel.ResponseCallback;
import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions;
import com.datastax.oss.protocol.internal.Frame;
import com.datastax.oss.protocol.internal.Message;
import com.datastax.oss.protocol.internal.ProtocolConstants;
import com.datastax.oss.protocol.internal.request.Prepare;
import com.datastax.oss.protocol.internal.request.Query;
import com.datastax.oss.protocol.internal.request.query.QueryOptions;
import com.datastax.oss.protocol.internal.response.result.Prepared;
import com.datastax.oss.protocol.internal.response.result.Rows;
import com.google.common.collect.Maps;
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.ScheduledFuture;
import java.net.InetAddress;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Handles the lifecyle of an admin request. */
public class AdminRequestHandler implements ResponseCallback {
  private static final Logger LOG = LoggerFactory.getLogger(AdminRequestHandler.class);

  public static AdminRequestHandler query(
      DriverChannel channel,
      String query,
      Map<String, Object> parameters,
      Duration timeout,
      int pageSize) {
    Query message =
        new Query(
            query,
            buildQueryOptions(pageSize, serialize(parameters, channel.protocolVersion()), null));
    return new AdminRequestHandler(channel, message, timeout);
  }

  public static AdminRequestHandler query(
      DriverChannel channel, String query, Duration timeout, int pageSize) {
    return query(channel, query, Collections.emptyMap(), timeout, pageSize);
  }

  public static AdminRequestHandler prepare(DriverChannel channel, String query, Duration timeout) {
    return new AdminRequestHandler(channel, new Prepare(query), timeout);
  }

  private final DriverChannel channel;
  private final Message message;
  private final Duration timeout;
  private final CompletableFuture<AdminResult> result = new CompletableFuture<>();

  // This is only ever accessed on the channel's event loop, so it doesn't need to be volatile
  private ScheduledFuture<?> timeoutFuture;

  private AdminRequestHandler(DriverChannel channel, Message message, Duration timeout) {
    this.channel = channel;
    this.message = message;
    this.timeout = timeout;
  }

  public CompletionStage<AdminResult> start() {
    if (LOG.isDebugEnabled()) {
      if (message instanceof Query) {
        Query query = (Query) this.message;
        LOG.debug("Executing query '{}' with values {}", query.query, query.options.namedValues);
      } else if (message instanceof Prepare) {
        LOG.debug("Preparing {}", ((Prepare) message).cqlQuery);
      }
    }
    channel.write(message, false, Frame.NO_PAYLOAD, this).addListener(this::onWriteComplete);
    return result;
  }

  private void onWriteComplete(Future<? super Void> future) {
    if (future.isSuccess()) {
      timeoutFuture =
          channel.eventLoop().schedule(this::fireTimeout, timeout.toNanos(), TimeUnit.NANOSECONDS);
      timeoutFuture.addListener(UncaughtExceptions::log);
    } else {
      result.completeExceptionally(future.cause());
    }
  }

  private void fireTimeout() {
    result.completeExceptionally(new DriverException("Query timed out"));
  }

  @Override
  public void onFailure(Throwable error) {
    if (timeoutFuture != null) {
      timeoutFuture.cancel(true);
    }
    result.completeExceptionally(error);
  }

  @Override
  public void onResponse(Frame responseFrame) {
    if (timeoutFuture != null) {
      timeoutFuture.cancel(true);
    }
    Message message = responseFrame.message;
    if (message instanceof Rows) {
      Rows rows = (Rows) message;
      ByteBuffer pagingState = rows.metadata.pagingState;
      AdminRequestHandler nextHandler = (pagingState == null) ? null : this.copy(pagingState);
      result.complete(new AdminResult(rows, nextHandler, channel.protocolVersion()));
    } else if (message instanceof Prepared) {
      // Internal prepares are only "reprepare on up" types of queries, where we only care about
      // success, not the actual result, so this is good enough:
{'fraction_non_alphanumeric': 0.09782507234201437, 'fraction_numerical': 0.0026136469709698495, 'mean_word_length': 4.283037475345168, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4235267', 'n_tokens_mistral': 2962, 'n_tokens_neox': 2899, 'n_words': 804}