Examples with AsynchronousSearchPersistenceModel org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceModel used on opensource projects

Example 1 with AsynchronousSearchPersistenceModel

use of org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceModel in project asynchronous-search by opensearch-project.

the class AsynchronousSearchPersistenceService method updateExpirationTime.

/**
 * Updates the expiration time field in index
 *
 * @param id                   asynchronous search id
 * @param expirationTimeMillis the new expiration time
 * @param user                 current user
 * @param listener             listener invoked with the response on completion of update request
 */
@SuppressWarnings("unchecked")
public void updateExpirationTime(String id, long expirationTimeMillis, User user, ActionListener<AsynchronousSearchPersistenceModel> listener) {
    if (indexExists() == false) {
        listener.onFailure(new ResourceNotFoundException(id));
        return;
    }
    UpdateRequest updateRequest = new UpdateRequest(ASYNC_SEARCH_RESPONSE_INDEX, id);
    updateRequest.retryOnConflict(5);
    if (user == null) {
        Map<String, Object> source = new HashMap<>();
        source.put(EXPIRATION_TIME_MILLIS, expirationTimeMillis);
        updateRequest.doc(source, XContentType.JSON);
    } else {
        String scriptCode = "if (ctx._source.user == null || ctx._source.user.backend_roles == null || " + "(params.backend_roles != null && params.backend_roles.containsAll(ctx._source.user.backend_roles))) " + "{ ctx._source.expiration_time_millis = params.expiration_time_millis } else { ctx.op = 'none' }";
        Map<String, Object> params = new HashMap<>();
        params.put(BACKEND_ROLES, user.getBackendRoles());
        params.put(EXPIRATION_TIME_MILLIS, expirationTimeMillis);
        Script conditionalUpdateScript = new Script(ScriptType.INLINE, "painless", scriptCode, params);
        updateRequest.script(conditionalUpdateScript);
    }
    updateRequest.fetchSource(FetchSourceContext.FETCH_SOURCE);
    client.update(updateRequest, ActionListener.wrap(updateResponse -> {
        switch(updateResponse.getResult()) {
            case NOOP:
                if (user != null) {
                    listener.onFailure(new OpenSearchSecurityException("User doesn't have necessary roles to access the asynchronous search with id " + id, RestStatus.FORBIDDEN));
                } else {
                    Map<String, Object> updatedSource = updateResponse.getGetResult().getSource();
                    listener.onResponse(new AsynchronousSearchPersistenceModel((long) updatedSource.get(START_TIME_MILLIS), (long) updatedSource.get(EXPIRATION_TIME_MILLIS), (String) updatedSource.get(RESPONSE), (String) updatedSource.get(ERROR), parseUser((Map<String, Object>) updatedSource.get(USER))));
                }
                break;
            case UPDATED:
                Map<String, Object> updatedSource = updateResponse.getGetResult().getSource();
                listener.onResponse(new AsynchronousSearchPersistenceModel((long) updatedSource.get(START_TIME_MILLIS), (long) updatedSource.get(EXPIRATION_TIME_MILLIS), (String) updatedSource.get(RESPONSE), (String) updatedSource.get(ERROR), parseUser((Map<String, Object>) updatedSource.get(USER))));
                break;
            case NOT_FOUND:
            case DELETED:
                logger.debug("Update Result [{}] for id [{}], expiration time requested, [{}]", updateResponse.getResult(), id, expirationTimeMillis);
                listener.onFailure(new ResourceNotFoundException(id));
                break;
        }
    }, exception -> {
        final Throwable cause = ExceptionsHelper.unwrapCause(exception);
        if (cause instanceof DocumentMissingException) {
            listener.onFailure(new ResourceNotFoundException(id));
        } else {
            logger.error(() -> new ParameterizedMessage("Exception occurred updating expiration time for asynchronous search [{}]", id), exception);
            listener.onFailure(cause instanceof Exception ? (Exception) cause : new NotSerializableExceptionWrapper(cause));
        }
    }));
}

Example 2 with AsynchronousSearchPersistenceModel

use of org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceModel in project asynchronous-search by opensearch-project.

the class AsynchronousSearchPersistenceServiceIT method testGetAndDeleteNonExistentId.

public void testGetAndDeleteNonExistentId() throws InterruptedException, IOException, ExecutionException {
    AsynchronousSearchPersistenceService persistenceService = getInstanceFromNode(AsynchronousSearchPersistenceService.class);
    TransportService transportService = getInstanceFromNode(TransportService.class);
    SearchResponse searchResponse = client().search(new SearchRequest(TEST_INDEX)).get();
    User user1 = TestClientUtils.randomUser();
    User user2 = TestClientUtils.randomUser();
    for (User originalUser : Arrays.asList(user1, null)) {
        AsynchronousSearchId asId = generateNewAsynchronousSearchId(transportService);
        AsynchronousSearchPersistenceModel model1 = new AsynchronousSearchPersistenceModel(System.currentTimeMillis(), System.currentTimeMillis() + new TimeValue(10, TimeUnit.DAYS).getMillis(), searchResponse, null, originalUser);
        CountDownLatch createLatch = new CountDownLatch(1);
        String id = AsynchronousSearchIdConverter.buildAsyncId(asId);
        persistenceService.storeResponse(id, model1, new LatchedActionListener<>(ActionListener.wrap(r -> assertSuccessfulResponseCreation(id, r), e -> {
            logger.debug("expect successful create, got", e);
            fail("Expected successful create, got " + e.getMessage());
        }), createLatch));
        createLatch.await();
        for (User currentuser : Arrays.asList(originalUser, user2)) {
            CountDownLatch latch = new CountDownLatch(2);
            // assert failure
            persistenceService.getResponse("id", currentuser, new LatchedActionListener<>(ActionListener.wrap((AsynchronousSearchPersistenceModel r) -> fail("Excepted resource_not_found_exception, got " + r), exception -> assertTrue("Expected resource_not_found expection, got " + exception.getClass().toString(), exception instanceof ResourceNotFoundException)), latch));
            // assert failure
            ActionListener<Boolean> wrap = ActionListener.wrap(r -> fail("Expected resource_not_found expection on delete, got acknowledgement " + r), ex -> assertTrue("Expected resource_not_found expection, got " + ex.getClass().toString(), ex instanceof ResourceNotFoundException));
            persistenceService.deleteResponse("id", currentuser, new LatchedActionListener<>(wrap, latch));
            latch.await();
        }
    }
}

Example 3 with AsynchronousSearchPersistenceModel

use of org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceModel in project asynchronous-search by opensearch-project.

the class AsynchronousSearchPersistenceServiceIT method testCreateAndGetAndDelete.

public void testCreateAndGetAndDelete() throws IOException, InterruptedException {
    AsynchronousSearchPersistenceService persistenceService = getInstanceFromNode(AsynchronousSearchPersistenceService.class);
    TransportService transportService = getInstanceFromNode(TransportService.class);
    AsynchronousSearchResponse asResponse = submitAndGetPersistedAsynchronousSearchResponse();
    AsynchronousSearchContextId asContextId = new AsynchronousSearchContextId(UUIDs.base64UUID(), randomInt(100));
    AsynchronousSearchId newAsynchronousSearchId = new AsynchronousSearchId(transportService.getLocalNode().getId(), 1, asContextId);
    String id = AsynchronousSearchIdConverter.buildAsyncId(newAsynchronousSearchId);
    User user1 = TestClientUtils.randomUser();
    User user2 = TestClientUtils.randomUser();
    for (User user : Arrays.asList(user1, null)) {
        AsynchronousSearchResponse newAsynchronousSearchResponse = new AsynchronousSearchResponse(id, AsynchronousSearchState.STORE_RESIDENT, asResponse.getStartTimeMillis(), asResponse.getExpirationTimeMillis(), asResponse.getSearchResponse(), asResponse.getError());
        createDoc(persistenceService, newAsynchronousSearchResponse, user);
        if (user != null) {
            CountDownLatch getLatch1 = new CountDownLatch(1);
            ActionListener<AsynchronousSearchPersistenceModel> getListener = ActionListener.wrap(r -> fail("Expected exception. Got " + r), e -> assertTrue(e instanceof OpenSearchSecurityException));
            persistenceService.getResponse(newAsynchronousSearchResponse.getId(), user2, new LatchedActionListener<>(getListener, getLatch1));
            getLatch1.await();
        }
        CountDownLatch getLatch2 = new CountDownLatch(1);
        persistenceService.getResponse(newAsynchronousSearchResponse.getId(), user, new LatchedActionListener<>(ActionListener.wrap(r -> assertEquals(new AsynchronousSearchPersistenceModel(asResponse.getStartTimeMillis(), asResponse.getExpirationTimeMillis(), asResponse.getSearchResponse(), null, user), r), e -> {
            logger.error("Expected get result got ", e);
            fail(e.getMessage());
        }), getLatch2));
        getLatch2.await();
        if (user != null) {
            CountDownLatch deleteLatch1 = new CountDownLatch(1);
            User diffUser = TestClientUtils.randomUser();
            ActionListener<Boolean> deleteListener = ActionListener.wrap(r -> fail("Expected exception on delete. Got acknowledgment" + r), e -> assertTrue(e instanceof OpenSearchSecurityException));
            persistenceService.deleteResponse(newAsynchronousSearchResponse.getId(), diffUser, new LatchedActionListener<>(deleteListener, deleteLatch1));
            deleteLatch1.await();
        }
        CountDownLatch deleteLatch2 = new CountDownLatch(1);
        ActionListener<Boolean> deleteListener = ActionListener.wrap(Assert::assertTrue, e -> {
            logger.debug(() -> new ParameterizedMessage("Delete failed unexpectedly "), e);
            fail("delete failed.expected success");
        });
        persistenceService.deleteResponse(newAsynchronousSearchResponse.getId(), user, new LatchedActionListener<>(deleteListener, deleteLatch2));
        deleteLatch2.await();
        // assert failure
        CountDownLatch getLatch3 = new CountDownLatch(2);
        ActionListener<AsynchronousSearchPersistenceModel> getListener = ActionListener.wrap((r) -> fail("Expected RNF, Got " + r), exception -> assertTrue(exception instanceof ResourceNotFoundException));
        persistenceService.getResponse(newAsynchronousSearchResponse.getId(), null, new LatchedActionListener<>(getListener, getLatch3));
        persistenceService.getResponse(newAsynchronousSearchResponse.getId(), user2, new LatchedActionListener<>(getListener, getLatch3));
        getLatch3.await();
    }
}

Example 4 with AsynchronousSearchPersistenceModel

use of org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceModel in project asynchronous-search by opensearch-project.

the class AsynchronousSearchPersistenceServiceIT method testDeleteExpiredResponse.

public void testDeleteExpiredResponse() throws InterruptedException, IOException {
    AsynchronousSearchPersistenceService persistenceService = getInstanceFromNode(AsynchronousSearchPersistenceService.class);
    AsynchronousSearchResponse asResponse = submitAndGetPersistedAsynchronousSearchResponse();
    CountDownLatch updateLatch = new CountDownLatch(1);
    long newExpirationTime = System.currentTimeMillis() + new TimeValue(100, TimeUnit.MILLISECONDS).getMillis();
    final AsynchronousSearchPersistenceModel newPersistenceModel = new AsynchronousSearchPersistenceModel(asResponse.getStartTimeMillis(), newExpirationTime, asResponse.getSearchResponse(), null, null);
    persistenceService.updateExpirationTime(asResponse.getId(), newExpirationTime, null, new LatchedActionListener<>(ActionListener.wrap(persistenceModel -> assertEquals(newPersistenceModel, persistenceModel), e -> {
        logger.debug("expect successful create, got", e);
        fail("Expected successful update result, got " + e.getMessage());
    }), updateLatch));
    updateLatch.await();
    CountDownLatch getLatch = new CountDownLatch(1);
    persistenceService.getResponse(asResponse.getId(), null, new LatchedActionListener<>(ActionListener.wrap(r -> assertEquals(newPersistenceModel, r), e -> {
        logger.debug("expect successful create, got", e);
        fail("Expected successful create, got " + e.getMessage());
    }), getLatch));
    getLatch.await();
    CountDownLatch deleteLatch = new CountDownLatch(1);
    persistenceService.deleteExpiredResponses(new LatchedActionListener<>(new ActionListener<AcknowledgedResponse>() {

        @Override
        public void onResponse(AcknowledgedResponse acknowledgedResponse) {
            assertTrue(acknowledgedResponse.isAcknowledged());
        }

        @Override
        public void onFailure(Exception e) {
            fail("Received exception while deleting expired response " + e.getMessage());
        }
    }, deleteLatch), System.currentTimeMillis());
    deleteLatch.await();
}

Example 5 with AsynchronousSearchPersistenceModel

use of org.opensearch.search.asynchronous.context.persistence.AsynchronousSearchPersistenceModel in project asynchronous-search by opensearch-project.

the class AsynchronousSearchPersistenceServiceIT method testUpdateExpiration.

public void testUpdateExpiration() throws InterruptedException, IOException {
    AsynchronousSearchPersistenceService persistenceService = getInstanceFromNode(AsynchronousSearchPersistenceService.class);
    ThreadPool threadPool1 = getInstanceFromNode(ThreadPool.class);
    User user1 = TestClientUtils.randomUser();
    User user2 = TestClientUtils.randomUser();
    for (User originalUser : Arrays.asList(user1, null)) {
        try (ThreadContext.StoredContext ctx = threadPool1.getThreadContext().stashContext()) {
            threadPool1.getThreadContext().putTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, getUserRolesString(originalUser));
            AsynchronousSearchResponse asResponse = submitAndGetPersistedAsynchronousSearchResponse();
            long newExpirationTime = System.currentTimeMillis() + new TimeValue(10, TimeUnit.DAYS).getMillis();
            final AsynchronousSearchPersistenceModel newPersistenceModel = new AsynchronousSearchPersistenceModel(asResponse.getStartTimeMillis(), newExpirationTime, asResponse.getSearchResponse(), null, originalUser);
            for (User currentUser : Arrays.asList(user2, user1, null)) {
                CountDownLatch updateLatch = new CountDownLatch(1);
                if (originalUser != null && currentUser != null && currentUser.equals(originalUser) == false) {
                    ActionListener<AsynchronousSearchPersistenceModel> updateListener = ActionListener.wrap(r -> fail("Expected security exception. Unauthorized update. Got " + r), e -> assertTrue(e instanceof OpenSearchSecurityException));
                    persistenceService.updateExpirationTime(asResponse.getId(), newExpirationTime, currentUser, new LatchedActionListener<>(updateListener, updateLatch));
                } else {
                    persistenceService.updateExpirationTime(asResponse.getId(), newExpirationTime, currentUser, new LatchedActionListener<>(ActionListener.wrap(persistenceModel -> assertEquals(newPersistenceModel, persistenceModel), e -> {
                        logger.debug("expect successful create, got", e);
                        fail("Expected successful create, got " + e.getMessage());
                    }), updateLatch));
                }
                updateLatch.await();
            }
            CountDownLatch getLatch = new CountDownLatch(1);
            persistenceService.getResponse(asResponse.getId(), originalUser, new LatchedActionListener<>(ActionListener.wrap(r -> assertEquals(newPersistenceModel, r), e -> {
                logger.debug("expect successful get result, got", e);
                fail("Expected successful get result, got " + e.getMessage());
            }), getLatch));
            getLatch.await();
        }
    }
}