Add the FeatureFlag entity (#2464)

* Add FeatureFlag entity

* Add converter

* Add loading cache

* Add more tests

* Fix NPE in cache

* small fixes

core/src/main/java/google/registry/model/common/FeatureFlag.java

@@ -0,0 +1,177 @@
+// Copyright 2024 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.common;
+
+import static com.google.api.client.util.Preconditions.checkState;
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static google.registry.config.RegistryConfig.getSingletonCacheRefreshDuration;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+
+import com.github.benmanes.caffeine.cache.CacheLoader;
+import com.github.benmanes.caffeine.cache.LoadingCache;
+import com.google.common.base.Joiner;
+import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Sets;
+import google.registry.model.Buildable;
+import google.registry.model.CacheUtils;
+import google.registry.model.ImmutableObject;
+import google.registry.persistence.VKey;
+import java.util.Map;
+import java.util.Set;
+import javax.persistence.Column;
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import org.joda.time.DateTime;
+
+@Entity
+public class FeatureFlag extends ImmutableObject implements Buildable {
+
+  /**
+   * The current status of the feature the flag represents.
+   *
+   * <p>Currently, there is no enforced ordering of these status values, but that may change in the
+   * future should new statuses be added to this enum that require it.
+   */
+  public enum FeatureStatus {
+    ACTIVE,
+    INACTIVE
+  }
+
+  /** The name of the flag/feature. */
+  @Id String featureName;
+
+  /** A map of times for each {@link FeatureStatus} the FeatureFlag should hold. */
+  @Column(nullable = false)
+  TimedTransitionProperty<FeatureStatus> status =
+      TimedTransitionProperty.withInitialValue(FeatureStatus.INACTIVE);
+
+  public static FeatureFlag get(String featureName) {
+    FeatureFlag maybeFeatureFlag = CACHE.get(featureName);
+    if (maybeFeatureFlag == null) {
+      throw new FeatureFlagNotFoundException(featureName);
+    } else {
+      return maybeFeatureFlag;
+    }
+  }
+
+  public static ImmutableSet<FeatureFlag> get(Set<String> featureNames) {
+    Map<String, FeatureFlag> featureFlags = CACHE.getAll(featureNames);
+    ImmutableSet<String> missingFlags =
+        Sets.difference(featureNames, featureFlags.keySet()).immutableCopy();
+    if (missingFlags.isEmpty()) {
+      return featureFlags.values().stream().collect(toImmutableSet());
+    } else {
+      throw new FeatureFlagNotFoundException(missingFlags);
+    }
+  }
+
+  /** A cache that loads the {@link FeatureFlag} for a given featureName. */
+  private static final LoadingCache<String, FeatureFlag> CACHE =
+      CacheUtils.newCacheBuilder(getSingletonCacheRefreshDuration())
+          .build(
+              new CacheLoader<>() {
+                @Override
+                public FeatureFlag load(final String featureName) {
+                  return tm().reTransact(() -> tm().loadByKeyIfPresent(createVKey(featureName)))
+                      .orElse(null);
+                }
+
+                @Override
+                public Map<? extends String, ? extends FeatureFlag> loadAll(
+                    Set<? extends String> featureFlagNames) {
+                  ImmutableMap<String, VKey<FeatureFlag>> keysMap =
+                      featureFlagNames.stream()
+                          .collect(
+                              toImmutableMap(featureName -> featureName, FeatureFlag::createVKey));
+                  Map<VKey<? extends FeatureFlag>, FeatureFlag> entities =
+                      tm().reTransact(() -> tm().loadByKeysIfPresent(keysMap.values()));
+                  return entities.values().stream()
+                      .collect(toImmutableMap(flag -> flag.featureName, flag -> flag));
+                }
+              });
+
+  public static VKey<FeatureFlag> createVKey(String featureName) {
+    return VKey.create(FeatureFlag.class, featureName);
+  }
+
+  @Override
+  public VKey<FeatureFlag> createVKey() {
+    return createVKey(featureName);
+  }
+
+  public String getFeatureName() {
+    return featureName;
+  }
+
+  public TimedTransitionProperty<FeatureStatus> getStatusMap() {
+    return status;
+  }
+
+  public FeatureStatus getStatus(DateTime time) {
+    return status.getValueAtTime(time);
+  }
+
+  @Override
+  public FeatureFlag.Builder asBuilder() {
+    return new FeatureFlag.Builder(clone(this));
+  }
+
+  /** A builder for constructing {@link FeatureFlag} objects, since they are immutable. */
+  public static class Builder extends Buildable.Builder<FeatureFlag> {
+
+    public Builder() {}
+
+    private Builder(FeatureFlag instance) {
+      super(instance);
+    }
+
+    @Override
+    public FeatureFlag build() {
+      checkArgument(
+          !Strings.isNullOrEmpty(getInstance().featureName),
+          "Feature name must not be null or empty");
+      getInstance().status.checkValidity();
+      return super.build();
+    }
+
+    public Builder setFeatureName(String featureName) {
+      checkState(getInstance().featureName == null, "Feature name can only be set once");
+      getInstance().featureName = featureName;
+      return this;
+    }
+
+    public Builder setStatus(ImmutableSortedMap<DateTime, FeatureStatus> statusMap) {
+      getInstance().status = TimedTransitionProperty.fromValueMap(statusMap);
+      return this;
+    }
+  }
+
+  /** Exception to throw when no FeatureFlag entity is found for given FeatureName string(s). */
+  public static class FeatureFlagNotFoundException extends RuntimeException {
+
+    FeatureFlagNotFoundException(ImmutableSet<String> featureNames) {
+      super("No feature flag object(s) found for " + Joiner.on(", ").join(featureNames));
+    }
+
+    FeatureFlagNotFoundException(String featureName) {
+      this(ImmutableSet.of(featureName));
+    }
+  }
+}

core/src/main/java/google/registry/model/tld/Tld.java

@@ -41,9 +41,9 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
-import com.google.common.collect.Maps;
 import com.google.common.collect.Ordering;
 import com.google.common.collect.Range;
+import com.google.common.collect.Sets;
 import com.google.common.net.InternetDomainName;
 import google.registry.model.Buildable;
 import google.registry.model.CacheUtils;
@@ -204,10 +204,7 @@ public class Tld extends ImmutableObject implements Buildable, UnsafeSerializabl
   public static ImmutableSet<Tld> get(Set<String> tlds) {
     Map<String, Tld> registries = CACHE.getAll(tlds);
     ImmutableSet<String> missingRegistries =
-        registries.entrySet().stream()
-            .filter(e -> e.getValue() == null)
-            .map(Map.Entry::getKey)
-            .collect(toImmutableSet());
+        Sets.difference(tlds, registries.keySet()).immutableCopy();
     if (missingRegistries.isEmpty()) {
       return registries.values().stream().collect(toImmutableSet());
     } else {
@@ -243,7 +240,8 @@ public class Tld extends ImmutableObject implements Buildable, UnsafeSerializabl
                       tlds.stream().collect(toImmutableMap(tld -> tld, Tld::createVKey));
                   Map<VKey<? extends Tld>, Tld> entities =
                       tm().reTransact(() -> tm().loadByKeysIfPresent(keysMap.values()));
-                  return Maps.transformEntries(keysMap, (k, v) -> entities.getOrDefault(v, null));
+                  return entities.values().stream()
+                      .collect(toImmutableMap(tld -> tld.tldStr, tld -> tld));
                 }
               });
 

core/src/main/java/google/registry/persistence/converter/FeatureFlagStatusConverter.java

@@ -0,0 +1,34 @@
+// Copyright 2024 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.persistence.converter;
+
+import google.registry.model.common.FeatureFlag.FeatureStatus;
+import javax.persistence.Converter;
+
+/** JPA converter for {@link google.registry.model.common.FeatureFlag} status transitions. */
+@Converter(autoApply = true)
+public class FeatureFlagStatusConverter
+    extends TimedTransitionPropertyConverterBase<FeatureStatus> {
+
+  @Override
+  protected String convertValueToString(FeatureStatus value) {
+    return value.toString();
+  }
+
+  @Override
+  protected FeatureStatus convertStringToValue(String string) {
+    return FeatureStatus.valueOf(string);
+  }
+}

core/src/main/resources/META-INF/persistence.xml

@@ -56,6 +56,7 @@
     <class>google.registry.model.contact.Contact</class>
     <class>google.registry.model.domain.Domain</class>
     <class>google.registry.model.domain.DomainHistory</class>
+    <class>google.registry.model.common.FeatureFlag</class>
     <class>google.registry.model.domain.GracePeriod</class>
     <class>google.registry.model.domain.GracePeriod$GracePeriodHistory</class>
     <class>google.registry.model.domain.secdns.DomainDsData</class>
@@ -97,6 +98,7 @@
     <class>google.registry.persistence.converter.CurrencyUnitConverter</class>
     <class>google.registry.persistence.converter.DateTimeConverter</class>
     <class>google.registry.persistence.converter.DurationConverter</class>
+    <class>google.registry.persistence.converter.FeatureFlagStatusConverter</class>
     <class>google.registry.persistence.converter.IdnTableEnumSetConverter</class>
     <class>google.registry.persistence.converter.InetAddressSetConverter</class>
     <class>google.registry.persistence.converter.LocalDateConverter</class>

core/src/test/java/google/registry/model/common/FeatureFlagTest.java

@@ -0,0 +1,207 @@
+// Copyright 2024 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.common;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.common.FeatureFlag.FeatureStatus.ACTIVE;
+import static google.registry.model.common.FeatureFlag.FeatureStatus.INACTIVE;
+import static google.registry.testing.DatabaseHelper.loadByEntity;
+import static google.registry.testing.DatabaseHelper.persistResource;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.joda.time.DateTimeZone.UTC;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
+import google.registry.model.EntityTestCase;
+import google.registry.model.common.FeatureFlag.FeatureFlagNotFoundException;
+import google.registry.model.common.FeatureFlag.FeatureStatus;
+import org.joda.time.DateTime;
+import org.junit.jupiter.api.Test;
+
+/** Unit tests for {@link FeatureFlag}. */
+public class FeatureFlagTest extends EntityTestCase {
+
+  public FeatureFlagTest() {
+    super(JpaEntityCoverageCheck.ENABLED);
+  }
+
+  @Test
+  void testSuccess_persistence() {
+    FeatureFlag featureFlag =
+        new FeatureFlag.Builder()
+            .setFeatureName("testFlag")
+            .setStatus(
+                ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                    .put(START_OF_TIME, INACTIVE)
+                    .put(DateTime.now(UTC).plusWeeks(8), ACTIVE)
+                    .build())
+            .build();
+    persistResource(featureFlag);
+    FeatureFlag flagFromDb = loadByEntity(featureFlag);
+    assertThat(featureFlag).isEqualTo(flagFromDb);
+  }
+
+  @Test
+  void testSuccess_getSingleFlag() {
+    FeatureFlag featureFlag =
+        new FeatureFlag.Builder()
+            .setFeatureName("testFlag")
+            .setStatus(
+                ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                    .put(START_OF_TIME, INACTIVE)
+                    .put(DateTime.now(UTC).plusWeeks(8), ACTIVE)
+                    .build())
+            .build();
+    persistResource(featureFlag);
+    assertThat(FeatureFlag.get("testFlag")).isEqualTo(featureFlag);
+  }
+
+  @Test
+  void testSuccess_getMultipleFlags() {
+    FeatureFlag featureFlag1 =
+        persistResource(
+            new FeatureFlag.Builder()
+                .setFeatureName("testFlag1")
+                .setStatus(
+                    ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                        .put(START_OF_TIME, INACTIVE)
+                        .put(DateTime.now(UTC).plusWeeks(8), ACTIVE)
+                        .build())
+                .build());
+    FeatureFlag featureFlag2 =
+        persistResource(
+            new FeatureFlag.Builder()
+                .setFeatureName("testFlag2")
+                .setStatus(
+                    ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                        .put(START_OF_TIME, INACTIVE)
+                        .put(DateTime.now(UTC).plusWeeks(3), INACTIVE)
+                        .build())
+                .build());
+    FeatureFlag featureFlag3 =
+        persistResource(
+            new FeatureFlag.Builder()
+                .setFeatureName("testFlag3")
+                .setStatus(
+                    ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                        .put(START_OF_TIME, INACTIVE)
+                        .build())
+                .build());
+    ImmutableSet<FeatureFlag> featureFlags =
+        FeatureFlag.get(ImmutableSet.of("testFlag1", "testFlag2", "testFlag3"));
+    assertThat(featureFlags.size()).isEqualTo(3);
+    assertThat(featureFlags).containsExactly(featureFlag1, featureFlag2, featureFlag3);
+  }
+
+  @Test
+  void testFailure_getMultipleFlagsOneMissing() {
+    persistResource(
+        new FeatureFlag.Builder()
+            .setFeatureName("testFlag1")
+            .setStatus(
+                ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                    .put(START_OF_TIME, INACTIVE)
+                    .put(DateTime.now(UTC).plusWeeks(8), ACTIVE)
+                    .build())
+            .build());
+    persistResource(
+        new FeatureFlag.Builder()
+            .setFeatureName("testFlag2")
+            .setStatus(
+                ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                    .put(START_OF_TIME, INACTIVE)
+                    .put(DateTime.now(UTC).plusWeeks(3), INACTIVE)
+                    .build())
+            .build());
+    FeatureFlagNotFoundException thrown =
+        assertThrows(
+            FeatureFlagNotFoundException.class,
+            () -> FeatureFlag.get(ImmutableSet.of("missingFlag", "testFlag1", "testFlag2")));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("No feature flag object(s) found for missingFlag");
+  }
+
+  @Test
+  void testFailure_featureFlagNotPresent() {
+    FeatureFlagNotFoundException thrown =
+        assertThrows(FeatureFlagNotFoundException.class, () -> FeatureFlag.get("fakeFlag"));
+    assertThat(thrown).hasMessageThat().isEqualTo("No feature flag object(s) found for fakeFlag");
+  }
+
+  @Test
+  void testFailure_resetFeatureName() {
+    FeatureFlag featureFlag =
+        new FeatureFlag.Builder()
+            .setFeatureName("testFlag")
+            .setStatus(
+                ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                    .put(START_OF_TIME, INACTIVE)
+                    .put(DateTime.now(UTC).plusWeeks(8), ACTIVE)
+                    .build())
+            .build();
+    IllegalStateException thrown =
+        assertThrows(
+            IllegalStateException.class,
+            () -> featureFlag.asBuilder().setFeatureName("differentName"));
+    assertThat(thrown).hasMessageThat().isEqualTo("Feature name can only be set once");
+  }
+
+  @Test
+  void testFailure_nullFeatureName() {
+    FeatureFlag.Builder featureFlagBuilder =
+        new FeatureFlag.Builder()
+            .setStatus(
+                ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                    .put(START_OF_TIME, INACTIVE)
+                    .put(DateTime.now(UTC).plusWeeks(8), ACTIVE)
+                    .build());
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, () -> featureFlagBuilder.build());
+    assertThat(thrown).hasMessageThat().isEqualTo("Feature name must not be null or empty");
+  }
+
+  @Test
+  void testFailure_emptyFeatureName() {
+    FeatureFlag.Builder featureFlagBuilder =
+        new FeatureFlag.Builder()
+            .setFeatureName("")
+            .setStatus(
+                ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                    .put(START_OF_TIME, INACTIVE)
+                    .put(DateTime.now(UTC).plusWeeks(8), ACTIVE)
+                    .build());
+    IllegalArgumentException thrown =
+        assertThrows(IllegalArgumentException.class, () -> featureFlagBuilder.build());
+    assertThat(thrown).hasMessageThat().isEqualTo("Feature name must not be null or empty");
+  }
+
+  @Test
+  void testFailure_invalidStatusMap() {
+    FeatureFlag.Builder featureFlagBuilder = new FeatureFlag.Builder().setFeatureName("testFlag");
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                featureFlagBuilder.setStatus(
+                    ImmutableSortedMap.<DateTime, FeatureStatus>naturalOrder()
+                        .put(DateTime.now(UTC).plusWeeks(8), ACTIVE)
+                        .build()));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("Must provide transition entry for the start of time (Unix Epoch)");
+  }
+}

core/src/test/java/google/registry/model/tld/TldTest.java

@@ -325,6 +325,16 @@ public final class TldTest extends EntityTestCase {
                 .values());
   }
 
+  @Test
+  void testFailure_testGetAllMissingTld() {
+    createTld("foo");
+    TldNotFoundException thrown =
+        assertThrows(
+            TldNotFoundException.class,
+            () -> assertThat(Tld.get(ImmutableSet.of("foo", "tld", "missing"))));
+    assertThat(thrown).hasMessageThat().isEqualTo("No TLD object(s) found for missing");
+  }
+
   @Test
   void testSetReservedLists() {
     ReservedList rl5 =

core/src/test/java/google/registry/schema/integration/SqlIntegrationTestSuite.java

@@ -23,6 +23,7 @@ import google.registry.bsa.persistence.BsaUnblockableDomainTest;
 import google.registry.model.billing.BillingBaseTest;
 import google.registry.model.common.CursorTest;
 import google.registry.model.common.DnsRefreshRequestTest;
+import google.registry.model.common.FeatureFlagTest;
 import google.registry.model.console.ConsoleEppActionHistoryTest;
 import google.registry.model.console.RegistrarPocUpdateHistoryTest;
 import google.registry.model.console.RegistrarUpdateHistoryTest;
@@ -103,6 +104,7 @@ import org.junit.runner.RunWith;
   DnsRefreshRequestTest.class,
   DomainSqlTest.class,
   DomainHistoryTest.class,
+  FeatureFlagTest.class,
   HostHistoryTest.class,
   LockTest.class,
   PollMessageTest.class,