Add support to filter on datasource for active tasks (#5998)

* Add support to filter on datasource for active tasks

* Added datasource filter to sql query for active tasks
* Fixed unit tests

* Address PR comments

common/src/main/java/io/druid/metadata/MetadataStorageActionHandler.java

@@ -120,7 +120,7 @@ public interface MetadataStorageActionHandler<EntryType, StatusType, LogType, Lo
    *
    * @return list of {@link TaskInfo}
    */
-  List<TaskInfo<EntryType>> getActiveTaskInfo();
+  List<TaskInfo<EntryType>> getActiveTaskInfo(@Nullable String dataSource);
 
   /**
    * Return createdDate and dataSource for the given id

indexing-service/src/main/java/io/druid/indexing/overlord/HeapMemoryTaskStorage.java

@@ -170,7 +170,7 @@ public class HeapMemoryTaskStorage implements TaskStorage
   }
 
   @Override
-  public List<TaskInfo<Task>> getActiveTaskInfo()
+  public List<TaskInfo<Task>> getActiveTaskInfo(@Nullable String dataSource)
   {
     giant.lock();
 

indexing-service/src/main/java/io/druid/indexing/overlord/MetadataTaskStorage.java

@@ -214,10 +214,10 @@ public class MetadataTaskStorage implements TaskStorage
   }
 
   @Override
-  public List<TaskInfo<Task>> getActiveTaskInfo()
+  public List<TaskInfo<Task>> getActiveTaskInfo(@Nullable String dataSource)
   {
     return ImmutableList.copyOf(
-        handler.getActiveTaskInfo()
+        handler.getActiveTaskInfo(dataSource)
     );
   }
 

indexing-service/src/main/java/io/druid/indexing/overlord/TaskStorage.java

@@ -127,9 +127,11 @@ public interface TaskStorage
    * Returns a list of currently running or pending tasks as stored in the storage facility as {@link TaskInfo}. No particular order
    * is guaranteed, but implementations are encouraged to return tasks in ascending order of creation.
    *
+   * @param datasource datasource
+   *
    * @return list of {@link TaskInfo}
    */
-  List<TaskInfo<Task>> getActiveTaskInfo();
+  List<TaskInfo<Task>> getActiveTaskInfo(@Nullable String dataSource);
 
   /**
    * Returns up to {@code maxTaskStatuses} {@link TaskInfo} objects of recently finished tasks as stored in the storage facility. No
@@ -137,7 +139,8 @@ public interface TaskStorage
    * No particular standard of "recent" is guaranteed, and in fact, this method is permitted to simply return nothing.
    *
    * @param maxTaskStatuses maxTaskStatuses
-   * @param duration duration
+   * @param duration        duration
+   * @param datasource      datasource
    *
    * @return list of {@link TaskInfo}
    */

indexing-service/src/main/java/io/druid/indexing/overlord/TaskStorageQueryAdapter.java

@@ -55,9 +55,9 @@ public class TaskStorageQueryAdapter
     return storage.getActiveTasks();
   }
 
-  public List<TaskInfo<Task>> getActiveTaskInfo()
+  public List<TaskInfo<Task>> getActiveTaskInfo(@Nullable String dataSource)
   {
-    return storage.getActiveTaskInfo();
+    return storage.getActiveTaskInfo(dataSource);
   }
 
   public List<TaskInfo<Task>> getRecentlyCompletedTaskInfo(

indexing-service/src/main/java/io/druid/indexing/overlord/http/OverlordResource.java

@@ -629,10 +629,10 @@ public class OverlordResource
       finalTaskList.addAll(completedTasks);
     }
 
-    List<TaskInfo<Task>> allActiveTaskInfo = Lists.newArrayList();
+    final List<TaskInfo<Task>> allActiveTaskInfo;
     final List<AnyTask> allActiveTasks = Lists.newArrayList();
     if (state == null || !"complete".equals(StringUtils.toLowerCase(state))) {
-      allActiveTaskInfo = taskStorageQueryAdapter.getActiveTaskInfo();
+      allActiveTaskInfo = taskStorageQueryAdapter.getActiveTaskInfo(dataSource);
       for (final TaskInfo<Task> task : allActiveTaskInfo) {
         allActiveTasks.add(
             new AnyTask(

indexing-service/src/test/java/io/druid/indexing/overlord/http/OverlordResourceTest.java

@@ -172,7 +172,7 @@ public class OverlordResourceTest
   public void testSecuredGetWaitingTask()
   {
     expectAuthorizationTokenCheck();
-    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo()).andStubReturn(
+    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo(null)).andStubReturn(
         ImmutableList.of(
             new TaskInfo(
                 "id_1",
@@ -278,7 +278,7 @@ public class OverlordResourceTest
             new MockTaskRunnerWorkItem(tasksIds.get(1), null)
         )
     );
-    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo()).andStubReturn(
+    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo(null)).andStubReturn(
         ImmutableList.of(
             new TaskInfo(
                 "id_1",
@@ -337,7 +337,7 @@ public class OverlordResourceTest
         )
     );
     //active tasks
-    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo()).andStubReturn(
+    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo(null)).andStubReturn(
         ImmutableList.of(
             new TaskInfo(
                 "id_1",
@@ -427,7 +427,7 @@ public class OverlordResourceTest
         )
     );
     //active tasks
-    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo()).andStubReturn(
+    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo("allow")).andStubReturn(
         ImmutableList.of(
             new TaskInfo(
                 "id_1",
@@ -491,7 +491,7 @@ public class OverlordResourceTest
   {
     expectAuthorizationTokenCheck();
     //active tasks
-    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo()).andStubReturn(
+    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo(null)).andStubReturn(
         ImmutableList.of(
             new TaskInfo(
                 "id_1",
@@ -549,7 +549,7 @@ public class OverlordResourceTest
   public void testGetTasksFilterRunningState()
   {
     expectAuthorizationTokenCheck();
-    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo()).andStubReturn(
+    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo("allow")).andStubReturn(
         ImmutableList.of(
             new TaskInfo(
                 "id_1",
@@ -615,7 +615,7 @@ public class OverlordResourceTest
             new MockTaskRunnerWorkItem(tasksIds.get(1), null)
         )
     );
-    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo()).andStubReturn(
+    EasyMock.expect(taskStorageQueryAdapter.getActiveTaskInfo(null)).andStubReturn(
         ImmutableList.of(
             new TaskInfo(
                 "id_1",

server/src/main/java/io/druid/metadata/DerbyMetadataStorageActionHandler.java

@@ -46,7 +46,7 @@ public class DerbyMetadataStorageActionHandler<EntryType, StatusType, LogType, L
 
   @Override
   protected Query<Map<String, Object>> createInactiveStatusesSinceQuery(
-      Handle handle, DateTime timestamp, @Nullable Integer maxNumStatuses, @Nullable String datasource
+      Handle handle, DateTime timestamp, @Nullable Integer maxNumStatuses, @Nullable String dataSource
   )
   {
     String sql = StringUtils.format(
@@ -59,7 +59,7 @@ public class DerbyMetadataStorageActionHandler<EntryType, StatusType, LogType, L
         + "FROM "
         + "  %s "
         + "WHERE "
-        + getWhereClauseForInactiveStatusesSinceQuery(datasource)
+        + getWhereClauseForInactiveStatusesSinceQuery(dataSource)
         + "ORDER BY created_date DESC",
         getEntryTable()
     );
@@ -72,8 +72,8 @@ public class DerbyMetadataStorageActionHandler<EntryType, StatusType, LogType, L
     if (maxNumStatuses != null) {
       query = query.bind("n", maxNumStatuses);
     }
-    if (datasource != null) {
+    if (dataSource != null) {
-      query = query.bind("ds", datasource);
+      query = query.bind("ds", dataSource);
     }
     return query;
   }

server/src/main/java/io/druid/metadata/MySQLMetadataStorageActionHandler.java

@@ -46,7 +46,7 @@ public class MySQLMetadataStorageActionHandler<EntryType, StatusType, LogType, L
 
   @Override
   protected Query<Map<String, Object>> createInactiveStatusesSinceQuery(
-      Handle handle, DateTime timestamp, @Nullable Integer maxNumStatuses, @Nullable String datasource
+      Handle handle, DateTime timestamp, @Nullable Integer maxNumStatuses, @Nullable String dataSource
   )
   {
     String sql = StringUtils.format(
@@ -59,7 +59,7 @@ public class MySQLMetadataStorageActionHandler<EntryType, StatusType, LogType, L
         + "FROM "
         + "  %s "
         + "WHERE "
-        + getWhereClauseForInactiveStatusesSinceQuery(datasource)
+        + getWhereClauseForInactiveStatusesSinceQuery(dataSource)
         + "ORDER BY created_date DESC",
         getEntryTable()
     );
@@ -72,8 +72,8 @@ public class MySQLMetadataStorageActionHandler<EntryType, StatusType, LogType, L
     if (maxNumStatuses != null) {
       query = query.bind("n", maxNumStatuses);
     }
-    if (datasource != null) {
+    if (dataSource != null) {
-      query = query.bind("ds", datasource);
+      query = query.bind("ds", dataSource);
     }
     return query;
   }

server/src/main/java/io/druid/metadata/PostgreSQLMetadataStorageActionHandler.java

@@ -46,7 +46,7 @@ public class PostgreSQLMetadataStorageActionHandler<EntryType, StatusType, LogTy
 
   @Override
   protected Query<Map<String, Object>> createInactiveStatusesSinceQuery(
-      Handle handle, DateTime timestamp, @Nullable Integer maxNumStatuses, @Nullable String datasource
+      Handle handle, DateTime timestamp, @Nullable Integer maxNumStatuses, @Nullable String dataSource
   )
   {
     String sql = StringUtils.format(
@@ -59,7 +59,7 @@ public class PostgreSQLMetadataStorageActionHandler<EntryType, StatusType, LogTy
         + "FROM "
         + "  %s "
         + "WHERE "
-        + getWhereClauseForInactiveStatusesSinceQuery(datasource)
+        + getWhereClauseForInactiveStatusesSinceQuery(dataSource)
         + "ORDER BY created_date DESC",
         getEntryTable()
     );
@@ -73,8 +73,8 @@ public class PostgreSQLMetadataStorageActionHandler<EntryType, StatusType, LogTy
     if (maxNumStatuses != null) {
       query = query.bind("n", maxNumStatuses);
     }
-    if (datasource != null) {
+    if (dataSource != null) {
-      query = query.bind("ds", datasource);
+      query = query.bind("ds", dataSource);
     }
     return query;
   }

server/src/main/java/io/druid/metadata/SQLMetadataStorageActionHandler.java

@@ -333,7 +333,7 @@ public abstract class SQLMetadataStorageActionHandler<EntryType, StatusType, Log
   public List<TaskInfo<EntryType>> getCompletedTaskInfo(
       DateTime timestamp,
       @Nullable Integer maxNumStatuses,
-      @Nullable String datasource
+      @Nullable String dataSource
   )
   {
     return getConnector().retryWithHandle(
@@ -342,7 +342,7 @@ public abstract class SQLMetadataStorageActionHandler<EntryType, StatusType, Log
               handle,
               timestamp,
               maxNumStatuses,
-              datasource
+              dataSource
           );
           return query.map(new TaskInfoMapper()).list();
         }
@@ -350,20 +350,52 @@ public abstract class SQLMetadataStorageActionHandler<EntryType, StatusType, Log
   }
 
   @Override
-  public List<TaskInfo<EntryType>> getActiveTaskInfo()
+  public List<TaskInfo<EntryType>> getActiveTaskInfo(@Nullable String dataSource)
   {
     return getConnector().retryWithHandle(
         handle -> {
-          return handle.createQuery(
+          final Query<Map<String, Object>> query = createActiveStatusesQuery(
-              StringUtils.format(
+              handle,
-                  "SELECT id, status_payload, payload, datasource, created_date FROM %s WHERE active = TRUE ORDER BY created_date",
+              dataSource
-                  entryTable
+          );
-              )
+          return query.map(new TaskInfoMapper()).list();
-          ).map(new TaskInfoMapper()).list();
         }
     );
   }
 
+  private Query<Map<String, Object>> createActiveStatusesQuery(Handle handle, @Nullable String dataSource)
+  {
+    String sql = StringUtils.format(
+        "SELECT "
+        + "  id, "
+        + "  status_payload, "
+        + "  payload, "
+        + "  datasource, "
+        + "  created_date "
+        + "FROM "
+        + "  %s "
+        + "WHERE "
+        + getWhereClauseForActiveStatusesQuery(dataSource)
+        + "ORDER BY created_date",
+        entryTable
+    );
+
+    Query<Map<String, Object>> query = handle.createQuery(sql);
+    if (dataSource != null) {
+      query = query.bind("ds", dataSource);
+    }
+    return query;
+  }
+
+  private String getWhereClauseForActiveStatusesQuery(String dataSource)
+  {
+    String sql = StringUtils.format("active = TRUE ");
+    if (dataSource != null) {
+      sql += " AND datasource = :ds ";
+    }
+    return sql;
+  }
+
   class TaskInfoMapper implements ResultSetMapper<TaskInfo<EntryType>>
   {
     @Override
@@ -401,7 +433,7 @@ public abstract class SQLMetadataStorageActionHandler<EntryType, StatusType, Log
       Handle handle,
       DateTime timestamp,
       @Nullable Integer maxNumStatuses,
-      @Nullable String datasource
+      @Nullable String dataSource
   );
 
   @Override

server/src/main/java/io/druid/metadata/SQLServerMetadataStorageActionHandler.java

@@ -46,7 +46,7 @@ public class SQLServerMetadataStorageActionHandler<EntryType, StatusType, LogTyp
 
   @Override
   protected Query<Map<String, Object>> createInactiveStatusesSinceQuery(
-      Handle handle, DateTime timestamp, @Nullable Integer maxNumStatuses, @Nullable String datasource
+      Handle handle, DateTime timestamp, @Nullable Integer maxNumStatuses, @Nullable String dataSource
   )
   {
     String sql = maxNumStatuses == null ? "SELECT " : "SELECT TOP (:n) ";
@@ -60,7 +60,7 @@ public class SQLServerMetadataStorageActionHandler<EntryType, StatusType, LogTyp
         + "FROM "
         + "  %s "
         + "WHERE "
-        + getWhereClauseForInactiveStatusesSinceQuery(datasource)
+        + getWhereClauseForInactiveStatusesSinceQuery(dataSource)
         + "ORDER BY created_date DESC",
         getEntryTable()
     );
@@ -70,8 +70,8 @@ public class SQLServerMetadataStorageActionHandler<EntryType, StatusType, LogTyp
     if (maxNumStatuses != null) {
       query = query.bind("n", maxNumStatuses);
     }
-    if (datasource != null) {
+    if (dataSource != null) {
-      query = query.bind("ds", datasource);
+      query = query.bind("ds", dataSource);
     }
     return query;
   }