From da54b4a5204b1983cb1ff0c258abd195db86badc Mon Sep 17 00:00:00 2001
From: Ruonan Li <rli@freewheel.tv>
Date: Thu, 30 May 2019 14:11:56 +0800
Subject: [PATCH] Expand Ajax API -- add  fetchProjectLogs and fetchflowgraph

---
 azkaban/remote.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/azkaban/remote.py b/azkaban/remote.py
index 7113398..7bcd515 100644
--- a/azkaban/remote.py
+++ b/azkaban/remote.py
@@ -341,6 +341,22 @@ def get_projects(self):
       params={'ajax': 'fetchallprojects'},
     ))
 
+  def get_project_logs(self, name):
+    """Get logs of a project.
+
+    :param name: Project name.
+
+    """
+    self._logger.debug('Getting logs of project %s.', name)
+    return _extract_json(self._request(
+      method='GET',
+      endpoint='manager',
+      params={
+        'ajax': 'fetchProjectLogs',
+        'project': name,
+      },
+    ))
+
   def create_project(self, name, description):
     """Create project.
 
@@ -737,6 +753,37 @@ def get_workflow_info(self, name, flow):
         # but sends a 200 empty response if the project doesn't exist
         raise AzkabanError('Project %s not found.', name)
 
+  def get_workflow_graph(self, name, flow):
+    """Get graph of jobs corresponding to a workflow.
+
+    :param name: Project name.
+    :param flow: Name of flow in project.
+
+    """
+    self._logger.debug(
+      'Fetching flow graph for workflow %s in project %s', flow, name
+    )
+    try:
+      res = self._request(
+        method='GET',
+        endpoint='manager',
+        params={
+          'ajax': 'fetchflowgraph',
+          'project': name,
+          'flow': flow,
+        },
+      )
+    except HTTPError:
+      # the Azkaban server throws a NullPointerException if the flow doesn't
+      # exist in the project, which causes a 500 response
+      raise AzkabanError('Worklow %s not found in project %s.', flow, name)
+    else:
+      try:
+        return _extract_json(res)
+      except ValueError:
+        # but sends a 200 empty response if the project doesn't exist
+        raise AzkabanError('Project %s not found.', name)
+
   def _refresh(self, password=None):
     """Refresh session ID.