Java tutorial
/* * 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.facebook.presto.metadata; import com.facebook.presto.metadata.ShardManagerDao.Utils; import com.facebook.presto.spi.PartitionKey; import com.facebook.presto.spi.TableHandle; import com.facebook.presto.split.NativePartitionKey; import com.google.common.base.Optional; import com.google.common.collect.ImmutableMultimap; import com.google.common.collect.Maps; import com.google.common.collect.Multimap; import org.skife.jdbi.v2.Handle; import org.skife.jdbi.v2.IDBI; import org.skife.jdbi.v2.TransactionStatus; import org.skife.jdbi.v2.VoidTransactionCallback; import javax.annotation.Nullable; import javax.inject.Inject; import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.Set; import static com.facebook.presto.util.SqlUtils.runIgnoringConstraintViolation; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; public class DatabaseShardManager implements ShardManager { private final IDBI dbi; private final ShardManagerDao dao; @Inject public DatabaseShardManager(@ForShardManager IDBI dbi) throws InterruptedException { this.dbi = dbi; this.dao = dbi.onDemand(ShardManagerDao.class); // keep retrying if database is unavailable when the server starts Utils.createShardTablesWithRetry(dao); } @Override public long allocateShard(TableHandle tableHandle) { checkNotNull(tableHandle, "tableHandle is null"); checkState(tableHandle instanceof NativeTableHandle, "can only allocate shards for native tables"); long tableId = ((NativeTableHandle) tableHandle).getTableId(); return dao.insertShard(tableId, false); } @Override public void commitShard(final long shardId, String nodeIdentifier) { final long nodeId = getOrCreateNodeId(nodeIdentifier); dbi.inTransaction(new VoidTransactionCallback() { @Override protected void execute(Handle handle, TransactionStatus status) { ShardManagerDao dao = handle.attach(ShardManagerDao.class); dao.commitShard(shardId); dao.insertShardNode(shardId, nodeId); } }); } @Override public void commitPartition(TableHandle tableHandle, final String partition, final List<? extends PartitionKey> partitionKeys, final Map<Long, String> shards) { checkNotNull(tableHandle, "tableHandle is null"); checkNotNull(partition, "partition is null"); checkNotNull(partitionKeys, "partitionKeys is null"); checkNotNull(shards, "shards is null"); checkState(tableHandle instanceof NativeTableHandle, "can only commit partitions for native tables"); final long tableId = ((NativeTableHandle) tableHandle).getTableId(); dbi.inTransaction(new VoidTransactionCallback() { @Override protected void execute(Handle handle, TransactionStatus status) { ShardManagerDao dao = handle.attach(ShardManagerDao.class); long partitionId = dao.insertPartition(tableId, partition); for (PartitionKey partitionKey : partitionKeys) { dao.insertPartitionKey(tableId, partition, partitionKey.getName(), partitionKey.getType().toString(), partitionKey.getValue()); } for (Map.Entry<Long, String> entry : shards.entrySet()) { long nodeId = getOrCreateNodeId(entry.getValue()); long shardId = entry.getKey(); dao.commitShard(shardId); dao.insertShardNode(shardId, nodeId); dao.insertPartitionShard(shardId, tableId, partitionId); } } }); } @Override public void disassociateShard(long shardId, @Nullable String nodeIdentifier) { dao.dropShardNode(shardId, nodeIdentifier); } @Override public void dropShard(final long shardId) { dbi.inTransaction(new VoidTransactionCallback() { @Override protected void execute(Handle handle, TransactionStatus status) { ShardManagerDao dao = handle.attach(ShardManagerDao.class); dao.deleteShardFromPartitionShards(shardId); dao.deleteShard(shardId); } }); } @Override public Set<TablePartition> getPartitions(TableHandle tableHandle) { checkNotNull(tableHandle, "tableHandle is null"); checkState(tableHandle instanceof NativeTableHandle, "can only commit partitions for native tables"); final long tableId = ((NativeTableHandle) tableHandle).getTableId(); return dao.getPartitions(tableId); } @Override public Multimap<String, ? extends PartitionKey> getAllPartitionKeys(TableHandle tableHandle) { checkNotNull(tableHandle, "tableHandle is null"); checkState(tableHandle instanceof NativeTableHandle, "can only commit partitions for native tables"); final long tableId = ((NativeTableHandle) tableHandle).getTableId(); Set<NativePartitionKey> partitionKeys = dao.getPartitionKeys(tableId); ImmutableMultimap.Builder<String, PartitionKey> builder = ImmutableMultimap.builder(); for (NativePartitionKey partitionKey : partitionKeys) { builder.put(partitionKey.getPartitionName(), partitionKey); } return builder.build(); } @Override public Multimap<Long, Entry<Long, String>> getCommittedPartitionShardNodes(TableHandle tableHandle) { checkNotNull(tableHandle, "tableHandle is null"); checkState(tableHandle instanceof NativeTableHandle, "tableHandle not a native table"); final long tableId = ((NativeTableHandle) tableHandle).getTableId(); ImmutableMultimap.Builder<Long, Entry<Long, String>> map = ImmutableMultimap.builder(); List<ShardNode> shardNodes = dao.getCommittedShardNodesByTableId(tableId); for (ShardNode shardNode : shardNodes) { map.put(shardNode.getPartitionId(), Maps.immutableEntry(shardNode.getShardId(), shardNode.getNodeIdentifier())); } return map.build(); } @Override public Multimap<Long, String> getCommittedShardNodesByTableId(TableHandle tableHandle) { checkNotNull(tableHandle, "tableHandle is null"); checkState(tableHandle instanceof NativeTableHandle, "tableHandle not a native table"); final long tableId = ((NativeTableHandle) tableHandle).getTableId(); ImmutableMultimap.Builder<Long, String> map = ImmutableMultimap.builder(); for (ShardNode shardNode : dao.getCommittedShardNodesByTableId(tableId)) { map.put(shardNode.getShardId(), shardNode.getNodeIdentifier()); } return map.build(); } @Override public Multimap<Long, String> getShardNodes(long tableId, String partitionName) { ImmutableMultimap.Builder<Long, String> map = ImmutableMultimap.builder(); for (ShardNode shardNode : dao.getAllShardNodes(tableId, partitionName)) { map.put(shardNode.getShardId(), shardNode.getNodeIdentifier()); } return map.build(); } @Override public Iterable<String> getAllNodesInUse() { return dao.getAllNodesInUse(); } @Override public void dropPartition(final TableHandle tableHandle, final String partitionName) { checkNotNull(tableHandle, "tableHandle is null"); checkState(tableHandle instanceof NativeTableHandle, "can only commit partitions for native tables"); final long tableId = ((NativeTableHandle) tableHandle).getTableId(); dbi.inTransaction(new VoidTransactionCallback() { @Override protected void execute(Handle handle, TransactionStatus status) throws Exception { ShardManagerDao dao = handle.attach(ShardManagerDao.class); List<Long> shardIds = dao.getAllShards(tableId, partitionName); for (Long shardId : shardIds) { dao.deleteShardFromPartitionShards(shardId); } dao.dropPartitionKeys(tableId, partitionName); dao.dropPartition(tableId, partitionName); } }); } @Override public Iterable<Long> getOrphanedShardIds(Optional<String> nodeIdentifier) { if (nodeIdentifier.isPresent()) { return dao.getOrphanedShards(nodeIdentifier.get()); } else { return dao.getAllOrphanedShards(); } } @Override public void dropOrphanedPartitions() { dao.dropAllOrphanedPartitions(); } private long getOrCreateNodeId(final String nodeIdentifier) { Long id = dao.getNodeId(nodeIdentifier); if (id != null) { return id; } // creating a node is idempotent runIgnoringConstraintViolation(new Runnable() { @Override public void run() { dao.insertNode(nodeIdentifier); } }); id = dao.getNodeId(nodeIdentifier); if (id == null) { throw new IllegalStateException("node does not exist after insert"); } return id; } }