All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.datastax.oss.simulacron.common.stubbing.PeerMetadataHandler Maven / Gradle / Ivy

There is a newer version: 0.12.0
Show newest version
/*
 * 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.simulacron.common.stubbing;

import static com.datastax.oss.protocol.internal.ProtocolConstants.DataType.ASCII;
import static com.datastax.oss.protocol.internal.ProtocolConstants.DataType.BOOLEAN;
import static com.datastax.oss.protocol.internal.ProtocolConstants.DataType.INET;
import static com.datastax.oss.protocol.internal.ProtocolConstants.DataType.UUID;

import com.datastax.oss.protocol.internal.Frame;
import com.datastax.oss.protocol.internal.request.Query;
import com.datastax.oss.protocol.internal.response.result.ColumnSpec;
import com.datastax.oss.protocol.internal.response.result.DefaultRows;
import com.datastax.oss.protocol.internal.response.result.RawType;
import com.datastax.oss.protocol.internal.response.result.Rows;
import com.datastax.oss.protocol.internal.response.result.RowsMetadata;
import com.datastax.oss.simulacron.common.cluster.AbstractCluster;
import com.datastax.oss.simulacron.common.cluster.AbstractNode;
import com.datastax.oss.simulacron.common.codec.Codec;
import com.datastax.oss.simulacron.common.codec.CodecUtils;
import com.datastax.oss.simulacron.common.codec.CqlMapper;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Queue;
import java.util.Set;
import java.util.UUID;
import java.util.function.Predicate;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class PeerMetadataHandler extends StubMapping implements InternalStubMapping {

  private static final List queries = new ArrayList<>();
  private static final List queryPatterns = new ArrayList<>();

  static final UUID schemaVersion = java.util.UUID.randomUUID();

  private static final String queryClusterName = "select cluster_name from system.local";
  private static final RowsMetadata queryClusterNameMetadata;

  static {
    CodecUtils.ColumnSpecBuilder systemLocal = CodecUtils.columnSpecBuilder("system", "local");
    List queryClusterNameSpecs =
        CodecUtils.columnSpecs(systemLocal.apply("cluster_name", CodecUtils.primitive(ASCII)));
    queryClusterNameMetadata = new RowsMetadata(queryClusterNameSpecs, null, new int[] {});
  }

  private static final Pattern queryPeers = Pattern.compile("SELECT (.*) FROM system\\.peers");
  private static final Pattern queryLocal =
      Pattern.compile("SELECT (.*) FROM system\\.local( WHERE key='local')*");
  private static final Pattern queryPeersWithAddr =
      Pattern.compile("SELECT \\* FROM system\\.peers WHERE peer='(.*)'");

  // query the java driver makes when refreshing node (i.e. after it comes back up)
  private static final String queryPeerWithNamedParam =
      "SELECT * FROM system.peers WHERE peer = :address";

  static {
    queries.add(queryClusterName);
    queries.add(queryPeerWithNamedParam);
  }

  static {
    queryPatterns.add(queryPeers);
    queryPatterns.add(queryLocal);
    queryPatterns.add(queryPeersWithAddr);
  }

  public PeerMetadataHandler() {}

  @Override
  public boolean matches(Frame frame) {
    if (frame.message instanceof Query) {
      Query query = (Query) frame.message;
      String queryStr = query.query;
      return queries.stream().anyMatch(q -> q.equalsIgnoreCase(queryStr))
          || queryPatternMatches(queryStr);
    }
    return false;
  }

  @Override
  public List getActions(AbstractNode node, Frame frame) {
    if (frame.message instanceof Query) {
      CqlMapper mapper = CqlMapper.forVersion(frame.protocolVersion);
      Query query = (Query) frame.message;

      if (query.query.equalsIgnoreCase(queryClusterName)) {
        return handleClusterNameQuery(node, mapper);
      } else {
        // if querying for particular peer, return information for only that peer.
        final Matcher peerAddrMatcher = queryPeersWithAddr.matcher(query.query);
        if (peerAddrMatcher.matches()) {
          return handlePeersQuery(
              node,
              mapper,
              n -> {
                InetAddress address;
                if (n.getAddress() instanceof InetSocketAddress) {
                  address = ((InetSocketAddress) n.getAddress()).getAddress();
                  String addrIp = address.getHostAddress();
                  return addrIp.equals(peerAddrMatcher.group(1));
                } else {
                  return false;
                }
              });
        } else if (query.query.equalsIgnoreCase(queryPeerWithNamedParam)) {
          ByteBuffer addressBuffer = query.options.namedValues.get("address");
          InetAddress address = mapper.inet.decode(addressBuffer);
          return handlePeersQuery(node, mapper, n -> n.inet().equals(address));
        }
        Matcher matcher = queryLocal.matcher(query.query);
        if (matcher.matches()) {
          return handleSystemLocalQuery(node, mapper);
        }
        matcher = queryPeers.matcher(query.query);
        if (matcher.matches()) {
          return handlePeersQuery(node, mapper, n -> n != node);
        }
      }
    }
    return Collections.emptyList();
  }

  private boolean queryPatternMatches(String query) {
    for (Pattern pattern : queryPatterns) {
      if (pattern.matcher(query).matches()) {
        return true;
      }
    }
    return false;
  }

  private Set resolveTokens(AbstractNode node) {
    String[] t = node.resolvePeerInfo("tokens", "0").split(",");
    return new LinkedHashSet<>(Arrays.asList(t));
  }

  private List handleSystemLocalQuery(AbstractNode node, CqlMapper mapper) {
    InetAddress address = resolveAddress(node);
    Codec> tokenCodec =
        mapper.codecFor(new RawType.RawSet(CodecUtils.primitive(ASCII)));

    List localRow =
        CodecUtils.row(
            CodecUtils.encodePeerInfo(node, mapper.ascii::encode, "key", "local"),
            CodecUtils.encodePeerInfo(node, mapper.ascii::encode, "bootstrapped", "COMPLETED"),
            mapper.inet.encode(node.resolvePeerInfo("rpc_address", address)),
            mapper.inet.encode(node.resolvePeerInfo("broadcast_address", address)),
            mapper.ascii.encode(node.getCluster().getName()),
            CodecUtils.encodePeerInfo(node, mapper.ascii::encode, "cql_version", "3.2.0"),
            mapper.ascii.encode(node.getDataCenter().getName()),
            mapper.inet.encode(address),
            CodecUtils.encodePeerInfo(
                node,
                mapper.ascii::encode,
                "partitioner",
                "org.apache.cassandra.dht.Murmur3Partitioner"),
            CodecUtils.encodePeerInfo(node, mapper.ascii::encode, "rack", "rack1"),
            mapper.ascii.encode(node.resolveCassandraVersion()),
            tokenCodec.encode(resolveTokens(node)),
            mapper.uuid.encode(schemaVersion));

    if (node.resolveDSEVersion() != null) {
      localRow.add(mapper.ascii.encode(node.resolveDSEVersion()));
      localRow.add(CodecUtils.encodePeerInfo(node, mapper.bool::encode, "graph", false));
    }

    Rows rows = new DefaultRows(buildSystemLocalRowsMetadata(node), CodecUtils.rows(localRow));
    MessageResponseAction action = new MessageResponseAction(rows);
    return Collections.singletonList(action);
  }

  private List handleClusterNameQuery(AbstractNode node, CqlMapper mapper) {
    Queue> clusterRow =
        CodecUtils.singletonRow(mapper.ascii.encode(node.getCluster().getName()));
    Rows rows = new DefaultRows(queryClusterNameMetadata, clusterRow);
    MessageResponseAction action = new MessageResponseAction(rows);
    return Collections.singletonList(action);
  }

  @SuppressWarnings("unchecked")
  private List handlePeersQuery(
      AbstractNode node, CqlMapper mapper, Predicate nodeFilter) {
    // For each node matching the filter, provide its peer information.
    Codec> tokenCodec =
        mapper.codecFor(new RawType.RawSet(CodecUtils.primitive(ASCII)));

    AbstractCluster cluster = node.getCluster();
    Stream stream = cluster.getNodes().stream();

    Queue> peerRows =
        new ArrayDeque<>(
            stream
                .filter(nodeFilter)
                .map(
                    n -> {
                      InetAddress address = resolveAddress(n);

                      List row =
                          CodecUtils.row(
                              mapper.inet.encode(n.resolvePeerInfo("peer", address)),
                              mapper.inet.encode(n.resolvePeerInfo("rpc_address", address)),
                              mapper.varchar.encode(
                                  n.resolvePeerInfo("data_center", n.getDataCenter().getName())),
                              CodecUtils.encodePeerInfo(n, mapper.varchar::encode, "rack", "rack1"),
                              mapper.varchar.encode(
                                  n.resolvePeerInfo(
                                      "release_version", n.resolveCassandraVersion())),
                              tokenCodec.encode(resolveTokens(n)),
                              mapper.inet.encode(n.resolvePeerInfo("listen_address", address)),
                              mapper.uuid.encode(n.resolvePeerInfo("host_id", schemaVersion)),
                              mapper.uuid.encode(
                                  n.resolvePeerInfo("schema_version", schemaVersion)));
                      if (node.resolveDSEVersion() != null) {
                        row.add(mapper.ascii.encode(n.resolveDSEVersion()));
                        row.add(CodecUtils.encodePeerInfo(n, mapper.bool::encode, "graph", false));
                      }
                      return row;
                    })
                .collect(Collectors.toList()));

    Rows rows = new DefaultRows(buildSystemPeersRowsMetadata(node), peerRows);
    MessageResponseAction action = new MessageResponseAction(rows);
    return Collections.singletonList(action);
  }

  private InetAddress resolveAddress(AbstractNode node) {
    InetAddress address;
    if (node.getAddress() instanceof InetSocketAddress) {
      address = ((InetSocketAddress) node.getAddress()).getAddress();
    } else {
      address = InetAddress.getLoopbackAddress();
    }
    return address;
  }

  private RowsMetadata buildSystemPeersRowsMetadata(AbstractNode node) {
    CodecUtils.ColumnSpecBuilder systemPeers = CodecUtils.columnSpecBuilder("system", "peers");
    List systemPeersSpecs =
        CodecUtils.columnSpecs(
            systemPeers.apply("peer", CodecUtils.primitive(INET)),
            systemPeers.apply("rpc_address", CodecUtils.primitive(INET)),
            systemPeers.apply("data_center", CodecUtils.primitive(ASCII)),
            systemPeers.apply("rack", CodecUtils.primitive(ASCII)),
            systemPeers.apply("release_version", CodecUtils.primitive(ASCII)),
            systemPeers.apply("tokens", new RawType.RawSet(CodecUtils.primitive(ASCII))),
            systemPeers.apply("listen_address", CodecUtils.primitive(INET)),
            systemPeers.apply("host_id", CodecUtils.primitive(UUID)),
            systemPeers.apply("schema_version", CodecUtils.primitive(UUID)));
    if (node.resolveDSEVersion() != null) {
      systemPeersSpecs.add(systemPeers.apply("dse_version", CodecUtils.primitive(ASCII)));
      systemPeersSpecs.add(systemPeers.apply("graph", CodecUtils.primitive(BOOLEAN)));
    }
    return new RowsMetadata(systemPeersSpecs, null, new int[] {0});
  }

  private RowsMetadata buildSystemLocalRowsMetadata(AbstractNode node) {
    CodecUtils.ColumnSpecBuilder systemLocal = CodecUtils.columnSpecBuilder("system", "local");
    List systemLocalSpecs =
        CodecUtils.columnSpecs(
            systemLocal.apply("key", CodecUtils.primitive(ASCII)),
            systemLocal.apply("bootstrapped", CodecUtils.primitive(ASCII)),
            systemLocal.apply("rpc_address", CodecUtils.primitive(INET)),
            systemLocal.apply("broadcast_address", CodecUtils.primitive(INET)),
            systemLocal.apply("cluster_name", CodecUtils.primitive(ASCII)),
            systemLocal.apply("cql_version", CodecUtils.primitive(ASCII)),
            systemLocal.apply("data_center", CodecUtils.primitive(ASCII)),
            systemLocal.apply("listen_address", CodecUtils.primitive(INET)),
            systemLocal.apply("partitioner", CodecUtils.primitive(ASCII)),
            systemLocal.apply("rack", CodecUtils.primitive(ASCII)),
            systemLocal.apply("release_version", CodecUtils.primitive(ASCII)),
            systemLocal.apply("tokens", new RawType.RawSet(CodecUtils.primitive(ASCII))),
            systemLocal.apply("schema_version", CodecUtils.primitive(UUID)));
    if (node.resolveDSEVersion() != null) {
      systemLocalSpecs.add(systemLocal.apply("dse_version", CodecUtils.primitive(ASCII)));
      systemLocalSpecs.add(systemLocal.apply("graph", CodecUtils.primitive(BOOLEAN)));
    }
    return new RowsMetadata(systemLocalSpecs, null, new int[] {0});
  }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy