diff --git a/src/main/java/org/opensearch/ubl/HeaderConstants.java b/src/main/java/org/opensearch/ubl/HeaderConstants.java index b848641..ba41e6f 100644 --- a/src/main/java/org/opensearch/ubl/HeaderConstants.java +++ b/src/main/java/org/opensearch/ubl/HeaderConstants.java @@ -8,8 +8,26 @@ package org.opensearch.ubl; -public class HeaderConstants { +public enum HeaderConstants { - public static final String EVENT_STORE_HEADER = "X-ubi-store"; + QUERY_ID_HEADER("X-ubl-query-id"), + EVENT_STORE_HEADER("X-ubl-store"), + USER_ID_HEADER("X-ubl-user-id"), + SESSION_ID_HEADER("X-ubl-session-id"); + + private final String header; + + private HeaderConstants(String header) { + this.header = header; + } + + public String getHeader() { + return header; + } + + @Override + public String toString() { + return header; + } } diff --git a/src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java b/src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java index a34a1bf..84fd8bd 100644 --- a/src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java +++ b/src/main/java/org/opensearch/ubl/UserBehaviorLoggingPlugin.java @@ -53,12 +53,20 @@ public class UserBehaviorLoggingPlugin extends Plugin implements ActionPlugin { @Override public Collection getRestHeaders() { - return List.of(new RestHeaderDefinition(HeaderConstants.EVENT_STORE_HEADER, false)); + return List.of( + new RestHeaderDefinition(HeaderConstants.EVENT_STORE_HEADER.getHeader(), false), + new RestHeaderDefinition(HeaderConstants.SESSION_ID_HEADER.getHeader(), false), + new RestHeaderDefinition(HeaderConstants.USER_ID_HEADER.getHeader(), false) + ); } @Override public Collection getTaskHeaders() { - return List.of(HeaderConstants.EVENT_STORE_HEADER); + return List.of( + HeaderConstants.EVENT_STORE_HEADER.getHeader(), + HeaderConstants.SESSION_ID_HEADER.getHeader(), + HeaderConstants.USER_ID_HEADER.getHeader() + ); } @Override diff --git a/src/main/java/org/opensearch/ubl/action/UserBehaviorLoggingActionFilter.java b/src/main/java/org/opensearch/ubl/action/UserBehaviorLoggingActionFilter.java index b56c65b..41e30fa 100644 --- a/src/main/java/org/opensearch/ubl/action/UserBehaviorLoggingActionFilter.java +++ b/src/main/java/org/opensearch/ubl/action/UserBehaviorLoggingActionFilter.java @@ -20,7 +20,6 @@ import org.opensearch.core.action.ActionResponse; import org.opensearch.tasks.Task; import org.opensearch.threadpool.ThreadPool; -import org.opensearch.transport.Header; import org.opensearch.ubl.HeaderConstants; import org.opensearch.ubl.backends.Backend; import org.opensearch.ubl.model.QueryRequest; @@ -64,19 +63,8 @@ public void app @Override public void onResponse(Response response) { - //LOGGER.info("Query ID header: " + task.getHeader("query-id")); - final long startTime = System.currentTimeMillis(); - String eventStore = task.getHeader(HeaderConstants.EVENT_STORE_HEADER); - - // If there is no event store header we'll use a "default" store. - if(eventStore == null || eventStore.trim().isEmpty()) { - eventStore = "default"; - } - - LOGGER.info("Using UBL event store: {}", eventStore); - // Get the search itself. final SearchRequest searchRequest = (SearchRequest) request; @@ -88,8 +76,11 @@ public void onResponse(Response response) { // Get all search hits from the response. if (response instanceof SearchResponse) { - // Create a UUID for this search request. - final String queryId = UUID.randomUUID().toString(); + // Get info from the headers. + final String queryId = getHeaderValue(HeaderConstants.QUERY_ID_HEADER, UUID.randomUUID().toString(), task); + final String eventStore = getHeaderValue(HeaderConstants.EVENT_STORE_HEADER, "default", task); + final String userId = getHeaderValue(HeaderConstants.USER_ID_HEADER, "", task); + final String sessionId = getHeaderValue(HeaderConstants.SESSION_ID_HEADER, "", task); // The query will be empty when there is no query, e.g. /_search final String query = searchRequest.source().toString(); @@ -98,19 +89,16 @@ public void onResponse(Response response) { final String queryResponseId = UUID.randomUUID().toString(); final List queryResponseHitIds = new LinkedList<>(); - final SearchResponse searchResponse = (SearchResponse) response; // Add each hit to the list of query responses. - searchResponse.getHits().forEach(hit -> { - queryResponseHitIds.add(String.valueOf(hit.docId())); - }); + searchResponse.getHits().forEach(hit -> queryResponseHitIds.add(String.valueOf(hit.docId()))); try { // Persist the query to the backend. backend.persistQuery(eventStore, - new QueryRequest(queryId, query), + new QueryRequest(queryId, query, userId, sessionId), new QueryResponse(queryId, queryResponseId, queryResponseHitIds)); } catch (Exception ex) { @@ -140,4 +128,16 @@ public void onFailure(Exception ex) { } + private String getHeaderValue(final HeaderConstants header, final String defaultValue, final Task task) { + + final String value = task.getHeader(header.getHeader()); + + if(value == null || value.trim().isEmpty()) { + return defaultValue; + } else { + return value; + } + + } + } diff --git a/src/main/java/org/opensearch/ubl/backends/OpenSearchBackend.java b/src/main/java/org/opensearch/ubl/backends/OpenSearchBackend.java index 7ea1e1c..1fba79d 100644 --- a/src/main/java/org/opensearch/ubl/backends/OpenSearchBackend.java +++ b/src/main/java/org/opensearch/ubl/backends/OpenSearchBackend.java @@ -84,6 +84,7 @@ public void delete(String storeName) { final String eventsIndexName = getEventsIndexName(storeName); final String queriesIndexName = getQueriesIndexName(storeName); final DeleteIndexRequest deleteIndexRequest = new DeleteIndexRequest(eventsIndexName, queriesIndexName); + client.admin().indices().delete(deleteIndexRequest); } @@ -102,17 +103,19 @@ public void persistEvent(String storeName, String eventJson) { } @Override - public void persistQuery(final String storeName, final QueryRequest queryRequest, QueryResponse queryResponse) throws Exception { + public void persistQuery(final String storeName, final QueryRequest queryRequest, QueryResponse queryResponse) { LOGGER.info("Writing query ID {} with response ID {}", queryRequest.getQueryId(), queryResponse.getQueryResponseId()); // What will be indexed - adheres to the queries-mapping.json final Map source = new HashMap<>(); source.put("timestamp", queryRequest.getTimestamp()); - source.put("queryId", queryRequest.getQueryId()); + source.put("query_id", queryRequest.getQueryId()); source.put("query", queryRequest.getQuery()); - source.put("queryResponseId", queryResponse.getQueryResponseId()); - source.put("queryResponseHitIds", queryResponse.getQueryResponseHitIds()); + source.put("query_response_id", queryResponse.getQueryResponseId()); + source.put("query_response_hit_ids", queryResponse.getQueryResponseHitIds()); + source.put("user_id", queryRequest.getUserId()); + source.put("session_id", queryRequest.getSessionId()); // Get the name of the queries. final String queriesIndexName = getQueriesIndexName(storeName); diff --git a/src/main/java/org/opensearch/ubl/model/QueryRequest.java b/src/main/java/org/opensearch/ubl/model/QueryRequest.java index a4ef2c5..7cd8f61 100644 --- a/src/main/java/org/opensearch/ubl/model/QueryRequest.java +++ b/src/main/java/org/opensearch/ubl/model/QueryRequest.java @@ -13,11 +13,15 @@ public class QueryRequest { private final long timestamp; private final String queryId; private final String query; + private final String userId; + private final String sessionId; - public QueryRequest(final String queryId, final String query) { + public QueryRequest(final String queryId, final String query, final String userId, final String sessionId) { this.timestamp = System.currentTimeMillis(); this.queryId = queryId; this.query = query; + this.userId = userId; + this.sessionId = sessionId; } public long getTimestamp() { @@ -32,4 +36,12 @@ public String getQuery() { return query; } + public String getUserId() { + return userId; + } + + public String getSessionId() { + return sessionId; + } + }