From b256e5b79d7dbe16800212e05f357c496edd6964 Mon Sep 17 00:00:00 2001
From: Rebeccah <rhunter@redhat.com>
Date: Fri, 21 May 2021 11:36:34 -0400
Subject: [PATCH] adding in a 3rd command that shows the associated templates,
 orgs, and invs the venv is tied to

---
 .../management/commands/export_custom_venv.py |  3 +-
 .../commands/get_custom_venv_associations.py  | 35 +++++++++++++++++++
 .../management/commands/list_custom_venvs.py  |  4 +--
 3 files changed, 38 insertions(+), 4 deletions(-)
 create mode 100644 awx/main/management/commands/get_custom_venv_associations.py

diff --git a/awx/main/management/commands/export_custom_venv.py b/awx/main/management/commands/export_custom_venv.py
index b63a7492af..fac39a052a 100644
--- a/awx/main/management/commands/export_custom_venv.py
+++ b/awx/main/management/commands/export_custom_venv.py
@@ -12,7 +12,6 @@ class Command(BaseCommand):
         parser.add_argument(
             'path',
             type=str,
-            nargs='?',
             default='',
             help='run this with a path to a virtual environment as an argument to see the pip freeze data',
         )
@@ -24,7 +23,7 @@ class Command(BaseCommand):
             if pip_data:
                 print(pip_data)
                 msg = [
-                    'To list all available custom virtual environments run:',
+                    'To list all (now deprecated) custom virtual environments run:',
                     'awx-manage list_custom_venvs',
                 ]
                 print('\n'.join(msg))
diff --git a/awx/main/management/commands/get_custom_venv_associations.py b/awx/main/management/commands/get_custom_venv_associations.py
new file mode 100644
index 0000000000..4ae877503f
--- /dev/null
+++ b/awx/main/management/commands/get_custom_venv_associations.py
@@ -0,0 +1,35 @@
+# Copyright (c) 2021 Ansible, Inc.
+# All Rights Reserved
+
+from django.core.management.base import BaseCommand
+from awx.main.models import Organization, InventorySource, UnifiedJobTemplate
+
+
+class Command(BaseCommand):
+    """Returns the pip freeze from the path passed in the argument"""
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            'path',
+            type=str,
+            nargs='?',
+            default='',
+            help='run this with a path to a virtual environment as an argument to see the pip freeze data',
+        )
+
+    def handle(self, *args, **options):
+        # look organiztions and unified job templates (which include JTs, workflows, and Inventory updates)
+        super(Command, self).__init__()
+        results = {}
+        if options.get('path'):
+            # sanity check here - is path in list
+            path = options.get('path')
+            orgs = [{"name": org.name, "id": org.id} for org in Organization.objects.filter(custom_virtualenv=path)]
+            ujts = [{"name": org.name, "id": org.id} for org in UnifiedJobTemplate.objects.filter(custom_virtualenv=path)]
+            invsrc = [{"name": inv.name, "id": inv.id} for inv in InventorySource.objects.filter(custom_virtualenv=path)]
+            results["Organizations"] = orgs
+            results["Unified Job Templates"] = ujts
+            results["Inventory Sources"] = invsrc
+            print(results)
+        else:
+            print("missing argument: please include a path argument following the command.")
diff --git a/awx/main/management/commands/list_custom_venvs.py b/awx/main/management/commands/list_custom_venvs.py
index bec712eaf9..b621ea3a3b 100644
--- a/awx/main/management/commands/list_custom_venvs.py
+++ b/awx/main/management/commands/list_custom_venvs.py
@@ -8,7 +8,7 @@ from django.conf import settings
 
 
 class Command(BaseCommand):
-    """Returns either a list of custom venv paths from the path passed in the argument"""
+    """Returns a list of custom venv paths from the path passed in the argument"""
 
     def handle(self, *args, **options):
         super(Command, self).__init__()
@@ -20,7 +20,7 @@ class Command(BaseCommand):
 
             msg = [
                 '',
-                'To export the contents of a virtual environment, ' 'run the following command while supplying the path as an argument:',
+                'To export the contents of a (deprecated) virtual environment, ' 'run the following command while supplying the path as an argument:',
                 'awx-manage export_custom_venv /path/to/venv',
             ]
             print('\n'.join(msg))