address feedback

startreedata · Jan 15, 2025 · 9f3463c · 9f3463c
1 parent cb23c39
commit 9f3463c
Show file tree

Hide file tree

Showing 11 changed files with 145 additions and 49 deletions.
diff --git a/thirdeye-integration-tests/src/test/java/ai/startree/thirdeye/HappyPathTest.java b/thirdeye-integration-tests/src/test/java/ai/startree/thirdeye/HappyPathTest.java
@@ -56,6 +56,7 @@
 import ai.startree.thirdeye.spi.api.NamespaceConfigurationApi;
 import ai.startree.thirdeye.spi.api.NotificationSchemesApi;
 import ai.startree.thirdeye.spi.api.PlanNodeApi;
+import ai.startree.thirdeye.spi.api.QuotasConfigurationApi;
 import ai.startree.thirdeye.spi.api.RcaInvestigationApi;
 import ai.startree.thirdeye.spi.api.SubscriptionGroupApi;
 import ai.startree.thirdeye.spi.api.TaskQuotasConfigurationApi;
@@ -741,7 +742,8 @@ public void testUpdateNamespaceConfiguration() {
     updatedCfg.setAuth(new AuthorizationConfigurationApi());
     updatedCfg.setId(namespaceConfigurationId);
     updatedCfg.setTemplateConfiguration(new TemplateConfigurationApi());
-    updatedCfg.setTaskQuotasConfiguration(new TaskQuotasConfigurationApi());
+    updatedCfg.setQuotasConfiguration(new QuotasConfigurationApi()
+        .setTaskQuotasConfiguration(new TaskQuotasConfigurationApi()));
     final Response response = request("api/workspace-configuration").put(
         Entity.json(updatedCfg));
     assertThat(response.getStatus()).isEqualTo(200);

diff --git a/...e/src/main/java/ai/startree/thirdeye/datalayer/bao/NamespaceConfigurationManagerImpl.java b/...e/src/main/java/ai/startree/thirdeye/datalayer/bao/NamespaceConfigurationManagerImpl.java
@@ -23,6 +23,7 @@
 import ai.startree.thirdeye.spi.datalayer.bao.NamespaceConfigurationManager;
 import ai.startree.thirdeye.spi.datalayer.dto.AuthorizationConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.NamespaceConfigurationDTO;
+import ai.startree.thirdeye.spi.datalayer.dto.QuotasConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.TaskQuotasConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.TemplateConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.TimeConfigurationDTO;
@@ -85,10 +86,15 @@ private boolean updateDefaults(final NamespaceConfigurationDTO existingNamespace
       existingNamespaceConfig.setTemplateConfiguration(defaultTemplateConfiguration());
       updated = true;
     }
-    if (force || existingNamespaceConfig.getTaskQuotasConfiguration() == null) {
-      existingNamespaceConfig.setTaskQuotasConfiguration(defaultTaskQuotasConfiguration());
+    if (force || existingNamespaceConfig.getQuotasConfiguration() == null) {
+      existingNamespaceConfig.setQuotasConfiguration(defaultQuotasConfiguration());
       updated = true;
     }
+    if (force || (existingNamespaceConfig.getQuotasConfiguration() != null &&
+        existingNamespaceConfig.getQuotasConfiguration().getTaskQuotasConfiguration() == null)) {
+      existingNamespaceConfig.getQuotasConfiguration()
+          .setTaskQuotasConfiguration(defaultTaskQuotasConfiguration());
+    }
     return updated;
   }
 
@@ -156,7 +162,7 @@ private NamespaceConfigurationDTO defaultNamespaceConfiguration(final String nam
     namespaceConfigurationDTO
         .setTimeConfiguration(defaultTimeConfiguration())
         .setTemplateConfiguration(defaultTemplateConfiguration())
-        .setTaskQuotasConfiguration(defaultTaskQuotasConfiguration())
+        .setQuotasConfiguration(defaultQuotasConfiguration())
         .setAuth(new AuthorizationConfigurationDTO().setNamespace(namespace));
     return namespaceConfigurationDTO;
   }
@@ -215,8 +221,15 @@ private TemplateConfigurationDTO defaultTemplateConfiguration() {
         .orElse(new TemplateConfigurationDTO());
   }
 
+  private QuotasConfigurationDTO defaultQuotasConfiguration() {
+    return optional(defaultNamespaceConfiguration.getQuotasConfiguration())
+        .orElse(new QuotasConfigurationDTO()
+            .setTaskQuotasConfiguration(new TaskQuotasConfigurationDTO()));
+  }
+
   private TaskQuotasConfigurationDTO defaultTaskQuotasConfiguration() {
-    return optional(defaultNamespaceConfiguration.getTaskQuotasConfiguration())
+    return optional(defaultNamespaceConfiguration.getQuotasConfiguration())
+        .map(QuotasConfigurationDTO::getTaskQuotasConfiguration)
         .orElse(new TaskQuotasConfigurationDTO());
   }
 }
diff --git a/thirdeye-persistence/src/test/java/ai/startree/thirdeye/datalayer/DatalayerTestUtils.java b/thirdeye-persistence/src/test/java/ai/startree/thirdeye/datalayer/DatalayerTestUtils.java
@@ -13,16 +13,15 @@
  */
 package ai.startree.thirdeye.datalayer;
 
-import static ai.startree.thirdeye.spi.Constants.DEFAULT_CHRONOLOGY;
 import static java.util.stream.Collectors.toSet;
 
-import ai.startree.thirdeye.spi.Constants;
 import ai.startree.thirdeye.spi.datalayer.Templatable;
 import ai.startree.thirdeye.spi.datalayer.dto.AbstractDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.AuthorizationConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.DatasetConfigDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.MetricConfigDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.NamespaceConfigurationDTO;
+import ai.startree.thirdeye.spi.datalayer.dto.QuotasConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.TaskQuotasConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.TemplateConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.TimeConfigurationDTO;
@@ -71,7 +70,8 @@ public static NamespaceConfigurationDTO buildNamespaceConfiguration(String names
     dto.setTimeConfiguration(new TimeConfigurationDTO());
     dto.setAuth(new AuthorizationConfigurationDTO().setNamespace(namespace));
     dto.setTemplateConfiguration(new TemplateConfigurationDTO());
-    dto.setTaskQuotasConfiguration(new TaskQuotasConfigurationDTO());
+    dto.setQuotasConfiguration(new QuotasConfigurationDTO()
+        .setTaskQuotasConfiguration(new TaskQuotasConfigurationDTO()));
     return dto;
   }
 }
diff --git a/...e/src/test/java/ai/startree/thirdeye/datalayer/bao/TestNamespaceConfigurationManager.java b/...e/src/test/java/ai/startree/thirdeye/datalayer/bao/TestNamespaceConfigurationManager.java
@@ -190,9 +190,9 @@ private void compareDtos(NamespaceConfigurationDTO dto1, NamespaceConfigurationD
         .isEqualTo(dto2.getTimeConfiguration().getMinimumOnboardingStartTime());
     assertThat(dto1.getTemplateConfiguration().getSqlLimitStatement())
         .isEqualTo(dto2.getTemplateConfiguration().getSqlLimitStatement());
-    assertThat(dto1.getTaskQuotasConfiguration().getDetectionTaskQuota())
-        .isEqualTo(dto2.getTaskQuotasConfiguration().getDetectionTaskQuota());
-    assertThat(dto1.getTaskQuotasConfiguration().getNotificationTaskQuota())
-        .isEqualTo(dto2.getTaskQuotasConfiguration().getNotificationTaskQuota());
+    assertThat(dto1.getQuotasConfiguration().getTaskQuotasConfiguration().getMaximumDetectionTasksPerMonth()).
+        isEqualTo(dto2.getQuotasConfiguration().getTaskQuotasConfiguration().getMaximumDetectionTasksPerMonth());
+    assertThat(dto1.getQuotasConfiguration().getTaskQuotasConfiguration().getMaximumNotificationTasksPerMonth()).
+        isEqualTo(dto2.getQuotasConfiguration().getTaskQuotasConfiguration().getMaximumNotificationTasksPerMonth());
   }
 }
diff --git a/...deye-server/src/main/java/ai/startree/thirdeye/service/NamespaceConfigurationService.java b/...deye-server/src/main/java/ai/startree/thirdeye/service/NamespaceConfigurationService.java
@@ -83,14 +83,14 @@ protected void validateUpdate(NamespaceConfigurationApi existing,
       NamespaceConfigurationApi updated) {
     if (!Objects.equals(existing.getId(), updated.getId()) ||
         !Objects.equals(existing.getAuth().getNamespace(), updated.getAuth().getNamespace()) ||
-        !Objects.equals(existing.getTaskQuotasConfiguration().getDetectionTaskQuota(),
-            updated.getTaskQuotasConfiguration().getDetectionTaskQuota()) ||
-        !Objects.equals(existing.getTaskQuotasConfiguration().getNotificationTaskQuota(),
-            updated.getTaskQuotasConfiguration().getNotificationTaskQuota())) {
+        !Objects.equals(existing.getQuotasConfiguration().getTaskQuotasConfiguration().getMaximumDetectionTasksPerMonth(),
+            updated.getQuotasConfiguration().getTaskQuotasConfiguration().getMaximumNotificationTasksPerMonth()) ||
+        !Objects.equals(existing.getQuotasConfiguration().getTaskQuotasConfiguration().getMaximumNotificationTasksPerMonth(),
+            updated.getQuotasConfiguration().getTaskQuotasConfiguration().getMaximumNotificationTasksPerMonth())) {
       throw badRequest(
           ThirdEyeStatus.ERR_NAMESPACE_CONFIGURATION_VALIDATION_FAILED,
           existing.namespace(),
-          "Updating Id, auth, or taskQuotasConfiguration is not allowed for Namespace Configuration");
+          "Updating Id, auth, or quotasConfiguration is not allowed for Namespace Configuration");
     }
   }
 }
diff --git a/thirdeye-spi/src/main/java/ai/startree/thirdeye/spi/api/NamespaceConfigurationApi.java b/thirdeye-spi/src/main/java/ai/startree/thirdeye/spi/api/NamespaceConfigurationApi.java
@@ -22,7 +22,7 @@ public class NamespaceConfigurationApi implements ThirdEyeCrudApi<NamespaceConfi
 
   private TemplateConfigurationApi templateConfiguration;
 
-  private TaskQuotasConfigurationApi taskQuotasConfiguration;
+  private QuotasConfigurationApi quotasConfiguration;
 
   public TimeConfigurationApi getTimeConfiguration() {
     return timeConfiguration;
@@ -62,13 +62,13 @@ public NamespaceConfigurationApi setTemplateConfiguration(
     return this;
   }
 
-  public TaskQuotasConfigurationApi getTaskQuotasConfiguration() {
-    return taskQuotasConfiguration;
+  public QuotasConfigurationApi getQuotasConfiguration() {
+    return quotasConfiguration;
   }
 
-  public NamespaceConfigurationApi setTaskQuotasConfiguration(
-      final TaskQuotasConfigurationApi taskQuotasConfiguration) {
-    this.taskQuotasConfiguration = taskQuotasConfiguration;
+  public NamespaceConfigurationApi setQuotasConfiguration(
+      final QuotasConfigurationApi quotasConfiguration) {
+    this.quotasConfiguration = quotasConfiguration;
     return this;
   }
 }
diff --git a/thirdeye-spi/src/main/java/ai/startree/thirdeye/spi/api/QuotasConfigurationApi.java b/thirdeye-spi/src/main/java/ai/startree/thirdeye/spi/api/QuotasConfigurationApi.java
@@ -0,0 +1,48 @@
+/*
+ * Copyright 2024 StarTree Inc
+ *
+ * Licensed under the StarTree Community License (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.startree.ai/legal/startree-community-license
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT * WARRANTIES OF ANY KIND,
+ * either express or implied.
+ * See the License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package ai.startree.thirdeye.spi.api;
+
+import java.util.Objects;
+
+public class QuotasConfigurationApi {
+
+  private TaskQuotasConfigurationApi taskQuotasConfiguration;
+
+  public TaskQuotasConfigurationApi getTaskQuotasConfiguration() {
+    return taskQuotasConfiguration;
+  }
+
+  public QuotasConfigurationApi setTaskQuotasConfiguration(
+      final TaskQuotasConfigurationApi taskQuotasConfiguration) {
+    this.taskQuotasConfiguration = taskQuotasConfiguration;
+    return this;
+  }
+
+  @Override
+  public boolean equals(final Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    final QuotasConfigurationApi that = (QuotasConfigurationApi) o;
+    return Objects.equals(taskQuotasConfiguration, that.taskQuotasConfiguration);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(taskQuotasConfiguration);
+  }
+}
diff --git a/thirdeye-spi/src/main/java/ai/startree/thirdeye/spi/api/TaskQuotasConfigurationApi.java b/thirdeye-spi/src/main/java/ai/startree/thirdeye/spi/api/TaskQuotasConfigurationApi.java
@@ -15,24 +15,26 @@
 
 public class TaskQuotasConfigurationApi {
 
-  private Long DetectionTaskQuota;
-  private Long NotificationTaskQuota;
+  private Long maximumDetectionTasksPerMonth;
+  private Long maximumNotificationTasksPerMonth;
 
-  public Long getDetectionTaskQuota() {
-    return DetectionTaskQuota;
+  public Long getMaximumDetectionTasksPerMonth() {
+    return maximumDetectionTasksPerMonth;
   }
 
-  public TaskQuotasConfigurationApi setDetectionTaskQuota(final Long detectionTaskQuota) {
-    this.DetectionTaskQuota = detectionTaskQuota;
+  public TaskQuotasConfigurationApi setMaximumDetectionTasksPerMonth(
+      final Long maximumDetectionTasksPerMonth) {
+    this.maximumDetectionTasksPerMonth = maximumDetectionTasksPerMonth;
     return this;
   }
 
-  public Long getNotificationTaskQuota() {
-    return NotificationTaskQuota;
+  public Long getMaximumNotificationTasksPerMonth() {
+    return maximumNotificationTasksPerMonth;
   }
 
-  public TaskQuotasConfigurationApi setNotificationTaskQuota(final Long notificationTaskQuota) {
-    this.NotificationTaskQuota = notificationTaskQuota;
+  public TaskQuotasConfigurationApi setMaximumNotificationTasksPerMonth(
+      final Long maximumNotificationTasksPerMonth) {
+    this.maximumNotificationTasksPerMonth = maximumNotificationTasksPerMonth;
     return this;
   }
 }
diff --git a/...e-spi/src/main/java/ai/startree/thirdeye/spi/datalayer/dto/NamespaceConfigurationDTO.java b/...e-spi/src/main/java/ai/startree/thirdeye/spi/datalayer/dto/NamespaceConfigurationDTO.java
@@ -13,6 +13,7 @@
  */
 package ai.startree.thirdeye.spi.datalayer.dto;
 
+import ai.startree.thirdeye.spi.api.QuotasConfigurationApi;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonInclude.Include;
@@ -25,7 +26,7 @@ public class NamespaceConfigurationDTO extends AbstractDTO {
 
   private TemplateConfigurationDTO templateConfiguration;
 
-  private TaskQuotasConfigurationDTO taskQuotasConfiguration;
+  private QuotasConfigurationDTO quotasConfiguration;
 
   public TimeConfigurationDTO getTimeConfiguration() {
     return timeConfiguration;
@@ -47,13 +48,13 @@ public NamespaceConfigurationDTO setTemplateConfiguration(
     return this;
   }
 
-  public TaskQuotasConfigurationDTO getTaskQuotasConfiguration() {
-    return taskQuotasConfiguration;
+  public QuotasConfigurationDTO getQuotasConfiguration() {
+    return quotasConfiguration;
   }
 
-  public NamespaceConfigurationDTO setTaskQuotasConfiguration(
-      final TaskQuotasConfigurationDTO taskQuotasConfiguration) {
-    this.taskQuotasConfiguration = taskQuotasConfiguration;
+  public NamespaceConfigurationDTO setQuotasConfiguration(
+      final QuotasConfigurationDTO quotasConfiguration) {
+    this.quotasConfiguration = quotasConfiguration;
     return this;
   }
 }
diff --git a/...deye-spi/src/main/java/ai/startree/thirdeye/spi/datalayer/dto/QuotasConfigurationDTO.java b/...deye-spi/src/main/java/ai/startree/thirdeye/spi/datalayer/dto/QuotasConfigurationDTO.java
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2024 StarTree Inc
+ *
+ * Licensed under the StarTree Community License (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.startree.ai/legal/startree-community-license
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT * WARRANTIES OF ANY KIND,
+ * either express or implied.
+ * See the License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package ai.startree.thirdeye.spi.datalayer.dto;
+
+public class QuotasConfigurationDTO {
+
+  private TaskQuotasConfigurationDTO taskQuotasConfiguration;
+
+  public TaskQuotasConfigurationDTO getTaskQuotasConfiguration() {
+    return taskQuotasConfiguration;
+  }
+
+  public QuotasConfigurationDTO setTaskQuotasConfiguration(
+      final TaskQuotasConfigurationDTO taskQuotasConfiguration) {
+    this.taskQuotasConfiguration = taskQuotasConfiguration;
+    return this;
+  }
+}
diff --git a/...-spi/src/main/java/ai/startree/thirdeye/spi/datalayer/dto/TaskQuotasConfigurationDTO.java b/...-spi/src/main/java/ai/startree/thirdeye/spi/datalayer/dto/TaskQuotasConfigurationDTO.java
@@ -18,27 +18,28 @@ public class TaskQuotasConfigurationDTO {
   /**
    * Monthly quota for number of detection task runs
    */
-  private Long DetectionTaskQuota;
+  private Long maximumDetectionTasksPerMonth;
   /**
    * Monthly quota for number of notification task runs
    */
-  private Long NotificationTaskQuota;
+  private Long maximumNotificationTasksPerMonth;
 
-  public Long getDetectionTaskQuota() {
-    return DetectionTaskQuota;
+  public Long getMaximumDetectionTasksPerMonth() {
+    return maximumDetectionTasksPerMonth;
   }
 
-  public TaskQuotasConfigurationDTO setDetectionTaskQuota(final Long detectionTaskQuota) {
-    this.DetectionTaskQuota = detectionTaskQuota;
+  public TaskQuotasConfigurationDTO setMaximumDetectionTasksPerMonth(
+      final Long maximumDetectionTasksPerMonth) {
+    this.maximumDetectionTasksPerMonth = maximumDetectionTasksPerMonth;
     return this;
   }
 
-  public Long getNotificationTaskQuota() {
-    return NotificationTaskQuota;
+  public Long getMaximumNotificationTasksPerMonth() {
+    return maximumNotificationTasksPerMonth;
   }
 
-  public TaskQuotasConfigurationDTO setNotificationTaskQuota(final Long notificationTaskQuota) {
-    this.NotificationTaskQuota = notificationTaskQuota;
+  public TaskQuotasConfigurationDTO setMaximumNotificationTasksPerMonth(final Long maximumNotificationTasksPerMonth) {
+    this.maximumNotificationTasksPerMonth = maximumNotificationTasksPerMonth;
     return this;
   }
 }