From e0f6af47001b38e0a615960551bdfe7328c9e2e8 Mon Sep 17 00:00:00 2001
From: Jeff Bradberry <jeff.bradberry@gmail.com>
Date: Mon, 23 Mar 2020 16:42:34 -0400
Subject: [PATCH] Enable the export of several relation field types

---
 awxkit/awxkit/api/pages/api.py  | 27 +++++++++++++++++++++++----
 awxkit/awxkit/api/pages/page.py |  2 +-
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/awxkit/awxkit/api/pages/api.py b/awxkit/awxkit/api/pages/api.py
index ec1e246707..92295dd8f6 100644
--- a/awxkit/awxkit/api/pages/api.py
+++ b/awxkit/awxkit/api/pages/api.py
@@ -21,6 +21,15 @@ EXPORTABLE_RESOURCES = [
 ]
 
 
+EXPORTABLE_RELATIONS = [
+    'Roles',
+    'NotificationTemplates',
+    'Labels',
+    'SurveySpec',
+    'WorkflowJobTemplateNodes',
+]
+
+
 NATURAL_KEYS = {
     'user': ('username',),
     'organization': ('name',),
@@ -35,6 +44,9 @@ NATURAL_KEYS = {
 
     # related resources
     'role': ('name', ':content_object'),
+    'notification_template': ('organization', 'name'),
+    'label': ('organization', 'name'),  # FIXME: label will need to be fully constructed from this
+    'workflow_job_template_node': ('workflow_job_template', 'identifier'),
 }
 
 
@@ -98,11 +110,18 @@ class ApiV2(base.Base):
 
         related = {}
         for k, related_endpoint in asset.related.items():
-            if k != 'roles':
+            if not related_endpoint:
+                continue
+            if k == 'object_roles':
+                continue
+            rel = related_endpoint._create()
+            if rel.__class__.__name__ not in EXPORTABLE_RELATIONS:
                 continue
             data = related_endpoint.get(all_pages=True)
             if 'results' in data:
                 related[k] = [get_natural_key(x) for x in data.results]
+            else:
+                related[k] = data.json
 
         related_fields = {'related': related} if related else {}
 
@@ -115,7 +134,7 @@ class ApiV2(base.Base):
         if value:
             from awxkit.cli.options import pk_or_name
 
-            pk = pk_or_name(self, resource, value)
+            pk = pk_or_name(self, resource, value)  # TODO: decide whether to support multiple
             results = endpoint.get(id=pk).results
         else:
             results = endpoint.get(all_pages=True).results
@@ -205,7 +224,7 @@ class ApiV2(base.Base):
             else:
                 pass  # admin role
 
-    def assign_related_assets(self, resource, assets):
+    def _assign_related_assets(self, resource, assets):
         for asset in assets:
             page = self._get_by_natural_key(asset['natural_key'])
             # FIXME: deal with `page is None` case
@@ -234,7 +253,7 @@ class ApiV2(base.Base):
             # FIXME: should we delete existing unpatched assets?
 
         # for resource, assets in data.items():
-        #     self.assign_related_assets(resource, assets)
+        #     self._assign_related_assets(resource, assets)
 
 
 page.register_page(resources.v2, ApiV2)
diff --git a/awxkit/awxkit/api/pages/page.py b/awxkit/awxkit/api/pages/page.py
index 7e714bbb05..256424a214 100644
--- a/awxkit/awxkit/api/pages/page.py
+++ b/awxkit/awxkit/api/pages/page.py
@@ -273,7 +273,7 @@ class Page(object):
     def get(self, all_pages=False, **query_parameters):
         r = self.connection.get(self.endpoint, query_parameters)
         page = self.page_identity(r)
-        if all_pages and page.next:
+        if all_pages and getattr(page, 'next', None):
             paged_results = [r.json()['results']]
             while page.next:
                 r = self.connection.get(self.next, query_parameters)