From ad272631bb521fb3551192327437919419dae2da Mon Sep 17 00:00:00 2001
From: Tyler Ouyang <tylerwowen@users.noreply.github.com>
Date: Thu, 25 Jul 2024 14:45:25 -0700
Subject: [PATCH] Allow hosts added as existing capacity when main environment
 is not found (#1672)

Hots added as existing capacities may not have an owning env. For those hosts, allow them to be added.
---
 .../db/DBEnvironDAOImplTest.java              | 26 ++++++++++++++++++-
 .../teletraan/resource/EnvCapacities.java     | 15 ++++++-----
 .../teletraan/resource/EnvCapacitiesTest.java |  6 ++---
 3 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/deploy-service/common/src/test/java/com/pinterest/deployservice/db/DBEnvironDAOImplTest.java b/deploy-service/common/src/test/java/com/pinterest/deployservice/db/DBEnvironDAOImplTest.java
index c3405ce71a..fea034fc5d 100644
--- a/deploy-service/common/src/test/java/com/pinterest/deployservice/db/DBEnvironDAOImplTest.java
+++ b/deploy-service/common/src/test/java/com/pinterest/deployservice/db/DBEnvironDAOImplTest.java
@@ -129,12 +129,36 @@ void testGetMainEnvByHostName_noHostAgent() throws Exception {
         sut.insert(expectedEnvBean);
 
         setUpHostBean(TEST_CLUSTER);
-        setUpHostBean(TEST_CLUSTER + "2");
+        setUpHostBean(TEST_CLUSTER + "3");
 
         EnvironBean actualEnvironBean = sut.getMainEnvByHostName(HOST_NAME);
         assertEnvironBean(expectedEnvBean, actualEnvironBean);
     }
 
+    @Test
+    void testGetMainEnvByHostId_noMainEnv() throws Exception {
+        EnvironBean envBean = EnvironBeanFixture.createRandomEnvironBean();
+        envBean.setCluster_name(null);
+        sut.insert(envBean);
+
+        setUpHostBean(TEST_CLUSTER);
+
+        EnvironBean actualEnvironBean = sut.getMainEnvByHostId(HOST_ID);
+        assertNull(actualEnvironBean);
+    }
+
+    @Test
+    void testGetMainEnvByHostName_noMainEnv() throws Exception {
+        EnvironBean envBean = EnvironBeanFixture.createRandomEnvironBean();
+        envBean.setCluster_name(null);
+        sut.insert(envBean);
+
+        setUpHostBean(TEST_CLUSTER);
+
+        EnvironBean actualEnvironBean = sut.getMainEnvByHostId(HOST_NAME);
+        assertNull(actualEnvironBean);
+    }
+
     private void setUpHostAgentBean() throws Exception {
         HostAgentBean hostAgentBean = new HostAgentBean();
         hostAgentBean.setHost_id(HOST_ID);
diff --git a/deploy-service/teletraanservice/src/main/java/com/pinterest/teletraan/resource/EnvCapacities.java b/deploy-service/teletraanservice/src/main/java/com/pinterest/teletraan/resource/EnvCapacities.java
index 16b06c6f29..521e5b8452 100644
--- a/deploy-service/teletraanservice/src/main/java/com/pinterest/teletraan/resource/EnvCapacities.java
+++ b/deploy-service/teletraanservice/src/main/java/com/pinterest/teletraan/resource/EnvCapacities.java
@@ -231,8 +231,7 @@ void authorize(
             CapacityType capacityType,
             List<String> capacities)
             throws Exception {
-        if (targetEnvironBean.getSystem_priority() != null
-                && targetEnvironBean.getSystem_priority() > 0) {
+        if (isSidecarEnvironment(targetEnvironBean)) {
             // Allow sidecars to add capacity
             return;
         }
@@ -272,12 +271,16 @@ private HashSet<AuthZResource> getCapacityMainEnvironments(
                 throw new InternalServerErrorException(e);
             }
 
-            if (envBean == null) {
-                throw new ForbiddenException(
-                        "Failed to get the main environment with capacity name: " + capacity);
+            if (envBean != null) {
+                resources.add(new AuthZResource(envBean.getEnv_name(), envBean.getStage_name()));
+            } else {
+                LOG.info("Failed to find main environment for capacity {}, skip authorization", capacity);
             }
-            resources.add(new AuthZResource(envBean.getEnv_name(), envBean.getStage_name()));
         }
         return resources;
     }
+
+    private boolean isSidecarEnvironment(EnvironBean environBean) {
+        return environBean.getSystem_priority() != null && environBean.getSystem_priority() > 0;
+    }
 }
diff --git a/deploy-service/teletraanservice/src/test/java/com/pinterest/teletraan/resource/EnvCapacitiesTest.java b/deploy-service/teletraanservice/src/test/java/com/pinterest/teletraan/resource/EnvCapacitiesTest.java
index 7cd11fba6c..8181c706a9 100644
--- a/deploy-service/teletraanservice/src/test/java/com/pinterest/teletraan/resource/EnvCapacitiesTest.java
+++ b/deploy-service/teletraanservice/src/test/java/com/pinterest/teletraan/resource/EnvCapacitiesTest.java
@@ -118,16 +118,14 @@ void authorizeSuccess(CapacityType type) throws Exception {
 
     @ParameterizedTest
     @MethodSource("capacityTypes")
-    void authorizeThrowsExceptionWhenNoMainEnvFound(CapacityType type) throws Exception {
+    void authorizeSucceedsWhenNoMainEnvFound(CapacityType type) throws Exception {
         EnvironBean envBean = EnvironBeanFixture.createRandomEnvironBean();
 
         for (String capacity : capacities) {
             when(environDAO.getMainEnvByHostName(capacity)).thenReturn(null);
             when(environDAO.getByCluster(capacity)).thenReturn(null);
         }
-        assertThrows(
-                ForbiddenException.class,
-                () -> sut.authorize(envBean, principal, type, capacities));
+        assertDoesNotThrow(() -> sut.authorize(envBean, principal, type, capacities));
     }
 
     @ParameterizedTest