com.datastax.oss.driver.internal.core.metadata.schema.queries.CassandraSchemaQueries Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of java-driver-core
The newest version!
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.metadata.schema.queries;

import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.internal.core.adminrequest.AdminRequestHandler;
import com.datastax.oss.driver.internal.core.adminrequest.AdminResult;
import com.datastax.oss.driver.internal.core.adminrequest.AdminRow;
import com.datastax.oss.driver.internal.core.channel.DriverChannel;
import com.datastax.oss.driver.internal.core.util.NanoTime;
import com.datastax.oss.driver.internal.core.util.concurrent.RunOrSchedule;
import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
import io.netty.util.concurrent.EventExecutor;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.function.Function;
import net.jcip.annotations.ThreadSafe;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@ThreadSafe
public abstract class CassandraSchemaQueries implements SchemaQueries {

  private static final Logger LOG = LoggerFactory.getLogger(CassandraSchemaQueries.class);

  private final DriverChannel channel;
  private final EventExecutor adminExecutor;
  private final Node node;
  private final String logPrefix;
  private final Duration timeout;
  private final int pageSize;
  private final KeyspaceFilter keyspaceFilter;
  // The future we return from execute, completes when all the queries are done.
  private final CompletableFuture schemaRowsFuture = new CompletableFuture<>();
  private final long startTimeNs = System.nanoTime();

  // All non-final fields are accessed exclusively on adminExecutor
  private CassandraSchemaRows.Builder schemaRowsBuilder;
  private int pendingQueries;

  protected CassandraSchemaQueries(
      DriverChannel channel, Node node, DriverExecutionProfile config, String logPrefix) {
    this.channel = channel;
    this.adminExecutor = channel.eventLoop();
    this.node = node;
    this.logPrefix = logPrefix;
    this.timeout = config.getDuration(DefaultDriverOption.METADATA_SCHEMA_REQUEST_TIMEOUT);
    this.pageSize = config.getInt(DefaultDriverOption.METADATA_SCHEMA_REQUEST_PAGE_SIZE);

    List refreshedKeyspaces =
        config.getStringList(
            DefaultDriverOption.METADATA_SCHEMA_REFRESHED_KEYSPACES, Collections.emptyList());
    assert refreshedKeyspaces != null; // per the default value
    this.keyspaceFilter = KeyspaceFilter.newInstance(logPrefix, refreshedKeyspaces);
  }

  protected abstract String selectKeyspacesQuery();

  protected abstract Optional selectVirtualKeyspacesQuery();

  protected abstract String selectTablesQuery();

  protected abstract Optional selectVirtualTablesQuery();

  protected abstract Optional selectViewsQuery();

  protected abstract Optional selectIndexesQuery();

  protected abstract String selectColumnsQuery();

  protected abstract Optional selectVirtualColumnsQuery();

  protected abstract String selectTypesQuery();

  protected abstract Optional selectFunctionsQuery();

  protected abstract Optional selectAggregatesQuery();

  protected abstract Optional selectEdgesQuery();

  protected abstract Optional selectVerticiesQuery();

  @Override
  public CompletionStage execute() {
    RunOrSchedule.on(adminExecutor, this::executeOnAdminExecutor);
    return schemaRowsFuture;
  }

  private void executeOnAdminExecutor() {
    assert adminExecutor.inEventLoop();

    schemaRowsBuilder = new CassandraSchemaRows.Builder(node, keyspaceFilter, logPrefix);
    String whereClause = keyspaceFilter.getWhereClause();

    query(selectKeyspacesQuery() + whereClause, schemaRowsBuilder::withKeyspaces);
    query(selectTypesQuery() + whereClause, schemaRowsBuilder::withTypes);
    query(selectTablesQuery() + whereClause, schemaRowsBuilder::withTables);
    query(selectColumnsQuery() + whereClause, schemaRowsBuilder::withColumns);
    selectIndexesQuery()
        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withIndexes));
    selectViewsQuery()
        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withViews));
    selectFunctionsQuery()
        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withFunctions));
    selectAggregatesQuery()
        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withAggregates));
    selectVirtualKeyspacesQuery()
        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withVirtualKeyspaces));
    selectVirtualTablesQuery()
        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withVirtualTables));
    selectVirtualColumnsQuery()
        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withVirtualColumns));
    selectEdgesQuery()
        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withEdges));
    selectVerticiesQuery()
        .ifPresent(select -> query(select + whereClause, schemaRowsBuilder::withVertices));
  }

  private void query(
      String queryString,
      Function, CassandraSchemaRows.Builder> builderUpdater) {
    assert adminExecutor.inEventLoop();

    pendingQueries += 1;
    query(queryString)
        .whenCompleteAsync(
            (result, error) -> handleResult(result, error, builderUpdater), adminExecutor);
  }

  @VisibleForTesting
  protected CompletionStage query(String query) {
    return AdminRequestHandler.query(channel, query, timeout, pageSize, logPrefix).start();
  }

  private void handleResult(
      AdminResult result,
      Throwable error,
      Function, CassandraSchemaRows.Builder> builderUpdater) {

    // If another query already failed, we've already propagated the failure so just ignore this one
    if (schemaRowsFuture.isCompletedExceptionally()) {
      return;
    }

    if (error != null) {
      schemaRowsFuture.completeExceptionally(error);
    } else {
      // Store the rows of the current page in the builder
      schemaRowsBuilder = builderUpdater.apply(result);
      if (result.hasNextPage()) {
        result
            .nextPage()
            .whenCompleteAsync(
                (nextResult, nextError) -> handleResult(nextResult, nextError, builderUpdater),
                adminExecutor);
      } else {
        pendingQueries -= 1;
        if (pendingQueries == 0) {
          LOG.debug(
              "[{}] Schema queries took {}", logPrefix, NanoTime.formatTimeSince(startTimeNs));
          schemaRowsFuture.complete(schemaRowsBuilder.build());
        }
      }
    }
  }
}