com.orientechnologies.orient.core.db.ODatabaseRecordThreadLocal Java Examples

private static <T> T withActiveDb(final ODatabase db, final Supplier<T> supplier) {
  @SuppressWarnings("resource")
  final ODatabaseDocumentInternal currentDb = ODatabaseRecordThreadLocal.instance().getIfDefined();
  if (db.equals(currentDb) || !(db instanceof ODatabaseDocumentInternal)) {
    return supplier.get();
  }
  try {
    ODatabaseRecordThreadLocal.instance().set((ODatabaseDocumentInternal) db);
    return supplier.get();
  }
  finally {
    if (currentDb != null) {
      ODatabaseRecordThreadLocal.instance().set(currentDb);
    }
    else {
      ODatabaseRecordThreadLocal.instance().remove();
    }
  }
}
 

@Override
public ODatabaseDocument get() {
    // lazy get: pool transaction will start not together with TransactionManager one, but as soon as
    // connection requested to avoid using connections of not used pools
    Preconditions.checkNotNull(pool, String.format("Pool %s not initialized", getType()));
    if (transaction.get() == null) {
        Preconditions.checkState(transactionManager.isTransactionActive(), String.format(
                "Can't obtain connection from pool %s: no transaction defined.", getType()));
        if (transactionManager.isExternalTransaction()) {
            // external mode: use already created connection
            transaction.set(ODatabaseRecordThreadLocal.instance().get());
            logger.trace("Pool {} use bound to thread connection (external mode)", getType());
        } else {
            // normal mode: create connection
            final ODatabaseDocument db = checkAndAcquireConnection();

            db.begin(transactionManager.getActiveTransactionType());
            transaction.set(db);
            logger.trace("Pool {} transaction started", getType());
        }
    }
    return (ODatabaseDocument) checkOpened(transaction.get()).activateOnCurrentThread();
}
 

protected void replaceStorage(final OPartitionedDatabasePool pool, final OStorage storage) {
  if (partitionsField != null) {
    ODatabaseDocumentInternal originalDb = ODatabaseRecordThreadLocal.instance().getIfDefined();
    try {
      // use reflection as workaround until public API is available
      for (Object partition : (Object[]) partitionsField.get(pool)) {
        for (ODatabaseDocumentTx db : (Iterable<ODatabaseDocumentTx>) partitionQueueField.get(partition)) {
          replaceStorage(db, storage);
        }
      }
    }
    catch (Exception | LinkageError e) {
      log.warn("Problem replacing storage for {}", storage.getName(), e);
    }
    finally {
      ODatabaseRecordThreadLocal.instance().set(originalDb);
    }
  }
}
 

@After
@Sudo
public void destroy() {
  ODatabaseDocument db = ODatabaseRecordThreadLocal.instance().get();

  notifications.forEach(notification -> {
    for (int i = 0; i < 3; i++) {
      try {
        db.begin();
        db.delete(notification.getDocument());
        db.commit();
        break;
      } catch (Exception e) {
        notification.reload();
      }
    }
  });
}
 

/**
 * Core orient field resolve method relies on bound connection, but it may be required to resolve
 * id outside of transaction. So we use orient method under transaction and do manual scan outside
 * of transaction.
 *
 * @param value object instance (non proxy)
 * @return object id field or null if not found
 */
private static Field findIdField(final Object value) {
    Field res = null;
    final Class<?> type = value.getClass();
    if (ODatabaseRecordThreadLocal.instance().isDefined()) {
        res = OObjectEntitySerializer.getIdField(type);
    } else {
        final String idField = OObjectSerializerHelper.getObjectIDFieldName(value);
        if (idField != null) {
            try {
                res = type.getDeclaredField(idField);
            } catch (NoSuchFieldException e) {
                LOGGER.warn(String
                        .format("Id field '%s' not found in class '%s'.", idField, type.getSimpleName()), e);
            }
        }
    }
    return res;
}
 

@After
@Sudo
public void destroy() {
  ODatabaseDocument db = ODatabaseRecordThreadLocal.instance().get();
  for (int i = 1; i <= 10; i++) {
    try {
      db.begin();
      db.delete(testNotification.getDocument());
      db.commit();
      break;
    } catch (Exception e) {
      if (i == 10) {
        throw new IllegalStateException(e);
      }
    }
  }
}
 

@Override
public Object execute(Object iThis, OIdentifiable iCurrentRecord, Object iCurrentResult, Object[] iParams,
    OCommandContext iContext) {

  String clazz = (String) iParams[0];
  String latField = (String) iParams[1];
  String lngField = (String) iParams[2];
  ODatabaseDocument databaseRecord = ODatabaseRecordThreadLocal.INSTANCE.get();
  Set<OIndex<?>> indexes = databaseRecord.getMetadata().getSchema().getClass(clazz).getInvolvedIndexes(latField, lngField);
  for (OIndex i : indexes) {
    if (OClass.INDEX_TYPE.SPATIAL.toString().equals(i.getInternal().getType())) {
      List<Object> params = new ArrayList<Object>();
      params.add(iParams[3]);
      params.add(iParams[4]);
      double distance = iParams.length > 5 ? ((Number) iParams[5]).doubleValue() : 0;
      Object ret = i.get(new OSpatialCompositeKey(params).setMaxDistance(distance));
      if (ret instanceof Collection) {
        if (context == null)
          context = new HashSet<Object>();
        context.addAll((Collection<?>) ret);
      }
      return ret;
    }
  }
  return null;
}
 

/**
 * @return result of execution
 */
public final V execute()
{
	ODatabaseDocument db = null;
	ODatabaseRecordThreadLocal orientDbThreadLocal = ODatabaseRecordThreadLocal.instance();
	ODatabaseDocument oldDb = orientDbThreadLocal.getIfDefined();
	if(oldDb!=null) orientDbThreadLocal.remove(); //Required to avoid stack of transactions
	try
	{
		db = getSettings().getDatabasePoolFactory().get(getDBUrl(), getUsername(), getPassword()).acquire();
		db.activateOnCurrentThread();
		return execute(db);
	} 
	finally
	{
		if(db!=null) db.close();
		if(oldDb!=null) orientDbThreadLocal.set((ODatabaseDocumentInternal)oldDb);
		else orientDbThreadLocal.remove();
	}
}
 

public void testOrientDbLifeCycle(String dbURL, boolean createDb, boolean dropDb) throws Exception
	{
		Orient.instance().startup();
		assertNotNull(ODatabaseRecordThreadLocal.instance());
		Orient.instance().removeShutdownHook();
		OServer server = OServerMain.create();
		server.startup(OrientDbTestWebApplication.class.getResource("db.config.xml").openStream());
		server.activate();
		if(createDb)
		{
			ODatabaseDocument dbToCreate = new ODatabaseDocumentTx(dbURL);
			if(!dbToCreate.exists()) dbToCreate.create();
			dbToCreate.close();
		}
		assertNotNull(ODatabaseRecordThreadLocal.instance());
		ODatabaseDocument db = new OPartitionedDatabasePoolFactory().get(dbURL, "admin", "admin").acquire();
		db.close();
		assertNotNull(ODatabaseRecordThreadLocal.instance());
		if(dropDb)
		{
			@SuppressWarnings("resource")
			ODatabaseDocument dbToDrop = new ODatabaseDocumentTx(dbURL);
			dbToDrop.open("admin", "admin");
			dbToDrop.drop();
		}
		server.shutdown();
		Orient.instance().shutdown();
//		Thread.sleep(50);
	}
 

/**
 * Returns the current DB connection. Normally this is whichever connection is active on the thread.
 * But when performing a commit we always return the connection which is being committed, so we can
 * track changes even when another connection is used to "fix" records during that commit.
 */
private ODatabaseDocumentInternal getCurrrentDb() {
  ODatabase db = commitDb.get();
  if (db == null) {
    db = ODatabaseRecordThreadLocal.instance().get();
  }
  return (ODatabaseDocumentInternal) db;
}
 

@Override
public void begin(final TxConfig config) {
    if (transaction.get() != null) {
        // transaction already in progress
        return;
    }
    if (config != null && config.isExternal()) {
        Preconditions.checkState(!ODatabaseRecordThreadLocal.instance().get().isClosed(),
                "Can't start external unit of work: connection bound to thread is closed");
    }
    transaction.set(MoreObjects.firstNonNull(config, defaultConfig));
    logger.trace("Transaction opened: {}", transaction.get());
}
 

@Test(enabled = false)
public void deInitDB() {
  if (remote) {
    process.destroy();
  } else {
    databaseDocumentTx.drop();
    ODatabaseRecordThreadLocal.INSTANCE.set(null);
  }
}
 

@Override
@SuppressWarnings("unchecked")
public void send(List<ODocument> notifications) {
  if (notifications == null || notifications.isEmpty()) {
    return;
  }
  ODatabaseDocument db = ODatabaseRecordThreadLocal.instance().get();

  IONotification notification = DAO.create(IONotification.class);
  IONotificationTransport transportWrapper = DAO.create(IONotificationTransport.class);

  notifications.forEach(notificationDoc -> {
    notification.fromStream(notificationDoc);
    transportWrapper.fromStream(notification.getTransport());
    ITransport transport = transportPool.acquire(transportWrapper.getAlias(), transportWrapper::createTransportService);

    for (int i = 1; i <= ATTEMPTS; i++) {
      try {
        if (i == 1) {
          handleSendingNotificationStatus(db, notification);
        }
        LOG.info("Send notification: {} {}", Thread.currentThread().getName(), notification.getDocument());
        transport.send(notificationDoc);
        handleSentNotificationStatus(db, notification);
        transportPool.release(transportWrapper.getAlias(), transport);
        break;
      } catch (Exception e) {
        if (i == ATTEMPTS) {
          handleFailedNotificationStatus(db, notification, e);
        }
      }
    }
  });
}
 

@After
@Sudo
public void destroy() {
  ODatabaseDocument db = ODatabaseRecordThreadLocal.instance().get();
  for (int i = 0; i < 10; i++) {
    try {
      db.delete(testNotification.getDocument());
      break;
    } catch (Exception e) {
      testNotification.reload();
    }
  }

}
 

@Before
@Sudo
public void init() {
  ONotificationScheduler.stopAll();

  ODatabaseDocument db = ODatabaseRecordThreadLocal.instance().get();

  ODocument mailTransport = notificationDAO.findTransportByAlias(TestDataModule.TRANSPORT_MAIL);

  if (mailTransport == null) {
    throw new IllegalStateException("There is no transport with alias: " + TestDataModule.TRANSPORT_MAIL);
  }

  OMail mail = OMailUtils.getOMailByName(TestDataModule.MAIL_TEST)
          .orElseThrow(IllegalStateException::new);

  notifications = new LinkedList<>();

  for (int i = 0; i < NOTIFICATIONS; i++) {
    db.begin();
    OPreparedMail preparedMail = new OPreparedMail(mail);
    IOMailNotification notification = DAO.create(IOMailNotification.class);
    notification.fromStream(new ODocument(IOMailNotification.CLASS_NAME));
    notification.setTransport(mailTransport);
    notification.setPreparedMail(preparedMail.getDocument());
    preparedMail.addRecipient("[email protected]");
    preparedMail.save();
    notification.save();
    notifications.add(notification);
    db.commit();
  }
}
 

@After
@Sudo
public void destroy() {
  ODatabaseDocument db = ODatabaseRecordThreadLocal.instance().get();
  for (int i = 0; i < 10; i++) {
    try {
      db.delete(notification.getDocument());
      break;
    } catch (Exception e) {
      notification.reload();
    }
  }
}
 

private LoadingCache<OClass, Collection<OProperty>> getCache()
{
	@SuppressWarnings("deprecation")
	int version = ODatabaseRecordThreadLocal.INSTANCE.get().getMetadata().getSchema().getVersion();
	if(version>currentSchemaVersion)
	{
		CACHE.invalidateAll();
		currentSchemaVersion=version;
	}
	return CACHE;
}
 

@Provides
public ODatabaseDocument getDatabaseRecord()
{
	ODatabaseDocument db = DefaultODatabaseThreadLocalFactory.castToODatabaseDocument(ODatabaseRecordThreadLocal.instance().get().getDatabaseOwner());
	if(db.isClosed())
	{
		ODatabaseRecordThreadLocal.instance().remove();
		db = DefaultODatabaseThreadLocalFactory.castToODatabaseDocument(ODatabaseRecordThreadLocal.instance().get().getDatabaseOwner());
	}
	return db;
}
 

@Override
public void onDetach(RequestCycle cycle) {
	ODatabaseDocument db = ODatabaseRecordThreadLocal.instance().getIfDefined();
	if(db!=null) {
		if(db.getTransaction().isActive()) db.commit(true);
		db.close();
		ODatabaseRecordThreadLocal.instance().remove();
	}
}
 

private OLocalPaginatedStorage storage() {
  if (storage == null || storage.isClosed()) {
    // use temp TX to get local storage; note we don't need a TX when reading write-ahead-log
    ODatabaseDocumentInternal currentDb = ODatabaseRecordThreadLocal.instance().getIfDefined();
    try (ODatabaseDocumentInternal db = databaseProvider.get().acquire()) {
      storage = (OLocalPaginatedStorage) db.getStorage().getUnderlying();
    }
    finally {
      ODatabaseRecordThreadLocal.instance().set(currentDb);
    }
  }
  return storage;
}
 

@Override
public void cleanup(final String context) {
  ODatabaseDocument database = ODatabaseRecordThreadLocal.instance().getIfDefined();
  if (database != null && database.getStatus() == STATUS.OPEN) {
    database.close();
    log.warn("{} left a database connection open. Any opened connections should also be explicitly closed.", context);
  }
}
 

@Test
public void toStringOutsideTxShouldNotLogSerializerError() {
  EntityMetadata metadata;

  try (ODatabaseDocumentTx db = database.getInstance().acquire()) {
    entityAdapter.register(db);

    // CREATE
    db.begin();
    TestEntity entity = new TestEntity();
    entity.primary = "Hello";
    entity.secondary = "World";
    entityAdapter.set(db, entity);
    db.commit();
  }

  try (ODatabaseDocumentTx db = database.getInstance().acquire()) {
    metadata = entityAdapter.get(db).getEntityMetadata();
  }

  ODatabaseRecordThreadLocal.INSTANCE.remove(); // remove reference to the originating database

  String toStringOutput = metadata.toString(); // attempt to dump the attached document's fields

  ArgumentCaptor<ILoggingEvent> events = ArgumentCaptor.forClass(ILoggingEvent.class);
  verify(mockAppender, atLeastOnce()).doAppend(events.capture());

  events.getAllValues().forEach(
      event -> assertThat(event.getFormattedMessage(), not(containsString("Error deserializing record"))));

  if (entityAdapter.isPartialEntity()) {
    assertThat(toStringOutput, containsString("{primary:Hello}"));
  }
  else {
    assertThat(toStringOutput, containsString("{primary:Hello,secondary:World}"));
  }
}
 

/**
 * Utility function for working around "ODatabaseException: Database instance is not set in current thread" issues.
 * The current database ThreadLocal is preserved and restored after calling the lambda.
 */
public static <T> T withOtherDatabase(Callable<T> function) {
  final ODatabaseDocumentInternal db = ODatabaseRecordThreadLocal.INSTANCE.getIfDefined();
  try {
    return function.call();
  }
  catch (Exception e) {
    Throwables.throwIfUnchecked(e);
    throw new RuntimeException(e);
  }
  finally {
    ODatabaseRecordThreadLocal.INSTANCE.set(db);
  }
}
 

@Override
public void signOut() {
	super.signOut();
	this.username=null;
	this.password=null;
	this.userModel.setObject(null);
	ODatabaseRecordThreadLocal.instance().remove();
}
 

@Override
public boolean isAllowed(ODocument iDocument,
		ORestrictedOperation iAllowOperation, boolean iReadOriginal) {
	database = ODatabaseRecordThreadLocal.instance().get();
	return super.isAllowed(iDocument, iAllowOperation, iReadOriginal);
}
 

@Override
public IRequestHandler onException(RequestCycle cycle, Exception ex) {
	ODatabaseDocument db = ODatabaseRecordThreadLocal.instance().getIfDefined();
	if(db!=null && !db.isClosed()) db.rollback();
	return null;
}
 

@Provides
public ODatabaseDocument getDatabaseRecord()
{
	return DefaultODatabaseThreadLocalFactory.castToODatabaseDocument(ODatabaseRecordThreadLocal.instance().get().getDatabaseOwner());
}
 

/**
 * Only allow detaching when we have no DB context or it is from a different DB.
 * If we have a valid context from the same DB then we don't need to detach yet.
 */
private boolean allowDetach() {
  ODatabase<?> db = ODatabaseRecordThreadLocal.INSTANCE.getIfDefined();
  return db == null || !getDbName().equals(db.getName());
}
 

public BpmnHook() {
  this.database = ODatabaseRecordThreadLocal.instance().get();
}
 

@Override
public void end(RequestCycle cycle) {
	ODatabaseDocument db = ODatabaseRecordThreadLocal.instance().getIfDefined();
	if(db!=null && db.getTransaction().isActive()) db.commit();
}