fix: Merge staging to 1284

Author: Alan Christie

.pylintrc

@@ -12,3 +12,4 @@ disable = C, F, R,
           fixme,
           import-error,
           imported-auth-user,
+	      new-db-field-with-default,

viewer/download_structures.py

@@ -43,6 +43,7 @@
     'single_sdf_file': (''),
     'metadata_info': (''),
     'smiles_info': (''),
+    'trans_matrix_info': (''),
     'extra_files': ('extra_files'),
     'readme': (''),
 }
@@ -66,6 +67,7 @@
         'smiles_info': {},
     },
     'metadata_info': None,
+    'trans_matrix_info': None,
 }
 
 
@@ -345,6 +347,35 @@ def _smiles_files_zip(zip_contents, ziparchive, download_path):
     os.remove(smiles_filename)
 
 
+def _trans_matrix_files_zip(ziparchive, target):
+    """Add transformation matrices to archive.
+
+    Note that this will always be the latest information - even for
+    preserved searches.
+    """
+    logger.info('+ Processing trans matrix files')
+
+    # grab the last set of files for this target
+    experiment_upload = target.experimentupload_set.order_by('commit_datetime').last()
+
+    trans_matrix_files = (
+        experiment_upload.neighbourhood_transforms,
+        experiment_upload.conformer_site_transforms,
+        experiment_upload.reference_structure_transforms,
+    )
+    for tmf in trans_matrix_files:
+        if tmf:
+            ziparchive.write(
+                Path(settings.MEDIA_ROOT).joinpath(str(tmf)),
+                os.path.join(
+                    _ZIP_FILEPATHS['trans_matrix_info'],
+                    Path(str(tmf)).name,
+                ),
+            )
+        else:
+            logger.info('File %s does not exist', Path(str(tmf)).name)
+
+
 def _extra_files_zip(ziparchive, target):
     """If an extra info folder exists at the target root level, then
     copy the contents to the output file as is.
@@ -507,6 +538,7 @@ def _create_structures_zip(target, zip_contents, file_url, original_search, host
             logger.warning('After _add_file_to_zip() errors=%s', errors)
 
         _extra_files_zip(ziparchive, target)
+        _trans_matrix_files_zip(ziparchive, target)
 
         _document_file_zip(ziparchive, download_path, original_search, host)
 
@@ -613,6 +645,10 @@ def _create_structures_dict(target, site_obvs, protein_params, other_params):
     if other_params['metadata_info'] is True:
         zip_contents['metadata_info'] = target.metadata.name
 
+    # Add the metadata file from the target
+    if other_params['trans_matrix_info'] is True:
+        zip_contents['trans_matrix_info'] = True
+
     return zip_contents
 
 
@@ -635,7 +671,13 @@ def get_download_params(request):
         'diff_file',
     ]
 
-    other_param_flags = ['sdf_info', 'single_sdf_file', 'metadata_info', 'smiles_info']
+    other_param_flags = [
+        'sdf_info',
+        'single_sdf_file',
+        'metadata_info',
+        'smiles_info',
+        'trans_matrix_info',
+    ]
 
     # protein_params = {'pdb_info': request.data['pdb_info'],
     #               'bound_info': request.data['bound_info'],

viewer/migrations/0031_compound_compound_code.py

@@ -0,0 +1,17 @@
+# Generated by Django 3.2.23 on 2024-01-19 14:38
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ('viewer', '0030_add_cm_reference_code'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='compound',
+            name='compound_code',
+            field=models.TextField(null=True),
+        ),
+    ]

viewer/migrations/0032_auto_20240119_1501.py

@@ -0,0 +1,38 @@
+# Generated by Django 3.2.23 on 2024-01-19 15:01
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ('viewer', '0031_compound_compound_code'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='experimentupload',
+            name='conformer_site_transforms',
+            field=models.FileField(
+                default='trans.yaml',
+                max_length=255,
+                upload_to='experiment-upload/',
+            ),
+            preserve_default=False,
+        ),
+        migrations.AddField(
+            model_name='experimentupload',
+            name='neighbourhood_transforms',
+            field=models.FileField(
+                default='trans.yaml', max_length=255, upload_to='experiment-upload/'
+            ),
+            preserve_default=False,
+        ),
+        migrations.AddField(
+            model_name='experimentupload',
+            name='reference_structure_transforms',
+            field=models.FileField(
+                default='trans.yaml', max_length=255, upload_to='experiment-upload/'
+            ),
+            preserve_default=False,
+        ),
+    ]

viewer/models.py

@@ -162,6 +162,15 @@ class ExperimentUpload(models.Model):
         help_text="Any message or task info associated with the upload."
         " Used for upload audit trail",
     )
+    neighbourhood_transforms = models.FileField(
+        upload_to="experiment-upload/", max_length=255
+    )
+    conformer_site_transforms = models.FileField(
+        upload_to="experiment-upload/", max_length=255
+    )
+    reference_structure_transforms = models.FileField(
+        upload_to="experiment-upload/", max_length=255
+    )
 
     def __str__(self) -> str:
         return f"{self.project}"
@@ -210,6 +219,7 @@ class Compound(models.Model):
 
     inchi = models.TextField(unique=False, db_index=True)
     smiles = models.CharField(max_length=255, db_index=True)
+    compound_code = models.TextField(null=True)
     current_identifier = models.CharField(
         max_length=255,
         db_index=True,

viewer/serializers.py

@@ -850,6 +850,7 @@ class DownloadStructuresSerializer(serializers.Serializer):
     smiles_info = serializers.BooleanField(default=False)
     static_link = serializers.BooleanField(default=False)
     file_url = serializers.CharField(max_length=200)
+    trans_matrix_info = serializers.BooleanField(default=False)
 
 
 # Start of Serializers for Squonk Jobs

viewer/target_loader.py

@@ -53,6 +53,11 @@
 # everything else
 METADATA_FILE = "meta_aligner.yaml"
 
+# transformation matrices
+TRANS_NEIGHBOURHOOD = "neighbourhood_transforms.yaml"
+TRANS_CONF_SITE = "conformer_site_transforms.yaml"
+TRANS_REF_STRUCT = "reference_structure_transforms.yaml"
+
 
 class UploadState(str, Enum):
     """Target loader progress state.
@@ -67,6 +72,7 @@ class UploadState(str, Enum):
     REPORTING = "REPORTING"
     SUCCESS = "SUCCESS"
     FAILED = "FAILED"
+    CANCELED = "CANCELED"
 
 
 class Level(str, Enum):
@@ -145,13 +151,14 @@ def log(self, level: Level, message: str) -> None:
         self.stack.append(UploadReportEntry(level=level, message=message))
         self._update_task(message)
 
-    def final(self, archive_name):
-        if self.upload_state == UploadState.PROCESSING:
-            self.upload_state = UploadState.SUCCESS
-            message = f"{archive_name} uploaded successfully."
+    def final(self, message, upload_state=None):
+        if upload_state:
+            self.upload_state = upload_state
         else:
-            self.upload_state = UploadState.FAILED
-            message = f"Uploading {archive_name} failed."
+            if self.upload_state == UploadState.PROCESSING:
+                self.upload_state = UploadState.SUCCESS
+            else:
+                self.upload_state = UploadState.FAILED
 
         self.stack.append(UploadReportEntry(message=message))
         self._update_task(self.json())
@@ -162,6 +169,7 @@ def json(self):
     def _update_task(self, message: str | list) -> None:
         if self.task:
             try:
+                logger.debug("taskstuff %s", dir(self.task))
                 self.task.update_state(
                     state=self.upload_state,
                     meta={
@@ -671,10 +679,12 @@ def process_compound(
         fields = {
             "smiles": smiles,
         }
+        defaults = {"compound_code": data.get("compound_code", None)}
 
         return ProcessedObject(
             model_class=Compound,
             fields=fields,
+            defaults=defaults,
             key=protein_name,
         )
 
@@ -1127,7 +1137,9 @@ def process_metadata(
         # moved this bit from init
         self.target, target_created = Target.objects.get_or_create(
             title=self.target_name,
-            display_name=self.target_name,
+            defaults={
+                "display_name": self.target_name,
+            },
         )
 
         # TODO: original target loader's function get_create_projects
@@ -1146,7 +1158,7 @@ def process_metadata(
             # remove uploaded file
             Path(self.bundle_path).unlink()
             msg = f"{self.bundle_name} already uploaded, skipping."
-            self.report.log(Level.INFO, msg)
+            self.report.final(msg, upload_state=UploadState.CANCELED)
             raise FileExistsError(msg)
 
         if project_created and committer.pk == settings.ANONYMOUS_USER:
@@ -1158,9 +1170,42 @@ def process_metadata(
         assert self.target
         self.target.project_id.add(self.project)
 
+        meta = self._load_yaml(Path(upload_root).joinpath(METADATA_FILE))
+
+        # collect top level info
+        self.version_number = meta["version_number"]
+        self.version_dir = meta["version_dir"]
+        self.previous_version_dirs = meta["previous_version_dirs"]
+
+        # check tranformation matrix files
+        (  # pylint: disable=unbalanced-tuple-unpacking
+            trans_neighbourhood,
+            trans_conf_site,
+            trans_ref_struct,
+        ) = self.validate_files(
+            obj_identifier="trans_matrices",
+            # since the paths are given if file as strings, I think I
+            # can get away with compiling them as strings here
+            file_struct={
+                TRANS_NEIGHBOURHOOD: f"{self.version_dir}/{TRANS_NEIGHBOURHOOD}",
+                TRANS_CONF_SITE: f"{self.version_dir}/{TRANS_CONF_SITE}",
+                TRANS_REF_STRUCT: f"{self.version_dir}/{TRANS_REF_STRUCT}",
+            },
+            required=(TRANS_NEIGHBOURHOOD, TRANS_CONF_SITE, TRANS_REF_STRUCT),
+        )
+
         self.experiment_upload.project = self.project
         self.experiment_upload.target = self.target
         self.experiment_upload.committer = committer
+        self.experiment_upload.neighbourhood_transforms = str(
+            self._get_final_path(trans_neighbourhood)
+        )
+        self.experiment_upload.conformer_site_transforms = str(
+            self._get_final_path(trans_conf_site)
+        )
+        self.experiment_upload.reference_structure_transforms = str(
+            self._get_final_path(trans_ref_struct)
+        )
         self.experiment_upload.save()
 
         (  # pylint: disable=unbalanced-tuple-unpacking
@@ -1171,13 +1216,6 @@ def process_metadata(
             blocks=("assemblies", "xtalforms"),
         )
 
-        meta = self._load_yaml(Path(upload_root).joinpath(METADATA_FILE))
-
-        # collect top level info
-        self.version_number = meta["version_number"]
-        self.version_dir = meta["version_dir"]
-        self.previous_version_dirs = meta["previous_version_dirs"]
-
         (  # pylint: disable=unbalanced-tuple-unpacking
             crystals,
             xtalforms,
@@ -1563,7 +1601,7 @@ def load_target(
                     )
         except IntegrityError as exc:
             logger.error(exc, exc_info=True)
-            target_loader.report.final(target_loader.data_bundle)
+            target_loader.report.final(f"Uploading {target_loader.data_bundle} failed")
             target_loader.experiment_upload.message = target_loader.report.json()
             raise IntegrityError from exc
 
@@ -1585,7 +1623,7 @@ def load_target(
 
         set_directory_permissions(target_loader.abs_final_path, 0o755)
 
-        target_loader.report.final(target_loader.data_bundle)
+        target_loader.report.final(f"{target_loader.data_bundle} uploaded successfully")
         target_loader.experiment_upload.message = target_loader.report.json()
 
         # logger.debug("%s", upload_report)

viewer/views.py

@@ -12,6 +12,7 @@
 
 import pandas as pd
 import pytz
+from celery import Celery
 from celery.result import AsyncResult
 from dateutil.parser import parse
 from django.conf import settings
@@ -1578,15 +1579,39 @@ def get(self, request, task_id, *args, **kwargs):
 
         # task_id is (will be) a UUID, but Celery expects a string
         task_id_str = str(task_id)
-        try:
-            result = AsyncResult(task_id_str)
-        except TimeoutError:
-            error = {'error': 'Task result query timed out'}
-            return Response(error, status=status.HTTP_408_REQUEST_TIMEOUT)
-
-        if result.state == 'PENDING':
-            error = {'error': 'Unknown task'}
-            return Response(error, status=status.HTTP_400_BAD_REQUEST)
+
+        celery_app = Celery("fragalysis")
+        inspect = celery_app.control.inspect()
+        ping = inspect.ping()
+
+        if ping:
+            active_tasks = inspect.active()
+
+            # active_tasks.values is a list of tasks for every worker
+            if task_id_str in [
+                k['id'] for worker in active_tasks.values() for k in worker
+            ]:
+                # celery confirms task exists
+                try:
+                    result = AsyncResult(task_id_str)
+                except TimeoutError:
+                    error = {'error': 'Task result query timed out'}
+                    return Response(error, status=status.HTTP_408_REQUEST_TIMEOUT)
+            else:
+                # no such task
+                error = {'error': 'Unknown task'}
+                return Response(error, status=status.HTTP_400_BAD_REQUEST)
+
+        else:
+            # task scheduler not running. This may be the case in local
+            # development, but this means there's really no way to
+            # validate whether the task exists
+            logger.warning('Celery not running!')
+            try:
+                result = AsyncResult(task_id_str)
+            except TimeoutError:
+                error = {'error': 'Task result query timed out'}
+                return Response(error, status=status.HTTP_408_REQUEST_TIMEOUT)
 
         # Extract messages (from task info)
         # Assuming the task has some info.