add dump operations

csst_common/pipeline.py

@@ -16,6 +16,7 @@ import sys
 import json
 import subprocess
 import os
+import copy
 import shutil
 import traceback
 import warnings
@@ -25,6 +26,7 @@ from astropy.io import fits
 
 import csst_dfs_client
 import csst_fs
+from csst_fs import s3_fs
 
 from .ccds import CCDS
 from .utils import retry
@@ -33,6 +35,8 @@ from .logger import get_logger
 from .status import CsstStatus, CsstResult
 from .fits import s3_options
 
+s3_options = csst_fs.s3_config.load_s3_options()
+
 
 def print_directory_tree(directory="."):
     for root, dirs, files in os.walk(directory):
@@ -180,22 +184,34 @@ class Pipeline:
             print("Failed to clean output directory!")
             print_directory_tree(d)
 
+    # file operations
     @staticmethod
     def move(file_src: str, file_dst: str) -> str:
         """Move file `file_src` to `file_dist`."""
         return shutil.move(file_src, file_dst)
 
     @staticmethod
-    def copy(file_src: str, file_dst: str) -> str:
-        """Move file `file_src` to `file_dist`."""
-        return shutil.copy(file_src, file_dst)
+    def copy_nas_file(file_src: str, file_dst: str) -> str:
+        """Copy NAS file `file_src` to `file_dist`."""
+        shutil.copy(file_src, file_dst)
+        assert os.path.exists(
+            file_dst
+        ), f"Failed to copy NAS file {file_src} to {file_dst}!"
+        return file_dst
 
-    @property
-    def summarize(self):
-        """Summarize this run."""
-        t_stop: Time = Time.now()
-        t_cost: float = (t_stop - self.t_start).value * 86400.0
-        self.logger.info(f"Total cost: {t_cost:.1f} sec")
+    @staticmethod
+    def dump_oss_file(rpath: str, lpath: str) -> str:
+        """Copy OSS file `file_src` to `file_dist`."""
+        s3_fs.get(rpath, lpath)
+        assert os.path.exists(lpath), f"Failed to dump OSS file {rpath} to {lpath}!"
+        return lpath
+
+    # @property
+    # def summarize(self):
+    #     """Summarize this run."""
+    #     t_stop: Time = Time.now()
+    #     t_cost: float = (t_stop - self.t_start).value * 86400.0
+    #     self.logger.info(f"Total cost: {t_cost:.1f} sec")
 
     def clean_output(self):
         """Clean output directory."""
@@ -209,40 +225,72 @@ class Pipeline:
         """Create new file in output directory."""
         return os.path.join(self.dir_output, file_name)
 
-    def download_oss_file(self, oss_file_path: str, dir_dst: str = None) -> str:
-        """Download an OSS file from OSS to output directory."""
-        if dir_dst is None:
-            dir_dst = self.dir_output
-        local_file_path = os.path.join(dir_dst, os.path.basename(oss_file_path))
-        csst_fs.s3_fs.get(oss_file_path, local_file_path, s3_options=s3_options)
-        assert os.path.exists(
-            local_file_path
-        ), f"Failed to download {oss_file_path} to {local_file_path}"
-        return local_file_path
-
-    def abspath(self, file_path: str) -> str:
-        """Return absolute path of `file_path`."""
-        if file_path.__contains__(":"):
-            # it's an OSS file path
-            assert self.use_oss, "USE_OSS must be True to use OSS file path!"
-            # download OSS file to output directory
-            local_file_path = self.download_oss_file(file_path)
-            # return local file path
-            return local_file_path
+    @staticmethod
+    def dump_file(remote_file_path: str, local_file_path: str = None) -> str:
+        """Copy file `remote_file_path` to `local_file_path`."""
+        is_oss = remote_file_path.__contains__(":")
+        if is_oss:
+            local_file_path = Pipeline.dump_oss_file(remote_file_path, local_file_path)
         else:
-            # it's a NAS file path
-            if file_path.startswith("CSST"):
-                # DFS
-                return os.path.join(self.dfs_root, file_path)
-            else:
-                # CCDS
-                return os.path.join(self.ccds_root, file_path)
+            local_file_path = Pipeline.copy_nas_file(remote_file_path, local_file_path)
+        return local_file_path
 
-    def download_dfs_file(self, file_path: str, dir_dst: str = None) -> str:
+    # abspath
+    def convert_to_abspath_for_dfs_recs(self, dfs_recs: list[dict]) -> list[dict]:
+        """Convert `file_path` to absolute path for DFS."""
+        dfs_recs_abs = copy.deepcopy(dfs_recs)
+        for rec in dfs_recs_abs:
+            rec["file_path"] = os.path.join(self.dfs_root, rec["file_path"])
+        return dfs_recs_abs
+
+    def convert_to_abspath_for_ccds_refs(self, ccds_refs: dict) -> dict:
+        """Convert `file_path` to absolute path for CCDS."""
+        ccds_refs_abs = copy.deepcopy(ccds_refs)
+        for ref_name, ref_path in ccds_refs_abs.items():
+            ccds_refs_abs[ref_name] = os.path.join(self.ccds_root, ref_path)
+        return ccds_refs_abs
+
+    def dump_dfs_recs(
+        self, dfs_recs_abs: list[dict], dir_dump: str = None
+    ) -> list[dict]:
+        """Copy DFS files to output directory."""
+        # set default dir_dump to output directory
+        if dir_dump is None:
+            dir_dump = os.path.join(self.dir_output, "dfs")
+            self.mkdir(dir_dump)
+        # dump data to dir_dump
+        dfs_recs_dump = copy.deepcopy(dfs_recs_abs)
+        for rec in dfs_recs_dump:
+            remote_file_path = rec["file_path"]
+            local_file_path = os.path.join(dir_dump, os.path.basename(remote_file_path))
+            # copy DFS file to local_file_path
+            self.dump_file(remote_file_path, local_file_path)
+            rec["file_path"] = local_file_path
+        return dfs_recs_dump
+
+    def dump_ccds_refs(self, refs: dict, dir_dump: str = None) -> dict:
+        """Copy raw file from CCDS to output directory."""
+        # set default dir_dump to output directory
+        if dir_dump is None:
+            dir_dump = os.path.join(self.dir_output, "ccds")
+            self.mkdir(dir_dump)
+        # dump data to dir_dump
+        ccds_refs_dump = copy.deepcopy(refs)
+        for ref_name, ref_path in ccds_refs_dump.items():
+            remote_file_path = ref_path
+            local_file_path = os.path.join(dir_dump, os.path.basename(remote_file_path))
+            # copy DFS file to local_file_path
+            self.dump_file(remote_file_path, local_file_path)
+            ccds_refs_dump[ref_name] = local_file_path
+        return ccds_refs_dump
+
+    def dump_dfs_file(self, file_path: str, dir_dst: str = None) -> str:
         """Copy DFS file to output directory."""
+        # by default, dump file to output directory
         if dir_dst is None:
             dir_dst = self.dir_output
 
+        # if use OSS,
         if self.use_oss:
             # download OSS file to dst directory
             return self.download_oss_file(file_path, dir_dst)
@@ -252,7 +300,7 @@ class Pipeline:
             self.copy(self.abspath(file_path), local_file_path)
             return local_file_path
 
-    def download_ccds_refs(self, refs: dict, dir_dst: str = None) -> dict:
+    def dump_ccds_refs(self, refs: dict, dir_dst: str = None) -> dict:
         """Copy raw file from CCDS to output directory."""
         if dir_dst is None:
             dir_dst = self.dir_output