/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.ambari.server.checks;

import java.io.File;
import java.io.IOException;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Scanner;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import javax.annotation.Nullable;

import jakarta.inject.Provider;
import jakarta.persistence.EntityManager;
import jakarta.persistence.Query;
import jakarta.persistence.TypedQuery;

import org.apache.ambari.server.AmbariException;
import org.apache.ambari.server.api.services.AmbariMetaInfo;
import org.apache.ambari.server.configuration.Configuration;
import org.apache.ambari.server.orm.DBAccessor;
import org.apache.ambari.server.orm.dao.AlertDefinitionDAO;
import org.apache.ambari.server.orm.dao.ClusterDAO;
import org.apache.ambari.server.orm.dao.ExecutionCommandDAO;
import org.apache.ambari.server.orm.dao.HostComponentDesiredStateDAO;
import org.apache.ambari.server.orm.dao.HostComponentStateDAO;
import org.apache.ambari.server.orm.dao.HostRoleCommandDAO;
import org.apache.ambari.server.orm.dao.MetainfoDAO;
import org.apache.ambari.server.orm.dao.StageDAO;
import org.apache.ambari.server.orm.entities.AlertDefinitionEntity;
import org.apache.ambari.server.orm.entities.ClusterConfigEntity;
import org.apache.ambari.server.orm.entities.HostComponentDesiredStateEntity;
import org.apache.ambari.server.orm.entities.HostComponentStateEntity;
import org.apache.ambari.server.orm.entities.MetainfoEntity;
import org.apache.ambari.server.state.ClientConfigFileDefinition;
import org.apache.ambari.server.state.Cluster;
import org.apache.ambari.server.state.Clusters;
import org.apache.ambari.server.state.ComponentInfo;
import org.apache.ambari.server.state.Host;
import org.apache.ambari.server.state.Service;
import org.apache.ambari.server.state.ServiceInfo;
import org.apache.ambari.server.state.State;
import org.apache.ambari.server.state.UpgradeState;
import org.apache.ambari.server.state.configgroup.ConfigGroup;
import org.apache.ambari.server.utils.VersionUtils;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Splitter;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Multimap;
import com.google.inject.Inject;
import com.google.inject.Injector;
import com.google.inject.persist.Transactional;

public class DatabaseConsistencyCheckHelper {

  private static final Logger LOG = LoggerFactory.getLogger(DatabaseConsistencyCheckHelper.class);

  @Inject
  private static Injector injector;

  private static MetainfoDAO metainfoDAO;
  private static AlertDefinitionDAO alertDefinitionDAO;
  private static Connection connection;
  private static AmbariMetaInfo ambariMetaInfo;
  private static DBAccessor dbAccessor;
  private static HostRoleCommandDAO hostRoleCommandDAO;
  private static ExecutionCommandDAO executionCommandDAO;
  private static StageDAO stageDAO;

  private static DatabaseConsistencyCheckResult checkResult = DatabaseConsistencyCheckResult.DB_CHECK_SUCCESS;
  public static final String GET_CONFIGS_SELECTED_MORE_THAN_ONCE_QUERY = "select c.cluster_name, cc.type_name from clusterconfig cc " +
      "join clusters c on cc.cluster_id=c.cluster_id " +
      "group by c.cluster_name, cc.type_name " +
      "having sum(cc.selected) > 1";

  /**
   * @return The result of the DB cheks run so far.
   */
  public static DatabaseConsistencyCheckResult getLastCheckResult() {
    return checkResult;
  }

  /**
   * Reset check result to {@link DatabaseConsistencyCheckResult#DB_CHECK_SUCCESS}.
   */
  public static void resetCheckResult() {
    checkResult = DatabaseConsistencyCheckResult.DB_CHECK_SUCCESS;
  }

  /**
   * Called internally to set the result of the DB checks. The new result is only recorded if it is more severe than
   * the existing result.
   *
    * @param newResult the new result
   */
  private static void setCheckResult(DatabaseConsistencyCheckResult newResult) {
    if (newResult.ordinal() > checkResult.ordinal()) {
      checkResult = newResult;
    }
  }

  /**
   * Called to indicate a warning during checks
   *
   * @param messageTemplate Message template (log4j format)
   * @param messageParams Message params
   */
  private static void warning(String messageTemplate, Object... messageParams) {
    LOG.warn(messageTemplate, messageParams);
    setCheckResult(DatabaseConsistencyCheckResult.DB_CHECK_WARNING);
  }

  /**
   * Called to indicate an error during checks
   *
   * @param messageTemplate Message template (log4j format)
   * @param messageParams Message params
   */
  private static void error(String messageTemplate, Object... messageParams) {
    LOG.error(messageTemplate, messageParams);
    setCheckResult(DatabaseConsistencyCheckResult.DB_CHECK_ERROR);
  }

  public static void setInjector(Injector injector) {
    DatabaseConsistencyCheckHelper.injector = injector;
    // Clean up: new injector means static fields should be reinitalized, though in real life it only occurs during testing
    closeConnection();
    connection = null;
    metainfoDAO = null;
    alertDefinitionDAO = null;
    ambariMetaInfo = null;
    dbAccessor = null;
  }

  public static void setConnection(Connection connection) {
    DatabaseConsistencyCheckHelper.connection = connection;
  }

  /*
    * method to close connection
    * */
  public static void closeConnection() {
    try {
      if (connection != null) {
        connection.close();
      }
    } catch (SQLException e) {
      LOG.error("Exception occurred during connection close procedure: ", e);
    }
  }

  public static DatabaseConsistencyCheckResult runAllDBChecks(boolean fixIssues) throws Throwable {
    LOG.info("******************************* Check database started *******************************");
    try {
      if (fixIssues) {
        fixHostComponentStatesCountEqualsHostComponentsDesiredStates();
        fixClusterConfigsNotMappedToAnyService();
        fixConfigGroupServiceNames();
        fixConfigGroupHostMappings();
        fixConfigGroupsForDeletedServices();
        fixConfigsSelectedMoreThanOnce();
        fixAlertsForDeletedServices();
      }
      checkSchemaName();
      checkMySQLEngine();
      checkForConfigsNotMappedToService();
      checkForConfigsSelectedMoreThanOnce();
      checkForHostsWithoutState();
      checkHostComponentStates();
      checkServiceConfigs();
      checkForLargeTables();
      checkConfigGroupsHasServiceName();
      checkConfigGroupHostMapping(true);
      checkConfigGroupsForDeletedServices(true);
      checkForStalealertdefs();
      LOG.info("******************************* Check database completed *******************************");
      return checkResult;
    }
    catch (Throwable ex) {
      LOG.error("An error occurred during database consistency check.", ex);
      throw ex;
    }
  }

  public static void checkDBVersionCompatible() throws AmbariException {
    LOG.info("Checking DB store version");

    if (metainfoDAO == null) {
      metainfoDAO = injector.getInstance(MetainfoDAO.class);
    }
    if (alertDefinitionDAO == null) {
      alertDefinitionDAO = injector.getInstance(AlertDefinitionDAO.class);
    }
    MetainfoEntity schemaVersionEntity = metainfoDAO.findByKey(Configuration.SERVER_VERSION_KEY);
    String schemaVersion = null;

    if (schemaVersionEntity != null) {
      schemaVersion = schemaVersionEntity.getMetainfoValue();
    }

    Configuration conf = injector.getInstance(Configuration.class);
    File versionFile = new File(conf.getServerVersionFilePath());
    if (!versionFile.exists()) {
      throw new AmbariException("Server version file does not exist.");
    }
    String serverVersion = null;
    try (Scanner scanner = new Scanner(versionFile)) {
      serverVersion = scanner.useDelimiter("\\Z").next();

    } catch (IOException ioe) {
      throw new AmbariException("Unable to read server version file.");
    }

    if (schemaVersionEntity==null || VersionUtils.compareVersions(schemaVersion, serverVersion, 3) != 0) {
      String error = "Current database store version is not compatible with " +
              "current server version"
              + ", serverVersion=" + serverVersion
              + ", schemaVersion=" + schemaVersion;
      LOG.error(error);
      throw new AmbariException(error);
    }

    LOG.info("DB store version is compatible");
  }

  /**
   * This method checks if ambari database has tables with too big size (according to limit).
   * First of all we are trying to get table size from schema information, but if it's not possible,
   * we will get tables rows count and compare it with row count limit.
   */
  static void checkForLargeTables() {
    LOG.info("Checking for tables with large physical size");
    ensureConnection();

    if (dbAccessor.getDbType() == DBAccessor.DbType.H2) {
      return;
    }

    DBAccessor.DbType dbType = dbAccessor.getDbType();
    String schemaName = dbAccessor.getDbSchema();

    String GET_TABLE_SIZE_IN_BYTES_POSTGRESQL = "SELECT pg_total_relation_size('%s') \"Table Size\"";
    String GET_TABLE_SIZE_IN_BYTES_MYSQL = "SELECT (data_length + index_length) \"Table Size\" FROM information_schema.TABLES WHERE table_schema = \"" + schemaName + "\" AND table_name =\"%s\"";
    String GET_TABLE_SIZE_IN_BYTES_ORACLE = "SELECT bytes \"Table Size\" FROM user_segments WHERE segment_type='TABLE' AND segment_name='%s'";
    String GET_ROW_COUNT_QUERY = "SELECT COUNT(*) FROM %s";

    Map<DBAccessor.DbType, String> tableSizeQueryMap = new HashMap<>();
    tableSizeQueryMap.put(DBAccessor.DbType.POSTGRES, GET_TABLE_SIZE_IN_BYTES_POSTGRESQL);
    tableSizeQueryMap.put(DBAccessor.DbType.MYSQL, GET_TABLE_SIZE_IN_BYTES_MYSQL);
    tableSizeQueryMap.put(DBAccessor.DbType.ORACLE, GET_TABLE_SIZE_IN_BYTES_ORACLE);

    List<String> tablesToCheck = Arrays.asList("host_role_command", "execution_command", "stage", "request", "alert_history");

    final double TABLE_SIZE_LIMIT_MB = 3000.0;
    final int TABLE_ROW_COUNT_LIMIT = 3000000;

    String findTableSizeQuery = tableSizeQueryMap.get(dbType);

    if (dbType == DBAccessor.DbType.ORACLE) {
      for (int i = 0;i < tablesToCheck.size(); i++) {
        tablesToCheck.set(i, tablesToCheck.get(i).toUpperCase());
      }
    }

    for (String tableName : tablesToCheck) {

      ResultSet rs = null;
      Statement statement = null;
      Double tableSizeInMB = null;
      Long tableSizeInBytes = null;
      int tableRowCount = -1;

      try {
        statement = connection.createStatement(ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_UPDATABLE);
        rs = statement.executeQuery(String.format(findTableSizeQuery, tableName));
        if (rs != null) {
          while (rs.next()) {
            tableSizeInBytes = rs.getLong(1);
            if (tableSizeInBytes != null) {
              tableSizeInMB = tableSizeInBytes / 1024.0 / 1024.0;
            }
          }
        }

        if (tableSizeInMB != null && tableSizeInMB > TABLE_SIZE_LIMIT_MB) {
          warning("The database table {} is currently {} MB (limit is {}) and may impact performance. It is recommended " +
                  "that you reduce its size by executing \"ambari-server db-purge-history\".",
                  tableName, tableSizeInMB, TABLE_SIZE_LIMIT_MB);
        } else if (tableSizeInMB != null && tableSizeInMB < TABLE_SIZE_LIMIT_MB) {
          LOG.info(String.format("The database table %s is currently %.3f MB and is within normal limits (%.3f)",
                  tableName, tableSizeInMB, TABLE_SIZE_LIMIT_MB));
        } else {
          throw new Exception();
        }
      } catch (Exception e) {
        LOG.error(String.format("Failed to get %s table size from database, will check row count: ", tableName), e);
        try {
          rs = statement.executeQuery(String.format(GET_ROW_COUNT_QUERY, tableName));
          if (rs != null) {
            while (rs.next()) {
              tableRowCount = rs.getInt(1);
            }
          }

          if (tableRowCount > TABLE_ROW_COUNT_LIMIT) {
            warning("The database table {} currently has {} rows (limit is {}) and may impact performance. It is " +
                    "recommended that you reduce its size by executing \"ambari-server db-purge-history\".",
                    tableName, tableRowCount, TABLE_ROW_COUNT_LIMIT);
          } else if (tableRowCount != -1 && tableRowCount < TABLE_ROW_COUNT_LIMIT) {
            LOG.info(String.format("The database table %s currently has %d rows and is within normal limits (%d)", tableName, tableRowCount, TABLE_ROW_COUNT_LIMIT));
          } else {
            warning("Unable to get size for table {}!", tableName);
          }
        } catch (SQLException ex) {
          warning(String.format("Failed to get %s row count: ", tableName), e);
        }
      } finally {
        if (rs != null) {
          try {
            rs.close();
          } catch (SQLException e) {
            LOG.error("Exception occurred during result set closing procedure: ", e);
          }
        }

        if (statement != null) {
          try {
            statement.close();
          } catch (SQLException e) {
            LOG.error("Exception occurred during statement closing procedure: ", e);
          }
        }
      }
    }

  }

  /**
   * This method checks if any config type in clusterconfig table, has more than
   * one versions selected. If config version is selected(in selected column =
   * 1), it means that this version of config is actual. So, if any config type
   * has more than one selected version it's a bug and we are showing error
   * message for user.
   */
  static void checkForConfigsSelectedMoreThanOnce() {
    LOG.info("Checking for more than 1 configuration of the same type being enabled.");

    String GET_CONFIGS_SELECTED_MORE_THAN_ONCE_QUERY = "select c.cluster_name, cc.type_name from clusterconfig cc "
        + "join clusters c on cc.cluster_id=c.cluster_id "
        + "group by c.cluster_name, cc.type_name " +
            "having sum(selected) > 1";

    Multimap<String, String> clusterConfigTypeMap = HashMultimap.create();
    ResultSet rs = null;
    Statement statement = null;

    ensureConnection();

    try {
      statement = connection.createStatement(ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_UPDATABLE);
      rs = statement.executeQuery(GET_CONFIGS_SELECTED_MORE_THAN_ONCE_QUERY);
      if (rs != null) {
        while (rs.next()) {
          clusterConfigTypeMap.put(rs.getString("cluster_name"), rs.getString("type_name"));
        }

        for (String clusterName : clusterConfigTypeMap.keySet()) {
          error("You have config(s), in cluster {}, that is(are) selected more than once in clusterconfig table: {}",
                  clusterName ,StringUtils.join(clusterConfigTypeMap.get(clusterName), ","));
        }
      }

    } catch (SQLException e) {
      warning("Exception occurred during check for config selected more than once procedure: ", e);
    } finally {
      if (rs != null) {
        try {
          rs.close();
        } catch (SQLException e) {
          LOG.error("Exception occurred during result set closing procedure: ", e);
        }
      }

      if (statement != null) {
        try {
          statement.close();
        } catch (SQLException e) {
          LOG.error("Exception occurred during statement closing procedure: ", e);
        }
      }
    }
  }

  /**
  * This method checks if all hosts from hosts table
  * has related host state info in hoststate table.
  * If not then we are showing error.
  * */
  static void checkForHostsWithoutState() {
    LOG.info("Checking for hosts without state");

    String GET_HOSTS_WITHOUT_STATUS_QUERY = "select host_name from hosts where host_id not in (select host_id from hoststate)";
    Set<String> hostsWithoutStatus = new HashSet<>();
    ResultSet rs = null;
    Statement statement = null;

    ensureConnection();

    try {
      statement = connection.createStatement(ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_UPDATABLE);
      rs = statement.executeQuery(GET_HOSTS_WITHOUT_STATUS_QUERY);
      if (rs != null) {
        while (rs.next()) {
          hostsWithoutStatus.add(rs.getString("host_name"));
        }

        if (!hostsWithoutStatus.isEmpty()) {
          warning("You have host(s) without state (in hoststate table): " + StringUtils.join(hostsWithoutStatus, ","));
        }
      }

    } catch (SQLException e) {
      warning("Exception occurred during check for host without state procedure: ", e);
    } finally {
      if (rs != null) {
        try {
          rs.close();
        } catch (SQLException e) {
          LOG.error("Exception occurred during result set closing procedure: ", e);
        }
      }

      if (statement != null) {
        try {
          statement.close();
        } catch (SQLException e) {
          LOG.error("Exception occurred during statement closing procedure: ", e);
        }
      }
    }
  }

  private static int runQuery(Statement statement, String query) {
    ResultSet rs = null;
    int result = 0;
    try {
      rs = statement.executeQuery(query);

      if (rs != null) {
        while (rs.next()) {
          result = rs.getInt(1);
        }
      }

    } catch (SQLException e) {
      warning("Exception occurred during topology request tables check: ", e);
    } finally {
      if (rs != null) {
        try {
          rs.close();
        } catch (SQLException e) {
          LOG.error("Exception occurred during result set closing procedure: ", e);
        }
      }
    }
    return result;
  }


  /**
  * This method checks if count of host component states equals count
  * of desired host component states. According to ambari logic these
  * two tables should have the same count of rows. If not then we are
  * showing error for user.
  *
  * One more, we are checking if all components has only one host
  * component state. If some component has more, it can cause issues
  * */
  static void checkHostComponentStates() {
    LOG.info("Checking host component states count equals host component desired states count");

    String GET_HOST_COMPONENT_STATE_COUNT_QUERY = "select count(*) from hostcomponentstate";
    String GET_HOST_COMPONENT_DESIRED_STATE_COUNT_QUERY = "select count(*) from hostcomponentdesiredstate";
    String GET_MERGED_TABLE_ROW_COUNT_QUERY = "select count(*) FROM hostcomponentstate hcs " +
            "JOIN hostcomponentdesiredstate hcds ON hcs.service_name=hcds.service_name AND hcs.component_name=hcds.component_name AND hcs.host_id=hcds.host_id";
    String GET_HOST_COMPONENT_STATE_DUPLICATES_QUERY = "select component_name, host_id from hostcomponentstate group by component_name, host_id having count(component_name) > 1";
    int hostComponentStateCount = 0;
    int hostComponentDesiredStateCount = 0;
    Map<String, String> hostComponentStateDuplicates = new HashMap<>();
    int mergedCount = 0;
    ResultSet rs = null;
    Statement statement = null;

    ensureConnection();

    try {
      statement = connection.createStatement(ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_UPDATABLE);

      rs = statement.executeQuery(GET_HOST_COMPONENT_STATE_COUNT_QUERY);
      if (rs != null) {
        while (rs.next()) {
          hostComponentStateCount = rs.getInt(1);
        }
      }

      rs = statement.executeQuery(GET_HOST_COMPONENT_DESIRED_STATE_COUNT_QUERY);
      if (rs != null) {
        while (rs.next()) {
          hostComponentDesiredStateCount = rs.getInt(1);
        }
      }

      rs = statement.executeQuery(GET_MERGED_TABLE_ROW_COUNT_QUERY);
      if (rs != null) {
        while (rs.next()) {
          mergedCount = rs.getInt(1);
        }
      }

      if (hostComponentStateCount != hostComponentDesiredStateCount || hostComponentStateCount != mergedCount) {
        warning("Your host component states (hostcomponentstate table) count not equals host component desired states (hostcomponentdesiredstate table) count!");
      }


      rs = statement.executeQuery(GET_HOST_COMPONENT_STATE_DUPLICATES_QUERY);
      if (rs != null) {
        while (rs.next()) {
          hostComponentStateDuplicates.put(rs.getString("component_name"), rs.getString("host_id"));
        }
      }

      for (Map.Entry<String, String> component : hostComponentStateDuplicates.entrySet()) {
        warning("Component {} on host with id {}, has more than one host component state (hostcomponentstate table)!", component.getKey(), component.getValue());
      }

    } catch (SQLException e) {
      warning("Exception occurred during check for same count of host component states and host component desired states: ", e);
    } finally {
      if (rs != null) {
        try {
          rs.close();
        } catch (SQLException e) {
          LOG.error("Exception occurred during result set closing procedure: ", e);
        }
      }

      if (statement != null) {
        try {
          statement.close();
        } catch (SQLException e) {
          LOG.error("Exception occurred during statement closing procedure: ", e);
        }
      }
    }

  }

  /**
  * Remove configs that are not mapped to any service.
  */
  @Transactional
  static void fixClusterConfigsNotMappedToAnyService() {
    LOG.info("Checking for configs not mapped to any Service");
    ClusterDAO clusterDAO = injector.getInstance(ClusterDAO.class);
    List<ClusterConfigEntity> notMappedClusterConfigs = getNotMappedClusterConfigsToService();

    for (ClusterConfigEntity clusterConfigEntity : notMappedClusterConfigs){
      if (clusterConfigEntity.isUnmapped()){
        continue; // skip clusterConfigs that did not leave after service deletion
      }
      List<String> types = new ArrayList<>();
      String type = clusterConfigEntity.getType();
      types.add(type);
      LOG.error("Removing config that is not mapped to any service", clusterConfigEntity);
      clusterDAO.removeConfig(clusterConfigEntity);
    }
  }


  /**
   * Find ClusterConfigs that are not mapped to Service
   * @return ClusterConfigs that are not mapped to Service
   */
  private static List<ClusterConfigEntity> getNotMappedClusterConfigsToService() {
    Provider<EntityManager> entityManagerProvider = injector.getProvider(EntityManager.class);
    EntityManager entityManager = entityManagerProvider.get();

    String queryName = "ClusterConfigEntity.findNotMappedClusterConfigsToService";
    TypedQuery<ClusterConfigEntity> query = entityManager.createNamedQuery(queryName, ClusterConfigEntity.class);

    return query.getResultList();
  }

  /**
   * Look for configs that are not mapped to any service.
   */
  static void checkForConfigsNotMappedToService() {
    LOG.info("Checking for configs that are not mapped to any service");
    List<ClusterConfigEntity> notMappedClasterConfigs = getNotMappedClusterConfigsToService();

    Set<String> nonMappedConfigs = new HashSet<>();
    for (ClusterConfigEntity clusterConfigEntity : notMappedClasterConfigs) {
      if (!clusterConfigEntity.isUnmapped()){ //this check does not report warning for configs from deleted services
        nonMappedConfigs.add(clusterConfigEntity.getType() + '-' + clusterConfigEntity.getTag());
      }
    }
    if (!nonMappedConfigs.isEmpty()){
      warning("You have config(s): {} that is(are) not mapped (in serviceconfigmapping table) to any service!", StringUtils.join(nonMappedConfigs, ","));
    }
  }

  /**
  * This method checks if count of host component states equals count
  * of desired host component states. According to ambari logic these
  * two tables should have the same count of rows. If not then we are
  * adding missed host components.
  * hard to predict root couse of inconsistency, so it can be dangerous
  */
  @Transactional
  static void fixHostComponentStatesCountEqualsHostComponentsDesiredStates() {
    LOG.info("Checking that there are the same number of actual and desired host components");

    HostComponentStateDAO hostComponentStateDAO = injector.getInstance(HostComponentStateDAO.class);
    HostComponentDesiredStateDAO hostComponentDesiredStateDAO = injector.getInstance(HostComponentDesiredStateDAO.class);

    List<HostComponentDesiredStateEntity> hostComponentDesiredStates = hostComponentDesiredStateDAO.findAll();
    List<HostComponentStateEntity> hostComponentStates = hostComponentStateDAO.findAll();

    Set<HostComponentDesiredStateEntity> missedHostComponentDesiredStates = new HashSet<>();
    missedHostComponentDesiredStates.addAll(hostComponentDesiredStates);
    Set<HostComponentStateEntity> missedHostComponentStates = new HashSet<>();
    missedHostComponentStates.addAll(hostComponentStates);

    for (Iterator<HostComponentStateEntity> stateIterator = missedHostComponentStates.iterator(); stateIterator.hasNext();){
      HostComponentStateEntity hostComponentStateEntity = stateIterator.next();
      for (Iterator<HostComponentDesiredStateEntity> desiredStateIterator = missedHostComponentDesiredStates.iterator(); desiredStateIterator.hasNext();) {
        HostComponentDesiredStateEntity hostComponentDesiredStateEntity = desiredStateIterator.next();
        if (hostComponentStateEntity.getComponentName().equals(hostComponentDesiredStateEntity.getComponentName()) &&
            hostComponentStateEntity.getServiceName().equals(hostComponentDesiredStateEntity.getServiceName()) &&
            hostComponentStateEntity.getHostId().equals(hostComponentDesiredStateEntity.getHostId())){
          desiredStateIterator.remove();
          stateIterator.remove();
          break;
        }
      }
    }

    for (HostComponentDesiredStateEntity hostComponentDesiredStateEntity : missedHostComponentDesiredStates) {
      HostComponentStateEntity stateEntity = new HostComponentStateEntity();
      stateEntity.setClusterId(hostComponentDesiredStateEntity.getClusterId());
      stateEntity.setComponentName(hostComponentDesiredStateEntity.getComponentName());
      stateEntity.setServiceName(hostComponentDesiredStateEntity.getServiceName());
      stateEntity.setVersion(State.UNKNOWN.toString());
      stateEntity.setHostEntity(hostComponentDesiredStateEntity.getHostEntity());
      stateEntity.setCurrentState(State.UNKNOWN);
      stateEntity.setUpgradeState(UpgradeState.NONE);
      stateEntity.setServiceComponentDesiredStateEntity(hostComponentDesiredStateEntity.getServiceComponentDesiredStateEntity());

      LOG.error("Trying to add missing record in hostcomponentstate: {}", stateEntity);
      hostComponentStateDAO.create(stateEntity);
    }

    for (HostComponentStateEntity missedHostComponentState : missedHostComponentStates) {

      HostComponentDesiredStateEntity stateEntity = new HostComponentDesiredStateEntity();
      stateEntity.setClusterId(missedHostComponentState.getClusterId());
      stateEntity.setComponentName(missedHostComponentState.getComponentName());
      stateEntity.setServiceName(missedHostComponentState.getServiceName());
      stateEntity.setHostEntity(missedHostComponentState.getHostEntity());
      stateEntity.setDesiredState(State.UNKNOWN);
      stateEntity.setServiceComponentDesiredStateEntity(missedHostComponentState.getServiceComponentDesiredStateEntity());

      LOG.error("Trying to add missing record in hostcomponentdesiredstate: {}", stateEntity);
      hostComponentDesiredStateDAO.create(stateEntity);
    }
  }

  /**
   * This makes the following checks for Postgres:
   * <ol>
   *   <li>Check if the connection's schema (first item on search path) is the one set in ambari.properties</li>
   *   <li>Check if the connection's schema is present in the DB</li>
   *   <li>Check if the ambari tables exist in the schema configured in ambari.properties</li>
   *   <li>Check if ambari tables don't exist in other shemas</li>
   * </ol>
   * The purpose of these checks is to avoid that tables and constraints in ambari's schema get confused with tables
   * and constraints in other schemas on the DB user's search path. This can happen after an improperly made DB restore
   * operation and can cause issues during upgrade.
   * we may have no permissions (e.g. external DB) to auto fix this problem
  **/
  static void checkSchemaName () {
    Configuration conf = injector.getInstance(Configuration.class);
    if(conf.getDatabaseType() == Configuration.DatabaseType.POSTGRES) {
      LOG.info("Ensuring that the schema set for Postgres is correct");

      ensureConnection();

      try (ResultSet schemaRs = connection.getMetaData().getSchemas();
           ResultSet searchPathRs = connection.createStatement().executeQuery("show search_path");
           ResultSet ambariTablesRs = connection.createStatement().executeQuery(
               "select table_schema from information_schema.tables where table_name = 'hostcomponentdesiredstate'")) {
        // Check if ambari's schema exists
        final boolean ambariSchemaExists = getResultSetColumn(schemaRs, "TABLE_SCHEM").contains(conf.getDatabaseSchema());
        if ( !ambariSchemaExists ) {
          warning("The schema [{}] defined for Ambari from ambari.properties has not been found in the database. " +
              "Storing Ambari tables under a different schema can lead to problems.", conf.getDatabaseSchema());
        }
        // Check if the right schema is first on the search path
        List<Object> searchPathResultColumn = getResultSetColumn(searchPathRs, "search_path");
        List<String> searchPath = searchPathResultColumn.isEmpty() ? ImmutableList.of() :
            ImmutableList.copyOf(Splitter.on(",").trimResults().split(String.valueOf(searchPathResultColumn.get(0))));
        String firstSearchPathItem = searchPath.isEmpty() ? null : searchPath.get(0);
        if (!Objects.equals(firstSearchPathItem, conf.getDatabaseSchema())) {
          warning("The schema [{}] defined for Ambari in ambari.properties is not first on the search path:" +
              " {}. This can lead to problems.", conf.getDatabaseSchema(), searchPath);
        }
        // Check schemas with Ambari tables.
        ArrayList<Object> schemasWithAmbariTables = getResultSetColumn(ambariTablesRs, "table_schema");
        if ( ambariSchemaExists && !schemasWithAmbariTables.contains(conf.getDatabaseSchema()) ) {
          warning("The schema [{}] defined for Ambari from ambari.properties does not contain the Ambari tables. " +
              "Storing Ambari tables under a different schema can lead to problems.", conf.getDatabaseSchema());
        }
        if ( schemasWithAmbariTables.size() > 1 ) {
          warning("Multiple schemas contain the Ambari tables: {}. This can lead to problems.", schemasWithAmbariTables);
        }
      }
      catch (SQLException e) {
        warning("Exception occurred during checking db schema name: ", e);
      }
    }
  }

  private static ArrayList<Object> getResultSetColumn(@Nullable ResultSet rs, String columnName) throws SQLException {
    ArrayList<Object> values = new ArrayList<>();
    if (null != rs) {
      while (rs.next()) {
        values.add(rs.getObject(columnName));
      }
    }
    return values;
  }

  /**
  * This method checks tables engine type to be innodb for MySQL.
  * it's too risky to autofix by migrating DB to innodb
  * */
  static void checkMySQLEngine () {
    Configuration conf = injector.getInstance(Configuration.class);
    if(conf.getDatabaseType()!=Configuration.DatabaseType.MYSQL) {
      return;
    }
