diff --git a/bin/datastore.py b/bin/datastore.py index 57b52af..3da14b1 100644 --- a/bin/datastore.py +++ b/bin/datastore.py @@ -438,6 +438,32 @@ class Datastore: self.maybe_commit() + def save_raw_smsrc (self, records, fileinfo, filepath, filedata = None): + + with self.conn.cursor() as cursor: + hash = self.add_file(filepath, cursor) + + # Start by deleting any online data we may have for this sequence + # NOTE: Do I need to do this? + #self.del_hash("*online*", cursor) + + # The shots should already exist, e.g., from a P1 import + + values = [ (json.dumps(record), fileinfo["sequence"], record["shot"]) for record in records ] + + qry = """ + UPDATE raw_shots + SET meta = jsonb_set(meta, '{smsrc}', %s::jsonb, true) + WHERE sequence = %s AND point = %s; + """ + + cursor.executemany(qry, values) + + if filedata is not None: + self.save_file_data(filepath, json.dumps(filedata), cursor) + + self.maybe_commit() + def save_file_data(self, path, filedata, cursor = None): """ diff --git a/bin/import_smsrc.py b/bin/import_smsrc.py new file mode 100755 index 0000000..f634f6e --- /dev/null +++ b/bin/import_smsrc.py @@ -0,0 +1,76 @@ +#!/usr/bin/python3 + +""" +Import SmartSource data. + +For each survey in configuration.surveys(), check for new +or modified final gun header files and (re-)import them into the +database. +""" + +import os +import sys +import pathlib +import re +import configuration +import smsrc +from datastore import Datastore + +if __name__ == '__main__': + + print("Reading configuration") + surveys = configuration.surveys() + + print("Connecting to database") + db = Datastore() + db.connect() + + print("Reading surveys") + for survey in surveys: + print(f'Survey: {survey["id"]} ({survey["schema"]})') + + db.set_survey(survey["schema"]) + + try: + raw_smsrc = survey["raw"]["smsrc"] + except KeyError: + print("No SmartSource data configuration") + continue + + flags = 0 + if "flags" in raw_smsrc: + configuration.rxflags(raw_smsrc["flags"]) + + pattern = raw_smsrc["pattern"] + rx = re.compile(pattern["regex"], flags) + + for fileprefix in raw_smsrc["paths"]: + print(f"Path prefix: {fileprefix}") + + for globspec in raw_smsrc["globs"]: + for filepath in pathlib.Path(fileprefix).glob(globspec): + filepath = str(filepath) + print(f"Found {filepath}") + + if not db.file_in_db(filepath): + print("Importing") + + match = rx.match(os.path.basename(filepath)) + if not match: + error_message = f"File path not matching the expected format! ({filepath} ~ {pattern['regex']})" + print(error_message, file=sys.stderr) + print("This file will be ignored!") + continue + + file_info = dict(zip(pattern["captures"], match.groups())) + + smsrc_records = smsrc.from_file(filepath) + + print("Saving") + + db.save_raw_smsrc(smsrc_records, file_info, filepath) + else: + print("Already in DB") + + print("Done") + diff --git a/bin/runner.sh b/bin/runner.sh index 879ba78..980ff79 100755 --- a/bin/runner.sh +++ b/bin/runner.sh @@ -95,6 +95,9 @@ run $BINDIR/import_final_p111.py print_log "Import final P1/90" run $BINDIR/import_final_p190.py +print_log "Import SmartSource data" +run $BINDIR/import_smsrc.py + if [[ -z "$RUNNER_NOEXPORT" ]]; then print_log "Export system data" run $BINDIR/system_exports.py diff --git a/bin/smsrc.py b/bin/smsrc.py new file mode 100644 index 0000000..e6a91c7 --- /dev/null +++ b/bin/smsrc.py @@ -0,0 +1,95 @@ +#!/usr/bin/python3 + +""" +SmartSource parsing functions. +""" + +import mmap +import struct +from collections import namedtuple + +def _str (v): + return str(v, 'ascii').strip() + +def _tstamp (v): + return str(v) # TODO + +def _f10 (v): + return float(v)/10 + +def _ignore (v): + return None + +st_smartsource_header = struct.Struct(">6s 4s 30s 10s 2s 1s 17s 1s 1s 2s 2s 2s 2s 2s 4s 6s 5s 5s 6s 4s 88s") + +fn_smartsource_header = ( + _str, int, _str, int, int, _str, _tstamp, int, int, int, int, int, int, int, int, int, + float, float, float, int, _str +) + +SmartsourceHeader = namedtuple("SmartsourceHeader", "header blk_siz line shot mask trg_mode time src_number num_subarray num_guns num_active num_delta num_auto num_nofire spread volume avg_delta std_delta baro_press manifold spare") + +st_smartsource_gun = struct.Struct(">1s 2s 1s 1s 1s 1s 1s 6s 6s 4s 4s 4s 4s 4s") + +fn_smartsource_gun = ( + int, int, int, _str, _str, lambda v: v=="Y", _str, + _f10, _f10, _f10, _f10, + int, int, int +) + +SmartsourceGun = namedtuple("SmartsourceGun", "string gun source mode detect autofire spare aim_point firetime delay depth pressure volume filltime") + +SmartSourceRecord = namedtuple("SmartSourceRecord", "header guns") + +def safe_apply (iter): + def safe_fn (fn, v): + try: + return fn(v) + except ValueError: + return None + return [safe_fn(v[0], v[1]) for v in iter] + +def _check_chunk_size(chunk, size): + return len(chunk) == size + +def from_file(path): + + records = [] + + with open(path, "rb") as fd: + with mmap.mmap(fd.fileno(), length=0, access=mmap.ACCESS_READ) as buffer: + + while True: + + offset = buffer.find(b"*SMSRC") + + if offset == -1: + break + + buffer = buffer[offset:] + record, length = read_smartsource(buffer) + + if record is not None: + records.append(record) + + if length != 0: + buffer = buffer[length:] + else: + buffer = buffer[1:] + + return records + +def read_smartsource(buffer): + length = 0 + header = st_smartsource_header.unpack_from(buffer, 0) + length += st_smartsource_header.size + header = SmartsourceHeader(*safe_apply(zip(fn_smartsource_header, header))) + record = dict(header._asdict()) + record["guns"] = [] + + for _ in range(header.num_guns): + gun = st_smartsource_gun.unpack_from(buffer, length) + record["guns"].append(SmartsourceGun(*safe_apply(zip(fn_smartsource_gun, gun)))) + length += st_smartsource_gun.size + + return (record, length)