reffits.py 7.28 KB
Newer Older
Wei Shoulin's avatar
Wei Shoulin committed
1
import logging
Xie Zhou's avatar
update    
Xie Zhou committed
2
import os
Wei Shoulin's avatar
Wei Shoulin committed
3
import time, datetime
Xie Zhou's avatar
update    
Xie Zhou committed
4
5
6
7
import shutil
from astropy.io import fits

from ..common.db import DBClient
Wei Shoulin's avatar
Wei Shoulin committed
8
from ..common.utils import *
Wei Shoulin's avatar
Wei Shoulin committed
9
10

log = logging.getLogger('csst')
Wei Shoulin's avatar
Wei Shoulin committed
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class RefFitsApi(object):
    REF_FITS_BIAS = "bias"
    REF_FITS_FLAT = "flat"
    REF_FITS_DARK = "dark"
    REF_FITS_SKY = "sky"
    REF_FITS_ARC = "arc"

    REF_FITS_TYPES = [REF_FITS_BIAS, REF_FITS_FLAT, REF_FITS_DARK, REF_FITS_SKY, REF_FITS_ARC]

    def __init__(self, sub_system = "ifs"):
        self.sub_system = sub_system
        self.root_dir = os.getenv("CSST_LOCAL_FILE_ROOT", "/opt/temp/csst")
        self.check_dir()
        self.db = DBClient()

    def check_dir(self):
        if not os.path.exists(self.root_dir):
            os.mkdir(self.root_dir)
            log.info("using [%s] as root directory", self.root_dir)
        if not os.path.exists(os.path.join(self.root_dir, "refs")):
            os.mkdir(os.path.join(self.root_dir, "refs"))

    def find(self, **kwargs):
        '''
        parameter kwargs:
Wei Shoulin's avatar
Wei Shoulin committed
36
37
            obs_time = [int],
            file_name = [str],
Wei Shoulin's avatar
Wei Shoulin committed
38
            ccd_num = [int],
Wei Shoulin's avatar
Wei Shoulin committed
39
40
            exp_time = (start, end),
            status = [int],
Wei Shoulin's avatar
Wei Shoulin committed
41
42
            ref_type = [str]

Wei Shoulin's avatar
Wei Shoulin committed
43
        return list of reference's files records
Wei Shoulin's avatar
Wei Shoulin committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
        '''
        obs_time = get_parameter(kwargs, "obs_time")
        file_name = get_parameter(kwargs, "file_name")
        exp_time = get_parameter(kwargs, "exp_time", (None, format_time_ms(time.time())))
        ccd_num = get_parameter(kwargs, "ccd_num")
        status = get_parameter(kwargs, "status")
        ref_type = get_parameter(kwargs, "ref_type")
        
        sql = []

        sql.append("select * from ifs_ref_fits where exp_time<='" + exp_time[1] + "'")

        if exp_time[0] is not None:
            sql.append(" and exp_time>='" + exp_time[0] + "'")
        if obs_time is not None:
            sql.append(" and obs_time=" + obs_time)
        if ccd_num is not None:
            sql.append(" and ccd_num=" + ccd_num)
        if ref_type is not None:
            sql.append(" and ref_type='" + ref_type + "'")           
        if status is not None:
            sql.append(" and status=" + status)
Xie Zhou's avatar
update    
Xie Zhou committed
66

Wei Shoulin's avatar
Wei Shoulin committed
67
68
69
70
71
72
73
74
        if file_name:
            sql = ["select * from ifs_ref_fits where filename='" + file_name + "'"]
        sql.append(" order by exp_time desc")
        _, r = self.db.select_many("".join(sql))

        return r

    def get(self, **kwargs):
Wei Shoulin's avatar
Wei Shoulin committed
75
76
        '''query database, return a record as dict

Xie Zhou's avatar
update    
Xie Zhou committed
77
        parameter kwargs:
Wei Shoulin's avatar
Wei Shoulin committed
78
79
80
81
82
83
84
85
            fits_id = [int] 

        return dict or None
        '''
        fits_id = get_parameter(kwargs, "fits_id", -1)
        r = self.db.select_one(
            "select * from ifs_ref_fits where id=?", (fits_id,))
        return r
Xie Zhou's avatar
update    
Xie Zhou committed
86

Wei Shoulin's avatar
Wei Shoulin committed
87
    def read(self, **kwargs):
Xie Zhou's avatar
update    
Xie Zhou committed
88
        '''
Wei Shoulin's avatar
Wei Shoulin committed
89
        parameter kwargs:
Wei Shoulin's avatar
Wei Shoulin committed
90
91
            fits_id = [int],
            file_path = [str], 
Wei Shoulin's avatar
Wei Shoulin committed
92
93
94
95
96
97
98
99
100
101
102
103
104
105
            chunk_size = [int]

        yield bytes of fits file
        '''
        fits_id = get_parameter(kwargs, "fits_id")
        file_path = get_parameter(kwargs, "file_path")

        if fits_id is None and file_path is None:
            raise Exception("fits_id or file_path need to be defined")

        if fits_id is not None:
            r = self.db.select_one(
                "select * from ifs_ref_fits where id=?", (fits_id))
            if r is not None:
Wei Shoulin's avatar
Wei Shoulin committed
106
                file_path = r["file_path"]
Wei Shoulin's avatar
Wei Shoulin committed
107
108

        if file_path is not None:
Wei Shoulin's avatar
Wei Shoulin committed
109
110
            chunk_size = get_parameter(kwargs, "chunk_size", 20480)
            return yield_file_bytes(os.path.join(self.root_dir, file_path), chunk_size)
Wei Shoulin's avatar
Wei Shoulin committed
111
112
113
114

    def update_status(self, **kwargs):
        '''
        parameter kwargs:
Wei Shoulin's avatar
Wei Shoulin committed
115
116
            fits_id = [int],
            status = [int]
Wei Shoulin's avatar
Wei Shoulin committed
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
        '''

        fits_id = get_parameter(kwargs, "fits_id")
        status = get_parameter(kwargs, "status")

        existed = self.db.exists(
            "select * from ifs_ref_fits where id=?",
            (fits_id,)
        )
        if existed:
            log.warning('%s not found' %(fits_id, ))
            return
        self.db.execute(
            'update ifs_ref_fits set status=? where id=?',
            (status, fits_id)
        )  
        self.db.end() 

    def import2db(self, **kwargs):
        '''
        parameter kwargs:
            file_path = [str]
            ref_type = [str]

            insert into database
        '''
       
Xie Zhou's avatar
update    
Xie Zhou committed
144
145
146
147
        file_path = get_parameter(kwargs, "file_path")

        if file_path is None:
            raise Exception("file_path need to be defined")
Wei Shoulin's avatar
Wei Shoulin committed
148
149
150
151
152
        
        file_full_path = os.path.join(self.root_dir, file_path)

        if not os.path.exists(file_full_path):
            raise Exception("%s not found"%(file_full_path))
Xie Zhou's avatar
update    
Xie Zhou committed
153

Wei Shoulin's avatar
Wei Shoulin committed
154
155
156
157
158
        file_name = os.path.basename(file_path)

        existed = self.db.exists(
            "select * from ifs_ref_fits where filename=?",
            (file_name,)
Xie Zhou's avatar
update    
Xie Zhou committed
159
        )
Wei Shoulin's avatar
Wei Shoulin committed
160
161
        if existed:
            log.warning('%s has already been imported' %(file_path, ))
Xie Zhou's avatar
update    
Xie Zhou committed
162
            return
Wei Shoulin's avatar
Wei Shoulin committed
163
164
165
166
167
168
169

        hu = fits.getheader(file_full_path)
        obs_time = hu['obst'] if 'obst' in hu else ''
        ccd_num = hu['ccd_num'] if 'ccd_num' in hu else 0
        exp_time = format_time_ms(time.time())

        ref_type = get_parameter(kwargs, "ref_type")
Xie Zhou's avatar
update    
Xie Zhou committed
170
        
Wei Shoulin's avatar
Wei Shoulin committed
171
172
173
174
175
176
177
178
179
180
181
        if ref_type is None:
            if 'flat' in file_name.lower():
                ref_type = 'flat'
            elif 'bias' in file_name.lower():
                ref_type = 'bias'
            elif 'hgar' in file_name.lower():
                ref_type = 'arc'
            elif 'sky' in file_name.lower():
                ref_type = 'sky'
            else:
                ref_type = ""
Xie Zhou's avatar
update    
Xie Zhou committed
182

Wei Shoulin's avatar
Wei Shoulin committed
183
        
Xie Zhou's avatar
update    
Xie Zhou committed
184
        self.db.execute(
Wei Shoulin's avatar
Wei Shoulin committed
185
186
187
            'INSERT INTO ifs_ref_fits (filename, obs_time, ccd_num, exp_time, file_path, ref_type, status, create_time) \
                VALUES(?,?,?,?,?,?,?,?)',
            (file_name, obs_time, ccd_num, exp_time, file_path, ref_type, 1, format_time_ms(time.time()))
Xie Zhou's avatar
update    
Xie Zhou committed
188
        )
Wei Shoulin's avatar
Wei Shoulin committed
189
190
191
192
193
194

        self.db.end()

        log.info("ref fits %s imported.", file_path)

    def write(self, **kwargs):
Wei Shoulin's avatar
Wei Shoulin committed
195
196
197
198
199
        ''' copy a local file to file storage, then reduce the header of fits file and insert a record into database
 
        parameter kwargs:
            file_path = [str]
        '''            
Wei Shoulin's avatar
Wei Shoulin committed
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
        file_path = get_parameter(kwargs, "file_path")

        new_file_dir = create_dir(os.path.join(self.root_dir, "refs"),
                self.sub_system, 
                "/".join([str(datetime.datetime.now().year),"%02d"%(datetime.datetime.now().month),"%02d"%(datetime.datetime.now().day)]))
        

        file_basename = os.path.basename(file_path)
        new_file_path = os.path.join(new_file_dir, file_basename)
        shutil.copyfile(file_path, new_file_path)

        self.import2db(file_path = new_file_path.replace(self.root_dir, '')[1:])
        

    def associate_raw(self, **kwargs):
        raw_fits_ids = get_parameter(kwargs, "raw_fits_ids")
        ref_fits_id = get_parameter(kwargs, "ref_fits_id")

        if raw_fits_ids is None or ref_fits_id is None:
            raise Exception("raw_fits_ids or ref_fits_id is None")

        sql = ['INSERT INTO ifs_raw_ref (fit_id, ref_id, create_time) values ']
        values = ["(%s,%s,%s)"%(i,ref_fits_id,format_time_ms(time.time())) for i in raw_fits_ids]

        self.db.execute(sql + ",".join())

        self.db.end()

        log.info("%s associate to %s imported.", raw_fits_ids, ref_fits_id)