1667

pending completion

Build # 1667

Build Type

Pull #257

travis-ci

Committed by

web-flow

Commit Message

<a href="https://github.com/LCOGT/banzai/commit/<a class=hub.com/LCOGT/banzai/commit/66c3e06c81cc1e61874d2c38f3c69a543c738131">66c3e06c8<a href="https://github.com/LCOGT/banzai/commit/66c3e06c81cc1e61874d2c38f3c69a543c738131">">Merge </a><a class="double-link" href="https://github.com/LCOGT/banzai/commit/<a class="double-link" href="https://github.com/LCOGT/banzai/commit/21dfc6b85419377ad239b4dba416dcecb849ee77">21dfc6b85</a>">21dfc6b85</a><a href="https://github.com/LCOGT/banzai/commit/66c3e06c81cc1e61874d2c38f3c69a543c738131"> into 1bf5aec56">1bf5aec56</a>

Pull Request Pull Request #257: Seek cursor to beginning of buffer before returning.

Run Details

27 of 27 new or added lines in 9 files covered. (100.0%)

1770 of 2818 relevant lines covered (62.81%)

0.63 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

42.4

/banzai/dbs.py

""" dbs.py: Database utility functions for banzai

    This is built around the SQLAlchemy ORM

Author
    Curtis McCully (cmccully@lcogt.net)

October 2015
"""
import os.path
import logging
import datetime
from dateutil.parser import parse

import numpy as np
import requests
from sqlalchemy import create_engine, pool, type_coerce, cast
from sqlalchemy.orm import sessionmaker
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean, CHAR, JSON, UniqueConstraint
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.sql.expression import true
from contextlib import contextmanager
from banzai.utils import file_utils

Base = declarative_base()

logger = logging.getLogger('banzai')


@contextmanager
def get_session(db_address):
    """
    Get a connection to the database.

    Returns
    -------
    session: SQLAlchemy Database Session
    """
    # Build a new engine for each session. This makes things thread safe.
    engine = create_engine(db_address, poolclass=pool.NullPool)
    Base.metadata.bind = engine

    # We don't use autoflush typically. I have run into issues where SQLAlchemy would try to flush
    # incomplete records causing a crash. None of the queries here are large, so it should be ok.
    db_session = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)
    session = db_session()
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()


class CalibrationImage(Base):
    """
    Master Calibration Image Database Record

    This defines the calimages table. We use this to keep track of the master calibration frames.
    Typically these are bias, darks, and flat field frames. These are indexed by dayobs to make it
    easy to find the closest calibration frame.
    """
    __tablename__ = 'calimages'
    id = Column(Integer, primary_key=True, autoincrement=True)
    type = Column(String(50), index=True)
    filename = Column(String(100), unique=True)
    filepath = Column(String(150))
    frameid = Column(Integer, nullable=True)
    dateobs = Column(DateTime, index=True)
    datecreated = Column(DateTime, index=True)
    instrument_id = Column(Integer, ForeignKey("instruments.id"), index=True)
    is_master = Column(Boolean)
    is_bad = Column(Boolean)
    good_until = Column(DateTime, default=datetime.datetime(3000, 1, 1))
    good_after = Column(DateTime, default=datetime.datetime(1000, 1, 1))
    attributes = Column(JSON)


class Instrument(Base):
    """
    Instrument Database Record

    This defines the instruments table.
    """
    __tablename__ = 'instruments'
    id = Column(Integer, primary_key=True, autoincrement=True)
    site = Column(String(15), ForeignKey('sites.id'), index=True, nullable=False)
    camera = Column(String(50), index=True, nullable=False)
    type = Column(String(100))
    name = Column(String(100), index=True, nullable=False)
    __table_args__ = (UniqueConstraint('site', 'camera', 'name', name='instrument_constraint'),)


class Site(Base):
    """
    Site Database Record

    This defines the sites table structure
    """
    __tablename__ = 'sites'
    id = Column(String(15), primary_key=True)
    timezone = Column(Integer)


class ProcessedImage(Base):
    __tablename__ = 'processedimages'
    id = Column(Integer, primary_key=True, autoincrement=True)
    filename = Column(String(100), index=True)
    frameid = Column(Integer, nullable=True)
    checksum = Column(CHAR(32), index=True, default='0'*32)
    success = Column(Boolean, default=False)
    tries = Column(Integer, default=0)


def parse_configdb(configdb_address):
    """
    Parse the contents of the configdb.

    Parameters
    ----------
    configdb_address : str
                      URL of the configdb, must be inside LCOGT VPN

    Returns
    -------
    sites : list of dicts
            each site dictionary contains a timezone.
    instruments : list of dicts
              each camera dictionary contains a site, instrument code, and camera type.
    """
    results = requests.get(configdb_address).json()['results']
    instruments = []
    sites = []
    # this will be removed when configdb is updated
    CAMERAS_FOR_INSTRUMENTS = {'nres01': 'fa09', 'nres02': 'fa17', 'nres03': 'fa13', 'nres04': 'fa18',
                               'floyds01': 'en06', 'floyds02': 'en12'}
    # end of hotfix
    for site in results:
        sites.append({'code': site['code'], 'timezone': site['timezone']})
        for enc in site['enclosure_set']:
            for tel in enc['telescope_set']:
                for ins in tel['instrument_set']:
                    sci_cam = ins.get('science_camera')
                    if sci_cam is not None:
                        instrument = {'site': site['code'],
                                      'camera': sci_cam['code'],
                                      'name': ins.get('code'),
                                      'type': sci_cam['camera_type']['code']}
                        # hotfix for configdb
                        if not instrument['name']:
                            instrument['name'] = instrument['camera']
                        if instrument['name'] in CAMERAS_FOR_INSTRUMENTS:
                            instrument['camera'] = CAMERAS_FOR_INSTRUMENTS[instrument['name']]
                        instruments.append(instrument)
    return sites, instruments


def add_instrument(instrument, db_address):
    with get_session(db_address=db_address) as db_session:
        equivalence_criteria = {'site': instrument['site'],
                                'camera': instrument['camera'],
                                'name': instrument['name']}
        record_attributes = {'site': instrument['site'],
                             'camera': instrument['camera'],
                             'name': instrument['name'],
                             'type': instrument['type']}

        instrument_record = add_or_update_record(db_session, Instrument, equivalence_criteria, record_attributes)
        db_session.commit()
    return instrument_record


def add_or_update_record(db_session, table_model, equivalence_criteria, record_attributes):
    """
    Add a record to the database if it does not exist or update the record if it does exist.

    Parameters
    ----------
    db_session : SQLAlchemy database session
                 session must be active

    table_model : SQLAlchemy Base
                  The class representation of the table of interest

    equivalence_criteria : dict
                           record attributes that need to match for the records to be considered
                           the same

    record_attributes : dict
                        record attributes that will be set/updated

    Returns
    -------
    record : SQLAlchemy Base
             The object representation of the added/updated record

    Notes
    -----
    The added/updated record is added to the database but not committed. You need to call
    db_session.commit() to write the changes to the database.
    """
    query = true()
    for key in equivalence_criteria.keys():
        query &= getattr(table_model, key) == equivalence_criteria[key]
    record = db_session.query(table_model).filter(query).first()
    if record is None:
        record = table_model(**record_attributes)
        db_session.add(record)
    for attribute in record_attributes:
        setattr(record, attribute, record_attributes[attribute])
    return record


class SiteMissingException(Exception):
    pass


def query_for_instrument(db_address, site, camera, name=None):
    # Short circuit
    if None in [site, camera]:
        return None
    with get_session(db_address=db_address) as db_session:
        criteria = (Instrument.site == site) & (Instrument.camera == camera)
        if name is not None:
            criteria &= Instrument.name == name
        instrument = db_session.query(Instrument).filter(criteria).order_by(Instrument.id.desc()).first()
    return instrument


def save_calibration_info(output_file, image, db_address):
    # Store the information into the calibration table
    # Check and see if the bias file is already in the database
    with get_session(db_address=db_address) as db_session:
        output_filename = os.path.basename(output_file)
        record_attributes = {'type': image.obstype.upper(),
                             'filename': output_filename,
                             'filepath': os.path.dirname(output_file),
                             'dateobs': image.dateobs,
                             'datecreated': image.datecreated,
                             'instrument_id': image.instrument.id,
                             'is_master': image.is_master,
                             'is_bad': image.is_bad,
                             'frameid': image.frame_id,
                             'attributes': {}}
        for attribute in image.grouping_criteria:
            record_attributes['attributes'][attribute] = str(getattr(image, attribute))

        add_or_update_record(db_session, CalibrationImage, {'filename': output_filename}, record_attributes)

        db_session.commit()


def get_processed_image(path, db_address):
    # TODO: add support for AWS path styles
    filename = os.path.basename(path)
    with get_session(db_address=db_address) as db_session:
        processed_image = add_or_update_record(db_session, ProcessedImage, {'filename': filename},
                                               {'filename': filename})
        db_session.commit()
    return processed_image


def commit_processed_image(processed_image, db_address):
    with get_session(db_address=db_address) as db_session:
        db_session.add(processed_image)
        db_session.commit()


def save_processed_image(path, md5, db_address):
    filename = os.path.basename(path)
    output_record = get_processed_image(filename, db_address)
    output_record.success = True
    output_record.checksum = md5
    commit_processed_image(output_record, db_address)


def get_timezone(site, db_address):
    with get_session(db_address=db_address) as db_session:
        site_list = db_session.query(Site).filter(Site.id == site).all()
    if len(site_list) == 0:
        raise SiteMissingException
    return site_list[0].timezone


def get_instruments_at_site(site, db_address):
    with get_session(db_address=db_address) as db_session:
        query = (Instrument.site == site)
        instruments = db_session.query(Instrument).filter(query).all()
    return instruments


def get_instrument_by_id(id, db_address):
    with get_session(db_address=db_address) as db_session:
        instrument = db_session.query(Instrument).filter(Instrument.id==id).first()
    return instrument


def cal_record_to_file_info(record):
    if record is None:
        return None
    # if the filepath in the DB is not set, make sure we set the path to None.
    if record.filepath is None:
        path = None
    else:
        path = os.path.join(record.filepath, record.filename)
    file_info = {'frameid': record.frameid,
                 'path': path,
                 'filename': record.filename}
    return file_info


def get_master_cal_record(image, calibration_type, master_selection_criteria, db_address,
                          use_only_older_calibrations=False):
    calibration_criteria = CalibrationImage.type == calibration_type.upper()
    calibration_criteria &= CalibrationImage.instrument_id == image.instrument.id
    calibration_criteria &= CalibrationImage.is_master.is_(True)
    calibration_criteria &= CalibrationImage.is_bad.is_(False)

    for criterion in master_selection_criteria:
        # We have to cast to strings according to the sqlalchemy docs for version 1.3:
        # https://docs.sqlalchemy.org/en/latest/core/type_basics.html?highlight=json#sqlalchemy.types.JSON
        calibration_criteria &= cast(CalibrationImage.attributes[criterion], String) ==\
                                type_coerce(str(getattr(image, criterion)), JSON)

    # During real-time reduction, we want to avoid using different master calibrations for the same block,
    # therefore we make sure the the calibration frame used was created before the block start time
    if use_only_older_calibrations and getattr(image, 'block_start') is not None:
        calibration_criteria &= CalibrationImage.datecreated < image.block_start

    calibration_criteria &= CalibrationImage.good_after <= image.dateobs
    calibration_criteria &= CalibrationImage.good_until >= image.dateobs

    with get_session(db_address=db_address) as db_session:
        calibration_images = db_session.query(CalibrationImage).filter(calibration_criteria).all()

    # Exit if no calibration file found
    if len(calibration_images) == 0:
        return None

    # Find the closest date
    date_deltas = np.abs(np.array([i.dateobs - image.dateobs for i in calibration_images]))
    closest_calibration_image = calibration_images[np.argmin(date_deltas)]

    return closest_calibration_image


def get_master_cal(image, calibration_type, master_selection_criteria, db_address,
                   use_only_older_calibrations=False):
    return cal_record_to_file_info(get_master_cal_record(image, calibration_type, master_selection_criteria, db_address,
                                                         use_only_older_calibrations=use_only_older_calibrations))


def get_individual_cal_records(instrument, calibration_type, min_date: str, max_date: str, db_address: str,
                               include_bad_frames: bool = False):
    calibration_criteria = CalibrationImage.instrument_id == instrument.id
    calibration_criteria &= CalibrationImage.type == calibration_type.upper()
    calibration_criteria &= CalibrationImage.dateobs >= parse(min_date).replace(tzinfo=None)
    calibration_criteria &= CalibrationImage.dateobs <= parse(max_date).replace(tzinfo=None)
    calibration_criteria &= CalibrationImage.is_master == False

    calibration_criteria &= CalibrationImage.is_master == False

    if not include_bad_frames:
        calibration_criteria &= CalibrationImage.is_bad == False

    with get_session(db_address=db_address) as db_session:
        image_records = db_session.query(CalibrationImage).filter(calibration_criteria).all()

    return image_records


def get_individual_cal_frames(instrument, calibration_type, min_date: str, max_date: str, db_address: str,
                              include_bad_frames: bool = False):
    image_records = get_individual_cal_records(instrument, calibration_type, min_date, max_date, db_address,
                                               include_bad_frames=include_bad_frames)
    return [cal_record_to_file_info(record) for record in image_records]


def mark_frame(filename, mark_as, db_address):
    set_is_bad_to = True if mark_as == "bad" else False
    logger.debug("Setting the is_bad parameter for {filename} to {set_is_bad_to}".format(
        filename=filename, set_is_bad_to=set_is_bad_to))
    with get_session(db_address=db_address) as db_session:
        # First check to make sure the image is in the database
        image = db_session.query(CalibrationImage).filter(CalibrationImage.filename == filename).first()
        if image is None:
            logger.error("Frame {filename} not found in database, exiting".format(filename=filename))
            return
        if image.is_bad is set_is_bad_to:
            logger.error("The is_bad parameter for {filename} is already set to {is_bad}, exiting".format(
                filename=filename, is_bad=image.is_bad))
            return
        equivalence_criteria = {'filename': filename}
        record_attributes = {'filename': filename,
                             'is_bad': set_is_bad_to}
        add_or_update_record(db_session, CalibrationImage, equivalence_criteria, record_attributes)
        db_session.commit()


def create_db(db_address):
    # Create an engine for the database
    engine = create_engine(db_address)

    # Create all tables in the engine
    # This only needs to be run once on initialization.
    Base.metadata.create_all(engine)


def populate_instrument_tables(db_address, configdb_address):
    """
    Populate the instrument table from the configdb

    Parameters
    ----------
    db_address : str
                 sqlalchemy address to the database of the form
                 mysql://username:password@localhost/test

    configdb_address : str
                       URL of the configdb

    Notes
    -----
    This only works inside the LCOGT VPN. This should be run at least when a new camera is
    added to the network.
    """
    sites, instruments = parse_configdb(configdb_address=configdb_address)
    with get_session(db_address=db_address) as db_session:
        for site in sites:
            add_or_update_record(db_session, Site, {'id': site['code']},
                                 {'id': site['code'], 'timezone': site['timezone']})

    for instrument in instruments:
        add_instrument(instrument, db_address)

1	""" dbs.py: Database utility functions for banzai
2
3	This is built around the SQLAlchemy ORM
4
5	Author
6	Curtis McCully (cmccully@lcogt.net)
7
8	October 2015
9	"""
10	import os.path	1✔
11	import logging	1✔
12	import datetime	1✔
13	from dateutil.parser import parse	1✔
14
15	import numpy as np	1✔
16	import requests	1✔
17	from sqlalchemy import create_engine, pool, type_coerce, cast	1✔
18	from sqlalchemy.orm import sessionmaker	1✔
19	from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean, CHAR, JSON, UniqueConstraint	1✔
20	from sqlalchemy.ext.declarative import declarative_base	1✔
21	from sqlalchemy.sql.expression import true	1✔
22	from contextlib import contextmanager	1✔
23	from banzai.utils import file_utils	1✔
24
25	Base = declarative_base()	1✔
26
27	logger = logging.getLogger('banzai')	1✔
28
29
30	@contextmanager	1✔
31	def get_session(db_address):
32	"""
33	Get a connection to the database.
34
35	Returns
36	-------
37	session: SQLAlchemy Database Session
38	"""
39	# Build a new engine for each session. This makes things thread safe.
40	engine = create_engine(db_address, poolclass=pool.NullPool)	1✔
41	Base.metadata.bind = engine	1✔
42
43	# We don't use autoflush typically. I have run into issues where SQLAlchemy would try to flush
44	# incomplete records causing a crash. None of the queries here are large, so it should be ok.
45	db_session = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)	1✔
46	session = db_session()	1✔
47	try:	1✔
48	yield session	1✔
49	session.commit()	1✔
50	except:	×
51	session.rollback()	×
52	raise	×
53	finally:
54	session.close()	1✔
55
56
57	class CalibrationImage(Base):	1✔
58	"""
59	Master Calibration Image Database Record
60
61	This defines the calimages table. We use this to keep track of the master calibration frames.
62	Typically these are bias, darks, and flat field frames. These are indexed by dayobs to make it
63	easy to find the closest calibration frame.
64	"""
65	__tablename__ = 'calimages'	1✔
66	id = Column(Integer, primary_key=True, autoincrement=True)	1✔
67	type = Column(String(50), index=True)	1✔
68	filename = Column(String(100), unique=True)	1✔
69	filepath = Column(String(150))	1✔
70	frameid = Column(Integer, nullable=True)	1✔
71	dateobs = Column(DateTime, index=True)	1✔
72	datecreated = Column(DateTime, index=True)	1✔
73	instrument_id = Column(Integer, ForeignKey("instruments.id"), index=True)	1✔
74	is_master = Column(Boolean)	1✔
75	is_bad = Column(Boolean)	1✔
76	good_until = Column(DateTime, default=datetime.datetime(3000, 1, 1))	1✔
77	good_after = Column(DateTime, default=datetime.datetime(1000, 1, 1))	1✔
78	attributes = Column(JSON)	1✔
79
80
81	class Instrument(Base):	1✔
82	"""
83	Instrument Database Record
84
85	This defines the instruments table.
86	"""
87	__tablename__ = 'instruments'	1✔
88	id = Column(Integer, primary_key=True, autoincrement=True)	1✔
89	site = Column(String(15), ForeignKey('sites.id'), index=True, nullable=False)	1✔
90	camera = Column(String(50), index=True, nullable=False)	1✔
91	type = Column(String(100))	1✔
92	name = Column(String(100), index=True, nullable=False)	1✔
93	__table_args__ = (UniqueConstraint('site', 'camera', 'name', name='instrument_constraint'),)	1✔
94
95
96	class Site(Base):	1✔
97	"""
98	Site Database Record
99
100	This defines the sites table structure
101	"""
102	__tablename__ = 'sites'	1✔
103	id = Column(String(15), primary_key=True)	1✔
104	timezone = Column(Integer)	1✔
105
106
107	class ProcessedImage(Base):	1✔
108	__tablename__ = 'processedimages'	1✔
109	id = Column(Integer, primary_key=True, autoincrement=True)	1✔
110	filename = Column(String(100), index=True)	1✔
111	frameid = Column(Integer, nullable=True)	1✔
112	checksum = Column(CHAR(32), index=True, default='0'*32)	1✔
113	success = Column(Boolean, default=False)	1✔
114	tries = Column(Integer, default=0)	1✔
115
116
117	def parse_configdb(configdb_address):	1✔
118	"""
119	Parse the contents of the configdb.
120
121	Parameters
122	----------
123	configdb_address : str
124	URL of the configdb, must be inside LCOGT VPN
125
126	Returns
127	-------
128	sites : list of dicts
129	each site dictionary contains a timezone.
130	instruments : list of dicts
131	each camera dictionary contains a site, instrument code, and camera type.
132	"""
133	results = requests.get(configdb_address).json()['results']	×
134	instruments = []	×
135	sites = []	×
136	# this will be removed when configdb is updated
137	CAMERAS_FOR_INSTRUMENTS = {'nres01': 'fa09', 'nres02': 'fa17', 'nres03': 'fa13', 'nres04': 'fa18',	×
138	'floyds01': 'en06', 'floyds02': 'en12'}
139	# end of hotfix
140	for site in results:	×
141	sites.append({'code': site['code'], 'timezone': site['timezone']})	×
142	for enc in site['enclosure_set']:	×
143	for tel in enc['telescope_set']:	×
144	for ins in tel['instrument_set']:	×
145	sci_cam = ins.get('science_camera')	×
146	if sci_cam is not None:	×
147	instrument = {'site': site['code'],	×
148	'camera': sci_cam['code'],
149	'name': ins.get('code'),
150	'type': sci_cam['camera_type']['code']}
151	# hotfix for configdb
152	if not instrument['name']:	×
153	instrument['name'] = instrument['camera']	×
154	if instrument['name'] in CAMERAS_FOR_INSTRUMENTS:	×
155	instrument['camera'] = CAMERAS_FOR_INSTRUMENTS[instrument['name']]	×
156	instruments.append(instrument)	×
157	return sites, instruments	×
158
159
160	def add_instrument(instrument, db_address):	1✔
161	with get_session(db_address=db_address) as db_session:	×
162	equivalence_criteria = {'site': instrument['site'],	×
163	'camera': instrument['camera'],
164	'name': instrument['name']}
165	record_attributes = {'site': instrument['site'],	×
166	'camera': instrument['camera'],
167	'name': instrument['name'],
168	'type': instrument['type']}
169
170	instrument_record = add_or_update_record(db_session, Instrument, equivalence_criteria, record_attributes)	×
171	db_session.commit()	×
172	return instrument_record	×
173
174
175	def add_or_update_record(db_session, table_model, equivalence_criteria, record_attributes):	1✔
176	"""
177	Add a record to the database if it does not exist or update the record if it does exist.
178
179	Parameters
180	----------
181	db_session : SQLAlchemy database session
182	session must be active
183
184	table_model : SQLAlchemy Base
185	The class representation of the table of interest
186
187	equivalence_criteria : dict
188	record attributes that need to match for the records to be considered
189	the same
190
191	record_attributes : dict
192	record attributes that will be set/updated
193
194	Returns
195	-------
196	record : SQLAlchemy Base
197	The object representation of the added/updated record
198
199	Notes
200	-----
201	The added/updated record is added to the database but not committed. You need to call
202	db_session.commit() to write the changes to the database.
203	"""
204	query = true()	1✔
205	for key in equivalence_criteria.keys():	1✔
206	query &= getattr(table_model, key) == equivalence_criteria[key]	1✔
207	record = db_session.query(table_model).filter(query).first()	1✔
208	if record is None:	1✔
209	record = table_model(**record_attributes)	1✔
210	db_session.add(record)	1✔
211	for attribute in record_attributes:	1✔
212	setattr(record, attribute, record_attributes[attribute])	1✔
213	return record	1✔
214
215
216	class SiteMissingException(Exception):	1✔
217	pass	1✔
218
219
220	def query_for_instrument(db_address, site, camera, name=None):	1✔
221	# Short circuit
222	if None in [site, camera]:	×
223	return None	×
224	with get_session(db_address=db_address) as db_session:	×
225	criteria = (Instrument.site == site) & (Instrument.camera == camera)	×
226	if name is not None:	×
227	criteria &= Instrument.name == name	×
228	instrument = db_session.query(Instrument).filter(criteria).order_by(Instrument.id.desc()).first()	×
229	return instrument	×
230
231
232	def save_calibration_info(output_file, image, db_address):	1✔
233	# Store the information into the calibration table
234	# Check and see if the bias file is already in the database
235	with get_session(db_address=db_address) as db_session:	×
236	output_filename = os.path.basename(output_file)	×
237	record_attributes = {'type': image.obstype.upper(),	×
238	'filename': output_filename,
239	'filepath': os.path.dirname(output_file),
240	'dateobs': image.dateobs,
241	'datecreated': image.datecreated,
242	'instrument_id': image.instrument.id,
243	'is_master': image.is_master,
244	'is_bad': image.is_bad,
245	'frameid': image.frame_id,
246	'attributes': {}}
247	for attribute in image.grouping_criteria:	×
248	record_attributes['attributes'][attribute] = str(getattr(image, attribute))	×
249
250	add_or_update_record(db_session, CalibrationImage, {'filename': output_filename}, record_attributes)	×
251
252	db_session.commit()	×
253
254
255	def get_processed_image(path, db_address):	1✔
256	# TODO: add support for AWS path styles
257	filename = os.path.basename(path)	×
258	with get_session(db_address=db_address) as db_session:	×
259	processed_image = add_or_update_record(db_session, ProcessedImage, {'filename': filename},	×
260	{'filename': filename})
261	db_session.commit()	×
262	return processed_image	×
263
264
265	def commit_processed_image(processed_image, db_address):	1✔
266	with get_session(db_address=db_address) as db_session:	×
267	db_session.add(processed_image)	×
268	db_session.commit()	×
269
270
271	def save_processed_image(path, md5, db_address):	1✔
272	filename = os.path.basename(path)	×
273	output_record = get_processed_image(filename, db_address)	×
274	output_record.success = True	×
275	output_record.checksum = md5	×
276	commit_processed_image(output_record, db_address)	×
277
278
279	def get_timezone(site, db_address):	1✔
280	with get_session(db_address=db_address) as db_session:	×
281	site_list = db_session.query(Site).filter(Site.id == site).all()	×
282	if len(site_list) == 0:	×
283	raise SiteMissingException	×
284	return site_list[0].timezone	×
285
286
287	def get_instruments_at_site(site, db_address):	1✔
288	with get_session(db_address=db_address) as db_session:	×
289	query = (Instrument.site == site)	×
290	instruments = db_session.query(Instrument).filter(query).all()	×
291	return instruments	×
292
293
294	def get_instrument_by_id(id, db_address):	1✔
295	with get_session(db_address=db_address) as db_session:	×
296	instrument = db_session.query(Instrument).filter(Instrument.id==id).first()	×
297	return instrument	×
298
299
300	def cal_record_to_file_info(record):	1✔
301	if record is None:	×
302	return None	×
303	# if the filepath in the DB is not set, make sure we set the path to None.
304	if record.filepath is None:	×
305	path = None	×
306	else:
307	path = os.path.join(record.filepath, record.filename)	×
308	file_info = {'frameid': record.frameid,	×
309	'path': path,
310	'filename': record.filename}
311	return file_info	×
312
313
314	def get_master_cal_record(image, calibration_type, master_selection_criteria, db_address,	1✔
315	use_only_older_calibrations=False):
316	calibration_criteria = CalibrationImage.type == calibration_type.upper()	×
317	calibration_criteria &= CalibrationImage.instrument_id == image.instrument.id	×
318	calibration_criteria &= CalibrationImage.is_master.is_(True)	×
319	calibration_criteria &= CalibrationImage.is_bad.is_(False)	×
320
321	for criterion in master_selection_criteria:	×
322	# We have to cast to strings according to the sqlalchemy docs for version 1.3:
323	# https://docs.sqlalchemy.org/en/latest/core/type_basics.html?highlight=json#sqlalchemy.types.JSON
324	calibration_criteria &= cast(CalibrationImage.attributes[criterion], String) ==\	×
325	type_coerce(str(getattr(image, criterion)), JSON)
326
327	# During real-time reduction, we want to avoid using different master calibrations for the same block,
328	# therefore we make sure the the calibration frame used was created before the block start time
329	if use_only_older_calibrations and getattr(image, 'block_start') is not None:	×
330	calibration_criteria &= CalibrationImage.datecreated < image.block_start	×
331
332	calibration_criteria &= CalibrationImage.good_after <= image.dateobs	×
333	calibration_criteria &= CalibrationImage.good_until >= image.dateobs	×
334
335	with get_session(db_address=db_address) as db_session:	×
336	calibration_images = db_session.query(CalibrationImage).filter(calibration_criteria).all()	×
337
338	# Exit if no calibration file found
339	if len(calibration_images) == 0:	×
340	return None	×
341
342	# Find the closest date
343	date_deltas = np.abs(np.array([i.dateobs - image.dateobs for i in calibration_images]))	×
344	closest_calibration_image = calibration_images[np.argmin(date_deltas)]	×
345
346	return closest_calibration_image	×
347
348
349	def get_master_cal(image, calibration_type, master_selection_criteria, db_address,	1✔
350	use_only_older_calibrations=False):
351	return cal_record_to_file_info(get_master_cal_record(image, calibration_type, master_selection_criteria, db_address,	×
352	use_only_older_calibrations=use_only_older_calibrations))
353
354
355	def get_individual_cal_records(instrument, calibration_type, min_date: str, max_date: str, db_address: str,	1✔
356	include_bad_frames: bool = False):
357	calibration_criteria = CalibrationImage.instrument_id == instrument.id	×
358	calibration_criteria &= CalibrationImage.type == calibration_type.upper()	×
359	calibration_criteria &= CalibrationImage.dateobs >= parse(min_date).replace(tzinfo=None)	×
360	calibration_criteria &= CalibrationImage.dateobs <= parse(max_date).replace(tzinfo=None)	×
361	calibration_criteria &= CalibrationImage.is_master == False	×
362
363	calibration_criteria &= CalibrationImage.is_master == False	×
364
365	if not include_bad_frames:	×
366	calibration_criteria &= CalibrationImage.is_bad == False	×
367
368	with get_session(db_address=db_address) as db_session:	×
369	image_records = db_session.query(CalibrationImage).filter(calibration_criteria).all()	×
370
371	return image_records	×
372
373
374	def get_individual_cal_frames(instrument, calibration_type, min_date: str, max_date: str, db_address: str,	1✔
375	include_bad_frames: bool = False):
376	image_records = get_individual_cal_records(instrument, calibration_type, min_date, max_date, db_address,	×
377	include_bad_frames=include_bad_frames)
378	return [cal_record_to_file_info(record) for record in image_records]	×
379
380
381	def mark_frame(filename, mark_as, db_address):	1✔
382	set_is_bad_to = True if mark_as == "bad" else False	×
383	logger.debug("Setting the is_bad parameter for {filename} to {set_is_bad_to}".format(	×
384	filename=filename, set_is_bad_to=set_is_bad_to))
385	with get_session(db_address=db_address) as db_session:	×
386	# First check to make sure the image is in the database
387	image = db_session.query(CalibrationImage).filter(CalibrationImage.filename == filename).first()	×
388	if image is None:	×
389	logger.error("Frame {filename} not found in database, exiting".format(filename=filename))	×
390	return	×
391	if image.is_bad is set_is_bad_to:	×
392	logger.error("The is_bad parameter for {filename} is already set to {is_bad}, exiting".format(	×
393	filename=filename, is_bad=image.is_bad))
394	return	×
395	equivalence_criteria = {'filename': filename}	×
396	record_attributes = {'filename': filename,	×
397	'is_bad': set_is_bad_to}
398	add_or_update_record(db_session, CalibrationImage, equivalence_criteria, record_attributes)	×
399	db_session.commit()	×
400
401
402	def create_db(db_address):	1✔
403	# Create an engine for the database
404	engine = create_engine(db_address)	1✔
405
406	# Create all tables in the engine
407	# This only needs to be run once on initialization.
408	Base.metadata.create_all(engine)	1✔
409
410
411	def populate_instrument_tables(db_address, configdb_address):	1✔
412	"""
413	Populate the instrument table from the configdb
414
415	Parameters
416	----------
417	db_address : str
418	sqlalchemy address to the database of the form
419	mysql://username:password@localhost/test
420
421	configdb_address : str
422	URL of the configdb
423
424	Notes
425	-----
426	This only works inside the LCOGT VPN. This should be run at least when a new camera is
427	added to the network.
428	"""
429	sites, instruments = parse_configdb(configdb_address=configdb_address)	×
430	with get_session(db_address=db_address) as db_session:	×
431	for site in sites:	×
432	add_or_update_record(db_session, Site, {'id': site['code']},	×
433	{'id': site['code'], 'timezone': site['timezone']})
434
435	for instrument in instruments:	×
436	add_instrument(instrument, db_address)	×

LCOGT / banzai / 1667

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous