Skip to content
This repository has been archived by the owner on Jun 6, 2024. It is now read-only.

Adding query_id to the response headers #43

Merged
merged 2 commits into from
Feb 19, 2024
Merged

Adding query_id to the response headers #43

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
90d7ed5
#9 Adding query_id to the response headers.
jzonthemtn Feb 17, 2024
120170e
Resolving merge conflicts.
jzonthemtn Feb 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,7 @@

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.ubl.action.UserBehaviorLoggingRestHandler;
import org.opensearch.ubl.action.UserBehaviorLoggingSearchFilter;
import org.opensearch.action.support.ActionFilter;
import org.opensearch.ubl.backends.Backend;
import org.opensearch.ubl.backends.OpenSearchBackend;
import org.opensearch.client.Client;
import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
import org.opensearch.cluster.node.DiscoveryNodes;
Expand All @@ -24,14 +20,18 @@
import org.opensearch.core.xcontent.NamedXContentRegistry;
import org.opensearch.env.Environment;
import org.opensearch.env.NodeEnvironment;
import org.opensearch.ubl.events.EventManager;
import org.opensearch.plugins.ActionPlugin;
import org.opensearch.plugins.Plugin;
import org.opensearch.repositories.RepositoriesService;
import org.opensearch.rest.RestController;
import org.opensearch.rest.RestHandler;
import org.opensearch.script.ScriptService;
import org.opensearch.threadpool.ThreadPool;
import org.opensearch.ubl.action.UserBehaviorLoggingActionFilter;
import org.opensearch.ubl.action.UserBehaviorLoggingRestHandler;
import org.opensearch.ubl.backends.Backend;
import org.opensearch.ubl.backends.OpenSearchBackend;
import org.opensearch.ubl.events.EventManager;
import org.opensearch.watcher.ResourceWatcherService;

import java.util.ArrayList;
Expand Down Expand Up @@ -98,7 +98,7 @@ public Collection<Object> createComponents(
) {

this.backend = new OpenSearchBackend(client);
this.userBehaviorLoggingFilter = new UserBehaviorLoggingSearchFilter(backend, environment.settings());
this.userBehaviorLoggingFilter = new UserBehaviorLoggingActionFilter(backend, environment.settings(), threadPool);

LOGGER.info("Creating scheduled task");

Expand Down
32 changes: 20 additions & 12 deletions ...tion/UserBehaviorLoggingSearchFilter.java → ...tion/UserBehaviorLoggingActionFilter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,25 +16,29 @@
import org.opensearch.action.support.ActionFilter;
import org.opensearch.action.support.ActionFilterChain;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.util.concurrent.ThreadContext;
import org.opensearch.core.action.ActionListener;
import org.opensearch.core.action.ActionResponse;
import org.opensearch.threadpool.ThreadPool;
import org.opensearch.ubl.model.QueryResponse;
import org.opensearch.ubl.backends.Backend;
import org.opensearch.ubl.model.QueryRequest;
import org.opensearch.tasks.Task;

import java.util.*;

public class UserBehaviorLoggingSearchFilter implements ActionFilter {
public class UserBehaviorLoggingActionFilter implements ActionFilter {

private static final Logger LOGGER = LogManager.getLogger(UserBehaviorLoggingSearchFilter.class);
private static final Logger LOGGER = LogManager.getLogger(UserBehaviorLoggingActionFilter.class);

private final Backend backend;
private final Settings settings;
private final ThreadPool threadPool;

public UserBehaviorLoggingSearchFilter(final Backend backend, final Settings settings) {
public UserBehaviorLoggingActionFilter(final Backend backend, final Settings settings, ThreadPool threadPool) {
this.backend = backend;
this.settings = settings;
this.threadPool = threadPool;
}

@Override
Expand All @@ -57,6 +61,8 @@ public <Request extends ActionRequest, Response extends ActionResponse> void app
@Override
public void onResponse(Response response) {

LOGGER.info("Query ID header: " + task.getHeader("query-id"));

final long startTime = System.currentTimeMillis();

// Get the search itself.
Expand All @@ -67,19 +73,19 @@ public void onResponse(Response response) {
//final Set<String> indicesToLog = new HashSet<>(Arrays.asList(settings.get(SettingsConstants.INDEX_NAMES).split(",")));
//if(indicesToLog.containsAll(indices)) {

// Create a UUID for this search request.
final String queryId = UUID.randomUUID().toString();
// Get all search hits from the response.
if (response instanceof SearchResponse) {

// The query will be empty when there is no query, e.g. /_search
final String query = searchRequest.source().toString();
// Create a UUID for this search request.
final String queryId = UUID.randomUUID().toString();

// Create a UUID for this search response.
final String queryResponseId = UUID.randomUUID().toString();
// The query will be empty when there is no query, e.g. /_search
final String query = searchRequest.source().toString();

final List<String> queryResponseHitIds = new LinkedList<>();
// Create a UUID for this search response.
final String queryResponseId = UUID.randomUUID().toString();

// Get all search hits from the response.
if (response instanceof SearchResponse) {
final List<String> queryResponseHitIds = new LinkedList<>();

final SearchResponse searchResponse = (SearchResponse) response;

Expand All @@ -101,6 +107,8 @@ public void onResponse(Response response) {
LOGGER.error("Unable to persist query.", ex);
}

threadPool.getThreadContext().addResponseHeader("query_id", queryId);

}

//}
Expand Down
Loading