From def2034883fea871b8c8775a14a78992dc8df6dc Mon Sep 17 00:00:00 2001
From: Ryan Petrello <rpetrell@redhat.com>
Date: Fri, 28 Apr 2017 11:11:16 -0400
Subject: [PATCH] add a new `vault_credential` relationship to Job/JobTemplate

additionally, properly assign vault credentials to Jobs and JobTemplates
as they're migrated to the new split model

see: #5807
see: #5878
---
 .../0040_v320_add_credentialtype_model.py     | 10 ++++++++++
 awx/main/migrations/_credentialtypes.py       | 20 +++++++++++++++++--
 awx/main/models/jobs.py                       |  8 ++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/awx/main/migrations/0040_v320_add_credentialtype_model.py b/awx/main/migrations/0040_v320_add_credentialtype_model.py
index 1042e68aaf..96c34246f2 100644
--- a/awx/main/migrations/0040_v320_add_credentialtype_model.py
+++ b/awx/main/migrations/0040_v320_add_credentialtype_model.py
@@ -50,6 +50,16 @@ class Migration(migrations.Migration):
             field=models.ForeignKey(related_name='credentials', to='main.CredentialType', null=True),
             preserve_default=False,
         ),
+        migrations.AddField(
+            model_name='job',
+            name='vault_credential',
+            field=models.ForeignKey(related_name='jobs_as_vault_credential+', on_delete=django.db.models.deletion.SET_NULL, default=None, blank=True, to='main.Credential', null=True),
+        ),
+        migrations.AddField(
+            model_name='jobtemplate',
+            name='vault_credential',
+            field=models.ForeignKey(related_name='jobtemplates_as_vault_credential+', on_delete=django.db.models.deletion.SET_NULL, default=None, blank=True, to='main.Credential', null=True),
+        ),
         migrations.AlterUniqueTogether(
             name='credential',
             unique_together=set([('organization', 'name', 'credential_type')]),
diff --git a/awx/main/migrations/_credentialtypes.py b/awx/main/migrations/_credentialtypes.py
index b387f33fd1..3e91dc3fd3 100644
--- a/awx/main/migrations/_credentialtypes.py
+++ b/awx/main/migrations/_credentialtypes.py
@@ -13,6 +13,8 @@ def migrate_to_v2_credentials(apps, schema_editor):
     try:
         utils.get_current_apps = lambda: apps
         for cred in apps.get_model('main', 'Credential').objects.all():
+            job_templates = cred.jobtemplates.all()
+            jobs = cred.jobs.all()
             data = {}
             if getattr(cred, 'vault_password', None):
                 data['vault_password'] = cred.vault_password
@@ -23,6 +25,15 @@ def migrate_to_v2_credentials(apps, schema_editor):
             for field in defined_fields:
                 if getattr(cred, field, None):
                     cred.inputs[field] = getattr(cred, field)
+            if cred.vault_password:
+                for jt in job_templates:
+                    jt.credential = None
+                    jt.vault_credential = cred
+                    jt.save()
+                for job in jobs:
+                    job.credential = None
+                    job.vault_credential = cred
+                    job.save()
             cred.save()
 
             #
@@ -45,11 +56,16 @@ def migrate_to_v2_credentials(apps, schema_editor):
                 new_cred.admin_role = None
                 new_cred.use_role = None
 
-                # TODO: Job Template assignments
-
                 if any([getattr(cred, field) for field in ssh_type.defined_fields]):
                     new_cred.save(force_insert=True)
 
+                    for jt in job_templates:
+                        jt.credential = new_cred
+                        jt.save()
+                    for job in jobs:
+                        job.credential = new_cred
+                        job.save()
+
                     # passwords must be decrypted and re-encrypted, because
                     # their encryption is based on the Credential's primary key
                     # (which has changed)
diff --git a/awx/main/models/jobs.py b/awx/main/models/jobs.py
index 3dc8b3dd1b..e9464d5ac2 100644
--- a/awx/main/models/jobs.py
+++ b/awx/main/models/jobs.py
@@ -88,6 +88,14 @@ class JobOptions(BaseModel):
         default=None,
         on_delete=models.SET_NULL,
     )
+    vault_credential = models.ForeignKey(
+        'Credential',
+        related_name='%(class)ss_as_vault_credential+',
+        blank=True,
+        null=True,
+        default=None,
+        on_delete=models.SET_NULL,
+    )
     cloud_credential = models.ForeignKey(
         'Credential',
         related_name='%(class)ss_as_cloud_credential+',