org.apache.beam.it.mongodb.MongoDBResourceManager Maven / Gradle / Ivy
Show all versions of beam-it-mongodb Show documentation
/*
* 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.beam.it.mongodb;
import static org.apache.beam.it.mongodb.MongoDBResourceManagerUtils.checkValidCollectionName;
import static org.apache.beam.it.mongodb.MongoDBResourceManagerUtils.generateDatabaseName;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoClient;
import com.mongodb.client.MongoClients;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import org.apache.beam.it.common.ResourceManager;
import org.apache.beam.it.testcontainers.TestContainerResourceManager;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.bson.Document;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testcontainers.containers.MongoDBContainer;
import org.testcontainers.utility.DockerImageName;
/**
* Client for managing MongoDB resources.
*
* The class supports one database and multiple collections per database object. A database is
* created when the first collection is created if one has not been created already.
*
*
The database name is formed using testId. The database name will be "{testId}-{ISO8601 time,
* microsecond precision}", with additional formatting.
*
*
The class is thread-safe.
*/
public class MongoDBResourceManager extends TestContainerResourceManager
implements ResourceManager {
private static final Logger LOG = LoggerFactory.getLogger(MongoDBResourceManager.class);
private static final String DEFAULT_MONGODB_CONTAINER_NAME = "mongo";
// A list of available MongoDB Docker image tags can be found at
// https://hub.docker.com/_/mongo/tags
private static final String DEFAULT_MONGODB_CONTAINER_TAG = "4.0.18";
// 27017 is the default port that MongoDB is configured to listen on
private static final int MONGODB_INTERNAL_PORT = 27017;
private final MongoClient mongoClient;
private final String databaseName;
private final String connectionString;
private final boolean usingStaticDatabase;
private MongoDBResourceManager(Builder builder) {
this(
/* mongoClient= */ null,
new MongoDBContainer(
DockerImageName.parse(builder.containerImageName).withTag(builder.containerImageTag)),
builder);
}
@VisibleForTesting
@SuppressWarnings("nullness")
MongoDBResourceManager(
@Nullable MongoClient mongoClient, MongoDBContainer container, Builder builder) {
super(container, builder);
this.usingStaticDatabase = builder.databaseName != null;
this.databaseName =
usingStaticDatabase ? builder.databaseName : generateDatabaseName(builder.testId);
this.connectionString =
String.format("mongodb://%s:%d", this.getHost(), this.getPort(MONGODB_INTERNAL_PORT));
this.mongoClient = mongoClient == null ? MongoClients.create(connectionString) : mongoClient;
}
public static Builder builder(String testId) {
return new Builder(testId);
}
/** Returns the URI connection string to the MongoDB Database. */
public synchronized String getUri() {
return connectionString;
}
/**
* Returns the name of the Database that this MongoDB manager will operate in.
*
* @return the name of the MongoDB Database.
*/
public synchronized String getDatabaseName() {
return databaseName;
}
private synchronized MongoDatabase getDatabase() {
try {
return mongoClient.getDatabase(databaseName);
} catch (Exception e) {
throw new MongoDBResourceManagerException(
"Error retrieving database " + databaseName + " from MongoDB.", e);
}
}
private synchronized boolean collectionExists(String collectionName) {
// Check collection name
checkValidCollectionName(databaseName, collectionName);
Iterable collectionNames = getDatabase().listCollectionNames();
for (String name : collectionNames) {
// The Collection already exists in the database, return false.
if (collectionName.equals(name)) {
return true;
}
}
return false;
}
/**
* Creates a collection in MongoDB for storing Documents.
*
* Note: Implementations may do database creation here, if one does not already exist.
*
* @param collectionName Collection name to associate with the given MongoDB instance.
* @return A boolean indicating whether the resource was created.
* @throws MongoDBResourceManagerException if there is an error creating the collection in
* MongoDB.
*/
public synchronized boolean createCollection(String collectionName)
throws MongoDBResourceManagerException {
LOG.info("Creating collection using collectionName '{}'.", collectionName);
try {
// Check to see if the Collection exists
if (collectionExists(collectionName)) {
return false;
}
// The Collection does not exist in the database, create it and return true.
getDatabase().createCollection(collectionName);
} catch (Exception e) {
throw new MongoDBResourceManagerException("Error creating collection.", e);
}
LOG.info("Successfully created collection {}.{}", databaseName, collectionName);
return true;
}
/**
* Helper method to retrieve a MongoCollection with the given name from the database and return
* it.
*
* @param collectionName The name of the MongoCollection.
* @param createCollection A boolean that specifies to create the Collection if it does not exist.
* @return A MongoCollection with the given name.
*/
private MongoCollection getMongoDBCollection(
String collectionName, boolean createCollection) {
if (!collectionExists(collectionName) && !createCollection) {
throw new MongoDBResourceManagerException(
"Collection " + collectionName + " does not exists in database " + databaseName);
}
return getDatabase().getCollection(collectionName);
}
/**
* Inserts the given Document into a collection.
*
* A database will be created here, if one does not already exist.
*
* @param collectionName The name of the collection to insert the document into.
* @param document The document to insert into the collection.
* @return A boolean indicating whether the Document was inserted successfully.
*/
public synchronized boolean insertDocument(String collectionName, Document document) {
return insertDocuments(collectionName, ImmutableList.of(document));
}
/**
* Inserts the given Documents into a collection.
*
*
Note: Implementations may do collection creation here, if one does not already exist.
*
* @param collectionName The name of the collection to insert the documents into.
* @param documents A list of documents to insert into the collection.
* @return A boolean indicating whether the Documents were inserted successfully.
* @throws MongoDBResourceManagerException if there is an error inserting the documents.
*/
public synchronized boolean insertDocuments(String collectionName, List documents)
throws MongoDBResourceManagerException {
LOG.info(
"Attempting to write {} documents to {}.{}.",
documents.size(),
databaseName,
collectionName);
try {
getMongoDBCollection(collectionName, /* createCollection= */ true).insertMany(documents);
} catch (Exception e) {
throw new MongoDBResourceManagerException("Error inserting documents.", e);
}
LOG.info(
"Successfully wrote {} documents to {}.{}", documents.size(), databaseName, collectionName);
return true;
}
/**
* Reads all the Documents in a collection.
*
* @param collectionName The name of the collection to read from.
* @return A List of all the Documents in the collection.
* @throws MongoDBResourceManagerException if there is an error reading the collection.
*/
public synchronized List readCollection(String collectionName)
throws MongoDBResourceManagerException {
LOG.info("Reading all documents from {}.{}", databaseName, collectionName);
FindIterable fetchRecords;
try {
fetchRecords = getMongoDBCollection(collectionName, /* createCollection= */ false).find();
} catch (Exception e) {
throw new MongoDBResourceManagerException("Error reading collection.", e);
}
List documents =
StreamSupport.stream(fetchRecords.spliterator(), false).collect(Collectors.toList());
LOG.info("Successfully loaded documents from {}.{}", databaseName, collectionName);
return documents;
}
/**
* Counts the number of Documents in a collection.
*
* @param collectionName The name of the collection to read from.
* @return The number of Documents in the collection.
* @throws MongoDBResourceManagerException if there is an error reading the collection.
*/
public synchronized long countCollection(String collectionName)
throws MongoDBResourceManagerException {
LOG.info("Counting all documents from {}.{}", databaseName, collectionName);
long numDocuments =
getMongoDBCollection(collectionName, /* createCollection= */ false).countDocuments();
LOG.info(
"Successfully counted {} documents from {}.{}", numDocuments, databaseName, collectionName);
return numDocuments;
}
@Override
public synchronized void cleanupAll() {
LOG.info("Attempting to cleanup MongoDB manager.");
boolean producedError = false;
// First, delete the database if it was not given as a static argument
try {
if (!usingStaticDatabase) {
mongoClient.getDatabase(databaseName).drop();
}
} catch (Exception e) {
LOG.error("Failed to delete MongoDB database {}.", databaseName, e);
producedError = true;
}
// Next, try to close the MongoDB client connection
try {
mongoClient.close();
} catch (Exception e) {
LOG.error("Failed to delete MongoDB client.", e);
producedError = true;
}
// Throw Exception at the end if there were any errors
if (producedError) {
throw new MongoDBResourceManagerException(
"Failed to delete resources. Check above for errors.");
}
super.cleanupAll();
LOG.info("MongoDB manager successfully cleaned up.");
}
/** Builder for {@link MongoDBResourceManager}. */
public static final class Builder
extends TestContainerResourceManager.Builder {
private @Nullable String databaseName;
private Builder(String testId) {
super(testId, DEFAULT_MONGODB_CONTAINER_NAME, DEFAULT_MONGODB_CONTAINER_TAG);
this.databaseName = null;
}
/**
* Sets the database name to that of a static database instance. Use this method only when
* attempting to operate on a pre-existing Mongo database.
*
* Note: if a database name is set, and a static MongoDB server is being used
* (useStaticContainer() is also called on the builder), then a database will be created on the
* static server if it does not exist, and it will not be removed when cleanupAll() is called on
* the MongoDBResourceManager.
*
* @param databaseName The database name.
* @return this builder object with the database name set.
*/
public Builder setDatabaseName(String databaseName) {
this.databaseName = databaseName;
return this;
}
@Override
public MongoDBResourceManager build() {
return new MongoDBResourceManager(this);
}
}
}