# -*- coding: utf-8 -*- """ Data extraction and transfer from miniSEED files """ from __future__ import (absolute_import, division, print_function, unicode_literals) from future.builtins import * # NOQA import re import os import ctypes import bisect from logging import getLogger from collections import namedtuple from io import BytesIO from obspy import read as mseed_read from obspy.core import UTCDateTime from obspy.core.stream import Stream from .msriterator import MSR_iterator logger = getLogger(__name__) class NoDataError(Exception): """ Error raised when no data is found """ pass class RequestLimitExceededError(Exception): """ Error raised when the amount of data exceeds the configured limit """ pass class ExtractedDataSegment(object): """ There are a few different forms that a chunk of extracted data can take, so we return a wrapped object that exposes a simple, consistent API for the handler to use. """ def write(self, wfile): """ Write the data to the given file-like object """ raise NotImplementedError() def get_num_bytes(self): """ Return the number of bytes in the segment """ raise NotImplementedError() def get_src_name(self): """ Return the name of the data source """ raise NotImplementedError() class MSRIDataSegment(ExtractedDataSegment): """ Segment of data from a MSR_iterator """ def __init__(self, msri, sample_rate, start_time, end_time, src_name): """ :param msri: A `MSR_iterator` :param sample_rate: Sample rate of the data :param start_time: A `UTCDateTime` giving the start of the requested data :param end_time: A `UTCDateTime` giving the end of the requested data :param src_name: Name of the data source for logging """ self.msri = msri self.sample_rate = sample_rate self.start_time = start_time self.end_time = end_time self.src_name = src_name def write(self, wfile): msrstart = self.msri.get_startepoch() msrend = self.msri.get_endepoch() reclen = self.msri.msr.contents.reclen sepoch = self.start_time.timestamp eepoch = self.end_time.timestamp # Process records that intersect with request time window if msrstart < eepoch and msrend > sepoch: # Trim record if coverage and partial overlap with request if self.sample_rate > 0 and (msrstart < self.start_time or msrend > self.end_time): logger.debug("Trimming record %s @ %s" % (self.src_name, self.msri.get_starttime())) tr = mseed_read(BytesIO(ctypes.string_at(self.msri.msr.contents.record, reclen)), format="MSEED")[0] tr.trim(self.start_time, self.end_time) st = Stream(traces=[tr]) st.write(wfile, format="MSEED") # Otherwise, write un-trimmed record else: # Construct to avoid copying the data, supposedly logger.debug("Writing full record %s @ %s" % (self.src_name, self.msri.get_starttime())) out = (ctypes.c_char * reclen).from_address( ctypes.addressof(self.msri.msr.contents.record.contents)) wfile.write(out.raw) def get_num_bytes(self): return self.msri.msr.contents.reclen def get_src_name(self): return self.src_name class FileDataSegment(ExtractedDataSegment): """ Segment of data that comes directly from a data file """ def __init__(self, filename, start_byte, num_bytes, src_name): """ :param filename: Name of data file :param start_byte: Return data starting from this offset :param num_bytes: Length of data to return :param src_name: Name of the data source for logging """ self.filename = filename self.start_byte = start_byte self.num_bytes = num_bytes self.src_name = src_name def write(self, wfile): with open(self.filename, "rb") as f: f.seek(self.start_byte) raw_data = f.read(self.num_bytes) wfile.write(raw_data) def get_num_bytes(self): return self.num_bytes def get_src_name(self): return self.src_name class MiniseedDataExtractor(object): """ Component for extracting, trimming, and validating data. """ def __init__(self, dp_replace=None, request_limit=0): """ :param dp_replace: optional tuple of (regex, replacement) indicating the location of data files :param request_limit: optional limit (in bytes) on how much data can be extracted at once """ if dp_replace: self.dp_replace_re = re.compile(dp_replace[0]) self.dp_replace_sub = dp_replace[1] else: self.dp_replace_re = None self.dp_replace_sub = None self.request_limit = request_limit def handle_trimming(self, stime, etime, NRow): """ Get the time & byte-offsets for the data in time range (stime, etime). This is done by finding the smallest section of the data in row that falls within the desired time range and is identified by the timeindex field of row. :returns: [(start time, start offset, trim_boolean), (end time, end offset, trim_boolean)] """ etime = UTCDateTime(NRow.requestend) row_stime = UTCDateTime(NRow.starttime) row_etime = UTCDateTime(NRow.endtime) # If we need a subset of the this block, trim it accordingly block_start = int(NRow.byteoffset) block_end = block_start + int(NRow.bytes) if stime > row_stime or etime < row_etime: tix = [x.split("=>") for x in NRow.timeindex.split(",")] if tix[-1][0] == 'latest': tix[-1] = [str(row_etime.timestamp), block_end] to_x = [float(x[0]) for x in tix] s_index = bisect.bisect_right(to_x, stime.timestamp) - 1 if s_index < 0: s_index = 0 e_index = bisect.bisect_right(to_x, etime.timestamp) off_start = int(tix[s_index][1]) if e_index >= len(tix): e_index = -1 off_end = int(tix[e_index][1]) return ([to_x[s_index], off_start, stime > row_stime], [to_x[e_index], off_end, etime < row_etime],) else: return ([row_stime.timestamp, block_start, False], [row_etime.timestamp, block_end, False]) def extract_data(self, index_rows): """ Perform the data extraction. :param index_rows: requested data, as produced by `HTTPServer_RequestHandler.fetch_index_rows` :yields: sequence of `ExtractedDataSegment`s """ # Pre-scan the index rows: # 1) Build processed list for extraction # 2) Check if the request is small enough to satisfy # Note: accumulated estimate of output bytes will be equal to or higher than actual output total_bytes = 0 request_rows = [] Request = namedtuple('Request', ['srcname', 'filename', 'starttime', 'endtime', 'triminfo', 'bytes', 'samplerate']) try: for NRow in index_rows: srcname = "_".join(NRow[:4]) filename = NRow.filename logger.debug("EXTRACT: src=%s, file=%s, bytes=%s, rate:%s" % (srcname, filename, NRow.bytes, NRow.samplerate)) starttime = UTCDateTime(NRow.requeststart) endtime = UTCDateTime(NRow.requestend) triminfo = self.handle_trimming(starttime, endtime, NRow) total_bytes += triminfo[1][1] - triminfo[0][1] if self.request_limit > 0 and total_bytes > self.request_limit: raise RequestLimitExceededError("Result exceeds limit of %d bytes" % self.request_limit) if self.dp_replace_re: filename = self.dp_replace_re.sub(self.dp_replace_sub, filename) if not os.path.exists(filename): raise Exception("Data file does not exist: %s" % filename) request_rows.append(Request(srcname=srcname, filename=filename, starttime=starttime, endtime=endtime, triminfo=triminfo, bytes=NRow.bytes, samplerate=NRow.samplerate)) logger.debug("EXTRACT: src=%s, file=%s, bytes=%s, rate:%s" % (srcname, filename, NRow.bytes, NRow.samplerate)) except Exception as err: import traceback traceback.print_exc() raise Exception("Error accessing data index: %s" % str(err)) # Error if request matches no data if total_bytes == 0: raise NoDataError() # Get & return the actual data for NRow in request_rows: logger.debug("Extracting %s (%s - %s) from %s" % (NRow.srcname, NRow.starttime, NRow.endtime, NRow.filename)) # Iterate through records in section if only part of the section is needed if NRow.triminfo[0][2] or NRow.triminfo[1][2]: for msri in MSR_iterator(filename=NRow.filename, startoffset=NRow.triminfo[0][1], dataflag=False): offset = msri.get_offset() # Done if we are beyond end offset if offset >= NRow.triminfo[1][1]: break yield MSRIDataSegment(msri, NRow.samplerate, NRow.starttime, NRow.endtime, NRow.srcname) # Check for passing end offset if (offset + msri.msr.contents.reclen) >= NRow.triminfo[1][1]: break # Otherwise, return the entire section else: yield FileDataSegment(NRow.filename, NRow.triminfo[0][1], NRow.bytes, NRow.srcname)