From ea6e78f1467e16d21b19211c06566914d0502731 Mon Sep 17 00:00:00 2001
From: Jordan <21129425+ItIsJordan@users.noreply.github.com>
Date: Fri, 26 Jan 2024 11:29:27 +0000
Subject: [PATCH] Add additional resource descriptions to index

Updates index to add file_description from additional resource objects within HEPSubmission/DataSubmission objects to the OpenSearch index.
---
 hepdata/ext/opensearch/document_enhancers.py | 29 ++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/hepdata/ext/opensearch/document_enhancers.py b/hepdata/ext/opensearch/document_enhancers.py
index daaaa347..f2cf8b80 100644
--- a/hepdata/ext/opensearch/document_enhancers.py
+++ b/hepdata/ext/opensearch/document_enhancers.py
@@ -34,6 +34,7 @@
 from hepdata.ext.opensearch.config.record_mapping import mapping as os_mapping
 from hepdata.modules.permissions.models import SubmissionParticipant
 from hepdata.modules.submission.api import get_latest_hepsubmission
+from hepdata.modules.submission.models import DataSubmission
 
 FORMATS = ['json', 'root', 'yaml', 'csv', 'yoda']
 
@@ -127,6 +128,19 @@ def add_data_keywords(doc):
     doc['data_keywords'] = dict(agg_keywords)
 
 
+def add_data_resources(doc):
+    """
+    Gets and adds additional resource descriptions for a DataSubmission object,
+    and adds it to the document object.
+
+    :param doc: The document object
+    :return:
+    """
+
+    submission = DataSubmission.query.filter_by(doi=doc["doi"]).one()
+    doc["resources"] = [s.file_description for s in submission.resources]
+
+
 def add_data_abstract(doc):
     """
     Adds the data abstract from its associated HEPSubmission to the document object
@@ -139,6 +153,19 @@ def add_data_abstract(doc):
     doc['data_abstract'] = submission.data_abstract
 
 
+def add_submission_resources(doc):
+    """
+    Gets and adds additional resource descriptions for a HEPSubmission object,
+    and adds it to the document object.
+
+    :param doc: The document object
+    :return:
+    """
+
+    submission = get_latest_hepsubmission(publication_recid=doc['recid'], overall_status='finished')
+    doc['resources'] = [s.file_description for s in submission.resources]
+
+
 def process_cmenergies(keywords):
     cmenergies = []
     if keywords['cmenergies']:
@@ -194,6 +221,7 @@ def enhance_data_document(doc):
     add_data_table_urls(doc)
     add_parent_publication(doc)
     add_data_keywords(doc)
+    add_data_resources(doc)
 
 
 def enhance_publication_document(doc):
@@ -205,3 +233,4 @@ def enhance_publication_document(doc):
     process_last_updates(doc)
     add_analyses(doc)
     add_parent_child_info(doc)
+    add_submission_resources(doc)