Implemented new REST API hpFlows to add high priority flows to Azkaban.

azkaban-common/src/main/java/azkaban/AzkabanCommonModule.java

@@ -32,6 +32,8 @@
 import azkaban.imagemgmt.converters.ImageRampupPlanConverter;
 import azkaban.imagemgmt.converters.ImageTypeConverter;
 import azkaban.imagemgmt.converters.ImageVersionConverter;
+import azkaban.imagemgmt.daos.HPFlowDao;
+import azkaban.imagemgmt.daos.HPFlowDaoImpl;
 import azkaban.imagemgmt.daos.ImageMgmtCommonDao;
 import azkaban.imagemgmt.daos.ImageMgmtCommonDaoImpl;
 import azkaban.imagemgmt.daos.ImageRampupDao;
@@ -104,7 +106,9 @@ protected void configure() {
               Constants.DEFAULT_AZKABAN_POLLING_INTERVAL_MS);
       return new OsCpuUtil(Math.max(1, (cpuLoadPeriodSec * 1000) / pollingIntervalMs));
     });
+    // Following bindings are needed for containerized flow execution
     bindImageManagementDependencies();
+    bindHPFlowManagementDependencies();
   }
 
   public Class<? extends Storage> resolveStorageClassType() {
@@ -200,4 +204,11 @@ private boolean isContainerizedDispatchMethodEnabled() {
         .getString(Constants.ConfigurationKeys.AZKABAN_EXECUTION_DISPATCH_METHOD,
             DispatchMethod.PUSH.name()));
   }
+
+  private void bindHPFlowManagementDependencies() {
+    if (!isContainerizedDispatchMethodEnabled()) {
+      return;
+    }
+    bind(HPFlowDao.class).to(HPFlowDaoImpl.class).in(Scopes.SINGLETON);
+  }
 }

azkaban-common/src/main/java/azkaban/imagemgmt/daos/HPFlowDao.java

@@ -0,0 +1,32 @@
+package azkaban.imagemgmt.daos;
+
+import azkaban.executor.ExecutableFlow;
+import java.util.List;
+
+
+/**
+ * Data access object (DAO) for accessing high priority flow metadata.
+ * This interface defines add/remove/get methods for high priority flows.
+ */
+public interface HPFlowDao {
+  /**
+   * isHPFlowOwner
+   * @param userId : owner of high priority flows.
+   * @return : True if user is owner, false otherwise.
+   */
+  boolean isHPFlowOwner(final String userId);
+
+  /**
+   * addHPFlows
+   * @param flowIds : list of flows.
+   * @param userId : userId of owner who is adding the HP flows.
+   */
+  int addHPFlows(final List<String> flowIds, final String userId);
+
+  /**
+   * isHPFlow
+   * @param flow : Given a flow, finds if it is a high priority flow.
+   * @return
+   */
+  boolean isHPFlow(final ExecutableFlow flow);
+}

azkaban-common/src/main/java/azkaban/imagemgmt/daos/HPFlowDaoImpl.java

@@ -0,0 +1,169 @@
+package azkaban.imagemgmt.daos;
+
+import azkaban.db.DatabaseOperator;
+import azkaban.db.SQLTransaction;
+import azkaban.executor.ExecutableFlow;
+import azkaban.imagemgmt.exception.ErrorCode;
+import azkaban.imagemgmt.exception.ImageMgmtDaoException;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Timestamp;
+import java.time.LocalDateTime;
+import java.util.List;
+import javax.inject.Inject;
+import javax.inject.Singleton;
+import org.apache.commons.dbutils.ResultSetHandler;
+import org.apache.log4j.Logger;
+
+import static azkaban.imagemgmt.utils.ErroCodeConstants.*;
+
+
+/**
+ * Dao implementation for managing high priority flows
+ */
+@Singleton
+public class HPFlowDaoImpl implements HPFlowDao {
+  private static final Logger log = Logger.getLogger(HPFlowDaoImpl.class);
+
+  private final DatabaseOperator databaseOperator;
+
+  // Flow ID is a fully qualified flow name with format,
+  // <project_name>.<flow_ame>
+  private static String INSERT_HP_FLOW_OWNER =
+      "INSERT into hp_flow_owners (name, created_on, created_by, "
+      + "modified_on, modified_by ) "
+      + "values (?, ?, ?, ?, ?)";
+
+  private static String INSERT_HP_FLOWS =
+      "INSERT into hp_flows (flow_id, created_on, created_by) "
+      + "values (?, ?, ?)";
+
+  @Inject
+  public HPFlowDaoImpl(final DatabaseOperator databaseOperator) {
+    this.databaseOperator = databaseOperator;
+  }
+
+  /**
+   * Returns if the provided user has high priority flow management access.
+   * @param userId : owner of high priority flows.
+   * @return
+   */
+  @Override
+  public boolean isHPFlowOwner(final String userId) {
+    final FetchHPFlowOwnership fetchHPFlowOwnership = new FetchHPFlowOwnership();
+    try {
+      final String returnedName = this.databaseOperator
+          .query(FetchHPFlowOwnership.FETCH_HP_OWNER_BY_NAME,
+              fetchHPFlowOwnership, userId.toLowerCase());
+      log.info(String.format(
+          "HPFlowDao : Returned user id %s for given user id %s", returnedName, userId));
+      return returnedName.equalsIgnoreCase(userId);
+    } catch (final SQLException e) {
+      log.error(FetchHPFlowOwnership.FETCH_HP_OWNER_BY_NAME + " failed.", e);
+      throw new ImageMgmtDaoException(ErrorCode.INTERNAL_SERVER_ERROR, String.format(
+          "Unable to fetch HP flow ownership for %s.", userId));
+    }
+  }
+
+  /**
+   * Adds a list of flows into high priority flows db.
+   * @param flowIds : list of flows.
+   * @param userId : userId of owner who is adding the HP flows.
+   * @return
+   */
+  @Override
+  public int addHPFlows(final List<String> flowIds, final String userId) {
+    // Create the transaction block
+    final SQLTransaction<Integer> insertAllFlows = transOperator -> {
+      // Insert each flow in its own row with same userId and created time.
+      final Timestamp currTimeStamp = Timestamp.valueOf(LocalDateTime.now());
+      int count = 0;
+      for (final String flowId : flowIds) {
+        transOperator.update(INSERT_HP_FLOWS, flowId, currTimeStamp, userId);
+        count++;
+      }
+      transOperator.getConnection().commit();
+      return count;
+    };
+
+    int count = 0;
+    try {
+      count = this.databaseOperator.transaction(insertAllFlows);
+      if (count < 1) {
+        log.error("Exception while adding HP flows");
+        throw new ImageMgmtDaoException(ErrorCode.BAD_REQUEST,
+            "Exception while adding HP flows");
+      }
+      log.info("Added the HP flows by user " + userId);
+      return count;
+    } catch (final SQLException e) {
+      log.error("Unable to add HP flows", e);
+      String errorMessage = "";
+      if (e.getErrorCode() == SQL_ERROR_CODE_DUPLICATE_ENTRY) {
+        errorMessage = "Reason: One or more flows already exists.";
+      } else {
+        errorMessage = "Reason: ";
+      }
+      throw new ImageMgmtDaoException(ErrorCode.BAD_REQUEST, "Exception while "
+          + "adding HP flows. " + errorMessage);
+    }
+  }
+
+  /**
+   * Checks if the provided flow is high priority.
+   * @param flow : Given a flow, finds if it is a high priority flow.
+   * @return
+   */
+  @Override
+  public boolean isHPFlow(final ExecutableFlow flow) {
+    if (null == flow) {
+      return false;
+    }
+    final FetchHPFlow fetchHPFlow = new FetchHPFlow();
+    final String FQFN = flow.getProjectName() + "." + flow.getFlowId();
+    try {
+      final String returnedFlowId = this.databaseOperator
+          .query(FetchHPFlow.SELECT_HP_FLOW, fetchHPFlow, FQFN);
+      log.info("HPFlowDao : Found high priority flow "
+          + returnedFlowId == null ? "" : returnedFlowId);
+      return returnedFlowId.equals(FQFN);
+    } catch (final SQLException e) {
+      log.error(FetchHPFlow.SELECT_HP_FLOW + " failed.", e);
+      throw new ImageMgmtDaoException(ErrorCode.INTERNAL_SERVER_ERROR, String.format(
+          "Unable to fetch high priority flow info for flow %s.", FQFN));
+    }
+  }
+
+  /**
+   * ResultSetHandler implementation class for fetching high priority flow ownership
+   */
+  public static class FetchHPFlowOwnership implements ResultSetHandler<String> {
+    private static final String FETCH_HP_OWNER_BY_NAME =
+        "SELECT name FROM hp_flow_owners where lower(name) = ?";
+
+    @Override
+    public String handle(final ResultSet rs) throws SQLException {
+      if (!rs.next()) {
+        return null;
+      }
+      return rs.getString("name");
+    }
+  }
+
+  /**
+   * ResultSetHandler implementation class for fetching high priority flow
+   */
+  public static class FetchHPFlow implements ResultSetHandler<String> {
+
+    private static String SELECT_HP_FLOW =
+        "SELECT count(*) from hp_flows where flow_id = ?";
+
+    @Override
+    public String handle(final ResultSet rs) throws SQLException {
+      if (!rs.next()) {
+        return null;
+      }
+      return rs.getString("flow_id");
+    }
+  }
+}

azkaban-common/src/main/java/azkaban/imagemgmt/dto/HPFlowDTO.java

@@ -0,0 +1,43 @@
+package azkaban.imagemgmt.dto;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import javax.validation.constraints.NotBlank;
+import org.codehaus.jackson.annotate.JsonProperty;
+
+
+/**
+ * This is Add HP Flow class for updating high priority flow metadata.
+ */
+public class HPFlowDTO extends BaseDTO {
+  // CSV of flow ids.
+  @JsonProperty("flowIds")
+  @NotBlank(message = "flowList cannot be empty.")
+  private String flowIds;
+
+  public void setFlowIds(final String flowIds) {
+    // Remove all whitespaces
+    this.flowIds = flowIds.replaceAll("\\s", "");
+  }
+
+  public String getFlowIds() {
+    return this.flowIds;
+  }
+
+  /**
+   * Converts the CSV into a list of flow IDs.
+   * @return list of flow IDs.
+   */
+  public List<String> getFlowIdList() {
+    if (this.flowIds.isEmpty()) {
+      return new ArrayList<>();
+    }
+    return Arrays.asList(flowIds.split(","));
+  }
+
+  @Override
+  public String toString() {
+    return "HPFlowDTO{" + "flowList='" + this.flowIds + '\'' + '}';
+  }
+}

azkaban-common/src/main/java/azkaban/imagemgmt/models/HPFlowOwnership.java

@@ -0,0 +1,17 @@
+package azkaban.imagemgmt.models;
+
+/**
+ * This class represents HP Flow managers' metadata
+ */
+public class HPFlowOwnership extends BaseModel {
+  // HP flow management owner name
+  private String name;
+
+  public String getName() {
+    return this.name;
+  }
+
+  public void setName(final String name) {
+    this.name = name;
+  }
+}

azkaban-common/src/main/java/azkaban/imagemgmt/permission/PermissionManager.java

@@ -36,4 +36,11 @@ public interface PermissionManager {
   public boolean hasPermission(final String imageTypeName, final String userId, final Type type)
       throws ImageMgmtException;
 
+  /**
+   * Checks if the user has permission to manage high priority flows.
+   * @param userId
+   * @return
+   * @throws ImageMgmtException
+   */
+  public boolean hasPermission(final String userId) throws ImageMgmtException;
 }

azkaban-common/src/main/java/azkaban/imagemgmt/permission/PermissionManagerImpl.java

@@ -15,6 +15,7 @@
  */
 package azkaban.imagemgmt.permission;
 
+import azkaban.imagemgmt.daos.HPFlowDao;
 import azkaban.imagemgmt.daos.ImageTypeDao;
 import azkaban.imagemgmt.exception.ErrorCode;
 import azkaban.imagemgmt.exception.ImageMgmtException;
@@ -27,6 +28,9 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import static java.util.Objects.*;
+
+
 /**
  * This class defines and manages the permission for accessing image management APIs. The
  * permissions are defined on image type. The permissions are created using user role and access
@@ -38,10 +42,13 @@ public class PermissionManagerImpl implements PermissionManager {
   private static final Logger log = LoggerFactory.getLogger(PermissionManagerImpl.class);
 
   private final ImageTypeDao imageTypeDao;
+  private final HPFlowDao hpFlowDao;
 
   @Inject
-  public PermissionManagerImpl(final ImageTypeDao imageTypeDao) {
+  public PermissionManagerImpl(final ImageTypeDao imageTypeDao,
+      final HPFlowDao hpFlowDao) {
     this.imageTypeDao = imageTypeDao;
+    this.hpFlowDao = hpFlowDao;
   }
 
   /**
@@ -76,4 +83,24 @@ public boolean hasPermission(final String imageTypeName, final String userId, fi
 
     return hasPermission;
   }
+
+  /**
+   * Checks if the user has permission to manage high priority flows.
+   * @param userId : must not be null.
+   * @return boolean
+   */
+  @Override
+  public boolean hasPermission(final String userId) {
+    requireNonNull(userId, "userId must not be null");
+    if (this.hpFlowDao.isHPFlowOwner(userId)) {
+      return true;
+    }
+
+    // Throw exception with FORBIDDEN error.
+    log.error(String.format("API access permission check failed. The user %s "
+        + "does not have access to manage high priority flows", userId));
+    throw new ImageMgmtException(ErrorCode.FORBIDDEN, String.format(
+        "API access permission check failed. The user %s does not have access "
+            + "to manage high priority flows", userId));
+  }
 }