LENS-753: Queue number for queries in submitted queue
Project: http://git-wip-us.apache.org/repos/asf/lens/repo Commit: http://git-wip-us.apache.org/repos/asf/lens/commit/4addd7b6 Tree: http://git-wip-us.apache.org/repos/asf/lens/tree/4addd7b6 Diff: http://git-wip-us.apache.org/repos/asf/lens/diff/4addd7b6 Branch: refs/heads/current-release-line Commit: 4addd7b62ed1db8fe1af498a6baf55ee35ad692a Parents: 3563aac Author: Akshay Goyal <[email protected]> Authored: Tue Sep 15 16:29:32 2015 +0530 Committer: Rajat Khandelwal <[email protected]> Committed: Tue Sep 15 16:29:32 2015 +0530 ---------------------------------------------------------------------- .../org/apache/lens/api/query/QueryStatus.java | 38 +++++++++++++++---- .../lens/cli/commands/LensQueryCommands.java | 18 +-------- .../server/api/driver/DriverQueryStatus.java | 7 ++-- .../server/api/query/FinishedLensQuery.java | 2 +- .../lens/server/api/query/QueryContext.java | 2 +- .../query/collect/ImmutableQueryCollection.java | 6 +++ .../server/query/QueryExecutionServiceImpl.java | 26 ++++++++----- .../DefaultEstimatedQueryCollection.java | 5 +++ .../query/collect/DefaultQueryCollection.java | 36 ++++++++++++++++-- .../ThreadSafeEstimatedQueryCollection.java | 5 +++ .../collect/ThreadSafeQueryCollection.java | 5 +++ .../collect/DefaultQueryCollectionTest.java | 34 ++++++++++++++++- .../server/query/collect/QueryCollectUtil.java | 39 +++++++++++++++++++- 13 files changed, 178 insertions(+), 45 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-api/src/main/java/org/apache/lens/api/query/QueryStatus.java ---------------------------------------------------------------------- diff --git a/lens-api/src/main/java/org/apache/lens/api/query/QueryStatus.java b/lens-api/src/main/java/org/apache/lens/api/query/QueryStatus.java index 9614caa..3c8531f 100644 --- a/lens-api/src/main/java/org/apache/lens/api/query/QueryStatus.java +++ b/lens-api/src/main/java/org/apache/lens/api/query/QueryStatus.java @@ -39,6 +39,8 @@ import lombok.*; * * @param progress * the progress + * @param queueNumber + * the queue number * @param status * the status * @param statusMessage @@ -121,6 +123,13 @@ public class QueryStatus implements Serializable { private double progress; /** + * Queue number of a query when it is in waiting state. + */ + @Getter + @Setter + private Integer queueNumber; + + /** * The status. */ @XmlElement @@ -165,20 +174,33 @@ public class QueryStatus implements Serializable { */ @Override public String toString() { - StringBuilder str = new StringBuilder(status.toString()).append(':').append(statusMessage); - if (status.equals(Status.RUNNING)) { - str.append(" - Progress:").append(progress).append(":").append(progressMessage); + + StringBuilder str = new StringBuilder(" Status : ").append(status.toString()).append("\n"); + if (statusMessage != null) { + str.append(" Message : ").append(statusMessage).append("\n"); + } + + str.append(" Progress : ").append(progress).append("\n"); + if (progressMessage != null) { + str.append(" Progress Message : ").append(progressMessage).append("\n"); } + + if (queueNumber != null) { + str.append(" Position in queue : ").append(queueNumber).append("\n"); + } + + if (errorMessage != null) { + str.append(" Error : ").append(errorMessage).append("\n"); + } + if (status.equals(Status.SUCCESSFUL)) { if (isResultSetAvailable) { - str.append(" - Result Available"); + str.append(" Result Available"); } else { - str.append(" - Result Not Available"); + str.append(" Result Not Available"); } } - if (status.equals(Status.FAILED)) { - str.append(" - Cause:").append(errorMessage); - } + return str.toString(); } http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-cli/src/main/java/org/apache/lens/cli/commands/LensQueryCommands.java ---------------------------------------------------------------------- diff --git a/lens-cli/src/main/java/org/apache/lens/cli/commands/LensQueryCommands.java b/lens-cli/src/main/java/org/apache/lens/cli/commands/LensQueryCommands.java index 44c0c62..006eaed 100644 --- a/lens-cli/src/main/java/org/apache/lens/cli/commands/LensQueryCommands.java +++ b/lens-cli/src/main/java/org/apache/lens/cli/commands/LensQueryCommands.java @@ -156,26 +156,10 @@ public class LensQueryCommands extends BaseLensCommand { public String getStatus( @CliOption(key = {"", "query_handle"}, mandatory = true, help = "<query_handle>") String qh) { QueryStatus status = getClient().getQueryStatus(new QueryHandle(UUID.fromString(qh))); - StringBuilder sb = new StringBuilder(); if (status == null) { return "Unable to find status for " + qh; } - sb.append("Status : ").append(status.getStatus()).append("\n"); - if (status.getStatusMessage() != null) { - sb.append("Message : ").append(status.getStatusMessage()).append("\n"); - } - if (status.getProgress() != 0) { - sb.append("Progress : ").append(status.getProgress()).append("\n"); - if (status.getProgressMessage() != null) { - sb.append("Progress Message : ").append(status.getProgressMessage()).append("\n"); - } - } - - if (status.getErrorMessage() != null) { - sb.append("Error : ").append(status.getErrorMessage()).append("\n"); - } - - return sb.toString(); + return status.toString(); } /** http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server-api/src/main/java/org/apache/lens/server/api/driver/DriverQueryStatus.java ---------------------------------------------------------------------- diff --git a/lens-server-api/src/main/java/org/apache/lens/server/api/driver/DriverQueryStatus.java b/lens-server-api/src/main/java/org/apache/lens/server/api/driver/DriverQueryStatus.java index f78b7c3..79d8bf1 100644 --- a/lens-server-api/src/main/java/org/apache/lens/server/api/driver/DriverQueryStatus.java +++ b/lens-server-api/src/main/java/org/apache/lens/server/api/driver/DriverQueryStatus.java @@ -167,7 +167,8 @@ public class DriverQueryStatus implements Serializable { break; } - return new QueryStatus(progress, qstate, statusMessage, isResultSetAvailable, progressMessage, errorMessage, null); + return new QueryStatus(progress, null, qstate, statusMessage, isResultSetAvailable, progressMessage, + errorMessage, null); } /** @@ -178,8 +179,8 @@ public class DriverQueryStatus implements Serializable { * @return the query status */ public static QueryStatus createQueryStatus(QueryStatus.Status state, DriverQueryStatus dstatus) { - return new QueryStatus(dstatus.progress, state, dstatus.statusMessage, dstatus.isResultSetAvailable, - dstatus.progressMessage, dstatus.errorMessage, null); + return new QueryStatus(dstatus.progress, null, state, dstatus.statusMessage, + dstatus.isResultSetAvailable, dstatus.progressMessage, dstatus.errorMessage, null); } /* http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server-api/src/main/java/org/apache/lens/server/api/query/FinishedLensQuery.java ---------------------------------------------------------------------- diff --git a/lens-server-api/src/main/java/org/apache/lens/server/api/query/FinishedLensQuery.java b/lens-server-api/src/main/java/org/apache/lens/server/api/query/FinishedLensQuery.java index 6cecf7e..d8c04db 100644 --- a/lens-server-api/src/main/java/org/apache/lens/server/api/query/FinishedLensQuery.java +++ b/lens-server-api/src/main/java/org/apache/lens/server/api/query/FinishedLensQuery.java @@ -204,7 +204,7 @@ public class FinishedLensQuery { qctx.setQueryHandle(QueryHandle.fromString(handle)); qctx.setLaunchTime(this.startTime); qctx.setEndTime(getEndTime()); - qctx.setStatusSkippingTransitionTest(new QueryStatus(0.0, QueryStatus.Status.valueOf(getStatus()), + qctx.setStatusSkippingTransitionTest(new QueryStatus(0.0, null, QueryStatus.Status.valueOf(getStatus()), getErrorMessage() == null ? "" : getErrorMessage(), getResult() != null, null, null, null)); qctx.getDriverStatus().setDriverStartTime(getDriverStartTime()); qctx.getDriverStatus().setDriverFinishTime(getDriverEndTime()); http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server-api/src/main/java/org/apache/lens/server/api/query/QueryContext.java ---------------------------------------------------------------------- diff --git a/lens-server-api/src/main/java/org/apache/lens/server/api/query/QueryContext.java b/lens-server-api/src/main/java/org/apache/lens/server/api/query/QueryContext.java index beaa72f..12de0a5 100644 --- a/lens-server-api/src/main/java/org/apache/lens/server/api/query/QueryContext.java +++ b/lens-server-api/src/main/java/org/apache/lens/server/api/query/QueryContext.java @@ -217,7 +217,7 @@ public class QueryContext extends AbstractQueryContext { super(userQuery, user, qconf, conf, drivers, mergeDriverConf); this.submissionTime = submissionTime; this.queryHandle = new QueryHandle(UUID.randomUUID()); - this.status = new QueryStatus(0.0f, Status.NEW, "Query just got created", false, null, null, null); + this.status = new QueryStatus(0.0f, null, Status.NEW, "Query just got created", false, null, null, null); this.priority = Priority.NORMAL; this.lensConf = qconf; this.conf = conf; http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server-api/src/main/java/org/apache/lens/server/api/query/collect/ImmutableQueryCollection.java ---------------------------------------------------------------------- diff --git a/lens-server-api/src/main/java/org/apache/lens/server/api/query/collect/ImmutableQueryCollection.java b/lens-server-api/src/main/java/org/apache/lens/server/api/query/collect/ImmutableQueryCollection.java index 344ff10..0dcf367 100644 --- a/lens-server-api/src/main/java/org/apache/lens/server/api/query/collect/ImmutableQueryCollection.java +++ b/lens-server-api/src/main/java/org/apache/lens/server/api/query/collect/ImmutableQueryCollection.java @@ -50,4 +50,10 @@ public interface ImmutableQueryCollection { * @return Count of existing queries */ int getQueriesCount(); + + /** + * + * @return Index of a query within collection + */ + Integer getQueryIndex(final QueryContext query); } http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server/src/main/java/org/apache/lens/server/query/QueryExecutionServiceImpl.java ---------------------------------------------------------------------- diff --git a/lens-server/src/main/java/org/apache/lens/server/query/QueryExecutionServiceImpl.java b/lens-server/src/main/java/org/apache/lens/server/query/QueryExecutionServiceImpl.java index 23cc748..9e27dd4 100644 --- a/lens-server/src/main/java/org/apache/lens/server/query/QueryExecutionServiceImpl.java +++ b/lens-server/src/main/java/org/apache/lens/server/query/QueryExecutionServiceImpl.java @@ -615,8 +615,8 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE checkEstimatedQueriesState(query); QueryStatus oldStatus = query.getStatus(); - QueryStatus newStatus = new QueryStatus(query.getStatus().getProgress(), QueryStatus.Status.LAUNCHED, - "Query is launched on driver", false, null, null, null); + QueryStatus newStatus = new QueryStatus(query.getStatus().getProgress(), null, + QueryStatus.Status.LAUNCHED, "Query is launched on driver", false, null, null, null); query.validateTransition(newStatus); // Check if we need to pass session's effective resources to selected driver @@ -717,7 +717,7 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE throws LensException { QueryStatus before = ctx.getStatus(); - ctx.setStatus(new QueryStatus(0.0f, FAILED, statusMsg, false, null, reason, lensErrorTO)); + ctx.setStatus(new QueryStatus(0.0f, null, FAILED, statusMsg, false, null, reason, lensErrorTO)); updateFinishedQuery(ctx, before); fireStatusChangeEvent(ctx, ctx.getStatus(), before); } @@ -731,7 +731,7 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE */ private void setCancelledStatus(QueryContext ctx, String statusMsg) throws LensException { QueryStatus before = ctx.getStatus(); - ctx.setStatus(new QueryStatus(0.0f, CANCELED, statusMsg, false, null, null, null)); + ctx.setStatus(new QueryStatus(0.0f, null, CANCELED, statusMsg, false, null, null, null)); updateFinishedQuery(ctx, before); fireStatusChangeEvent(ctx, ctx.getStatus(), before); } @@ -765,7 +765,7 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE void setSuccessState(QueryContext ctx) throws LensException { QueryStatus before = ctx.getStatus(); - ctx.setStatus(new QueryStatus(1.0f, SUCCESSFUL, "Query is successful!", ctx + ctx.setStatus(new QueryStatus(1.0f, null, SUCCESSFUL, "Query is successful!", ctx .isResultAvailableInDriver(), null, null, null)); updateFinishedQuery(ctx, before); fireStatusChangeEvent(ctx, ctx.getStatus(), before); @@ -807,6 +807,14 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE fireStatusChangeEvent(ctx, ctx.getStatus(), before); } } + if (ctx.queued()) { + Integer queryIndex = waitingQueries.getQueryIndex(ctx); + // Query index could be null when the query status is queued but + // query is present in priorityblocking queue for processing + if (queryIndex != null) { + ctx.getStatus().setQueueNumber(queryIndex); + } + } } } } @@ -946,7 +954,7 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE resultSets.remove(finished.getCtx().getQueryHandle()); } fireStatusChangeEvent(finished.getCtx(), - new QueryStatus(1f, CLOSED, "Query purged", false, null, null, null), finished.getCtx().getStatus()); + new QueryStatus(1f, null, CLOSED, "Query purged", false, null, null, null), finished.getCtx().getStatus()); log.info("Query purged: {}", finished.getCtx().getQueryHandle()); } catch (LensException e) { @@ -1006,8 +1014,8 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE this.launchedQueries = new ThreadSafeEstimatedQueryCollection(new DefaultEstimatedQueryCollection(new DefaultQueryCollection())); - this.waitingQueries - = new ThreadSafeEstimatedQueryCollection(new DefaultEstimatedQueryCollection(new DefaultQueryCollection())); + this.waitingQueries = new ThreadSafeEstimatedQueryCollection(new DefaultEstimatedQueryCollection( + new DefaultQueryCollection(new TreeSet<QueryContext>(new QueryContextPriorityComparator())))); ImmutableSet<QueryLaunchingConstraint> queryConstraints = getImplementations( QUERY_LAUNCHING_CONSTRAINT_FACTORIES_KEY, hiveConf); @@ -1696,7 +1704,7 @@ public class QueryExecutionServiceImpl extends BaseLensService implements QueryE private QueryHandle submitQuery(final QueryContext ctx) throws LensException { QueryStatus before = ctx.getStatus(); - ctx.setStatus(new QueryStatus(0.0, QUEUED, "Query is queued", false, null, null, null)); + ctx.setStatus(new QueryStatus(0.0, null, QUEUED, "Query is queued", false, null, null, null)); queuedQueries.add(ctx); log.debug("Added to Queued Queries:{}", ctx.getQueryHandleString()); allQueries.put(ctx.getQueryHandle(), ctx); http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server/src/main/java/org/apache/lens/server/query/collect/DefaultEstimatedQueryCollection.java ---------------------------------------------------------------------- diff --git a/lens-server/src/main/java/org/apache/lens/server/query/collect/DefaultEstimatedQueryCollection.java b/lens-server/src/main/java/org/apache/lens/server/query/collect/DefaultEstimatedQueryCollection.java index e3505bb..908b86f 100644 --- a/lens-server/src/main/java/org/apache/lens/server/query/collect/DefaultEstimatedQueryCollection.java +++ b/lens-server/src/main/java/org/apache/lens/server/query/collect/DefaultEstimatedQueryCollection.java @@ -138,6 +138,11 @@ public class DefaultEstimatedQueryCollection implements EstimatedQueryCollection return this.queries.getQueriesCount(); } + @Override + public Integer getQueryIndex(QueryContext query) { + return this.queries.getQueryIndex(query); + } + @VisibleForTesting void checkState(final QueryContext query) { Preconditions.checkState(query.getSelectedDriver() != null); http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server/src/main/java/org/apache/lens/server/query/collect/DefaultQueryCollection.java ---------------------------------------------------------------------- diff --git a/lens-server/src/main/java/org/apache/lens/server/query/collect/DefaultQueryCollection.java b/lens-server/src/main/java/org/apache/lens/server/query/collect/DefaultQueryCollection.java index f9e7701..844237a 100644 --- a/lens-server/src/main/java/org/apache/lens/server/query/collect/DefaultQueryCollection.java +++ b/lens-server/src/main/java/org/apache/lens/server/query/collect/DefaultQueryCollection.java @@ -26,7 +26,6 @@ import org.apache.commons.collections.CollectionUtils; import org.apache.commons.collections.map.MultiValueMap; import com.google.common.collect.Sets; -import lombok.NoArgsConstructor; import lombok.NonNull; import lombok.extern.slf4j.Slf4j; @@ -36,17 +35,28 @@ import lombok.extern.slf4j.Slf4j; * * @see QueryCollection */ -@NoArgsConstructor @Slf4j public class DefaultQueryCollection implements QueryCollection { - private final Set<QueryContext> queries = Sets.newLinkedHashSet(); + private final Set<QueryContext> queries; private final MultiValueMap queriesByUser = MultiValueMap.decorate(new HashMap(), LinkedHashSet.class); + public DefaultQueryCollection() { + this.queries = Sets.newLinkedHashSet(); + } + public DefaultQueryCollection(@NonNull final Set<QueryContext> queries) { + this(); addAll(queries); } + public DefaultQueryCollection(final TreeSet<QueryContext> treeSet) { + this.queries = treeSet; + for (QueryContext query : treeSet) { + queriesByUser.put(query.getSubmittedUser(), query); + } + } + @Override public boolean add(final QueryContext query) { @@ -96,6 +106,26 @@ public class DefaultQueryCollection implements QueryCollection { return queries.size(); } + + /** + * Since the collection is a linkedHashSet, the order of queries is always maintained. + * @param query + * @return + */ + @Override + public Integer getQueryIndex(QueryContext query) { + Iterator iterator = queries.iterator(); + int index = 1; + while (iterator.hasNext()) { + QueryContext queuedQuery = (QueryContext) iterator.next(); + if (queuedQuery.getQueryHandle().equals(query.getQueryHandle())) { + return index; + } + index += 1; + } + return null; + } + private Collection<QueryContext> getQueriesCollectionForUser(final String user) { final Collection<QueryContext> userQueries = queriesByUser.getCollection(user); http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server/src/main/java/org/apache/lens/server/query/collect/ThreadSafeEstimatedQueryCollection.java ---------------------------------------------------------------------- diff --git a/lens-server/src/main/java/org/apache/lens/server/query/collect/ThreadSafeEstimatedQueryCollection.java b/lens-server/src/main/java/org/apache/lens/server/query/collect/ThreadSafeEstimatedQueryCollection.java index cdbd2ad..5d24379 100644 --- a/lens-server/src/main/java/org/apache/lens/server/query/collect/ThreadSafeEstimatedQueryCollection.java +++ b/lens-server/src/main/java/org/apache/lens/server/query/collect/ThreadSafeEstimatedQueryCollection.java @@ -70,6 +70,11 @@ public class ThreadSafeEstimatedQueryCollection implements EstimatedQueryCollect } @Override + public synchronized Integer getQueryIndex(QueryContext query) { + return this.estimatedQueries.getQueryIndex(query); + } + + @Override public synchronized boolean add(QueryContext query) { return this.estimatedQueries.add(query); } http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server/src/main/java/org/apache/lens/server/query/collect/ThreadSafeQueryCollection.java ---------------------------------------------------------------------- diff --git a/lens-server/src/main/java/org/apache/lens/server/query/collect/ThreadSafeQueryCollection.java b/lens-server/src/main/java/org/apache/lens/server/query/collect/ThreadSafeQueryCollection.java index 7b43a38..a49c7de 100644 --- a/lens-server/src/main/java/org/apache/lens/server/query/collect/ThreadSafeQueryCollection.java +++ b/lens-server/src/main/java/org/apache/lens/server/query/collect/ThreadSafeQueryCollection.java @@ -72,4 +72,9 @@ public class ThreadSafeQueryCollection implements QueryCollection { public synchronized int getQueriesCount() { return this.queries.getQueriesCount(); } + + @Override + public synchronized Integer getQueryIndex(QueryContext query) { + return this.queries.getQueryIndex(query); + } } http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server/src/test/java/org/apache/lens/server/query/collect/DefaultQueryCollectionTest.java ---------------------------------------------------------------------- diff --git a/lens-server/src/test/java/org/apache/lens/server/query/collect/DefaultQueryCollectionTest.java b/lens-server/src/test/java/org/apache/lens/server/query/collect/DefaultQueryCollectionTest.java index 7a81e83..d4b6bdf 100644 --- a/lens-server/src/test/java/org/apache/lens/server/query/collect/DefaultQueryCollectionTest.java +++ b/lens-server/src/test/java/org/apache/lens/server/query/collect/DefaultQueryCollectionTest.java @@ -29,17 +29,25 @@ import java.util.Set; import org.apache.lens.server.api.query.QueryContext; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; public class DefaultQueryCollectionTest { private static final String MOCK_USER = "MockUserEmail"; + private static final String MOCK_HANDLE = "0-0-0-0-"; + + @DataProvider + public Object[][] dpQueryCosts() { + return new Object[][]{{new double[]{20.0, 50.0, 10.0, 80.0, 40.0, }, }, }; + + } /* Note: Since verification of addition/removal required calling get methods, hence methods getQueriesCount and getQueries(user) are indirectly getting tested in these tests */ @Test - public void testAddMethodAddsQueriesToAllViews(){ + public void testAddMethodAddsQueriesToAllViews() { /* Initialization */ final int noOfQueriesUsedInTest = 2; @@ -75,6 +83,30 @@ public class DefaultQueryCollectionTest { assertEquals(queries.getQueries(MOCK_USER).size(), 0); } + @Test(dataProvider = "dpQueryCosts") + public void testRemoveMethodMustChangeQueryIndices(final double[] queryCosts) { + + /* Initialization */ + int numberOfQueries = queryCosts.length; + QueryCollection collection = createQueriesTreeSetWithQueryHandleAndCostStubbing(queryCosts, MOCK_HANDLE); + + QueryContext completedQuery = getMockedQueryFromQueries(collection.getQueries(), MOCK_HANDLE, 1); + QueryContext queuedQuery = getMockedQueryFromQueries(collection.getQueries(), MOCK_HANDLE, 5); + + /* Verification 1: Verifies that all queries were added into the collection*/ + assertEquals(collection.getQueriesCount(), numberOfQueries); + + /* Execution */ + collection.remove(completedQuery); + + /* Verification 2: Verifies that queries were removed from the collection */ + assertEquals(collection.getQueriesCount(), numberOfQueries - 1); + + /* Verification 3: Verifies that query index is decreased after removal of queries which were present before + them in the queries list */ + assertEquals(collection.getQueryIndex(queuedQuery).intValue(), 2); + } + @Test public void testGetQueriesMustReturnCopyOfUnderlyingCollection() { http://git-wip-us.apache.org/repos/asf/lens/blob/4addd7b6/lens-server/src/test/java/org/apache/lens/server/query/collect/QueryCollectUtil.java ---------------------------------------------------------------------- diff --git a/lens-server/src/test/java/org/apache/lens/server/query/collect/QueryCollectUtil.java b/lens-server/src/test/java/org/apache/lens/server/query/collect/QueryCollectUtil.java index 51fcf00..7789a4c 100644 --- a/lens-server/src/test/java/org/apache/lens/server/query/collect/QueryCollectUtil.java +++ b/lens-server/src/test/java/org/apache/lens/server/query/collect/QueryCollectUtil.java @@ -22,14 +22,19 @@ package org.apache.lens.server.query.collect; import static java.lang.reflect.Modifier.isPublic; import static java.lang.reflect.Modifier.isSynchronized; +import org.apache.lens.api.query.QueryHandle; +import org.apache.lens.server.api.query.QueryContext; +import org.apache.lens.server.api.query.cost.FactPartitionBasedQueryCost; +import org.apache.lens.server.query.QueryContextPriorityComparator; + import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.testng.Assert.assertTrue; import java.lang.reflect.Method; +import java.util.Iterator; import java.util.Set; - -import org.apache.lens.server.api.query.QueryContext; +import java.util.TreeSet; import com.google.common.collect.Sets; @@ -77,6 +82,36 @@ public class QueryCollectUtil { return new DefaultQueryCollection(mockQueries); } + public static QueryCollection createQueriesTreeSetWithQueryHandleAndCostStubbing(final double[] queryCosts, + final String handlePrefix) { + + TreeSet<QueryContext> mockQueries = new TreeSet<>(new QueryContextPriorityComparator()); + + for (int index = 1; index <= queryCosts.length; ++index) { + mockQueries.add(createQueryInstanceWithQueryHandleAndCostStubbing(handlePrefix, index, queryCosts[index - 1])); + } + return new DefaultQueryCollection(mockQueries); + } + + public static QueryContext createQueryInstanceWithQueryHandleAndCostStubbing(String handlePrefix, int index, + double queryCost) { + QueryContext mockQuery = mock(QueryContext.class); + when(mockQuery.getQueryHandle()).thenReturn(QueryHandle.fromString(handlePrefix + index)); + when(mockQuery.getSelectedDriverQueryCost()).thenReturn(new FactPartitionBasedQueryCost(queryCost)); + return mockQuery; + } + + public static QueryContext getMockedQueryFromQueries(Set<QueryContext> queries, String mockHandle, int index) { + Iterator iterator = queries.iterator(); + while (iterator.hasNext()) { + QueryContext queuedQuery = (QueryContext) iterator.next(); + if (queuedQuery.getQueryHandle().equals(QueryHandle.fromString(mockHandle + index))) { + return queuedQuery; + } + } + return null; + } + public static QueryCollection stubMockQueryAndCreateQueriesInstance(final QueryContext mockQuery, final String mockUser) {
