From 6d70651611550cb373d4136e16c72f1b1be9b307 Mon Sep 17 00:00:00 2001
From: Jeff Bradberry <jeff.bradberry@gmail.com>
Date: Thu, 7 Feb 2019 16:40:07 -0500
Subject: [PATCH] Update the inventory_import management command

to respect the new Organization.max_hosts limit.
---
 .../management/commands/inventory_import.py   | 35 +++++++++++++++++--
 awx/main/models/inventory.py                  |  4 +++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/awx/main/management/commands/inventory_import.py b/awx/main/management/commands/inventory_import.py
index 3224edf5dd..0974dcb7fa 100644
--- a/awx/main/management/commands/inventory_import.py
+++ b/awx/main/management/commands/inventory_import.py
@@ -904,10 +904,27 @@ class Command(BaseCommand):
                 logger.error(LICENSE_MESSAGE % d)
             raise CommandError('License count exceeded!')
 
+    def check_org_host_limit(self):
+        license_info = get_licenser().validate()
+        if license_info.get('license_type', 'UNLICENSED') == 'open':
+            return
+
+        org = self.inventory.organization
+        if org is None or org.max_hosts == 0:
+            return
+
+        active_count = Host.objects.org_active_count(org.id)
+        if active_count > org.max_hosts:
+            raise CommandError('Host limit for organization exceeded!')
+
     def mark_license_failure(self, save=True):
         self.inventory_update.license_error = True
         self.inventory_update.save(update_fields=['license_error'])
 
+    def mark_org_limits_failure(self, save=True):
+        self.inventory_update.org_host_limit_error = True
+        self.inventory_update.save(update_fields=['org_host_limit_error'])
+
     def handle(self, *args, **options):
         self.verbosity = int(options.get('verbosity', 1))
         self.set_logging_level()
@@ -961,6 +978,13 @@ class Command(BaseCommand):
             self.mark_license_failure(save=True)
             raise e
 
+        try:
+            # Check the per-org host limits
+            self.check_org_host_limit()
+        except CommandError as e:
+            self.mark_org_limits_failure(save=True)
+            raise e
+
         status, tb, exc = 'error', '', None
         try:
             if settings.SQL_DEBUG:
@@ -1032,9 +1056,17 @@ class Command(BaseCommand):
                         # If the license is not valid, a CommandError will be thrown, 
                         # and inventory update will be marked as invalid.
                         # with transaction.atomic() will roll back the changes.
+                        license_fail = True
                         self.check_license()
+
+                        # Check the per-org host limits
+                        license_fail = False
+                        self.check_org_host_limit()
                 except CommandError as e:
-                    self.mark_license_failure()
+                    if license_fail:
+                        self.mark_license_failure()
+                    else:
+                        self.mark_org_limits_failure()
                     raise e
 
                 if settings.SQL_DEBUG:
@@ -1062,7 +1094,6 @@ class Command(BaseCommand):
             else:
                 tb = traceback.format_exc()
                 exc = e
-            transaction.rollback()
 
         if self.invoked_from_dispatcher is False:
             with ignore_inventory_computed_fields():
diff --git a/awx/main/models/inventory.py b/awx/main/models/inventory.py
index b098d2e4a9..df04d5b1ad 100644
--- a/awx/main/models/inventory.py
+++ b/awx/main/models/inventory.py
@@ -1653,6 +1653,10 @@ class InventoryUpdate(UnifiedJob, InventorySourceOptions, JobNotificationMixin,
         default=False,
         editable=False,
     )
+    org_host_limit_error = models.BooleanField(
+        default=False,
+        editable=False,
+    )
     source_project_update = models.ForeignKey(
         'ProjectUpdate',
         related_name='scm_inventory_updates',