Пример #1
0
/**
 * @author Guillaume Scheibel <[email protected]>
 * @author Sanne Grinovero <[email protected]>
 */
public class MongoDBTestHelper implements GridDialectTestHelper {

  private static final Log log = LoggerFactory.getLogger();

  static {
    // Read host and port from environment variable
    // Maven's surefire plugin set it to the string 'null'
    String mongoHostName = System.getenv("MONGODB_HOSTNAME");
    if (isNotNull(mongoHostName)) {
      System.getProperties().setProperty(OgmProperties.HOST, mongoHostName);
    }
    String mongoPort = System.getenv("MONGODB_PORT");
    if (isNotNull(mongoPort)) {
      System.getProperties().setProperty(OgmProperties.PORT, mongoPort);
    }
  }

  private static boolean isNotNull(String mongoHostName) {
    return mongoHostName != null && mongoHostName.length() > 0 && !"null".equals(mongoHostName);
  }

  @Override
  public long getNumberOfEntities(Session session) {
    return getNumberOfEntities(session.getSessionFactory());
  }

  @Override
  public long getNumberOfEntities(SessionFactory sessionFactory) {
    MongoDBDatastoreProvider provider = MongoDBTestHelper.getProvider(sessionFactory);
    DB db = provider.getDatabase();
    int count = 0;

    for (String collectionName : getEntityCollections(sessionFactory)) {
      count += db.getCollection(collectionName).count();
    }

    return count;
  }

  private boolean isSystemCollection(String collectionName) {
    return collectionName.startsWith("system.");
  }

  @Override
  public long getNumberOfAssociations(Session session) {
    return getNumberOfAssociations(session.getSessionFactory());
  }

  @Override
  public long getNumberOfAssociations(SessionFactory sessionFactory) {
    long associationCount = getNumberOfAssociationsFromGlobalCollection(sessionFactory);
    associationCount += getNumberOfAssociationsFromDedicatedCollections(sessionFactory);
    associationCount += getNumberOfEmbeddedAssociations(sessionFactory);

    return associationCount;
  }

  public long getNumberOfAssociationsFromGlobalCollection(SessionFactory sessionFactory) {
    DB db = getProvider(sessionFactory).getDatabase();
    return db.getCollection(MongoDBConfiguration.DEFAULT_ASSOCIATION_STORE).count();
  }

  public long getNumberOfAssociationsFromDedicatedCollections(SessionFactory sessionFactory) {
    DB db = getProvider(sessionFactory).getDatabase();

    Set<String> associationCollections = getDedicatedAssociationCollections(sessionFactory);
    long associationCount = 0;
    for (String collectionName : associationCollections) {
      associationCount += db.getCollection(collectionName).count();
    }

    return associationCount;
  }

  // TODO Use aggregation framework for a more efficient solution; Given that there will only be a
  // few
  // test collections/entities, that's good enough for now
  public long getNumberOfEmbeddedAssociations(SessionFactory sessionFactory) {
    DB db = getProvider(sessionFactory).getDatabase();
    long associationCount = 0;

    for (String entityCollection : getEntityCollections(sessionFactory)) {
      DBCursor entities = db.getCollection(entityCollection).find();

      while (entities.hasNext()) {
        DBObject entity = entities.next();
        associationCount += getNumberOfEmbeddedAssociations(entity);
      }
    }

    return associationCount;
  }

  private int getNumberOfEmbeddedAssociations(DBObject entity) {
    int numberOfReferences = 0;

    for (String fieldName : entity.keySet()) {
      Object field = entity.get(fieldName);
      if (isAssociation(field)) {
        numberOfReferences++;
      }
    }

    return numberOfReferences;
  }

  private boolean isAssociation(Object field) {
    return (field instanceof List);
  }

  private Set<String> getEntityCollections(SessionFactory sessionFactory) {
    DB db = MongoDBTestHelper.getProvider(sessionFactory).getDatabase();
    Set<String> names = new HashSet<String>();

    for (String collectionName : db.getCollectionNames()) {
      if (!isSystemCollection(collectionName)
          && !isDedicatedAssociationCollection(collectionName)
          && !isGlobalAssociationCollection(collectionName)) {
        names.add(collectionName);
      }
    }

    return names;
  }

  private Set<String> getDedicatedAssociationCollections(SessionFactory sessionFactory) {
    DB db = MongoDBTestHelper.getProvider(sessionFactory).getDatabase();
    Set<String> names = new HashSet<String>();

    for (String collectionName : db.getCollectionNames()) {
      if (isDedicatedAssociationCollection(collectionName)) {
        names.add(collectionName);
      }
    }

    return names;
  }

  private boolean isDedicatedAssociationCollection(String collectionName) {
    return collectionName.startsWith(MongoDBDialect.ASSOCIATIONS_COLLECTION_PREFIX);
  }

  private boolean isGlobalAssociationCollection(String collectionName) {
    return collectionName.equals(MongoDBConfiguration.DEFAULT_ASSOCIATION_STORE);
  }

  @Override
  @SuppressWarnings("unchecked")
  public Map<String, Object> extractEntityTuple(Session session, EntityKey key) {
    MongoDBDatastoreProvider provider = MongoDBTestHelper.getProvider(session.getSessionFactory());
    DBObject finder = new BasicDBObject(MongoDBDialect.ID_FIELDNAME, key.getColumnValues()[0]);
    DBObject result = provider.getDatabase().getCollection(key.getTable()).findOne(finder);
    replaceIdentifierColumnName(result, key);
    return result.toMap();
  }

  /**
   * The MongoDB dialect replaces the name of the column identifier, so when the tuple is extracted
   * from the db we replace the column name of the identifier with the original one. We are assuming
   * the identifier is not embedded and is a single property.
   */
  private void replaceIdentifierColumnName(DBObject result, EntityKey key) {
    Object idValue = result.get(MongoDBDialect.ID_FIELDNAME);
    result.removeField(MongoDBDialect.ID_FIELDNAME);
    result.put(key.getColumnNames()[0], idValue);
  }

  @Override
  public boolean backendSupportsTransactions() {
    return false;
  }

  private static MongoDBDatastoreProvider getProvider(SessionFactory sessionFactory) {
    DatastoreProvider provider =
        ((SessionFactoryImplementor) sessionFactory)
            .getServiceRegistry()
            .getService(DatastoreProvider.class);
    if (!(MongoDBDatastoreProvider.class.isInstance(provider))) {
      throw new RuntimeException("Not testing with MongoDB, cannot extract underlying cache");
    }
    return MongoDBDatastoreProvider.class.cast(provider);
  }

  @Override
  public void dropSchemaAndDatabase(SessionFactory sessionFactory) {
    MongoDBDatastoreProvider provider = getProvider(sessionFactory);
    try {
      provider.getDatabase().dropDatabase();
    } catch (MongoException ex) {
      throw log.unableToDropDatabase(ex, provider.getDatabase().getName());
    }
  }

  @Override
  public Map<String, String> getEnvironmentProperties() {
    // read variables from the System properties set in the static initializer
    Map<String, String> envProps = new HashMap<String, String>(2);
    copyFromSystemPropertiesToLocalEnvironment(OgmProperties.HOST, envProps);
    copyFromSystemPropertiesToLocalEnvironment(OgmProperties.PORT, envProps);
    copyFromSystemPropertiesToLocalEnvironment(OgmProperties.USERNAME, envProps);
    copyFromSystemPropertiesToLocalEnvironment(OgmProperties.PASSWORD, envProps);
    return envProps;
  }

  private void copyFromSystemPropertiesToLocalEnvironment(
      String environmentVariableName, Map<String, String> envProps) {
    String value = System.getProperties().getProperty(environmentVariableName);
    if (value != null && value.length() > 0) {
      envProps.put(environmentVariableName, value);
    }
  }

  @Override
  public long getNumberOfAssociations(SessionFactory sessionFactory, AssociationStorageType type) {
    switch (type) {
      case ASSOCIATION_DOCUMENT:
        return getNumberOfAssociationsFromGlobalCollection(sessionFactory);
      case IN_ENTITY:
        return getNumberOfEmbeddedAssociations(sessionFactory);
      default:
        throw new IllegalArgumentException("Unexpected association storaget type " + type);
    }
  }

  @Override
  public GridDialect getGridDialect(DatastoreProvider datastoreProvider) {
    return new MongoDBDialect((MongoDBDatastoreProvider) datastoreProvider);
  }

  public static void assertDbObject(
      OgmSessionFactory sessionFactory,
      String collection,
      String queryDbObject,
      String expectedDbObject) {
    assertDbObject(sessionFactory, collection, queryDbObject, null, expectedDbObject);
  }

  public static void assertDbObject(
      OgmSessionFactory sessionFactory,
      String collection,
      String queryDbObject,
      String projectionDbObject,
      String expectedDbObject) {
    DBObject finder = (DBObject) JSON.parse(queryDbObject);
    DBObject fields = projectionDbObject != null ? (DBObject) JSON.parse(projectionDbObject) : null;

    MongoDBDatastoreProvider provider = MongoDBTestHelper.getProvider(sessionFactory);
    DBObject actual = provider.getDatabase().getCollection(collection).findOne(finder, fields);

    assertJsonEquals(expectedDbObject, actual.toString());
  }

  public static Map<String, DBObject> getIndexes(
      OgmSessionFactory sessionFactory, String collection) {
    MongoDBDatastoreProvider provider = MongoDBTestHelper.getProvider(sessionFactory);
    List<DBObject> indexes = provider.getDatabase().getCollection(collection).getIndexInfo();
    Map<String, DBObject> indexMap = new HashMap<>();
    for (DBObject index : indexes) {
      indexMap.put(index.get("name").toString(), index);
    }
    return indexMap;
  }

  public static void dropIndexes(OgmSessionFactory sessionFactory, String collection) {
    MongoDBDatastoreProvider provider = MongoDBTestHelper.getProvider(sessionFactory);
    provider.getDatabase().getCollection(collection).dropIndexes();
  }

  public static void assertJsonEquals(String expectedJson, String actualJson) {
    try {
      JSONCompareResult result =
          JSONCompare.compareJSON(expectedJson, actualJson, JSONCompareMode.NON_EXTENSIBLE);

      if (result.failed()) {
        throw new AssertionError(result.getMessage() + "; Actual: " + actualJson);
      }
    } catch (JSONException e) {
      Exceptions.<RuntimeException>sneakyThrow(e);
    }
  }

  @Override
  public Class<? extends DatastoreConfiguration<?>> getDatastoreConfigurationType() {
    return MongoDB.class;
  }
}
/**
 * Configuration for {@link MongoDBDatastoreProvider}.
 *
 * @author Guillaume Scheibel <*****@*****.**>
 * @author Gunnar Morling
 */
public class MongoDBConfiguration extends DocumentStoreConfiguration {

  public static final String DEFAULT_ASSOCIATION_STORE = "Associations";

  /**
   * The default value used to set the timeout during the connection to the MongoDB instance This
   * value is set in milliseconds.
   *
   * @see MongoDBProperties#TIMEOUT
   */
  private static final int DEFAULT_TIMEOUT = 5000;

  private static final int DEFAULT_PORT = 27017;

  private static final Log log = LoggerFactory.getLogger();

  private static final TimeoutValidator TIMEOUT_VALIDATOR = new TimeoutValidator();

  private final int timeout;
  private final WriteConcern writeConcern;
  private final ReadPreference readPreference;

  /**
   * Creates a new {@link MongoDBConfiguration}.
   *
   * @param configurationValues configuration values given via {@code persistence.xml} etc.
   * @param globalOptions global settings given via an option configurator
   */
  public MongoDBConfiguration(
      ConfigurationPropertyReader propertyReader, OptionsContext globalOptions) {
    super(propertyReader, DEFAULT_PORT);

    this.timeout =
        propertyReader
            .property(MongoDBProperties.TIMEOUT, int.class)
            .withDefault(DEFAULT_TIMEOUT)
            .withValidator(TIMEOUT_VALIDATOR)
            .getValue();

    this.writeConcern = globalOptions.getUnique(WriteConcernOption.class);
    this.readPreference = globalOptions.getUnique(ReadPreferenceOption.class);
  }

  /**
   * Create a {@link MongoClientOptions} using the {@link MongoDBConfiguration}.
   *
   * @return the {@link MongoClientOptions} corresponding to the {@link MongoDBConfiguration}
   */
  public MongoClientOptions buildOptions() {
    MongoClientOptions.Builder optionsBuilder = new MongoClientOptions.Builder();

    optionsBuilder.connectTimeout(timeout);
    optionsBuilder.writeConcern(writeConcern);
    optionsBuilder.readPreference(readPreference);

    return optionsBuilder.build();
  }

  private static class TimeoutValidator implements PropertyValidator<Integer> {

    @Override
    public void validate(Integer value) throws HibernateException {
      if (value < 0) {
        throw log.mongoDBTimeOutIllegalValue(value);
      }
    }
  }
}