models.py

from django.db import models
from django_filters import rest_framework as filters
from urllib.parse import urlsplit
from django.contrib.auth.models import User
from django.contrib.postgres.fields import ArrayField
from celery.result import AsyncResult
from ldvspec.celery import app as spec_app


class DataLocation(models.Model):
    name = models.CharField(max_length=50, primary_key=True)
    uri = models.CharField(max_length=200)

    @staticmethod
    def insert_location_from_string(location_string):
        """
        Insert a datalocation from a srm string (e.g. srm://srm-url:4321)
        Data Location names are: Sara, Juelich and Poznan
        :param str location_string: SRM url
        :return: DataLocation object
        """
        # netloc will become srm-url:4321
        _, netloc, *_ = urlsplit(location_string)
        if "sara" in netloc.lower():
            loc_name = "Sara"
        elif "fz-juelich" in netloc.lower():
            loc_name = "Juelich"
        elif "psnc" in netloc.lower():
            loc_name = "Poznan"
        else:
            loc_name = "Unknown"

        dataloc = DataLocation(name=loc_name, uri=netloc)
        dataloc.save()
        return dataloc


class DataProduct(models.Model):
    obs_id = models.CharField(verbose_name="OBS_ID", max_length=15)
    oid_source = models.CharField(verbose_name='OBS_ID_SOURCE', max_length=15)
    dataproduct_source = models.CharField(max_length=20)
    dataproduct_type = models.CharField(max_length=50)
    project = models.CharField(max_length=50)
    location = models.ForeignKey(DataLocation, on_delete=models.DO_NOTHING)
    activity = models.CharField(max_length=50)
    surl = models.CharField(max_length=200)
    filesize = models.PositiveBigIntegerField()
    additional_meta = models.JSONField()

    @staticmethod
    def insert_dataproduct(obs_id,
                           oid_source,
                           dataproduct_source,
                           dataproduct_type,
                           project,
                           activity,
                           surl,
                           filesize,
                           additional_meta):
        scheme, netloc, *_ = urlsplit(surl)

        dp = DataProduct(obs_id=obs_id,
                         oid_source=oid_source,
                         dataproduct_source=dataproduct_source,
                         dataproduct_type=dataproduct_type,
                         project=project,
                         location=DataLocation.insert_location_from_string('://'.join((scheme, netloc))),
                         activity=activity,
                         filesize=filesize,
                         additional_meta=additional_meta,
                         surl=surl
                         )
        dp.save()
        return dp


class DataProductFilter(models.Model):
    field = models.CharField(max_length=100)
    name = models.CharField(max_length=20)
    lookup_type = models.CharField(max_length=100)


class ATDBProcessingSite(models.Model):
    name = models.CharField(primary_key=True)
    url = models.URLField()


class WorkSpecification(models.Model):
    created_on = models.DateTimeField(auto_now_add=True)
    created_by = models.ForeignKey(User, on_delete=models.DO_NOTHING, null=True)
    filters = models.JSONField(null=True)
    inputs = models.JSONField(null=True)
    related_tasks = ArrayField(models.IntegerField(), null=True)
    is_ready = models.BooleanField(default=False)
    is_defined = models.BooleanField(default=False)
    async_task_result = models.CharField(max_length=100, null=True)
    processing_site = models.ForeignKey(ATDBProcessingSite, null=True, on_delete=models.DO_NOTHING)

    def save(self, force_insert=False, force_update=False, using=None,
             update_fields=None):
        super(WorkSpecification, self).save(force_insert=force_insert, force_update=force_update, using=using,
                                            update_fields=update_fields)
        if self.async_task_result is None:
            from lofardata.tasks import define_work_specification
            res: AsyncResult = define_work_specification.delay(self.pk)
            self.async_task_result = res.id
            super(WorkSpecification, self).save(update_fields=['async_task_result'])