Add option to only include flag data files relevant for a system

5 files changed, 90 insertions, 4 deletions

diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/systemflags/v1/SystemFlagsDataArchive.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/systemflags/v1/SystemFlagsDataArchive.java
index 8f91a8127bd..57dafcd3291 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/systemflags/v1/SystemFlagsDataArchive.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/systemflags/v1/SystemFlagsDataArchive.java
@@ -12,6 +12,7 @@ import com.yahoo.vespa.flags.FetchVector;
 import com.yahoo.vespa.flags.FlagId;
 import com.yahoo.vespa.flags.json.DimensionHelper;
 import com.yahoo.vespa.flags.json.FlagData;
+import com.yahoo.vespa.hosted.controller.api.integration.zone.ZoneRegistry;
 
 import java.io.BufferedInputStream;
 import java.io.IOException;
@@ -71,7 +72,7 @@ public class SystemFlagsDataArchive {
                 if (!entry.isDirectory() && name.startsWith("flags/")) {
                     Path filePath = Paths.get(name);
                     String rawData = new String(zipIn.readAllBytes(), StandardCharsets.UTF_8);
-                    addFile(builder, rawData, filePath);
+                    addFile(builder, rawData, filePath, Set.of());
                 }
             }
             return builder.build();
@@ -80,7 +81,14 @@ public class SystemFlagsDataArchive {
         }
     }
 
-    public static SystemFlagsDataArchive fromDirectory(Path directory) {
+    public static SystemFlagsDataArchive fromDirectoryAndSystem(Path directory, ZoneRegistry systemDefinition) {
+        return fromDirectory(directory, systemDefinition);
+    }
+
+    public static SystemFlagsDataArchive fromDirectory(Path directory) { return fromDirectory(directory, null); }
+
+    private static SystemFlagsDataArchive fromDirectory(Path directory, ZoneRegistry systemDefinition) {
+        Set<String> filenamesForSystem = getFilenamesForSystem(systemDefinition);
         Path root = directory.toAbsolutePath();
         Path flagsDirectory = directory.resolve("flags");
         if (!Files.isDirectory(flagsDirectory)) {
@@ -93,7 +101,7 @@ public class SystemFlagsDataArchive {
                 if (!Files.isDirectory(absolutePath) &&
                         relativePath.startsWith("flags")) {
                     String rawData = uncheck(() -> Files.readString(absolutePath, StandardCharsets.UTF_8));
-                    addFile(builder, rawData, relativePath);
+                    addFile(builder, rawData, relativePath, filenamesForSystem);
                 }
             });
             return builder.build();
@@ -102,6 +110,7 @@ public class SystemFlagsDataArchive {
         }
     }
 
+
     public void toZip(OutputStream out) {
         ZipOutputStream zipOut = new ZipOutputStream(out);
         files.forEach((flagId, fileMap) -> {
@@ -152,11 +161,21 @@ public class SystemFlagsDataArchive {
         });
     }
 
-    private static void addFile(Builder builder, String rawData, Path filePath) {
+    private static Set<String> getFilenamesForSystem(ZoneRegistry systemDefinition) {
+        if (systemDefinition == null) return Set.of();
+        return FlagsTarget.getAllTargetsInSystem(systemDefinition).stream()
+                .flatMap(target -> target.flagDataFilesPrioritized().stream())
+                .collect(Collectors.toSet());
+    }
+
+    private static void addFile(Builder builder, String rawData, Path filePath, Set<String> filenamesForSystem) {
         String filename = filePath.getFileName().toString();
         if (filename.startsWith(".")) {
             return; // Ignore files starting with '.'
         }
+        if (!filenamesForSystem.isEmpty() && !filenamesForSystem.contains(filename)) {
+            return; // Ignore files irrelevant for system
+        }
         if (!filename.endsWith(".json")) {
             throw new IllegalArgumentException(String.format("Only JSON files are allowed in 'flags/' directory (found '%s')", filePath.toString()));
         }
diff --git a/controller-api/src/test/java/com/yahoo/vespa/hosted/controller/api/systemflags/v1/SystemFlagsDataArchiveTest.java b/controller-api/src/test/java/com/yahoo/vespa/hosted/controller/api/systemflags/v1/SystemFlagsDataArchiveTest.java
index 771e42e85f9..d29657035c4 100644
--- a/controller-api/src/test/java/com/yahoo/vespa/hosted/controller/api/systemflags/v1/SystemFlagsDataArchiveTest.java
+++ b/controller-api/src/test/java/com/yahoo/vespa/hosted/controller/api/systemflags/v1/SystemFlagsDataArchiveTest.java
@@ -2,16 +2,20 @@
 
 package com.yahoo.vespa.hosted.controller.api.systemflags.v1;
 
+import com.yahoo.config.provision.CloudName;
 import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.SystemName;
+import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.config.provision.zone.ZoneId;
+import com.yahoo.config.provision.zone.ZoneList;
 import com.yahoo.text.JSON;
 import com.yahoo.vespa.athenz.api.AthenzService;
 import com.yahoo.vespa.flags.FetchVector;
 import com.yahoo.vespa.flags.FlagId;
 import com.yahoo.vespa.flags.RawFlag;
 import com.yahoo.vespa.flags.json.FlagData;
+import com.yahoo.vespa.hosted.controller.api.integration.zone.ZoneRegistry;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
@@ -26,6 +30,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.net.URI;
+import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.List;
 import java.util.Map;
@@ -35,6 +40,9 @@ import static java.util.stream.Collectors.toList;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 /**
  * @author bjorncs
@@ -52,6 +60,7 @@ public class SystemFlagsDataArchiveTest {
     public final ExpectedException expectedException = ExpectedException.none();
 
     private static final FlagsTarget mainControllerTarget = FlagsTarget.forController(SYSTEM);
+    private static final FlagsTarget cdControllerTarget = FlagsTarget.forController(SystemName.cd);
     private static final FlagsTarget prodUsWestCfgTarget = createConfigserverTarget(Environment.prod, "us-west-1");
     private static final FlagsTarget prodUsEast3CfgTarget = createConfigserverTarget(Environment.prod, "us-east-3");
     private static final FlagsTarget devUsEast1CfgTarget = createConfigserverTarget(Environment.dev, "us-east-1");
@@ -260,6 +269,30 @@ public class SystemFlagsDataArchiveTest {
         }
     }
 
+    @Test
+    public void ignores_files_not_related_to_specified_system_definition() {
+        ZoneRegistry registry = createZoneRegistryMock();
+        Path testDirectory = Paths.get("src/test/resources/system-flags-for-multiple-systems/");
+        var archive = SystemFlagsDataArchive.fromDirectoryAndSystem(testDirectory, registry);
+        assertFlagDataHasValue(archive, MY_TEST_FLAG, cdControllerTarget, "default"); // Would be 'cd.controller' if files for CD system were included
+        assertFlagDataHasValue(archive, MY_TEST_FLAG, mainControllerTarget, "default");
+        assertFlagDataHasValue(archive, MY_TEST_FLAG, prodUsWestCfgTarget, "main.prod.us-west-1");
+    }
+
+    @SuppressWarnings("unchecked") // workaround for mocking a method for generic return type
+    private static ZoneRegistry createZoneRegistryMock() {
+        // Cannot use the standard registry mock as it's located in controller-server module
+        ZoneRegistry registryMock = mock(ZoneRegistry.class);
+        when(registryMock.system()).thenReturn(SystemName.main);
+        when(registryMock.getConfigServerVipUri(any())).thenReturn(URI.create("http://localhost:8080/"));
+        when(registryMock.getConfigServerHttpsIdentity(any())).thenReturn(new AthenzService("domain", "servicename"));
+        ZoneList zoneListMock = mock(ZoneList.class);
+        when(zoneListMock.reachable()).thenReturn(zoneListMock);
+        when(zoneListMock.zones()).thenReturn((List)List.of(new SimpleZone("prod.us-west-1"), new SimpleZone("prod.us-east-3")));
+        when(registryMock.zones()).thenReturn(zoneListMock);
+        return registryMock;
+    }
+
     private static void assertArchiveReturnsCorrectTestFlagDataForTarget(SystemFlagsDataArchive archive) {
         assertFlagDataHasValue(archive, MY_TEST_FLAG, mainControllerTarget, "main.controller");
         assertFlagDataHasValue(archive, MY_TEST_FLAG, prodUsWestCfgTarget, "main.prod.us-west-1");
@@ -286,4 +319,14 @@ public class SystemFlagsDataArchiveTest {
                 .collect(toList());
     }
 
+    private static class SimpleZone implements ZoneApi {
+        final ZoneId zoneId;
+        SimpleZone(String zoneId) { this.zoneId = ZoneId.from(zoneId); }
+
+        @Override public SystemName getSystemName() { return  SystemName.main; }
+        @Override public ZoneId getId() { return zoneId; }
+        @Override public CloudName getCloudName() { throw new UnsupportedOperationException(); }
+        @Override public String getCloudNativeRegionName() { throw new UnsupportedOperationException(); }
+    }
+
 }
\ No newline at end of file
diff --git a/controller-api/src/test/resources/system-flags-for-multiple-systems/flags/my-test-flag/cd.controller.json b/controller-api/src/test/resources/system-flags-for-multiple-systems/flags/my-test-flag/cd.controller.json
new file mode 100644
index 00000000000..ce3cdd43889
--- /dev/null
+++ b/controller-api/src/test/resources/system-flags-for-multiple-systems/flags/my-test-flag/cd.controller.json
@@ -0,0 +1,8 @@
+{
+  "id" : "my-test-flag",
+  "rules" : [
+    {
+      "value" : "cd.controller"
+    }
+  ]
+}
\ No newline at end of file
diff --git a/controller-api/src/test/resources/system-flags-for-multiple-systems/flags/my-test-flag/default.json b/controller-api/src/test/resources/system-flags-for-multiple-systems/flags/my-test-flag/default.json
new file mode 100644
index 00000000000..5924eb860c0
--- /dev/null
+++ b/controller-api/src/test/resources/system-flags-for-multiple-systems/flags/my-test-flag/default.json
@@ -0,0 +1,8 @@
+{
+  "id" : "my-test-flag",
+  "rules" : [
+    {
+      "value" : "default"
+    }
+  ]
+}
\ No newline at end of file
diff --git a/controller-api/src/test/resources/system-flags-for-multiple-systems/flags/my-test-flag/main.prod.us-west-1.json b/controller-api/src/test/resources/system-flags-for-multiple-systems/flags/my-test-flag/main.prod.us-west-1.json
new file mode 100644
index 00000000000..45989773df8
--- /dev/null
+++ b/controller-api/src/test/resources/system-flags-for-multiple-systems/flags/my-test-flag/main.prod.us-west-1.json
@@ -0,0 +1,8 @@
+{
+  "id" : "my-test-flag",
+  "rules" : [
+    {
+      "value" : "main.prod.us-west-1"
+    }
+  ]
+}
\ No newline at end of file