Examples with MongoDatabase com.mongodb.client.MongoDatabase used on opensource projects

Example 16 with MongoDatabase

use of com.mongodb.client.MongoDatabase in project drill by apache.

the class MongoGroupScan method init.

@SuppressWarnings({ "rawtypes" })
private void init() throws IOException {
    List<String> h = storagePluginConfig.getHosts();
    List<ServerAddress> addresses = Lists.newArrayList();
    for (String host : h) {
        addresses.add(new ServerAddress(host));
    }
    MongoClient client = storagePlugin.getClient();
    chunksMapping = Maps.newHashMap();
    chunksInverseMapping = Maps.newLinkedHashMap();
    if (isShardedCluster(client)) {
        MongoDatabase db = client.getDatabase(CONFIG);
        MongoCollection<Document> chunksCollection = db.getCollection(CHUNKS);
        Document filter = new Document();
        filter.put(NS, this.scanSpec.getDbName() + "." + this.scanSpec.getCollectionName());
        Document projection = new Document();
        projection.put(SHARD, select);
        projection.put(MIN, select);
        projection.put(MAX, select);
        FindIterable<Document> chunkCursor = chunksCollection.find(filter).projection(projection);
        MongoCursor<Document> iterator = chunkCursor.iterator();
        MongoCollection<Document> shardsCollection = db.getCollection(SHARDS);
        projection = new Document();
        projection.put(HOST, select);
        boolean hasChunks = false;
        while (iterator.hasNext()) {
            Document chunkObj = iterator.next();
            String shardName = (String) chunkObj.get(SHARD);
            String chunkId = (String) chunkObj.get(ID);
            filter = new Document(ID, shardName);
            FindIterable<Document> hostCursor = shardsCollection.find(filter).projection(projection);
            MongoCursor<Document> hostIterator = hostCursor.iterator();
            while (hostIterator.hasNext()) {
                Document hostObj = hostIterator.next();
                String hostEntry = (String) hostObj.get(HOST);
                String[] tagAndHost = StringUtils.split(hostEntry, '/');
                String[] hosts = tagAndHost.length > 1 ? StringUtils.split(tagAndHost[1], ',') : StringUtils.split(tagAndHost[0], ',');
                List<String> chunkHosts = Arrays.asList(hosts);
                Set<ServerAddress> addressList = getPreferredHosts(storagePlugin.getClient(addresses), chunkHosts);
                if (addressList == null) {
                    addressList = Sets.newHashSet();
                    for (String host : chunkHosts) {
                        addressList.add(new ServerAddress(host));
                    }
                }
                chunksMapping.put(chunkId, addressList);
                ServerAddress address = addressList.iterator().next();
                List<ChunkInfo> chunkList = chunksInverseMapping.get(address.getHost());
                if (chunkList == null) {
                    chunkList = Lists.newArrayList();
                    chunksInverseMapping.put(address.getHost(), chunkList);
                }
                List<String> chunkHostsList = new ArrayList<String>();
                for (ServerAddress serverAddr : addressList) {
                    chunkHostsList.add(serverAddr.toString());
                }
                ChunkInfo chunkInfo = new ChunkInfo(chunkHostsList, chunkId);
                Document minMap = (Document) chunkObj.get(MIN);
                Map<String, Object> minFilters = Maps.newHashMap();
                Set keySet = minMap.keySet();
                for (Object keyObj : keySet) {
                    Object object = minMap.get(keyObj);
                    if (!(object instanceof MinKey)) {
                        minFilters.put(keyObj.toString(), object);
                    }
                }
                chunkInfo.setMinFilters(minFilters);
                Map<String, Object> maxFilters = Maps.newHashMap();
                Map maxMap = (Document) chunkObj.get(MAX);
                keySet = maxMap.keySet();
                for (Object keyObj : keySet) {
                    Object object = maxMap.get(keyObj);
                    if (!(object instanceof MaxKey)) {
                        maxFilters.put(keyObj.toString(), object);
                    }
                }
                chunkInfo.setMaxFilters(maxFilters);
                chunkList.add(chunkInfo);
            }
            hasChunks = true;
        }
        // unsharded collection and it will be stored in the primary shard of that database.
        if (!hasChunks) {
            handleUnshardedCollection(getPrimaryShardInfo(client));
        }
    } else {
        handleUnshardedCollection(storagePluginConfig.getHosts());
    }
}

Example 17 with MongoDatabase

use of com.mongodb.client.MongoDatabase in project drill by apache.

the class MongoGroupScan method getPrimaryShardInfo.

private List<String> getPrimaryShardInfo(MongoClient client) {
    MongoDatabase database = storagePlugin.getClient().getDatabase(CONFIG);
    //Identify the primary shard of the queried database.
    MongoCollection<Document> collection = database.getCollection(DATABASES);
    Bson filter = new Document(ID, this.scanSpec.getDbName());
    Bson projection = new Document(PRIMARY, select);
    Document document = collection.find(filter).projection(projection).first();
    Preconditions.checkNotNull(document);
    String shardName = document.getString(PRIMARY);
    Preconditions.checkNotNull(shardName);
    //Identify the host(s) on which this shard resides.
    MongoCollection<Document> shardsCol = database.getCollection(SHARDS);
    filter = new Document(ID, shardName);
    projection = new Document(HOST, select);
    Document hostInfo = shardsCol.find(filter).projection(projection).first();
    Preconditions.checkNotNull(hostInfo);
    String hostEntry = hostInfo.getString(HOST);
    Preconditions.checkNotNull(hostEntry);
    String[] tagAndHost = StringUtils.split(hostEntry, '/');
    String[] hosts = tagAndHost.length > 1 ? StringUtils.split(tagAndHost[1], ',') : StringUtils.split(tagAndHost[0], ',');
    return Lists.newArrayList(hosts);
}

Example 18 with MongoDatabase

use of com.mongodb.client.MongoDatabase in project drill by apache.

the class MongoGroupScan method isShardedCluster.

private boolean isShardedCluster(MongoClient client) {
    MongoDatabase db = client.getDatabase(scanSpec.getDbName());
    String msg = db.runCommand(new Document("isMaster", 1)).getString("msg");
    return msg == null ? false : msg.equals("isdbgrid");
}

Example 19 with MongoDatabase

use of com.mongodb.client.MongoDatabase in project drill by apache.

the class MongoGroupScan method getPreferredHosts.

@SuppressWarnings("unchecked")
private Set<ServerAddress> getPreferredHosts(MongoClient client, List<String> hosts) {
    Set<ServerAddress> addressList = Sets.newHashSet();
    MongoDatabase db = client.getDatabase(scanSpec.getDbName());
    ReadPreference readPreference = client.getReadPreference();
    Document command = db.runCommand(new Document("isMaster", 1));
    final String primaryHost = command.getString("primary");
    final List<String> hostsList = (List<String>) command.get("hosts");
    switch(readPreference.getName().toUpperCase()) {
        case "PRIMARY":
        case "PRIMARYPREFERRED":
            if (primaryHost == null) {
                return null;
            }
            addressList.add(new ServerAddress(primaryHost));
            return addressList;
        case "SECONDARY":
        case "SECONDARYPREFERRED":
            if (primaryHost == null || hostsList == null) {
                return null;
            }
            hostsList.remove(primaryHost);
            for (String host : hostsList) {
                addressList.add(new ServerAddress(host));
            }
            return addressList;
        case "NEAREST":
            if (hostsList == null) {
                return null;
            }
            for (String host : hostsList) {
                addressList.add(new ServerAddress(host));
            }
            return addressList;
        default:
            return null;
    }
}

Example 20 with MongoDatabase

use of com.mongodb.client.MongoDatabase in project drill by apache.

the class MongoTestSuit method createDbAndCollections.

private static void createDbAndCollections(String dbName, String collectionName, String indexFieldName) {
    MongoDatabase db = mongoClient.getDatabase(dbName);
    MongoCollection<Document> mongoCollection = db.getCollection(collectionName);
    if (mongoCollection == null) {
        db.createCollection(collectionName);
        mongoCollection = db.getCollection(collectionName);
    }
    IndexOptions indexOptions = new IndexOptions().unique(true).background(false).name(indexFieldName);
    Bson keys = Indexes.ascending(indexFieldName);
    mongoCollection.createIndex(keys, indexOptions);
}