gistfile1.txt

# -*- coding: utf-8 -*-
# Generated by Django 1.9.8 on 2017-03-29 18:34
from __future__ import unicode_literals

from django.db import migrations, models

import uuid
from django.db.models.signals import post_save
from django.db.models.query import Q
from core.hooks.allocation_source import (
    listen_for_allocation_threshold_met,
    listen_for_instance_allocation_changes,
    listen_for_allocation_source_created_or_renewed,
    listen_for_allocation_source_compute_allowed_changed
)
from jetstream.models import TASAPIDriver


def toggle_signals(event_model, on=True):
    if on:
        post_save.connect(listen_for_allocation_threshold_met, sender=event_model)
        post_save.connect(listen_for_instance_allocation_changes, sender=event_model)
        post_save.connect(listen_for_allocation_source_created_or_renewed,sender=event_model)
        post_save.connect(listen_for_allocation_source_compute_allowed_changed,sender=event_model)

    else:
        post_save.disconnect(listen_for_allocation_threshold_met, sender=event_model)
        post_save.disconnect(listen_for_instance_allocation_changes, sender=event_model)
        post_save.disconnect(listen_for_allocation_source_created_or_renewed, sender=event_model)
        post_save.disconnect(listen_for_allocation_source_compute_allowed_changed, sender=event_model)


def add_events_for_old_allocationsource(apps,schema_editor):
    EventTable = apps.get_model('core','EventTable')
    AllocationSourceTable = apps.get_model('core','AllocationSource')

    api = TASAPIDriver()
    projects = api.get_all_projects()

    # switch off signals
    toggle_signals(EventTable, on=False)

    all_source_names = {}

    for project in projects:
        for allocation in project['allocations']:

            payload = {
                "start_date": allocation['start'],
                "end_date": allocation['end'],
                "compute_allocated": allocation['computeAllocated'],
                "allocation_source_name": project['chargeCode'],
            }

            created_event_key = 'sn=%s,si=%s,ev=%s,dc=jetstream,dc=atmosphere' % (
                project['chargeCode'], allocation['id'], 'allocation_source_created_or_renewed')
            created_event_uuid = uuid.uuid5(uuid.NAMESPACE_X500, str(created_event_key))

            aso = AllocationSourceTable.objects.filter(name=payload["allocation_source_name"])
            if not aso:
                continue

            try:
                e = EventTable(
                    uuid=created_event_uuid,
                    name="allocation_source_created_or_renewed",
                    payload=payload,
                    entity_id=payload['allocation_source_name'],
                    timestamp=payload['start_date'])
                e.save()
            except:
                pass

            compute_event_key = 'ca=%s,sn=%s,si=%s,ev=%s,dc=jetstream,dc=atmosphere' % (
                allocation['computeAllocated'], project['chargeCode'], allocation['id'],
                'allocation_source_compute_allowed_changed')
            compute_event_uuid = uuid.uuid5(uuid.NAMESPACE_X500, str(compute_event_key))

            try:
                e = EventTable(
                    uuid=compute_event_uuid,
                    name="allocation_source_compute_allowed_changed",
                    payload=payload,
                    entity_id=payload['allocation_source_name'],
                    timestamp=payload['start_date'])
                e.save()
            except:
                pass

    #switch signals back on
    toggle_signals(EventTable)

class Migration(migrations.Migration):

    dependencies = [
        ('core', '0078_unique_allocationsource_names_part2'),
    ]

    operations = [
        migrations.RunPython(
            add_events_for_old_allocationsource,
            reverse_code=migrations.RunPython.noop),
    ]