from collections import OrderedDict
import pymongo
_HISTORY_COUNTER_COLLECTION = "history_counters"
_JOB_ID_COUNTER_NAME = "jobid"
def _jobid_counter(db):
counters = db[_HISTORY_COUNTER_COLLECTION]
counters.create_index([("counter_name", pymongo.ASCENDING)], unique=True)
return counters
[docs]
def get_jobid(db):
"""
Ask MongoDB for a valid jobid.
All processing jobs should have a call to this function at the beginning
of the job script. Job ids are allocated by atomically incrementing a
named counter in the ``history_counters`` collection. The first id
allocated for a new database is 1.
:param db: database handle
:type db: top level database handle returned by a call to MongoClient.database
"""
counters = _jobid_counter(db)
counter = counters.find_one_and_update(
{"counter_name": _JOB_ID_COUNTER_NAME},
{"$inc": {"value": 1}},
upsert=True,
return_document=pymongo.ReturnDocument.AFTER,
)
return counter["value"]
[docs]
def bootstrap_history_jobid_counter(db):
"""Raise the job-id counter to the largest legacy history job id.
This is an explicit, idempotent deployment step for databases created
before job ids were allocated from ``history_counters``. Runtime job-id
allocation does not call this function or scan the history collection.
:param db: database handle
:type db: top level database handle returned by a call to
MongoClient.database
:return: the resulting counter high-water mark
"""
legacy_max_document = db.history.find_one(
{"jobid": {"$type": ["int", "long"]}},
projection={"_id": False, "jobid": True},
sort=[("jobid", pymongo.DESCENDING)],
)
legacy_max = legacy_max_document["jobid"] if legacy_max_document else 0
counters = _jobid_counter(db)
current_value = {"$ifNull": ["$value", 0]}
counter = counters.find_one_and_update(
{"counter_name": _JOB_ID_COUNTER_NAME},
[
{
"$set": {
"counter_name": _JOB_ID_COUNTER_NAME,
"value": {"$max": [current_value, legacy_max]},
}
}
],
upsert=True,
return_document=pymongo.ReturnDocument.AFTER,
)
return counter["value"]
def _reserve_requested_jobid(db, requested_jobid):
"""Atomically reserve an explicit job id or the next available value."""
counters = _jobid_counter(db)
next_jobid = {"$add": [{"$ifNull": ["$value", 0]}, 1]}
previous = counters.find_one_and_update(
{"counter_name": _JOB_ID_COUNTER_NAME},
[
{
"$set": {
"counter_name": _JOB_ID_COUNTER_NAME,
"value": {
"$cond": [
{"$gt": [requested_jobid, next_jobid]},
requested_jobid,
next_jobid,
]
},
}
}
],
upsert=True,
return_document=pymongo.ReturnDocument.BEFORE,
)
previous_value = previous["value"] if previous is not None else 0
allocated_jobid = previous_value + 1
if requested_jobid > allocated_jobid:
return requested_jobid, True
return allocated_jobid, False
def _antelope_pf_to_dict(pf):
"""Return one complete AntelopePf node without changing the source."""
result = OrderedDict()
for simple_key in pf.keys():
result[simple_key] = pf.get(simple_key)
for table_key in pf.tbl_keys():
result[table_key] = pf.get_tbl(table_key)
for branch_key in pf.arr_keys():
result[branch_key] = _antelope_pf_to_dict(pf.get_branch(branch_key))
return result
[docs]
def pfbranch_to_dict(pf, key):
"""
Recursive function to convert a single branch in an AntelopePf to a python dict.
This function utilizes recursion to follow a chain of arbitrary length of
branches defined in an AntelopePf object. Result is a dict with a chain of
dicts of the same length. i.e. if AntelopePf has 3 levels of branches
the dict will have a 3 levels of associative arrays keyed by the same
branch names as the Arr items in the original Pf file. Note this should
be called from the top level one branch at a time. i.e. for the parent
AntelopePf this function should be called once for each returned key by
pf.arr_keys().
Note that at each level Tbl& sections of the original pf are parsed to
be converted to lists of strings with each line of the Tbl section being
one string in the list.
:param pf: is an AntelopePf. Recursive calls use get_branch outputs
that return one of these.
:param key: key used to access the branch requested
:type key: string
:return: python dict translation of AntelopePf branch structure
:raise: RunTime errors are possible from the ccore methods that are called.
"""
return _antelope_pf_to_dict(pf.get_branch(key))
[docs]
class basic_history_data:
"""
This is a pure data object that if it were written in C could be
defined as a struct. It holds the data used to define the
parameters for a given algorithm.
"""
def __init__(self, job):
self.jobid = job
self.algorithm = "UNDEFINED"
self.param_type = "NONE"
self.params = {} # dict with unspecified content dumped to collection
[docs]
def load_algorithm_args(self, alg, argdict):
"""
Loads parameters defined to a set of function arguments.
Simple algorithms without a lot of parameters often simply need a
set of argument values. Here we require this to be defined by
a set of key:value pairs that map to dict. We also consider this
the lowest common denominator for a parameter definition so make
it a part of the base class.
:param alg: This should be a string defining the algorithm being registered.
:param argdict: This should be a dict of key:value pairs defining input
parameters. For algorithms defined at the top level by a python
function this should match the names of parameters in the arg list.
For C++ functions wrapped with pybind11 it should match the arg
keys defined in the wrappers. The key string will be used for
key:value pair in BSON written to MongoDB.
"""
self.algorithm = alg
self.param_type = "dict"
self.params = argdict
[docs]
class pf_history_data(basic_history_data):
"""
Loads history data container with data from an AntelopePf object.
mspasspy.ccore.utility defines the AntelopePf object that is an option for parameter
inputs. The file structure is identical to the
Antelope Pf file syntax. The API to an AntelopePF is not, however,
the same as the python bindings in Antelope as it handles Tbl
and Arr sections completely differently more in line with alternatives
like YAML. This method converts the data in an AntelopePf to a python
dict that can be dumped directly to MongoDB with pymongo's insert
methods. Converting the MongoDB document back to a pf structure requires
the inverse operator that does not exist, but should eventually be
created if this approach sees extensive use.
"""
def __init__(self, job, alg, pf):
"""
Basic constructor for this subclass.
This constructor applies the construction is initialization model
of oop. The AntelopePf pointed to by pf is parsed in this constuctor
to file and set the params dict and other attributes.
:param job: jobid (integer) normally should be preceded by call to get_jobid function.
:param alg: string defining a name assigned to the algorithm field
:param pf: AntelopePf object to be parsed and posted.
"""
self.jobid = job
self.algorithm = alg
self.param_type = "AntelopePf"
self.params = _antelope_pf_to_dict(pf)
[docs]
class HistoryLogger:
"""
Base class for generic, global history/provenance preservation in MsPASS.
The main concept of this object that a pymongo script to run a processing
job would create this object or one of it's children to preserve the
global run parameters for the a processing sequence. We limit that to
mean a sequence of processing algorithms that have a set of predefined
parameters that control their behaviour. The global parameters are
preserved in a special collection in MongoDB we give the (fixed) name
of "history". Processing steps are written as an ordered ``steps``
array. Legacy documents that keyed steps by algorithm name remain
readable, but that legacy representation is no longer written.
"""
def __init__(self, db, job=0):
"""
Basic constructor.
This is currently the only constructor for this object. It creates
a handle to MongoDB and sets a unique integer key called jobid.
Calling this constructor will guaranetee the jobid will be unique.
:param db: is a top level handle to a MongoDB server created by
calling the database method of a MongoClient instance.
:param job: job can be used to manually set the jobid. We use
an atomic counter comparable to lastid in the Antelope/Datascope
database. Hence if the input value of job is less than the next
value allocated by the counter, the jobid is silently set to that
allocated value.
(default is 0 which automatically allocates from the counter)
Users can get the actual value set from the jobid variable after
successful creation of this object.
:raises TypeError: if job is not a non-boolean integer.
"""
if isinstance(job, bool) or not isinstance(job, int):
raise TypeError("HistoryLogger job must be a non-boolean integer")
self.history_collection = db.history
# Check the input job id for validity and use get_jobid if needed
if job == 0:
self.jobid = get_jobid(db)
else:
self.jobid, requested_jobid_was_used = _reserve_requested_jobid(db, job)
if not requested_jobid_was_used:
print(
"HistoryLogger(Warning): input jobid=",
job,
" was invalid. Set jobid=",
self.jobid,
)
self.history_chain = [] # create empty container for history record
[docs]
def register(self, alg, partype, params):
"""
Register an algorithm's signature to preserve processing history.
Each algorithm in a processing chains should be registered by this
mechanism before starting a mspass processing chain. The
register method should be called in the order in which the algorithms are
applied.
:param alg: is the name of the algorithm that will be run. Assumed to be
a string.
:param partype: defines the format of the data defining input parameters
to this algorithm (Must be either 'dict' or 'AntelopePf')
:param params: is the actual input data. Actual type of this data
this arg references will depend up partype. partype defines
the type of the object expect (dict in this case means a python
dict object)
:raise: Throws a RuntimeError with a message if partype is not on the
list of supported parameter types
"""
if partype == "dict":
bhd = basic_history_data(self.jobid)
bhd.load_algorithm_args(alg, params)
self.history_chain.append(bhd)
elif partype == "AntelopePf":
pfhis = pf_history_data(self.jobid, alg, params)
self.history_chain.append(pfhis)
else:
raise RuntimeError(
"HistoryLogger (Warning): Unsupported parameter type=" + partype
)
[docs]
def save(self):
"""
Save the contents to the history collection.
The document stores processing steps in registration order in a
``steps`` array. Algorithm names are values, not document keys, so
repeated invocations are preserved.
"""
doc = {
"jobid": self.jobid,
"steps": [
{
"algorithm": d.algorithm,
"param_type": d.param_type,
"params": d.params,
}
for d in self.history_chain
],
}
if hasattr(self, "_loaded_document_id"):
self.history_collection.replace_one({"_id": self._loaded_document_id}, doc)
else:
self.history_collection.insert_one(doc)
[docs]
@classmethod
def load(cls, db, jobid):
"""
Load a saved processing history.
Both the ordered ``steps`` representation and legacy documents with
one top-level key per algorithm are accepted. The returned object can
be saved to migrate a legacy record to the ordered representation.
:param db: database handle
:param jobid: identifier of the history document to load
:return: a populated :class:`HistoryLogger`, or ``None`` when the job
does not exist
"""
history_collection = db.history
doc = history_collection.find_one({"jobid": jobid})
if doc is None:
return None
if isinstance(doc.get("steps"), list):
steps = doc["steps"]
else:
steps = [step for key, step in doc.items() if key not in {"_id", "jobid"}]
result = cls.__new__(cls)
result.history_collection = history_collection
result.jobid = doc["jobid"]
result.history_chain = []
if "_id" in doc:
result._loaded_document_id = doc["_id"]
for step in steps:
history_data = basic_history_data(result.jobid)
history_data.algorithm = step["algorithm"]
history_data.param_type = step["param_type"]
history_data.params = step["params"]
result.history_chain.append(history_data)
return result
