CASSANDRA-21513: Add virtual tables for data placements and partition locations

src/java/org/apache/cassandra/db/virtual/DataPlacementsTable.java

@@ -0,0 +1,212 @@
+/*
+ * 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 org.apache.cassandra.db.virtual;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import com.google.common.base.Preconditions;
+
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.marshal.BytesType;
+import org.apache.cassandra.db.marshal.CompositeType;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.db.marshal.SetType;
+import org.apache.cassandra.db.marshal.UTF8Type;
+import org.apache.cassandra.dht.IPartitioner;
+import org.apache.cassandra.dht.LocalPartitioner;
+import org.apache.cassandra.dht.Range;
+import org.apache.cassandra.dht.ReversedLongLocalPartitioner;
+import org.apache.cassandra.dht.Token;
+import org.apache.cassandra.locator.EndpointsForRange;
+import org.apache.cassandra.locator.Replica;
+import org.apache.cassandra.schema.KeyspaceMetadata;
+import org.apache.cassandra.schema.ReplicationParams;
+import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.tcm.ClusterMetadata;
+import org.apache.cassandra.tcm.membership.NodeId;
+import org.apache.cassandra.tcm.ownership.DataPlacement;
+import org.apache.cassandra.tcm.ownership.DataPlacements;
+import org.apache.cassandra.tcm.ownership.VersionedEndpoints;
+
+final class DataPlacementsTable extends AbstractVirtualTable
+{
+    static final String TABLE_COMMENT = "Data placement information showing read and write endpoints per range";
+
+    static final String TABLE_NAME = "data_placements";
+
+    // Partition key columns
+    static final String KEYSPACE_NAME = "keyspace_name";
+    static final String TABLE_NAME_COLUMN = "table_name";
+    static final String RANGE_START = "range_start";
+    static final String RANGE_END = "range_end";
+
+    // Regular columns
+    static final String RANGE_START_BYTES = "range_start_bytes";
+    static final String RANGE_END_BYTES = "range_end_bytes";
+    static final String TOKEN_TYPE = "token_type";
+    static final String READ_ENDPOINTS = "read_endpoints";
+    static final String WRITE_ENDPOINTS = "write_endpoints";
+    static final String READ_REPLICAS = "read_replicas";
+    static final String WRITE_REPLICAS = "write_replicas";
+
+    DataPlacementsTable(String keyspace)
+    {
+        super(buildTableMetadata(keyspace));
+    }
+
+    @Override
+    public DataSet data()
+    {
+        SimpleDataSet result = new SimpleDataSet(metadata());
+
+        ClusterMetadata metadata = ClusterMetadata.current();
+        DataPlacements placements = metadata.placements;
+
+        for (KeyspaceMetadata ksm : metadata.schema.getKeyspaces())
+        {
+            ReplicationParams params = ksm.params.replication;
+            DataPlacement placement = placements.get(params);
+
+            for (TableMetadata table : ksm.tables)
+                addPlacementRows(result, ksm.name, table.name, params, placement, metadata, null, null);
+        }
+
+        return result;
+    }
+
+    @Override
+    public DataSet data(DecoratedKey partitionKey)
+    {
+        SimpleDataSet result = new SimpleDataSet(metadata());
+
+        ByteBuffer keyBytes = partitionKey.getKey();
+
+        Preconditions.checkArgument(metadata().partitionKeyType instanceof CompositeType,
+                                    "Expected CompositeType partition key, got %s", metadata().partitionKeyType.getClass());
+
+        ByteBuffer[] components = ((CompositeType) metadata().partitionKeyType).split(keyBytes);
+        Preconditions.checkArgument(components.length == 1 || components.length == 2 || components.length == 4,
+                                    "Expected 1, 2, or 4 partition key components (keyspace[, table[, range_start, range_end]]), got %d",
+                                    components.length);
+
+        String keyspaceName = UTF8Type.instance.compose(components[0]);
+        String tableName = components.length >= 2 ? UTF8Type.instance.compose(components[1]) : null;
+        String rangeStartFilter = components.length == 4 ? UTF8Type.instance.compose(components[2]) : null;
+        String rangeEndFilter = components.length == 4 ? UTF8Type.instance.compose(components[3]) : null;
+
+        ClusterMetadata metadata = ClusterMetadata.current();
+        KeyspaceMetadata ksm = metadata.schema.getKeyspaces().getNullable(keyspaceName);
+        if (ksm == null)
+            return result;
+
+        ReplicationParams params = ksm.params.replication;
+        DataPlacement placement = metadata.placements.get(params);
+        Preconditions.checkState(placement != null, "Placements should never be null for %s", keyspaceName);
+
+        if (tableName != null)
+        {
+            TableMetadata table = ksm.getTableOrViewNullable(tableName);
+            Preconditions.checkArgument(table != null, "Couldn't find table %s", tableName);
+            addPlacementRows(result, ksm.name, table.name, params, placement, metadata, rangeStartFilter, rangeEndFilter);
+        }
+        else
+        {
+            for (TableMetadata table : ksm.tables)
+                addPlacementRows(result, ksm.name, table.name, params, placement, metadata, null, null);
+        }
+
+        return result;
+    }
+
+    private void addPlacementRows(SimpleDataSet result, String keyspaceName, String tableName, ReplicationParams params, DataPlacement placement, ClusterMetadata metadata, String rangeStartFilter, String rangeEndFilter)
+    {
+        Preconditions.checkState(placement.reads.ranges.equals(placement.writes.ranges),
+                                 "Read ranges (%s) are not the same as write ranges (%s)",
+                                 placement.reads.ranges, placement.writes.ranges);
+
+        List<Range<Token>> ranges = placement.reads.ranges;
+        for (int i = 0; i < ranges.size(); i++)
+        {
+            Range<Token> range = ranges.get(i);
+            if (rangeStartFilter != null && !range.left.toString().equals(rangeStartFilter))
+                continue;
+            if (rangeEndFilter != null && !range.right.toString().equals(rangeEndFilter))
+                continue;
+
+            VersionedEndpoints.ForRange readEndpoints = placement.reads.forRange(range);
+            VersionedEndpoints.ForRange writeEndpoints = placement.writes.forRange(range);
+
+            Set<String> readEndpointSet = toEndpointStrings(readEndpoints.get());
+            Set<String> writeEndpointSet = toEndpointStrings(writeEndpoints.get());
+
+            Set<Integer> readReplicaSet = toNodeIds(readEndpoints.get(), metadata);
+            Set<Integer> writeReplicaSet = toNodeIds(writeEndpoints.get(), metadata);
+
+            IPartitioner partitioner = params.isMeta() ? ReversedLongLocalPartitioner.instance : metadata.partitioner;
+            result.row(keyspaceName, tableName, range.left.toString(), range.right.toString())
+                  .column(TOKEN_TYPE, range.left.getClass().getCanonicalName())
+                  .column(RANGE_START_BYTES, partitioner.getTokenFactory().toByteArray(range.left))
+                  .column(RANGE_END_BYTES, partitioner.getTokenFactory().toByteArray(range.right))
+                  .column(READ_ENDPOINTS, readEndpointSet)
+                  .column(WRITE_ENDPOINTS, writeEndpointSet)
+                  .column(READ_REPLICAS, readReplicaSet)
+                  .column(WRITE_REPLICAS, writeReplicaSet);
+        }
+    }
+
+    static Set<String> toEndpointStrings(EndpointsForRange endpoints)
+    {
+        return endpoints.stream()
+                       .map(Replica::endpoint)
+                       .map(Object::toString)
+                       .collect(Collectors.toSet());
+    }
+
+    static Set<Integer> toNodeIds(EndpointsForRange endpoints, ClusterMetadata metadata)
+    {
+        return endpoints.stream()
+                       .map(Replica::endpoint)
+                       .map(metadata.directory::peerId)
+                       .map(NodeId::id)
+                       .collect(Collectors.toSet());
+    }
+
+    private static TableMetadata buildTableMetadata(String keyspace)
+    {
+        return TableMetadata.builder(keyspace, TABLE_NAME)
+                            .comment(TABLE_COMMENT)
+                            .kind(TableMetadata.Kind.VIRTUAL)
+                            .partitioner(new LocalPartitioner(CompositeType.getInstance(UTF8Type.instance, UTF8Type.instance, UTF8Type.instance, UTF8Type.instance)))
+                            .addPartitionKeyColumn(KEYSPACE_NAME, UTF8Type.instance)
+                            .addPartitionKeyColumn(TABLE_NAME_COLUMN, UTF8Type.instance)
+                            .addPartitionKeyColumn(RANGE_START, UTF8Type.instance)
+                            .addPartitionKeyColumn(RANGE_END, UTF8Type.instance)
+                            .addRegularColumn(TOKEN_TYPE, UTF8Type.instance)
+                            .addRegularColumn(RANGE_START_BYTES, BytesType.instance)
+                            .addRegularColumn(RANGE_END_BYTES, BytesType.instance)
+                            .addRegularColumn(READ_ENDPOINTS, SetType.getInstance(UTF8Type.instance, false))
+                            .addRegularColumn(WRITE_ENDPOINTS, SetType.getInstance(UTF8Type.instance, false))
+                            .addRegularColumn(READ_REPLICAS, SetType.getInstance(Int32Type.instance, false))
+                            .addRegularColumn(WRITE_REPLICAS, SetType.getInstance(Int32Type.instance, false))
+                            .build();
+    }
+}

src/java/org/apache/cassandra/db/virtual/PartitionLocationTable.java

@@ -0,0 +1,163 @@
+/*
+ * 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 org.apache.cassandra.db.virtual;
+
+import java.nio.ByteBuffer;
+import java.util.Set;
+
+import com.google.common.base.Preconditions;
+
+import org.apache.cassandra.db.DecoratedKey;
+import org.apache.cassandra.db.marshal.BytesType;
+import org.apache.cassandra.db.marshal.CompositeType;
+import org.apache.cassandra.db.marshal.Int32Type;
+import org.apache.cassandra.db.marshal.SetType;
+import org.apache.cassandra.db.marshal.UTF8Type;
+import org.apache.cassandra.dht.IPartitioner;
+import org.apache.cassandra.dht.LocalPartitioner;
+import org.apache.cassandra.dht.Range;
+import org.apache.cassandra.dht.ReversedLongLocalPartitioner;
+import org.apache.cassandra.dht.Token;
+import org.apache.cassandra.exceptions.InvalidRequestException;
+import org.apache.cassandra.schema.KeyspaceMetadata;
+import org.apache.cassandra.schema.ReplicationParams;
+import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.tcm.ClusterMetadata;
+import org.apache.cassandra.tcm.ownership.DataPlacement;
+import org.apache.cassandra.tcm.ownership.VersionedEndpoints;
+
+import static com.google.common.base.Preconditions.*;
+
+final class PartitionLocationTable extends AbstractVirtualTable
+{
+    static final String TABLE_NAME = "partition_location";
+    static final String TABLE_COMMENT = "shows the token range and replicas (read and write) for a given partition";
+
+    // Partition keys
+    static final String COLUMN_KEYSPACE_NAME = "keyspace_name";
+    static final String COLUMN_TABLE_NAME = "table_name";
+    static final String COLUMN_KEY = "key";
+
+    // Regular columns
+    static final String COLUMN_TOKEN = "tkn"; // can't use "token" because this is reserved in CQL
+    static final String COLUMN_TOKEN_BYTES = "tkn_bytes";
+    static final String COLUMN_RANGE_START = "range_start";
+    static final String COLUMN_RANGE_END = "range_end";
+    static final String COLUMN_RANGE_START_BYTES = "range_start_bytes";
+    static final String COLUMN_RANGE_END_BYTES = "range_end_bytes";
+    static final String COLUMN_READ_ENDPOINTS = "read_endpoints";
+    static final String COLUMN_WRITE_ENDPOINTS = "write_endpoints";
+    static final String COLUMN_READ_REPLICAS = "read_replicas";
+    static final String COLUMN_WRITE_REPLICAS = "write_replicas";
+
+    PartitionLocationTable(String keyspace)
+    {
+        super(buildTableMetadata(keyspace));
+    }
+
+    @Override
+    public DataSet data()
+    {
+        throw new InvalidRequestException("Partition location table requires a partition key, for example: " +
+                                          "SELECT * FROM system_views.partition_location WHERE keyspace_name = 'ks' AND table_name = 'tbl' AND key = '1:a'");
+    }
+
+    @Override
+    public DataSet data(DecoratedKey partitionKey)
+    {
+        SimpleDataSet result = new SimpleDataSet(metadata());
+
+        // Partition key is (keyspace_name, table_name, key)
+        ByteBuffer keyBytes = partitionKey.getKey();
+
+        checkArgument(metadata().partitionKeyType instanceof CompositeType,
+                                    "PartitionLocationTable partition key type must be CompositeType, got %s",
+                                    metadata().partitionKeyType.getClass().getName());
+
+        ByteBuffer[] components = ((CompositeType) metadata().partitionKeyType).split(keyBytes);
+        checkArgument(components.length == 3,
+                                    "PartitionLocationTable partition key must have exactly 3 components: keyspace_name, table_name, key; got %d",
+                                    components.length);
+
+        String keyspaceName = UTF8Type.instance.compose(components[0]);
+        String tableName = UTF8Type.instance.compose(components[1]);
+        String keyString = UTF8Type.instance.compose(components[2]);
+
+        ClusterMetadata metadata = ClusterMetadata.current();
+        KeyspaceMetadata ksm = checkNotNull(metadata.schema.getKeyspaceMetadata(keyspaceName),
+                                       "Keyspace %s is not found in metadata", keyspaceName);
+        TableMetadata table = checkNotNull(ksm.getTableOrViewNullable(tableName),
+                "Table %s is not found in metadata (within keyspace %s)", tableName, keyspaceName);
+
+        ByteBuffer partitionKeyBytes = table.partitionKeyType.fromString(keyString);
+
+        DecoratedKey dk = table.partitioner.decorateKey(partitionKeyBytes);
+        Token token = dk.getToken();
+
+        ReplicationParams replicationParams = ksm.params.replication;
+        DataPlacement placement = metadata.placements.get(replicationParams);
+
+        VersionedEndpoints.ForRange readEndpoints = placement.reads.forRange(token);
+        VersionedEndpoints.ForRange writeEndpoints = placement.writes.forRange(token);
+
+        Range<Token> range = readEndpoints.get().range();
+
+        Set<String> readEndpointSet = DataPlacementsTable.toEndpointStrings(readEndpoints.get());
+        Set<String> writeEndpointSet = DataPlacementsTable.toEndpointStrings(writeEndpoints.get());
+
+        Set<Integer> readReplicaSet = DataPlacementsTable.toNodeIds(readEndpoints.get(), metadata);
+        Set<Integer> writeReplicaSet = DataPlacementsTable.toNodeIds(writeEndpoints.get(), metadata);
+
+        IPartitioner partitioner = replicationParams.isMeta() ? ReversedLongLocalPartitioner.instance : metadata.partitioner;
+        result.row(keyspaceName, tableName, keyString)
+              .column(COLUMN_TOKEN, token.toString())
+              .column(COLUMN_TOKEN_BYTES, partitioner.getTokenFactory().toByteArray(token))
+              .column(COLUMN_RANGE_START, range.left.toString())
+              .column(COLUMN_RANGE_END, range.right.toString())
+              .column(COLUMN_RANGE_START_BYTES, partitioner.getTokenFactory().toByteArray(range.left))
+              .column(COLUMN_RANGE_END_BYTES, partitioner.getTokenFactory().toByteArray(range.right))
+              .column(COLUMN_READ_ENDPOINTS, readEndpointSet)
+              .column(COLUMN_WRITE_ENDPOINTS, writeEndpointSet)
+              .column(COLUMN_READ_REPLICAS, readReplicaSet)
+              .column(COLUMN_WRITE_REPLICAS, writeReplicaSet);
+
+        return result;
+    }
+
+    private static TableMetadata buildTableMetadata(String keyspace) {
+        return TableMetadata.builder(keyspace, TABLE_NAME)
+                            .comment(TABLE_COMMENT)
+                            .kind(TableMetadata.Kind.VIRTUAL)
+                            .partitioner(new LocalPartitioner(CompositeType.getInstance(UTF8Type.instance, UTF8Type.instance, UTF8Type.instance)))
+                            .addPartitionKeyColumn(COLUMN_KEYSPACE_NAME, UTF8Type.instance)
+                            .addPartitionKeyColumn(COLUMN_TABLE_NAME, UTF8Type.instance)
+                            .addPartitionKeyColumn(COLUMN_KEY, UTF8Type.instance)
+                            .addRegularColumn(COLUMN_TOKEN, UTF8Type.instance)
+                            .addRegularColumn(COLUMN_TOKEN_BYTES, BytesType.instance)
+                            .addRegularColumn(COLUMN_RANGE_START, UTF8Type.instance)
+                            .addRegularColumn(COLUMN_RANGE_END, UTF8Type.instance)
+                            .addRegularColumn(COLUMN_RANGE_START_BYTES, BytesType.instance)
+                            .addRegularColumn(COLUMN_RANGE_END_BYTES, BytesType.instance)
+                            .addRegularColumn(COLUMN_READ_ENDPOINTS, SetType.getInstance(UTF8Type.instance, false))
+                            .addRegularColumn(COLUMN_WRITE_ENDPOINTS, SetType.getInstance(UTF8Type.instance, false))
+                            .addRegularColumn(COLUMN_READ_REPLICAS, SetType.getInstance(Int32Type.instance, false))
+                            .addRegularColumn(COLUMN_WRITE_REPLICAS, SetType.getInstance(Int32Type.instance, false))
+                            .build();
+    }
+}

src/java/org/apache/cassandra/db/virtual/SystemViewsKeyspace.java

@@ -69,6 +69,8 @@ private SystemViewsKeyspace()
                     .add(new LocalTable(VIRTUAL_VIEWS))
                     .add(new ClusterMetadataLogTable(VIRTUAL_VIEWS))
                     .add(new ClusterMetadataDirectoryTable(VIRTUAL_VIEWS))
+                    .add(new DataPlacementsTable(VIRTUAL_VIEWS))
+                    .add(new PartitionLocationTable(VIRTUAL_VIEWS))
                     .addAll(LocalRepairTables.getAll(VIRTUAL_VIEWS))
                     .addAll(CIDRFilteringMetricsTable.getAll(VIRTUAL_VIEWS))
                     .addAll(StorageAttachedIndexTables.getAll(VIRTUAL_VIEWS))