From e92c8cfdccee15e19e4c47a55749414d978713cb Mon Sep 17 00:00:00 2001
From: Jeff Bradberry <jeff.bradberry@gmail.com>
Date: Mon, 30 Mar 2020 16:07:12 -0400
Subject: [PATCH] Remove the natural key lookup dict and make each Page
 responsible instead

---
 awxkit/awxkit/api/pages/api.py                | 55 ++-----------------
 awxkit/awxkit/api/pages/credentials.py        |  3 +
 awxkit/awxkit/api/pages/inventory.py          |  1 +
 awxkit/awxkit/api/pages/job_templates.py      |  1 +
 awxkit/awxkit/api/pages/labels.py             |  1 +
 .../api/pages/notification_templates.py       |  1 +
 awxkit/awxkit/api/pages/organizations.py      |  2 +
 awxkit/awxkit/api/pages/page.py               | 21 +++++++
 awxkit/awxkit/api/pages/projects.py           |  1 +
 awxkit/awxkit/api/pages/roles.py              | 14 ++++-
 awxkit/awxkit/api/pages/schedules.py          |  2 +-
 awxkit/awxkit/api/pages/teams.py              |  1 +
 awxkit/awxkit/api/pages/users.py              |  2 +
 .../api/pages/workflow_job_template_nodes.py  |  1 +
 .../api/pages/workflow_job_templates.py       |  1 +
 awxkit/awxkit/exceptions.py                   |  5 ++
 16 files changed, 61 insertions(+), 51 deletions(-)

diff --git a/awxkit/awxkit/api/pages/api.py b/awxkit/awxkit/api/pages/api.py
index 8b45213459..a79e49d57e 100644
--- a/awxkit/awxkit/api/pages/api.py
+++ b/awxkit/awxkit/api/pages/api.py
@@ -35,50 +35,6 @@ EXPORTABLE_DEPENDENT_OBJECTS = [
 ]
 
 
-NATURAL_KEYS = {
-    'user': ('username',),
-    'organization': ('name',),
-    'team': ('organization', 'name'),
-    'credential_type': ('name', 'kind'),
-    'credential': ('organization', 'name', 'credential_type'),
-    'notification_template': ('organization', 'name'),
-    'project': ('organization', 'name'),
-    'inventory': ('organization', 'name'),
-    'job_template': ('organization', 'name'),
-    'workflow_job_template': ('organization', 'name'),
-
-    # 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'),
-    'schedule': ('unified_job_template', 'name'),
-}
-
-
-def get_natural_key(pg):
-    natural_key = {'type': pg['type']}
-    lookup = NATURAL_KEYS.get(pg['type'], ())
-
-    for key in lookup or ():
-        if key.startswith(':'):
-            # treat it like a special-case related object
-            related_objs = [
-                related for name, related in pg.related.items()
-                if name not in ('users', 'teams')
-            ]
-            if related_objs:
-                natural_key[key[1:]] = get_natural_key(related_objs[0].get())
-        elif key in pg.related:
-            natural_key[key] = get_natural_key(pg.related[key].get())
-        elif key in pg:
-            natural_key[key] = pg[key]
-
-    if not natural_key:
-        return None
-    return natural_key
-
-
 def freeze(key):
     if key is None:
         return None
@@ -113,10 +69,11 @@ class ApiV2(base.Base):
                 key: asset.json[key] for key in options
                 if key in asset.json and key not in asset.related and key != 'id'
             }
-            fields['natural_key'] = get_natural_key(asset)
+            fields['natural_key'] = asset.get_natural_key()
 
             fk_fields = {
-                key: get_natural_key(asset.related[key].get()) for key in options
+                # FIXME: use caching by url
+                key: asset.related[key].get().get_natural_key() for key in options
                 if key in asset.related
             }
 
@@ -124,7 +81,7 @@ class ApiV2(base.Base):
             for key, related_endpoint in asset.related.items():
                 if key in asset.json or not related_endpoint:
                     continue
-                if key == 'object_roles':
+                if key == 'object_roles':  # FIXME
                     continue
                 rel = related_endpoint._create()
 
@@ -139,7 +96,7 @@ class ApiV2(base.Base):
                 if 'results' in data:
                     related_options = self._get_options(related_endpoint)
                     related[key] = [
-                        get_natural_key(x) if by_natural_key else self._serialize_asset(x, related_options)
+                        x.get_natural_key() if by_natural_key else self._serialize_asset(x, related_options)
                         for x in data.results
                     ]
                 else:
@@ -190,7 +147,7 @@ class ApiV2(base.Base):
             yield page_resource[page_cls]
 
     def _register_page(self, page):
-        natural_key = freeze(get_natural_key(page))
+        natural_key = freeze(page.get_natural_key())
         # FIXME: we need to keep a reference for the case where we
         # don't have a natural key, so we can delete
         if natural_key is not None:
diff --git a/awxkit/awxkit/api/pages/credentials.py b/awxkit/awxkit/api/pages/credentials.py
index 8a5e7e0eca..88f89dd3b3 100644
--- a/awxkit/awxkit/api/pages/credentials.py
+++ b/awxkit/awxkit/api/pages/credentials.py
@@ -149,6 +149,8 @@ def get_payload_field_and_value_from_kwargs_or_config_cred(
 
 class CredentialType(HasCreate, base.Base):
 
+    NATURAL_KEY = ('name', 'kind')
+
     def silent_delete(self):
         if not self.managed_by_tower:
             return super(CredentialType, self).silent_delete()
@@ -204,6 +206,7 @@ class Credential(HasCopy, HasCreate, base.Base):
 
     dependencies = [CredentialType]
     optional_dependencies = [Organization, User, Team]
+    NATURAL_KEY = ('organization', 'name', 'credential_type')
 
     def payload(
             self,
diff --git a/awxkit/awxkit/api/pages/inventory.py b/awxkit/awxkit/api/pages/inventory.py
index 2b057153d3..6d89fcba3a 100644
--- a/awxkit/awxkit/api/pages/inventory.py
+++ b/awxkit/awxkit/api/pages/inventory.py
@@ -32,6 +32,7 @@ log = logging.getLogger(__name__)
 class Inventory(HasCopy, HasCreate, HasInstanceGroups, HasVariables, base.Base):
 
     dependencies = [Organization]
+    NATURAL_KEY = ('organization', 'name')
 
     def print_ini(self):
         """Print an ini version of the inventory"""
diff --git a/awxkit/awxkit/api/pages/job_templates.py b/awxkit/awxkit/api/pages/job_templates.py
index ad59bfb742..cd45fc0c87 100644
--- a/awxkit/awxkit/api/pages/job_templates.py
+++ b/awxkit/awxkit/api/pages/job_templates.py
@@ -24,6 +24,7 @@ class JobTemplate(
         UnifiedJobTemplate):
 
     optional_dependencies = [Inventory, Credential, Project]
+    NATURAL_KEY = ('organization', 'name')
 
     def launch(self, payload={}):
         """Launch the job_template using related->launch endpoint."""
diff --git a/awxkit/awxkit/api/pages/labels.py b/awxkit/awxkit/api/pages/labels.py
index f545ef3776..a76b6920a5 100644
--- a/awxkit/awxkit/api/pages/labels.py
+++ b/awxkit/awxkit/api/pages/labels.py
@@ -9,6 +9,7 @@ from . import page
 class Label(HasCreate, base.Base):
 
     dependencies = [Organization]
+    NATURAL_KEY = ('organization', 'name')
 
     def silent_delete(self):
         """Label pages do not support DELETE requests. Here, we override the base page object
diff --git a/awxkit/awxkit/api/pages/notification_templates.py b/awxkit/awxkit/api/pages/notification_templates.py
index f69bb7b8a7..ff192d4433 100644
--- a/awxkit/awxkit/api/pages/notification_templates.py
+++ b/awxkit/awxkit/api/pages/notification_templates.py
@@ -24,6 +24,7 @@ notification_types = (
 class NotificationTemplate(HasCopy, HasCreate, base.Base):
 
     dependencies = [Organization]
+    NATURAL_KEY = ('organization', 'name')
 
     def test(self):
         """Create test notification"""
diff --git a/awxkit/awxkit/api/pages/organizations.py b/awxkit/awxkit/api/pages/organizations.py
index cdc24083b8..413ecf4961 100644
--- a/awxkit/awxkit/api/pages/organizations.py
+++ b/awxkit/awxkit/api/pages/organizations.py
@@ -8,6 +8,8 @@ from . import page
 
 class Organization(HasCreate, HasInstanceGroups, HasNotifications, base.Base):
 
+    NATURAL_KEY = ('name',)
+
     def add_admin(self, user):
         if isinstance(user, page.Page):
             user = user.json
diff --git a/awxkit/awxkit/api/pages/page.py b/awxkit/awxkit/api/pages/page.py
index 256424a214..9effa8ef64 100644
--- a/awxkit/awxkit/api/pages/page.py
+++ b/awxkit/awxkit/api/pages/page.py
@@ -317,6 +317,24 @@ class Page(object):
         page_cls = get_registered_page(endpoint)
         return page_cls(self.connection, endpoint=endpoint).get(**kw)
 
+    def get_natural_key(self):
+        if not getattr(self, 'NATURAL_KEY', None):
+            raise exc.NoNaturalKey(
+                "Page does not have a natural key: {}".format(getattr(self, 'endpoint', repr(self.__class__)))
+            )
+        natural_key = {}
+        for key in self.NATURAL_KEY:
+            if key in self.related:
+                # FIXME: use caching by url
+                natural_key[key] = self.related[key].get().get_natural_key()
+            elif key in self:
+                natural_key[key] = self[key]
+        if not natural_key:
+            return None
+
+        natural_key['type'] = self['type']
+        return natural_key
+
 
 _exception_map = {http.NO_CONTENT: exc.NoContent,
                   http.NOT_FOUND: exc.NotFound,
@@ -376,6 +394,9 @@ class PageList(object):
     def create(self, *a, **kw):
         return self.__item_class__(self.connection).create(*a, **kw)
 
+    def get_natural_key(self):
+        raise exc.NoNaturalKey
+
 
 class TentativePage(str):
 
diff --git a/awxkit/awxkit/api/pages/projects.py b/awxkit/awxkit/api/pages/projects.py
index 584c151f78..e40191260c 100644
--- a/awxkit/awxkit/api/pages/projects.py
+++ b/awxkit/awxkit/api/pages/projects.py
@@ -14,6 +14,7 @@ class Project(HasCopy, HasCreate, HasNotifications, UnifiedJobTemplate):
 
     optional_dependencies = [Credential, Organization]
     optional_schedule_fields = tuple()
+    NATURAL_KEY = ('organization', 'name')
 
     def payload(self, organization, scm_type='git', **kwargs):
         payload = PseudoNamespace(
diff --git a/awxkit/awxkit/api/pages/roles.py b/awxkit/awxkit/api/pages/roles.py
index f8d6e39f70..a08d670840 100644
--- a/awxkit/awxkit/api/pages/roles.py
+++ b/awxkit/awxkit/api/pages/roles.py
@@ -5,7 +5,19 @@ from . import page
 
 class Role(base.Base):
 
-    pass
+    NATURAL_KEY = ('name',)
+
+    def get_natural_key(self):
+        natural_key = super(Role, self).get_natural_key()
+        related_objs = [
+            related for name, related in self.related.items()
+            if name not in ('users', 'teams')
+        ]
+        if related_objs:
+            # FIXME: use caching by url
+            natural_key['content_object'] = related_objs[0].get().get_natural_key()
+
+        return natural_key
 
 
 page.register_page(resources.role, Role)
diff --git a/awxkit/awxkit/api/pages/schedules.py b/awxkit/awxkit/api/pages/schedules.py
index b0b21645eb..8603b2ad5b 100644
--- a/awxkit/awxkit/api/pages/schedules.py
+++ b/awxkit/awxkit/api/pages/schedules.py
@@ -8,7 +8,7 @@ from . import base
 
 class Schedule(UnifiedJob):
 
-    pass
+    NATURAL_KEY = ('unified_job_template', 'name')
 
 
 page.register_page([resources.schedule,
diff --git a/awxkit/awxkit/api/pages/teams.py b/awxkit/awxkit/api/pages/teams.py
index fc1e9de3f8..cb5577b5b2 100644
--- a/awxkit/awxkit/api/pages/teams.py
+++ b/awxkit/awxkit/api/pages/teams.py
@@ -11,6 +11,7 @@ from . import page
 class Team(HasCreate, base.Base):
 
     dependencies = [Organization]
+    NATURAL_KEY = ('organization', 'name')
 
     def add_user(self, user):
         if isinstance(user, page.Page):
diff --git a/awxkit/awxkit/api/pages/users.py b/awxkit/awxkit/api/pages/users.py
index 4039ef2e9d..22ab78dd11 100644
--- a/awxkit/awxkit/api/pages/users.py
+++ b/awxkit/awxkit/api/pages/users.py
@@ -9,6 +9,8 @@ from . import page
 
 class User(HasCreate, base.Base):
 
+    NATURAL_KEY = ('username',)
+
     def payload(self, **kwargs):
         payload = PseudoNamespace(
             username=kwargs.get('username') or 'User-{}'.format(
diff --git a/awxkit/awxkit/api/pages/workflow_job_template_nodes.py b/awxkit/awxkit/api/pages/workflow_job_template_nodes.py
index 01e0b41e0a..1b61754928 100644
--- a/awxkit/awxkit/api/pages/workflow_job_template_nodes.py
+++ b/awxkit/awxkit/api/pages/workflow_job_template_nodes.py
@@ -10,6 +10,7 @@ from . import page
 class WorkflowJobTemplateNode(HasCreate, base.Base):
 
     dependencies = [WorkflowJobTemplate, UnifiedJobTemplate]
+    NATURAL_KEY = ('workflow_job_template', 'identifier')
 
     def payload(self, workflow_job_template, unified_job_template, **kwargs):
         if not unified_job_template:
diff --git a/awxkit/awxkit/api/pages/workflow_job_templates.py b/awxkit/awxkit/api/pages/workflow_job_templates.py
index 6a28891e04..17f3b56342 100644
--- a/awxkit/awxkit/api/pages/workflow_job_templates.py
+++ b/awxkit/awxkit/api/pages/workflow_job_templates.py
@@ -13,6 +13,7 @@ from . import page
 class WorkflowJobTemplate(HasCopy, HasCreate, HasNotifications, HasSurvey, UnifiedJobTemplate):
 
     optional_dependencies = [Organization]
+    NATURAL_KEY = ('organization', 'name')
 
     def launch(self, payload={}):
         """Launch using related->launch endpoint."""
diff --git a/awxkit/awxkit/exceptions.py b/awxkit/awxkit/exceptions.py
index c1ff01719a..cf26097cef 100644
--- a/awxkit/awxkit/exceptions.py
+++ b/awxkit/awxkit/exceptions.py
@@ -101,3 +101,8 @@ class UnexpectedAWXState(Common):
 class IsMigrating(Common):
 
     pass
+
+
+class NoNaturalKey(Common):
+
+    pass