Move the import logic into methods on the ApiV2 class

awxkit/awxkit/api/pages/api.py

@@ -1,6 +1,10 @@
+import itertools
+
 from awxkit.api.resources import resources
+import awxkit.exceptions as exc
 from . import base
 from . import page
+from ..mixins import has_create
 
 
 EXPORTABLE_RESOURCES = [
@@ -34,29 +38,35 @@ NATURAL_KEYS = {
 }
 
 
-def get_natural_key(page):
+def get_natural_key(pg):
-    natural_key = {'type': page['type']}
+    natural_key = {'type': pg['type']}
-    lookup = NATURAL_KEYS.get(page['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 page.related.items()
+                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 page.related:
+        elif key in pg.related:
-            natural_key[key] = get_natural_key(page.related[key].get())
+            natural_key[key] = get_natural_key(pg.related[key].get())
-        elif key in page:
+        elif key in pg:
-            natural_key[key] = page[key]
+            natural_key[key] = pg[key]
 
     if not natural_key:
         return None
     return natural_key
 
 
+def freeze(key):
+    if key is None:
+        return None
+    return frozenset((k, freeze(v) if isinstance(v, dict) else v) for k, v in key.items())
+
+
 class Api(base.Base):
 
     pass
@@ -114,6 +124,97 @@ class ApiV2(base.Base):
         assets = (self._serialize_asset(asset, options) for asset in results)
         return [asset for asset in assets if asset is not None]
 
+    def _dependent_resources(self, data):
+        page_resource = {getattr(self, resource)._create().__item_class__: resource
+                         for resource in self.json}
+        data_pages = [getattr(self, resource)._create().__item_class__ for resource in data]
+
+        for page_cls in itertools.chain(*has_create.page_creation_order(*data_pages)):
+            yield page_resource[page_cls]
+
+    def _register_page(self, page):
+        natural_key = freeze(get_natural_key(page))
+        if natural_key is not None:
+            if getattr(self, '_natural_key', None) is None:
+                self._natural_key = {}
+
+            self._natural_key[natural_key] = page
+
+    def _register_existing_assets(self, resource):
+        endpoint = getattr(self, resource)
+        options = self._get_options(endpoint)
+        if getattr(self, '_options', None) is None:
+            self._options = {}
+        self._options[resource] = options
+
+        results = endpoint.get(all_pages=True).results
+        for pg in results:
+            self._register_page(pg)
+
+    def _get_by_natural_key(self, key, fetch=True):
+        frozen_key = freeze(key)
+        if frozen_key is not None and frozen_key not in self._natural_key and fetch:
+            pass  # FIXME
+
+        return self._natural_key.get(frozen_key)
+
+    def _create_assets(self, data, resource):
+        if resource not in data or resource not in EXPORTABLE_RESOURCES:
+            return
+
+        endpoint = getattr(self, resource)
+        options = self._options[resource]
+        assets = data[resource]
+        for asset in assets:
+            post_data = {}
+            for field, value in asset.items():
+                if field not in options:
+                    continue
+                if options[field]['type'] == 'id':
+                    page = self._get_by_natural_key(value)
+                    post_data[field] = page['id'] if page is not None else None
+                else:
+                    post_data[field] = value
+
+            page = self._get_by_natural_key(asset['natural_key'], fetch=False)
+            if page is None:
+                if resource == 'users':
+                    # We should only impose a default password if the resource doesn't exist.
+                    post_data.setdefault('password', 'abc123')
+                page = endpoint.post(post_data)
+            else:
+                page = page.put(post_data)
+
+            self._register_page(page)
+
+    def _assign_related(self, page, name, related_set):
+        pass  # FIXME
+
+    def _assign_roles(self, page, roles):
+        role_endpoint = page.json['related']['roles']
+        for role in roles:
+            if 'content_object' not in role:
+                continue  # admin role
+            obj_page = self._get_by_natural_key(role['content_object'])
+            if obj_page is not None:
+                role_page = obj_page.get_object_role(role['name'], by_name=True)
+                try:
+                    role_endpoint.post({'id': role_page['id']})
+                except exc.NoContent:  # desired exception on successful (dis)association
+                    pass
+            else:
+                pass  # admin role
+
+    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
+            for name, S in asset.get('related', {}).items():
+                if name == 'roles':
+                    self._assign_roles(page, S)
+                else:
+                    self._assign_related(page, name, S)
+
     def export_assets(self, **kwargs):
         # If no resource kwargs are explicitly used, export everything.
         all_resources = all(kwargs.get(resource) is None for resource in EXPORTABLE_RESOURCES)
@@ -126,8 +227,14 @@ class ApiV2(base.Base):
 
         return data
 
-    def import_assets(self):
+    def import_assets(self, data):
-        pass
+        for resource in self._dependent_resources(data):
+            self._register_existing_assets(resource)
+            self._create_assets(data, resource)
+            # FIXME: should we delete existing unpatched assets?
+
+        # for resource, assets in data.items():
+        #     self.assign_related_assets(resource, assets)
 
 
 page.register_page(resources.v2, ApiV2)

awxkit/awxkit/cli/resource.py

@@ -1,14 +1,11 @@
-import itertools
 import json
 import os
 import sys
 
 from awxkit import api, config
-import awxkit.exceptions as exc
 from awxkit.utils import to_str
-from awxkit.api.mixins import has_create
 from awxkit.api.pages import Page
-from awxkit.api.pages.api import EXPORTABLE_RESOURCES, get_natural_key
+from awxkit.api.pages.api import EXPORTABLE_RESOURCES
 from awxkit.cli.format import FORMATTERS, format_response, add_authentication_arguments
 from awxkit.cli.utils import CustomRegistryMeta, cprint
 
@@ -129,125 +126,19 @@ class Config(CustomCommand):
         }
 
 
-def freeze(key):
-    if key is None:
-        return None
-    return frozenset((k, freeze(v) if isinstance(v, dict) else v) for k, v in key.items())
-
-
 class Import(CustomCommand):
     name = 'import'
     help_text = 'import resources into Tower'
 
-    def __init__(self, *args, **kwargs):
-        super(Import, self).__init__(*args, **kwargs)
-        self._natural_key = {}
-        self._options = {}
-
-    def get_by_natural_key(self, key, fetch=True):
-        frozen_key = freeze(key)
-        if frozen_key is not None and frozen_key not in self._natural_key and fetch:
-            pass
-
-        return self._natural_key.get(frozen_key)
-
-    def create_assets(self, data, resource):
-        if resource not in data or resource not in EXPORTABLE_RESOURCES:
-            return
-        cprint("importing {}".format(resource), 'red', file=self.client.stderr)
-
-        endpoint = getattr(self.v2, resource)
-        options = self._options[resource]
-        assets = data[resource]
-        for asset in assets:
-            post_data = {}
-            for field, value in asset.items():
-                if field in ('related', 'natural_key'):
-                    continue
-                if options[field]['type'] == 'id':
-                    page = self.get_by_natural_key(value)
-                    post_data[field] = page['id'] if page is not None else None
-                else:
-                    post_data[field] = value
-
-            page = self.get_by_natural_key(asset['natural_key'], fetch=False)
-            if page is None:
-                if resource == 'users':
-                    # We should only impose a default password if the resource doesn't exist.
-                    post_data.setdefault('password', 'password')
-                page = endpoint.post(post_data)
-            else:
-                page = page.put(post_data)
-
-            self.register_page(page)
-
-    def register_page(self, page):
-        natural_key = freeze(get_natural_key(page))
-        if natural_key is not None:
-            self._natural_key[natural_key] = page
-
-    def register_existing_assets(self, resource):
-        endpoint = getattr(self.v2, resource)
-        options = endpoint.options().json['actions']['POST']
-        self._options[resource] = options
-
-        results = endpoint.get(all_pages=True).results
-        for page in results:
-            self.register_page(page)
-
-    def assign_roles(self, page, roles):
-        role_endpoint = page.json['related']['roles']
-        for role in roles:
-            if 'content_object' not in role:
-                continue  # admin role
-            obj_page = self.get_by_natural_key(role['content_object'])
-            if obj_page is not None:
-                role_page = obj_page.get_object_role(role['name'], by_name=True)
-                try:
-                    role_endpoint.post({'id': role_page['id']})
-                except exc.NoContent:  # desired exception on successful (dis)association
-                    pass
-            else:
-                pass  # admin role
-
-    def assign_related(self, page, name, related_set):
-        pass
-
-    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
-            for name, S in asset.get('related', {}).items():
-                if name == 'roles':
-                    self.assign_roles(page, S)
-                else:
-                    self.assign_related(page, name, S)
-
-    def dependent_resources(self, data):
-        page_resource = {getattr(self.v2, resource)._create().__item_class__: resource
-                         for resource in self.v2.json}
-        data_pages = [getattr(self.v2, resource)._create().__item_class__ for resource in data]
-
-        for page_cls in itertools.chain(*has_create.page_creation_order(*data_pages)):
-            yield page_resource[page_cls]
-
     def handle(self, client, parser):
         if client.help:
             parser.print_help()
             raise SystemExit()
 
         data = json.load(sys.stdin)
+
         client.authenticate()
-        self.client = client
+        client.v2.import_assets(data)
-        self.v2 = client.v2
-
-        for resource in self.dependent_resources(data):
-            self.register_existing_assets(resource)
-            self.create_assets(data, resource)
-            # FIXME: should we delete existing unpatched assets?
-
-        # for resource, assets in data.items():
-        #     self.assign_related_assets(resource, assets)
 
         return {}