MetadataResource: Fix handling of includeDisabled. (#3042)

server/src/main/java/io/druid/server/http/MetadataResource.java

@@ -22,8 +22,9 @@ package io.druid.server.http;
 import com.google.common.base.Function;
 import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
-import com.google.common.collect.Lists;
+import com.google.common.collect.Sets;
 import com.google.inject.Inject;
 import com.metamx.common.Pair;
 import com.sun.jersey.spi.container.ResourceFilters;
@@ -51,11 +52,10 @@ import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
 import java.io.IOException;
-import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  */
@@ -82,71 +82,18 @@ public class MetadataResource
   @Path("/datasources")
   @Produces(MediaType.APPLICATION_JSON)
   public Response getDatabaseDataSources(
-      @QueryParam("full") String full,
+      @QueryParam("full") final String full,
-      @QueryParam("includeDisabled") String includeDisabled,
+      @QueryParam("includeDisabled") final String includeDisabled,
       @Context final HttpServletRequest req
   )
   {
-    Response.ResponseBuilder builder = Response.status(Response.Status.OK);
+    final Set<String> dataSourceNamesPreAuth;
-
-    final Collection<DruidDataSource> druidDataSources;
-    if (authConfig.isEnabled()) {
-      // This is an experimental feature, see - https://github.com/druid-io/druid/pull/2424
-      final Map<Pair<Resource, Action>, Access> resourceAccessMap = new HashMap<>();
-      final AuthorizationInfo authorizationInfo = (AuthorizationInfo) req.getAttribute(AuthConfig.DRUID_AUTH_TOKEN);
-      if (includeDisabled != null) {
-        return builder.entity(
-            Collections2.filter(
-                metadataSegmentManager.getAllDatasourceNames(),
-                new Predicate<String>()
-                {
-                  @Override
-                  public boolean apply(String input)
-                  {
-                    Resource resource = new Resource(input, ResourceType.DATASOURCE);
-                    Action action = Action.READ;
-                    Pair<Resource, Action> key = new Pair<>(resource, action);
-                    if (resourceAccessMap.containsKey(key)) {
-                      return resourceAccessMap.get(key).isAllowed();
-                    } else {
-                      Access access = authorizationInfo.isAuthorized(key.lhs, key.rhs);
-                      resourceAccessMap.put(key, access);
-                      return access.isAllowed();
-                    }
-                  }
-                }
-            )).build();
-      } else {
-        druidDataSources =
-            Collections2.filter(
-                metadataSegmentManager.getInventory(),
-                new Predicate<DruidDataSource>()
-                {
-                  @Override
-                  public boolean apply(DruidDataSource input)
-                  {
-                    Resource resource = new Resource(input.getName(), ResourceType.DATASOURCE);
-                    Action action = Action.READ;
-                    Pair<Resource, Action> key = new Pair<>(resource, action);
-                    if (resourceAccessMap.containsKey(key)) {
-                      return resourceAccessMap.get(key).isAllowed();
-                    } else {
-                      Access access = authorizationInfo.isAuthorized(key.lhs, key.rhs);
-                      resourceAccessMap.put(key, access);
-                      return access.isAllowed();
-                    }
-                  }
-                }
-            );
-      }
-    } else {
-      druidDataSources = metadataSegmentManager.getInventory();
-    }
-
     if (includeDisabled != null) {
-      return builder.entity(
+      dataSourceNamesPreAuth = Sets.newTreeSet(metadataSegmentManager.getAllDatasourceNames());
-          Collections2.transform(
+    } else {
-              druidDataSources,
+      dataSourceNamesPreAuth = Sets.newTreeSet(
+          Iterables.transform(
+              metadataSegmentManager.getInventory(),
               new Function<DruidDataSource, String>()
               {
                 @Override
@@ -156,29 +103,60 @@ public class MetadataResource
                 }
               }
           )
-      ).build();
+      );
-    }
-    if (full != null) {
-      return builder.entity(druidDataSources).build();
     }
 
-    List<String> dataSourceNames = Lists.newArrayList(
+    final Set<String> dataSourceNamesPostAuth;
-        Iterables.transform(
+
-            druidDataSources,
+    if (authConfig.isEnabled()) {
-            new Function<DruidDataSource, String>()
+      // This is an experimental feature, see - https://github.com/druid-io/druid/pull/2424
-            {
+      final Map<Pair<Resource, Action>, Access> resourceAccessMap = new HashMap<>();
-              @Override
+      final AuthorizationInfo authorizationInfo = (AuthorizationInfo) req.getAttribute(AuthConfig.DRUID_AUTH_TOKEN);
-              public String apply(DruidDataSource dataSource)
+      dataSourceNamesPostAuth = ImmutableSet.copyOf(
+          Sets.filter(
+              dataSourceNamesPreAuth,
+              new Predicate<String>()
               {
-                return dataSource.getName();
+                @Override
+                public boolean apply(String input)
+                {
+                  Resource resource = new Resource(input, ResourceType.DATASOURCE);
+                  Action action = Action.READ;
+                  Pair<Resource, Action> key = new Pair<>(resource, action);
+                  if (resourceAccessMap.containsKey(key)) {
+                    return resourceAccessMap.get(key).isAllowed();
+                  } else {
+                    Access access = authorizationInfo.isAuthorized(key.lhs, key.rhs);
+                    resourceAccessMap.put(key, access);
+                    return access.isAllowed();
+                  }
+                }
               }
-            }
+          )
-        )
+      );
-    );
+    } else {
+      dataSourceNamesPostAuth = dataSourceNamesPreAuth;
+    }
 
-    Collections.sort(dataSourceNames);
+    // Cannot do both includeDisabled and full, let includeDisabled take priority
-
+    // Always use dataSourceNamesPostAuth to determine the set of returned dataSources
-    return builder.entity(dataSourceNames).build();
+    if (full != null && includeDisabled == null) {
+      return Response.ok().entity(
+          Collections2.filter(
+              metadataSegmentManager.getInventory(),
+              new Predicate<DruidDataSource>()
+              {
+                @Override
+                public boolean apply(DruidDataSource input)
+                {
+                  return dataSourceNamesPostAuth.contains(input.getName());
+                }
+              }
+          )
+      ).build();
+    } else {
+      return Response.ok().entity(dataSourceNamesPostAuth).build();
+    }
   }
 
   @GET