[persistence] Add TaskQuotasConfiguration to workspace configuration

thirdeye-integration-tests/src/test/java/ai/startree/thirdeye/HappyPathTest.java

@@ -58,6 +58,8 @@
 import ai.startree.thirdeye.spi.api.PlanNodeApi;
 import ai.startree.thirdeye.spi.api.RcaInvestigationApi;
 import ai.startree.thirdeye.spi.api.SubscriptionGroupApi;
+import ai.startree.thirdeye.spi.api.TaskQuotasConfigurationApi;
+import ai.startree.thirdeye.spi.api.TemplateConfigurationApi;
 import ai.startree.thirdeye.spi.api.TimeConfigurationApi;
 import ai.startree.thirdeye.spi.detection.AnomalyCause;
 import ai.startree.thirdeye.spi.detection.AnomalyFeedbackType;
@@ -738,6 +740,8 @@ public void testUpdateNamespaceConfiguration() {
             .setMinimumOnboardingStartTime(996684800000L));
     updatedCfg.setAuth(new AuthorizationConfigurationApi());
     updatedCfg.setId(namespaceConfigurationId);
+    updatedCfg.setTemplateConfiguration(new TemplateConfigurationApi());
+    updatedCfg.setTaskQuotasConfiguration(new TaskQuotasConfigurationApi());
     final Response response = request("api/workspace-configuration").put(
         Entity.json(updatedCfg));
     assertThat(response.getStatus()).isEqualTo(200);

thirdeye-persistence/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.TaskQuotasConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.TemplateConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.TimeConfigurationDTO;
 import com.google.inject.Inject;
@@ -84,6 +85,10 @@ private boolean updateDefaults(final NamespaceConfigurationDTO existingNamespace
       existingNamespaceConfig.setTemplateConfiguration(defaultTemplateConfiguration());
       updated = true;
     }
+    if (force || existingNamespaceConfig.getTaskQuotasConfiguration() == null) {
+      existingNamespaceConfig.setTaskQuotasConfiguration(defaultTaskQuotasConfiguration());
+      updated = true;
+    }
     return updated;
   }
 
@@ -151,6 +156,7 @@ private NamespaceConfigurationDTO defaultNamespaceConfiguration(final String nam
     namespaceConfigurationDTO
         .setTimeConfiguration(defaultTimeConfiguration())
         .setTemplateConfiguration(defaultTemplateConfiguration())
+        .setTaskQuotasConfiguration(defaultTaskQuotasConfiguration())
         .setAuth(new AuthorizationConfigurationDTO().setNamespace(namespace));
     return namespaceConfigurationDTO;
   }
@@ -208,4 +214,9 @@ private TemplateConfigurationDTO defaultTemplateConfiguration() {
     return optional(defaultNamespaceConfiguration.getTemplateConfiguration())
         .orElse(new TemplateConfigurationDTO());
   }
+
+  private TaskQuotasConfigurationDTO defaultTaskQuotasConfiguration() {
+    return optional(defaultNamespaceConfiguration.getTaskQuotasConfiguration())
+        .orElse(new TaskQuotasConfigurationDTO());
+  }
 }

thirdeye-persistence/src/test/java/ai/startree/thirdeye/datalayer/DatalayerTestUtils.java

@@ -23,6 +23,8 @@
 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.TaskQuotasConfigurationDTO;
+import ai.startree.thirdeye.spi.datalayer.dto.TemplateConfigurationDTO;
 import ai.startree.thirdeye.spi.datalayer.dto.TimeConfigurationDTO;
 import ai.startree.thirdeye.spi.metric.MetricType;
 import ai.startree.thirdeye.spi.util.SpiUtils;
@@ -66,12 +68,10 @@ public static DatasetConfigDTO getTestDatasetConfig(String collection) {
 
   public static NamespaceConfigurationDTO buildNamespaceConfiguration(String namespace) {
     final NamespaceConfigurationDTO dto = new NamespaceConfigurationDTO();
-    dto.setTimeConfiguration(new TimeConfigurationDTO()
-        .setDateTimePattern(Constants.NOTIFICATIONS_DEFAULT_DATE_PATTERN)
-        .setTimezone(DEFAULT_CHRONOLOGY.getZone())
-        .setMinimumOnboardingStartTime(946684800000L));
-    dto.setAuth(new AuthorizationConfigurationDTO()
-        .setNamespace(namespace));
+    dto.setTimeConfiguration(new TimeConfigurationDTO());
+    dto.setAuth(new AuthorizationConfigurationDTO().setNamespace(namespace));
+    dto.setTemplateConfiguration(new TemplateConfigurationDTO());
+    dto.setTaskQuotasConfiguration(new TaskQuotasConfigurationDTO());
     return dto;
   }
 }

thirdeye-persistence/src/test/java/ai/startree/thirdeye/datalayer/bao/TestNamespaceConfigurationManager.java

@@ -188,5 +188,11 @@ private void compareDtos(NamespaceConfigurationDTO dto1, NamespaceConfigurationD
         .isEqualTo(dto2.getTimeConfiguration().getTimezone());
     assertThat(dto1.getTimeConfiguration().getMinimumOnboardingStartTime())
         .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());
   }
 }

thirdeye-server/src/main/java/ai/startree/thirdeye/service/NamespaceConfigurationService.java

@@ -82,11 +82,15 @@ protected NamespaceConfigurationApi toApi(final NamespaceConfigurationDTO dto) {
   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.getAuth().getNamespace(), updated.getAuth().getNamespace()) ||
+        !Objects.equals(existing.getTaskQuotasConfiguration().getDetectionTaskQuota(),
+            updated.getTaskQuotasConfiguration().getDetectionTaskQuota()) ||
+        !Objects.equals(existing.getTaskQuotasConfiguration().getNotificationTaskQuota(),
+            updated.getTaskQuotasConfiguration().getNotificationTaskQuota())) {
       throw badRequest(
           ThirdEyeStatus.ERR_NAMESPACE_CONFIGURATION_VALIDATION_FAILED,
           existing.namespace(),
-          "Updating Id or auth is not allowed for Namespace Configuration");
+          "Updating Id, auth, or taskQuotasConfiguration is not allowed for Namespace Configuration");
     }
   }
 }

thirdeye-spi/src/main/java/ai/startree/thirdeye/spi/api/NamespaceConfigurationApi.java

@@ -22,6 +22,7 @@ public class NamespaceConfigurationApi implements ThirdEyeCrudApi<NamespaceConfi
 
   private TemplateConfigurationApi templateConfiguration;
 
+  private TaskQuotasConfigurationApi taskQuotasConfiguration;
 
   public TimeConfigurationApi getTimeConfiguration() {
     return timeConfiguration;
@@ -60,4 +61,14 @@ public NamespaceConfigurationApi setTemplateConfiguration(
     this.templateConfiguration = templateConfiguration;
     return this;
   }
+
+  public TaskQuotasConfigurationApi getTaskQuotasConfiguration() {
+    return taskQuotasConfiguration;
+  }
+
+  public NamespaceConfigurationApi setTaskQuotasConfiguration(
+      final TaskQuotasConfigurationApi taskQuotasConfiguration) {
+    this.taskQuotasConfiguration = taskQuotasConfiguration;
+    return this;
+  }
 }

thirdeye-spi/src/main/java/ai/startree/thirdeye/spi/api/TaskQuotasConfigurationApi.java

@@ -0,0 +1,38 @@
+/*
+ * 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;
+
+public class TaskQuotasConfigurationApi {
+
+  private Long DetectionTaskQuota;
+  private Long NotificationTaskQuota;
+
+  public Long getDetectionTaskQuota() {
+    return DetectionTaskQuota;
+  }
+
+  public TaskQuotasConfigurationApi setDetectionTaskQuota(final Long detectionTaskQuota) {
+    this.DetectionTaskQuota = detectionTaskQuota;
+    return this;
+  }
+
+  public Long getNotificationTaskQuota() {
+    return NotificationTaskQuota;
+  }
+
+  public TaskQuotasConfigurationApi setNotificationTaskQuota(final Long notificationTaskQuota) {
+    this.NotificationTaskQuota = notificationTaskQuota;
+    return this;
+  }
+}

thirdeye-spi/src/main/java/ai/startree/thirdeye/spi/datalayer/dto/NamespaceConfigurationDTO.java

@@ -23,7 +23,9 @@ public class NamespaceConfigurationDTO extends AbstractDTO {
 
   TimeConfigurationDTO timeConfiguration;
 
-  private TemplateConfigurationDTO templateConfiguration; 
+  private TemplateConfigurationDTO templateConfiguration;
+
+  private TaskQuotasConfigurationDTO taskQuotasConfiguration;
 
   public TimeConfigurationDTO getTimeConfiguration() {
     return timeConfiguration;
@@ -44,4 +46,14 @@ public NamespaceConfigurationDTO setTemplateConfiguration(
     this.templateConfiguration = templateConfiguration;
     return this;
   }
+
+  public TaskQuotasConfigurationDTO getTaskQuotasConfiguration() {
+    return taskQuotasConfiguration;
+  }
+
+  public NamespaceConfigurationDTO setTaskQuotasConfiguration(
+      final TaskQuotasConfigurationDTO taskQuotasConfiguration) {
+    this.taskQuotasConfiguration = taskQuotasConfiguration;
+    return this;
+  }
 }

thirdeye-spi/src/main/java/ai/startree/thirdeye/spi/datalayer/dto/TaskQuotasConfigurationDTO.java

@@ -0,0 +1,44 @@
+/*
+ * 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 TaskQuotasConfigurationDTO {
+
+  /**
+   * Monthly quota for number of detection task runs
+   */
+  private Long DetectionTaskQuota;
+  /**
+   * Monthly quota for number of notification task runs
+   */
+  private Long NotificationTaskQuota;
+
+  public Long getDetectionTaskQuota() {
+    return DetectionTaskQuota;
+  }
+
+  public TaskQuotasConfigurationDTO setDetectionTaskQuota(final Long detectionTaskQuota) {
+    this.DetectionTaskQuota = detectionTaskQuota;
+    return this;
+  }
+
+  public Long getNotificationTaskQuota() {
+    return NotificationTaskQuota;
+  }
+
+  public TaskQuotasConfigurationDTO setNotificationTaskQuota(final Long notificationTaskQuota) {
+    this.NotificationTaskQuota = notificationTaskQuota;
+    return this;
+  }
+}