From ed912c8abd11de63881653f10b913c5d380fbb6b Mon Sep 17 00:00:00 2001
From: GCHQDev404 <gchqdev404@users.noreply.github.com>
Date: Tue, 6 Jul 2021 07:17:28 -0700
Subject: [PATCH] gh-2455 FederatedStore GetTraits Audit. Refactoring
 FederatedGetTraitsHandler to use FederatedChain. Fixing graphSizeCount.

---
 .../impl/FederatedGetTraitsHandler.java        | 18 ++++--------------
 1 file changed, 4 insertions(+), 14 deletions(-)

diff --git a/store-implementation/federated-store/src/main/java/uk/gov/gchq/gaffer/federatedstore/operation/handler/impl/FederatedGetTraitsHandler.java b/store-implementation/federated-store/src/main/java/uk/gov/gchq/gaffer/federatedstore/operation/handler/impl/FederatedGetTraitsHandler.java
index b94eb134862..0e69afa744f 100644
--- a/store-implementation/federated-store/src/main/java/uk/gov/gchq/gaffer/federatedstore/operation/handler/impl/FederatedGetTraitsHandler.java
+++ b/store-implementation/federated-store/src/main/java/uk/gov/gchq/gaffer/federatedstore/operation/handler/impl/FederatedGetTraitsHandler.java
@@ -18,8 +18,8 @@
 
 import uk.gov.gchq.gaffer.commonutil.iterable.CloseableIterable;
 import uk.gov.gchq.gaffer.commonutil.stream.Streams;
+import uk.gov.gchq.gaffer.federatedstore.FederatedStore;
 import uk.gov.gchq.gaffer.federatedstore.operation.FederatedOperationChain;
-import uk.gov.gchq.gaffer.federatedstore.operation.GetAllGraphIds;
 import uk.gov.gchq.gaffer.operation.OperationChain;
 import uk.gov.gchq.gaffer.operation.OperationException;
 import uk.gov.gchq.gaffer.store.Context;
@@ -36,6 +36,7 @@
 
 import static java.util.Objects.isNull;
 import static java.util.Objects.nonNull;
+import static uk.gov.gchq.gaffer.federatedstore.FederatedStoreConstants.KEY_OPERATION_OPTIONS_GRAPH_IDS;
 
 /**
  * returns a set of {@link StoreTrait} that are common for all visible graphs.
@@ -48,13 +49,10 @@ public class FederatedGetTraitsHandler implements OutputOperationHandler<GetTrai
     @Override
     public Set<StoreTrait> doOperation(final GetTraits operation, final Context context, final Store store) throws OperationException {
         try {
-            HashMap<String, String> options = isNull(operation.getOptions()) ? new HashMap<>() : new HashMap<>(operation.getOptions());
-            int graphIdsSize = getGraphIdsSize(options, store, context);
-
             FederatedOperationChain<Void, StoreTrait> wrappedFedChain = new FederatedOperationChain.Builder<Void, StoreTrait>()
                     .operationChain(OperationChain.wrap(operation))
                     //deep copy options
-                    .options(options)
+                    .options(isNull(operation.getOptions()) ? new HashMap<>() : new HashMap<>(operation.getOptions()))
                     .build();
 
             final CloseableIterable<StoreTrait> concatResults = store.execute(wrappedFedChain, context);
@@ -65,6 +63,7 @@ public Set<StoreTrait> doOperation(final GetTraits operation, final Context cont
                         //.flatMap(Collection::stream)
                         .collect(Collectors.toMap(t -> t, ignore -> 1, (existing, replacement) -> existing + replacement));
 
+                long graphIdsSize = ((FederatedStore) store).getGraphs(context.getUser(), operation.getOption(KEY_OPERATION_OPTIONS_GRAPH_IDS), operation).stream().count();
                 rtn.values().removeIf(v -> v < graphIdsSize);
             } else {
                 rtn = Collections.EMPTY_MAP;
@@ -75,13 +74,4 @@ public Set<StoreTrait> doOperation(final GetTraits operation, final Context cont
             throw new OperationException("Error getting federated traits.", e);
         }
     }
-
-    private int getGraphIdsSize(final HashMap<String, String> options, final Store store, final Context context) throws OperationException {
-        int graphIdsSize = 0;
-        Iterable<? extends String> execute = store.execute(new GetAllGraphIds.Builder().options(new HashMap<>(options)).build(), context);
-        for (final String ignore : execute) {
-            graphIdsSize++;
-        }
-        return graphIdsSize;
-    }
 }