Examples with AirflowDAG io.hops.hopsworks.common.util.templates.airflow.AirflowDAG used on opensource projects

Example 1 with AirflowDAG

use of io.hops.hopsworks.common.util.templates.airflow.AirflowDAG in project hopsworks by logicalclocks.

the class AirflowService method composeDAG.

@POST
@Path("/dag")
@Produces(MediaType.APPLICATION_JSON)
@AllowedProjectRoles({ AllowedProjectRoles.DATA_OWNER, AllowedProjectRoles.DATA_SCIENTIST })
@JWTRequired(acceptedTokens = { Audience.API }, allowedUserRoles = { "HOPS_ADMIN", "HOPS_USER" })
@ApiOperation(value = "Generate an Airflow Python DAG file from a DAG definition")
public Response composeDAG(AirflowDagDTO dagDefinition, @Context SecurityContext sc) throws AirflowException {
    Users user = jwtHelper.getUserPrincipal(sc);
    AirflowDAG dag = AirflowDagDTO.toAirflowDagTemplate(dagDefinition, user, project);
    Map<String, Object> dataModel = new HashMap<>(4);
    dataModel.put(AirflowJobLaunchOperator.class.getSimpleName(), AirflowJobLaunchOperator.class);
    dataModel.put(AirflowJobSuccessSensor.class.getSimpleName(), AirflowJobSuccessSensor.class);
    dataModel.put("dag", dag);
    java.nio.file.Path outputFile = Paths.get(airflowJWTManager.getProjectDagDirectory(project.getId()).toString(), dagDefinition.getName() + ".py");
    try (FileWriter out = new FileWriter(outputFile.toFile(), false)) {
        templateEngine.template(AirflowDAG.TEMPLATE_NAME, dataModel, out);
    } catch (IOException | TemplateException ex) {
        throw new AirflowException(RESTCodes.AirflowErrorCode.DAG_NOT_TEMPLATED, Level.SEVERE, "Could not template DAG file for Project " + project.getName(), ex.getMessage(), ex);
    }
    return Response.ok().build();
}

Example 2 with AirflowDAG

use of io.hops.hopsworks.common.util.templates.airflow.AirflowDAG in project hopsworks by logicalclocks.

the class AirflowDagDTO method toAirflowDagTemplate.

public static AirflowDAG toAirflowDagTemplate(AirflowDagDTO dagDefinition, Users owner, Project project) {
    AirflowDAG dag = new AirflowDAG(dagDefinition.getName(), owner.getUsername(), project.getName());
    dag.setScheduleInterval(dagDefinition.getScheduleInterval());
    if (dagDefinition.getApiKey() != null) {
        dag.setApiKey(dagDefinition.getApiKey());
    }
    List<AirflowOperator> operators = new ArrayList<>(dagDefinition.getOperators().size());
    for (AirflowOperatorDTO op : dagDefinition.getOperators()) {
        AirflowOperator operator = null;
        switch(op.getName()) {
            case AirflowJobLaunchOperator.NAME:
                operator = new AirflowJobLaunchOperator(project.getName(), op.getId(), op.getJobName(), op.getJobArgs());
                ((AirflowJobLaunchOperator) operator).setWait(op.isWait());
                break;
            case AirflowJobSuccessSensor.NAME:
                operator = new AirflowJobSuccessSensor(project.getName(), op.getId(), op.getJobName());
                break;
            default:
        }
        if (operator != null) {
            if (op.getDependsOn() != null && !op.getDependsOn().isEmpty()) {
                StringBuilder sb = new StringBuilder();
                if (op.getDependsOn().size() == 1) {
                    sb.append(AirflowOperator.sanitizeId(op.getDependsOn().get(0)));
                } else {
                    sb.append("[");
                    for (String dependency : op.getDependsOn()) {
                        sb.append(AirflowOperator.sanitizeId(dependency)).append(",");
                    }
                    sb.append("]");
                }
                operator.setUpstream(sb.toString());
            }
            operators.add(operator);
        }
    }
    dag.setOperators(operators);
    return dag;
}