def _GetCursor(self, spec, timestamp, limit):
"""Create a mongo cursor based on the timestamp restriction."""
if timestamp == self.NEWEST_TIMESTAMP or timestamp is None:
collection = self.latest_collection
elif timestamp == self.ALL_TIMESTAMPS:
collection = self.versioned_collection
elif isinstance(timestamp, tuple):
collection = self.versioned_collection
start, end = timestamp
spec = {
"$and": [
dict(timestamp={"$gte": int(start)}),
dict(timestamp={"$lte": int(end)}), spec
]
}
else:
raise data_store.Error("Undefined timestamp specification.")
cursor = collection.find(spec).sort("timestamp", pymongo.DESCENDING)
if limit:
cursor = cursor.limit(limit)
return cursor
def RecreatePathing(self, pathing=None):
if not pathing:
pathing = config_lib.CONFIG.Get("Datastore.pathing")
try:
self.path_regexes = [re.compile(path) for path in pathing]
except re.error:
raise data_store.Error("Invalid regular expression in Datastore.pathing")
def CheckResponseStatus(response):
"""Catch error conditions from the response and raise them."""
# Common case exit early.
if response.status == rdfvalue.DataStoreResponse.Status.OK:
return response
elif (response.status ==
rdfvalue.DataStoreResponse.Status.AUTHORIZATION_DENIED):
raise access_control.UnauthorizedAccess(response.status_desc,
response.failed_subject)
elif response.status == rdfvalue.DataStoreResponse.Status.TIMEOUT_ERROR:
raise data_store.TimeoutError(response.status_desc)
elif response.status == rdfvalue.DataStoreResponse.Status.DATA_STORE_ERROR:
raise data_store.Error(response.status_desc)
raise data_store.Error("Unknown error %s" % response.status_desc)
def MultiSet(self, subject, values, timestamp=None, replace=True,
sync=True, to_delete=None, token=None):
"""Set multiple attributes' values for this subject in one operation."""
self.security_manager.CheckDataStoreAccess(token, [subject], "w")
if timestamp is None:
timestamp = time.time() * 1e6
# Prepare a mongo bulk insert for all the values.
documents = []
subject = utils.SmartUnicode(subject)
to_delete = set(to_delete or [])
latest = {}
# Build a document for each unique timestamp.
for attribute, sequence in values.items():
for value in sequence:
if isinstance(value, tuple):
value, entry_timestamp = value
else:
entry_timestamp = timestamp
if entry_timestamp is None:
entry_timestamp = timestamp
attribute = utils.SmartUnicode(attribute)
prefix = attribute.split(":", 1)[0]
document = dict(subject=subject, timestamp=int(entry_timestamp),
predicate=attribute, prefix=prefix)
_Encode(document, value)
documents.append(document)
latest[attribute] = document
# Replacing means to delete all versions of the attribute first.
if replace:
to_delete.add(attribute)
if to_delete:
self.DeleteAttributes(subject, to_delete, token=token)
# Just write using bulk insert mode.
if documents:
try:
self.versioned_collection.insert(documents, w=1 if sync else 0)
except errors.PyMongoError as e:
logging.error("Mongo Error %s", e)
raise data_store.Error(utils.SmartUnicode(e))
# Maintain the latest documents in the latest collection.
for attribute, document in latest.items():
document.pop("_id", None)
self.latest_collection.update(
dict(subject=subject, predicate=attribute, prefix=prefix),
document, upsert=True, w=1 if sync else 0)
def _TimestampToFilter(self, timestamp):
if timestamp == data_store.DataStore.ALL_TIMESTAMPS:
return None
if timestamp is None or timestamp == data_store.DataStore.NEWEST_TIMESTAMP:
# Latest value only
return row_filters.CellsColumnLimitFilter(1)
if isinstance(timestamp, tuple):
return row_filters.TimestampRangeFilter(
self._TimestampRangeFromTuple(timestamp))
raise data_store.Error("Invalid timestamp specification: %s." % timestamp)
def _MultiSet(self, subject, document, spec):
if document:
# First try to update the data
try:
result = self.collection.update(spec, {"$set": document},
upsert=False,
safe=True)
except errors.PyMongoError as e:
logging.error("Mongo Error %s", e)
raise data_store.Error(utils.SmartUnicode(e))
# If the document does not already exist, just save a new one
if not result["updatedExisting"]:
document["_id"] = URNEncode(EscapeKey(subject))
self.collection.save(document)
