Core sampler: add ability to accept either a project code or an SDM name

• Core sampler now accepts either a project code or an SDM name

apps/cli/utilities/core_sampler/core_sampler/core_sampler.py

@@ -27,7 +27,10 @@ The core sampler outputs an SQL file you can use to load the core sample into a
 import argparse
 
 import psycopg2 as pg
-import psycopg2.extras as extras
+from psycopg2 import extras
+
+# pylint: disable=C0103, E0402, R0201
+
 from pycapo import CapoConfig
 
 from .database import PGTable
@@ -38,7 +41,6 @@ from .row_writer import (
     UniquifyingRowWriter,
 )
 
-
 # stolen shamelessly from aat_wrest
 class MDDBConnector:
     """Use this connection to interrogate this science product locator"""
@@ -89,16 +91,33 @@ class CoreSampler:
         self.writer = TopologicallySortingRowWriter(UniquifyingRowWriter(PostgresCopyRowWriter()))
         self.visited = set()
 
-    def sample(self, project_code: str):
+    def sample_project(self, project_code: str):
         """
-        Sample the database.
+        Get project metadata from the archive database.
+
+        :param project_code: project code of interest
+        :return:
         """
+
         # the first time through, we select from the projects table and get that row
         projects = self.table("projects")
         requested = projects.fetch({"project_code": project_code})
         self.save(requested)
         self.writer.close()
 
+    def sample_eb(self, exec_block_id: int):
+        """
+        Pull execution block metadata from the archive database.
+
+        :param exec_block_id: execution_block_id of interest
+        :return:
+        """
+        ebs = self.table("execution_blocks")
+
+        requested = ebs.fetch({"execution_block_id": exec_block_id})
+        self.save(requested)
+        self.writer.close()
+
     def save(self, rows: "RowSet"):
         """
         Save some rows, and then go and fetch their related rows and save them too, recursively.
@@ -138,10 +157,25 @@ class CoreSampler:
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("project_code", type=str, help="Project code to start core sampling from")
-
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "-p", "--project_code", type=str, nargs=1, help="Project code from which to start core sampling", action="store"
+    )
+    group.add_argument(
+        "-e",
+        "--exec_block_id",
+        type=int,
+        nargs=1,
+        help="execution_block_id of SDM for which to pull a core sample",
+        action="store",
+    )
     ns = parser.parse_args()
-    CoreSampler(MDDBConnector()).sample(ns.project_code)
+
+    if ns.project_code:
+        CoreSampler(MDDBConnector()).sample_project(ns.project_code[0])
+    elif ns.exec_block_id:
+        CoreSampler(MDDBConnector()).sample_eb(int(ns.exec_block_id[0]))
+
     return 0