tweaking wresters

apps/cli/utilities/aat_wrest/aat_wrest/metadata_wrester.py

@@ -366,7 +366,22 @@ class WrestWorkflowMetadata:
         make_json = {}
         try:
             cursor = self.conn.cursor()
-            cursor.execute(query, {"spl": self.spl[0]})
+            spl = self.spl[0]
+            if self.spl[0].isdigit():
+                spl_query = """
+                SELECT science_product_locator
+                FROM science_products_product_groups
+                WHERE product_group_id = %(pg_id)s;                       
+                """
+                cursor.execute(spl_query, {"pg_id": self.spl[0]})
+                data = cursor.fetchall()
+                if data:
+                    spl = data[0][0]
+                else:
+                    self.logger.error(f"ERROR: Failed to determine SPL from product group id, got {data}")
+                    return make_json
+
+            cursor.execute(query, {"spl": spl})
             data = cursor.fetchall()
             if data:
                 make_json = json.dumps(
@@ -379,7 +394,7 @@ class WrestWorkflowMetadata:
             else:
                 self.logger.error(
                     f"ERROR: aat-wrest query returned no results!"
-                    f" The database appears to be missing information for spl id {self.spl[0]}!"
+                    f" The database appears to be missing information for spl id {spl}!"
                 )
         finally:
             self.conn.close()
@@ -399,25 +414,25 @@ class WrestWorkflowMetadata:
                 product_group_id = self.spl[0]
             else:
                 # We were given an SPL, need to get the associated product group ID
-                prod_id_query = f"""
+                prod_id_query = """
                 SELECT product_group_id
                 FROM science_products_product_groups
-                WHERE science_product_locator = '{self.spl[0]}';
+                WHERE science_product_locator = %(spl)s;
                 """
-                cursor.execute(prod_id_query)
+                cursor.execute(prod_id_query, {"spl": self.spl[0]})
                 data = cursor.fetchall()
                 if data:
                     product_group_id = data[0][0]
                 else:
                     self.logger.error(f"ERROR: Failed to fetch product group id from SPL, got {data}")
 
-            files_query = f"""
+            files_query = """
             WITH product_locators AS (
                 SELECT sp.science_product_locator,
                        ap.ancillary_product_locator
                 FROM   science_products_product_groups sp
                 JOIN   ancillary_products ap ON sp.product_group_id = ap.product_group_id
-                WHERE  sp.product_group_id = '{product_group_id}'
+                WHERE  sp.product_group_id = %(pg_id)s
             ), filegroups AS (
                 SELECT ap2.filegroup_id
                 FROM ancillary_products ap2
@@ -432,7 +447,7 @@ class WrestWorkflowMetadata:
             JOIN filegroups fg ON f.filegroup = fg.filegroup_id
             """
             if product_group_id != "":
-                cursor.execute(files_query)
+                cursor.execute(files_query, {"pg_id": product_group_id})
                 data = cursor.fetchall()
                 if data:
                     # This should have all the ancillary and science product names
@@ -442,6 +457,19 @@ class WrestWorkflowMetadata:
                     self.logger.error(f"ERROR: Failed to fetch products from product ID {product_group_id}")
             else:
                 self.logger.error(f"ERROR: Failed to get product group ID")
+
+            input_group_query = """
+            SELECT DISTINCT s.science_product_locator
+            FROM product_groups pg2
+            JOIN science_products_product_groups sppg on pg2.product_group_id = sppg.product_group_id
+            JOIN science_products_product_groups s on pg2.parent_product_group_id = s.product_group_id
+            WHERE sppg.product_group_id = %(pg_id)s;
+            """
+            cursor.execute(input_group_query, {"pg_id": product_group_id})
+            input_data = cursor.fetchall()
+            if input_data:
+                input_info = json.dumps({"input_group_locator": input_data[0][0]})
+                make_json = {**make_json, **input_info}
         finally:
             self.conn.close()
         return make_json

apps/cli/utilities/aat_wrest/aat_wrest/wrest.py

@@ -160,7 +160,10 @@ def determine_wrester(connection: MDDBConnector, args: argparse.Namespace):
     elif args.curator:
         data = WrestWorkflowMetadata(connection, spl=args.curator).wrest_curator()
     elif args.curator_products:
-        data = WrestWorkflowMetadata(connection, spl=args.curator_products).wrest_curator_products()
+        general_data = WrestWorkflowMetadata(connection, spl=args.curator_products).wrest_curator()
+        connection2 = MDDBConnector()
+        product_data = WrestWorkflowMetadata(connection2, spl=args.curator_products).wrest_curator_products()
+        data = {**general_data, **product_data}
     else:
         data = None