ingest.py

import os, sys
import argparse
import logging
from astropy.io import fits
import datetime
import shutil

from csst_dfs_api_local.common.db import DBClient
from csst_dfs_commons.utils.fits import get_header_value

log = logging.getLogger('csst-dfs-api-local')

def ingest():
    
    parser = argparse.ArgumentParser(prog=f"{sys.argv[0]}", description="ingest the local files")
    parser.add_argument('-i','--infile', dest="infile", help="a file or a directory")
    parser.add_argument('-m', '--copyfiles', dest="copyfiles", action='store_true', default=False, help="whether copy files after import")
    args = parser.parse_args(sys.argv[1:])
    
    import_root_dir = args.infile
    if import_root_dir is None or (not os.path.isfile(import_root_dir) and not os.path.isdir(import_root_dir)):
        parser.print_help()
        sys.exit(0)

    db = DBClient()
    if os.path.isfile(import_root_dir):
        log.info(f"prepare import {import_root_dir}")
        ingest_one(import_root_dir, db, args.copyfiles)
    if os.path.isdir(import_root_dir):
        for (path, _, file_names) in os.walk(import_root_dir):
            for filename in file_names:
                if filename.find(".fits") > 0:
                    file_full_path = os.path.join(path, filename)
                    log.info(f"prepare import {file_full_path}")
                    ingest_one(file_full_path, db, args.copyfiles)

    db.close()

def ingest_one(file_path, db, copyfiles):
    dest_root_dir = os.getenv("CSST_LOCAL_FILE_ROOT", "/opt/temp/csst")

    hdul = fits.open(file_path)
    fits_header = hdul[0].header
    img_header = hdul[1].header

    obs_id = fits_header["OBSID"]
    exp_start_time = f"{fits_header['DATE-OBS']} {fits_header['TIME-OBS']}"
    exp_time = fits_header['EXPTIME']
    
    module_id = fits_header["INSTRUME"]
    obs_type = fits_header["FILETYPE"]
    qc0_status = -1
    prc_status = -1
    time_now = datetime.datetime.now()
    create_time = time_now.strftime('%Y-%m-%d %H:%M:%S')

    facility_status_id = 0
    module_status_id = 0

    existed = db.exists("select * from t_observation where obs_id=?", (obs_id,))
    if not existed:
        db.execute("insert into t_observation \
            (obs_id,obs_time,exp_time,module_id,obs_type,facility_status_id, module_status_id, qc0_status, prc_status,create_time) \
            values (?,?,?,?,?,?,?,?,?,?)",
        (obs_id,exp_start_time,exp_time,module_id,obs_type,facility_status_id,module_status_id,qc0_status, prc_status,create_time))
        db.end()
    #level0
    detector = fits_header["DETECTOR"]
    filename = get_header_value("FILENAME", fits_header, os.path.basename(file_path))
    
    existed = db.exists(
            "select * from msc_level0_data where filename=?",
            (filename,)
        )
    if existed:
        log.warning('%s has already been imported' %(file_path, ))
        db.end()
        return    

    detector_status_id = 0

    file_full_path = file_path

    if copyfiles:
        file_dir = f"{dest_root_dir}/{module_id}/{obs_type.upper()}/{fits_header['EXPSTART']}/{obs_id}/MS"
        if not os.path.exists(file_dir):
            os.makedirs(file_dir)
        file_full_path = f"{file_dir}/{filename}.fits"

    level0_id = f"{obs_id}{detector}"  

    c = db.execute("insert into msc_level0_data \
        (level0_id, obs_id, detector_no, obs_type, obs_time, exp_time,detector_status_id, filename, file_path,qc0_status, prc_status,create_time) \
        values (?,?,?,?,?,?,?,?,?,?,?,?)",
        (level0_id, obs_id, detector, obs_type, exp_start_time, exp_time, detector_status_id, filename, file_full_path, qc0_status, prc_status,create_time))
    db.end()
    level0_id_id = db.last_row_id()
    #level0-header
    ra_obj = fits_header["RA_OBJ"]
    dec_obj = fits_header["DEC_OBJ"]
    crpix1 = img_header["CRPIX1"]
    crpix2 = img_header["CRPIX2"]
    crval1 = img_header["CRVAL1"]
    crval1 = img_header["CRVAL2"]
    ctype1 = img_header["CTYPE1"]
    ctype2 = img_header["CTYPE2"]
    cd1_1 = img_header["CD1_1"]
    cd1_2 = img_header["CD1_2"]
    cd2_1 = img_header["CD2_1"]
    cd2_2 = img_header["CD2_2"]
    db.execute("delete from msc_level0_header where id=?",(level0_id_id,))    
    db.execute("insert into msc_level0_header \
        (id, ra_obj, dec_obj, crpix1, crpix2, crval1, crval1, \
        ctype1, ctype2, cd1_1, cd1_2, cd2_1, cd2_2, create_time) \
        values (?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
        (level0_id_id, ra_obj, dec_obj, crpix1, crpix2, crval1, crval1, ctype1, ctype2, cd1_1, cd1_2, cd2_1, cd2_2, create_time))
    
    if copyfiles:
        #copy files
        shutil.copyfile(file_path, file_full_path)

    db.end()

    print(f"{file_path} imported")

if __name__ == "__main__":
    ingest()