DAGPrinter { public DAG execute(List<ProcessInfo> processInfos, String dagName) { CommonNodeMaintainer nodeMaintainer = new CommonNodeMaintainer(); DAGNodeCollection nc = nodeMaintainer.getPnc(); Set<String> printedNodeNames = new HashSet<String>(); if (dagName == null || dagName.trim().isEmpty()) { LOGGER.error(EMPTYERROR); throw new MetadataException(EMPTYERROR); } String pid = processInfos.get(0).getProcessId().toString(); String owner =System.getProperty("user.name"); final String prefixDAG = "\n#!/usr/bin/env python\n"+ "from airflow.operators import BashOperator,BranchPythonOperator,DummyOperator,HiveOperator\n"+ "import subprocess\n"+ "from datetime import datetime, timedelta\n"+ "from airflow import DAG\n"+ "import logging\n" + "import os\n" + "import ast\n" + "args = {'owner': '"+ owner+"','start_date': datetime(2015, 10, 1, 5, 40, 0), 'depends_on_past': False}\n" + "\n" + "dag = DAG(dag_id='dag_"+ processInfos.get(0).getBusDomainId().toString()+"_"+ processInfos.get(0).getProcessTypeId().toString() + "_" + pid+"', default_args=args)"+"\n"+ "dict = {}\n" + "logger = logging.getLogger(__name__)\n"; StringBuilder credentials = new StringBuilder(); credentials.append(isSecurityEnabled(pid, "security")); StringBuilder airflowDAG = new StringBuilder(); StringBuilder stepDAG = new StringBuilder(); LOGGER.info("Starting dag generation for " + dagName); Map<Integer, ProcessInfo> tempProcessInfos = new HashMap<Integer, ProcessInfo>(); for (ProcessInfo processInfo : processInfos) { tempProcessInfos.put(processInfo.getProcessId(), processInfo); } for (ProcessInfo processInfo : processInfos) { LOGGER.info("processing " + processInfo); DAGNodeCollection nodeCollection = uniqNodeCollectionTreeMap.get(processInfo.getProcessId()); if (nodeCollection == null) { nodeCollection = new DAGNodeCollection(processInfo); } String[] children = processInfo.getNextProcessIds().split(","); for (String child : children) { LOGGER.debug("Analyzing child: " + child); Integer childKey = new Integer(child); DAGNodeCollection childNodeCollection = uniqNodeCollectionTreeMap.get(childKey); if (childNodeCollection == null) { childNodeCollection = new DAGNodeCollection(tempProcessInfos.get(childKey)); } childNodeCollection.addParent(nodeCollection, nodeMaintainer); nodeCollection.addChild(childNodeCollection, nodeMaintainer); uniqNodeCollectionTreeMap.put(childKey, childNodeCollection); } if (nodeCollection.getHaltStepNode().getToNode().getId().equals(tempProcessInfos.get(nodeCollection.getId()).getParentProcessId())) { nodeCollection.getHaltStepNode().setToNode(nc.getHaltJobNode()); } if (tempProcessInfos.get(nodeCollection.getId()).getParentProcessId() == 0) { DAGNode firstNode = null; if (nodeCollection.getDAGForkNode() != null) { LOGGER.info("First node in the dag is a fork"); firstNode = nodeCollection.getDAGForkNode(); } else { firstNode = nodeCollection.getChildren().get(0).getInitStepNode(); } LOGGER.info("firstNode=" + firstNode.getName()); nc.getInitJobNode().setToNode(firstNode); } uniqNodeCollectionTreeMap.put(processInfo.getProcessId(), nodeCollection); } LOGGER.info("uniqNodeCollectionTreeMap=" + uniqNodeCollectionTreeMap); Map<String,DAGNode> dagNodeMap = new HashMap<String, DAGNode>(); dagNodeMap.put(nc.getInitJobNode().getName(), nc.getInitJobNode()); LOGGER.info("nodeMaintainer.restartNodes=" + nodeMaintainer.getRestartNodes()); for (Integer id : uniqNodeCollectionTreeMap.keySet()) { if (tempProcessInfos.get(id).getParentProcessId() == 0) { uniqNodeCollectionTreeMap.get(id).setDAGTaskNode(null); uniqNodeCollectionTreeMap.get(id).setTermStepNode(null); uniqNodeCollectionTreeMap.get(id).setHaltStepNode(null); uniqNodeCollectionTreeMap.get(id).setInitStepNode(null); nc.setIdForProcessNodes(id); } LOGGER.info("Adding collection for: " + id); dagNodeMap.putAll(uniqNodeCollectionTreeMap.get(id).getDAGNodes()); stepDAG.append(uniqNodeCollectionTreeMap.get(id).toXML(printedNodeNames)); } dagNodeMap.put(nc.getHaltJobNode().getName(), nc.getHaltJobNode()); dagNodeMap.put(nc.getTermJobNode().getName(), nc.getTermJobNode()); LOGGER.info("oozieNodeMap size=" + dagNodeMap.size()); airflowDAG.append(prefixDAG); airflowDAG.append(credentials.toString()); airflowDAG.append(nc.getInitJobNode().getDAG()); airflowDAG.append(stepDAG); airflowDAG.append(nc.getHaltJobNode().getDAG()); airflowDAG.append(nc.getTermJobNode().getDAG()); String postfixXml =""; String homeDir = System.getProperty("user.home"); try { InputStream fis = new FileInputStream(homeDir+"/defFile.txt"); InputStreamReader isr = new InputStreamReader(fis, Charset.forName("UTF-8")); BufferedReader br = new BufferedReader(isr); String line; Boolean flag = true; line = br.readLine(); while ( (line = br.readLine()) != null) { if(flag && line.contains("dag_halt_step_")) continue; else flag = false; postfixXml = postfixXml + line+"\n"; } File file = new File(homeDir+"/defFile.txt"); file.delete(); }catch (IOException e){ System.out.println("e = " + e); } airflowDAG.append(postfixXml); LOGGER.info("Complete !"); DAG dag = new DAG(); dag.setDAG(airflowDAG); return dag; }  DAG execute(List<ProcessInfo> processInfos, String dagName); String isSecurityEnabled(String pid, String configGroup); String getCredProperties(String pid, String configGroup); DAG execInfo(List<ProcessInfo> processInfos, String dagName);  }
@Test @Ignore public void testExecute1() throws Exception { List<ProcessInfo> processInfos = getProcessBeans3(); String workflowXML = new DAGPrinter().execute(processInfos, "test-workflow").getDAG().toString(); LOGGER.info("workflowxml is "+workflowXML); }
@Test @Ignore public void testExecute2() throws Exception { List<ProcessInfo> processInfos = getProcessBeans2(); String airflowDAG = new DAGPrinter().execute(processInfos, "test-dag").getDAG().toString(); PrintWriter out = new PrintWriter("unitTest2.dag.py"); out.println(airflowDAG); out.close(); }