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)