Introduce host extractor

commit-id:ef160c01

deploy-service/common/src/main/java/com/pinterest/deployservice/dao/HostAgentDAO.java

@@ -1,5 +1,5 @@
-/*
- * Copyright 2020 Pinterest, Inc.
+/**
+ * Copyright (c) 2020-2024 Pinterest, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,14 +15,12 @@
  */
 package com.pinterest.deployservice.dao;
 
+import com.pinterest.deployservice.bean.EnvironBean;
+import com.pinterest.deployservice.bean.HostAgentBean;
 import java.sql.SQLException;
 import java.util.List;
 
-import com.pinterest.deployservice.bean.HostAgentBean;
-
-/**
- * A collection of methods to help hosts and groups mapping
- */
+/** A collection of methods to help hosts and groups mapping */
 public interface HostAgentDAO {
     void insert(HostAgentBean hostAgentBean) throws Exception;
 
@@ -36,11 +34,24 @@ public interface HostAgentDAO {
 
     List<HostAgentBean> getStaleHosts(long lastUpdateBefore) throws SQLException;
 
-    List<HostAgentBean> getStaleHosts(long lastUpdateAfter, long lastUpdateBefore) throws SQLException;
+    List<HostAgentBean> getStaleHosts(long lastUpdateAfter, long lastUpdateBefore)
+            throws SQLException;
 
     List<HostAgentBean> getStaleEnvHosts(long lastUpdateBefore) throws Exception;
 
-    List<HostAgentBean> getHostsByAgent(String agentVersion, long pageIndex, int pageSize) throws Exception;
+    List<HostAgentBean> getHostsByAgent(String agentVersion, long pageIndex, int pageSize)
+            throws Exception;
 
     long getDistinctHostsCount() throws SQLException;
+
+    /**
+     * Retrieves the main environment ID for the specified host ID.
+     *
+     * <p>The main environment is where the cluster that the host belongs to is created.
+     *
+     * @param hostId The ID of the host.
+     * @return The bean represents the main environment for the specified host ID.
+     * @throws SQLException if an error occurs while retrieving the main environment ID.
+     */
+    EnvironBean getMainEnvByHostId(String hostId) throws SQLException;
 }

deploy-service/common/src/main/java/com/pinterest/deployservice/db/DBHostAgentDAOImpl.java

@@ -1,5 +1,5 @@
-/*
- * Copyright 2020 Pinterest, Inc.
+/**
+ * Copyright (c) 2020-2024 Pinterest, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,30 +15,40 @@
  */
 package com.pinterest.deployservice.db;
 
+import com.pinterest.deployservice.bean.EnvironBean;
+import com.pinterest.deployservice.bean.HostAgentBean;
+import com.pinterest.deployservice.bean.SetClause;
+import com.pinterest.deployservice.dao.HostAgentDAO;
 import java.sql.SQLException;
 import java.util.List;
-
 import org.apache.commons.dbcp.BasicDataSource;
 import org.apache.commons.dbutils.QueryRunner;
 import org.apache.commons.dbutils.ResultSetHandler;
 import org.apache.commons.dbutils.handlers.BeanHandler;
 import org.apache.commons.dbutils.handlers.BeanListHandler;
 
-import com.pinterest.deployservice.bean.HostAgentBean;
-import com.pinterest.deployservice.bean.SetClause;
-import com.pinterest.deployservice.dao.HostAgentDAO;
-
 public class DBHostAgentDAOImpl implements HostAgentDAO {
-    private static final String INSERT_HOST_TEMPLATE = "INSERT INTO hosts_and_agents SET %s ON DUPLICATE KEY UPDATE %s";
-    private static final String UPDATE_HOST_BY_ID = "UPDATE hosts_and_agents SET %s WHERE host_id=?";
+    private static final String INSERT_HOST_TEMPLATE =
+            "INSERT INTO hosts_and_agents SET %s ON DUPLICATE KEY UPDATE %s";
+    private static final String UPDATE_HOST_BY_ID =
+            "UPDATE hosts_and_agents SET %s WHERE host_id=?";
     private static final String DELETE_HOST_BY_ID = "DELETE FROM hosts_and_agents WHERE host_id=?";
-    private static final String GET_HOST_BY_NAME = "SELECT * FROM hosts_and_agents WHERE host_name=?";
-    private static final String GET_HOST_BY_HOSTID = "SELECT * FROM hosts_and_agents WHERE host_id=?";
-    private static final String GET_HOSTS_BY_LAST_UPDATE = "SELECT DISTINCT * FROM hosts_and_agents WHERE last_update<?";
-    private static final String GET_HOSTS_BY_LAST_UPDATES = "SELECT DISTINCT * FROM hosts_and_agents WHERE last_update>? AND last_update<?";
-    private static final String GET_STALE_ENV_HOST = "SELECT DISTINCT hosts_and_agents.* FROM hosts_and_agents INNER JOIN hosts_and_envs ON hosts_and_agents.host_name=hosts_and_envs.host_name WHERE hosts_and_agents.last_update<?";
-    private static final String GET_HOSTS_BY_AGENT = "SELECT * FROM hosts_statuses WHERE agent_version=? ORDER BY host_id LIMIT ?,?";
-    private static final String GET_DISTINCT_HOSTS_COUNT = "SELECT COUNT(DISTINCT host_id) FROM hosts_and_agents";
+    private static final String GET_HOST_BY_NAME =
+            "SELECT * FROM hosts_and_agents WHERE host_name=?";
+    private static final String GET_HOST_BY_HOSTID =
+            "SELECT * FROM hosts_and_agents WHERE host_id=?";
+    private static final String GET_HOSTS_BY_LAST_UPDATE =
+            "SELECT DISTINCT * FROM hosts_and_agents WHERE last_update<?";
+    private static final String GET_HOSTS_BY_LAST_UPDATES =
+            "SELECT DISTINCT * FROM hosts_and_agents WHERE last_update>? AND last_update<?";
+    private static final String GET_STALE_ENV_HOST =
+            "SELECT DISTINCT hosts_and_agents.* FROM hosts_and_agents INNER JOIN hosts_and_envs ON hosts_and_agents.host_name=hosts_and_envs.host_name WHERE hosts_and_agents.last_update<?";
+    private static final String GET_HOSTS_BY_AGENT =
+            "SELECT * FROM hosts_statuses WHERE agent_version=? ORDER BY host_id LIMIT ?,?";
+    private static final String GET_DISTINCT_HOSTS_COUNT =
+            "SELECT COUNT(DISTINCT host_id) FROM hosts_and_agents";
+    private static final String GET_MAIN_ENV_BY_HOSTID =
+            "SELECT e.* FROM hosts_and_agents ha JOIN environs e ON ha.auto_scaling_group = e.cluster_name WHERE ha.host_id = ?";
 
     private BasicDataSource dataSource;
 
@@ -49,7 +59,9 @@ public DBHostAgentDAOImpl(BasicDataSource dataSource) {
     @Override
     public void insert(HostAgentBean hostAgentBean) throws Exception {
         SetClause setClause = hostAgentBean.genSetClause();
-        String clause = String.format(INSERT_HOST_TEMPLATE, setClause.getClause(), HostAgentBean.UPDATE_CLAUSE);
+        String clause =
+                String.format(
+                        INSERT_HOST_TEMPLATE, setClause.getClause(), HostAgentBean.UPDATE_CLAUSE);
         new QueryRunner(dataSource).update(clause, setClause.getValueArray());
     }
 
@@ -60,6 +72,7 @@ public void update(String id, HostAgentBean bean) throws Exception {
         setClause.addValue(id);
         new QueryRunner(dataSource).update(clause, setClause.getValueArray());
     }
+
     @Override
     public void delete(String id) throws Exception {
         new QueryRunner(dataSource).update(DELETE_HOST_BY_ID, id);
@@ -84,9 +97,11 @@ public List<HostAgentBean> getStaleHosts(long lastUpdateBefore) throws SQLExcept
     }
 
     @Override
-    public List<HostAgentBean> getStaleHosts(long lastUpdateAfter, long lastUpdateBefore) throws SQLException {
+    public List<HostAgentBean> getStaleHosts(long lastUpdateAfter, long lastUpdateBefore)
+            throws SQLException {
         ResultSetHandler<List<HostAgentBean>> h = new BeanListHandler<>(HostAgentBean.class);
-        return new QueryRunner(dataSource).query(GET_HOSTS_BY_LAST_UPDATES, h, lastUpdateAfter, lastUpdateBefore);
+        return new QueryRunner(dataSource)
+                .query(GET_HOSTS_BY_LAST_UPDATES, h, lastUpdateAfter, lastUpdateBefore);
     }
 
     @Override
@@ -96,15 +111,26 @@ public List<HostAgentBean> getStaleEnvHosts(long after) throws Exception {
     }
 
     @Override
-    public List<HostAgentBean> getHostsByAgent(String agentVersion, long pageIndex, int pageSize) throws Exception {
+    public List<HostAgentBean> getHostsByAgent(String agentVersion, long pageIndex, int pageSize)
+            throws Exception {
         ResultSetHandler<List<HostAgentBean>> h = new BeanListHandler<>(HostAgentBean.class);
-        return new QueryRunner(dataSource).query(GET_HOSTS_BY_AGENT, h, agentVersion, (pageIndex - 1) * pageSize, pageSize);
+        return new QueryRunner(dataSource)
+                .query(GET_HOSTS_BY_AGENT, h, agentVersion, (pageIndex - 1) * pageSize, pageSize);
     }
 
     @Override
     public long getDistinctHostsCount() throws SQLException {
-        Long n = new QueryRunner(dataSource).query(GET_DISTINCT_HOSTS_COUNT,
-                SingleResultSetHandlerFactory.<Long>newObjectHandler());
+        Long n =
+                new QueryRunner(dataSource)
+                        .query(
+                                GET_DISTINCT_HOSTS_COUNT,
+                                SingleResultSetHandlerFactory.<Long>newObjectHandler());
         return n == null ? 0 : n;
     }
+
+    @Override
+    public EnvironBean getMainEnvByHostId(String hostId) throws SQLException {
+        ResultSetHandler<EnvironBean> h = new BeanHandler<>(EnvironBean.class);
+        return new QueryRunner(dataSource).query(GET_MAIN_ENV_BY_HOSTID, h, hostId);
+    }
 }

deploy-service/common/src/test/java/com/pinterest/deployservice/db/DBHostAgentDAOImplTest.java

@@ -0,0 +1,89 @@
+/**
+ * Copyright (c) 2024 Pinterest, Inc.
+ *
+ * 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 com.pinterest.deployservice.db;
+
+import static org.junit.Assert.assertNull;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import com.pinterest.deployservice.bean.EnvironBean;
+import com.pinterest.deployservice.bean.HostAgentBean;
+import com.pinterest.deployservice.dao.EnvironDAO;
+import com.pinterest.deployservice.dao.HostAgentDAO;
+import com.pinterest.deployservice.fixture.EnvironBeanFixture;
+import org.apache.commons.dbcp.BasicDataSource;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class DBHostAgentDAOImplTest {
+    private static final String HOST_ID = "host123";
+    private static final String TEST_CLUSTER = "test-cluster";
+    private static BasicDataSource dataSource;
+    private HostAgentDAO sut;
+
+    @BeforeAll
+    static void setUpAll() throws Exception {
+        dataSource = DBUtils.createTestDataSource();
+    }
+
+    @BeforeEach
+    void setUp() {
+        sut = new DBHostAgentDAOImpl(dataSource);
+    }
+
+    @AfterEach
+    void tearDown() throws Exception {
+        DBUtils.truncateAllTables(dataSource);
+    }
+
+    @Test
+    void testGetMainEnvByHostId_happyPath() throws Exception {
+        EnvironDAO environDAO = new DBEnvironDAOImpl(dataSource);
+        EnvironBean expectedEnvBean = EnvironBeanFixture.createRandomEnvironBean();
+        expectedEnvBean.setCluster_name(TEST_CLUSTER);
+        environDAO.insert(expectedEnvBean);
+
+        HostAgentBean hostAgentBean = new HostAgentBean();
+        hostAgentBean.setHost_id(HOST_ID);
+        hostAgentBean.setAuto_scaling_group(TEST_CLUSTER);
+        sut.insert(hostAgentBean);
+
+        EnvironBean actualEnvironBean = sut.getMainEnvByHostId(HOST_ID);
+        assertEquals(expectedEnvBean.getEnv_name(), actualEnvironBean.getEnv_name());
+        assertEquals(expectedEnvBean.getStage_name(), actualEnvironBean.getStage_name());
+
+        EnvironBean nullEnvironBean = sut.getMainEnvByHostId("random-host-id");
+        assertNull(nullEnvironBean);
+    }
+
+    @Test
+    void testGetMainEnvByHostId_noHost() throws Exception {
+        EnvironBean actualEnvironBean = sut.getMainEnvByHostId(HOST_ID);
+        assertNull(actualEnvironBean);
+    }
+
+    @Test
+    void testGetMainEnvByHostId_noEnv() throws Exception {
+        HostAgentBean hostAgentBean = new HostAgentBean();
+        hostAgentBean.setHost_id(HOST_ID);
+        hostAgentBean.setAuto_scaling_group(TEST_CLUSTER);
+        sut.insert(hostAgentBean);
+
+        EnvironBean actualEnvironBean = sut.getMainEnvByHostId(HOST_ID);
+        assertNull(actualEnvironBean);
+    }
+}