def query(
self,
criteria: Optional[Dict] = None,
properties: Union[Dict, List, None] = None,
sort: Optional[Dict[str, Union[Sort, int]]] = None,
skip: int = 0,
limit: int = 0,
) -> Iterator[Dict]:
"""
Queries the Store for a set of documents
Args:
criteria: PyMongo filter for documents to search in
properties: properties to return in grouped documents
sort: Dictionary of sort order for fields. Keys are field names and
values are 1 for ascending or -1 for descending.
skip: number documents to skip
limit: limit on total number of documents returned
"""
criteria = criteria if criteria else {}
if properties is not None:
if isinstance(properties, list):
properties = {p: 1 for p in properties}
substitute(properties, self.reverse_aliases)
lazy_substitute(criteria, self.reverse_aliases)
for d in self.store.query(properties=properties,
criteria=criteria,
sort=sort,
limit=limit,
skip=skip):
substitute(d, self.aliases)
yield d
def remove_docs(self, criteria: Dict):
"""
Remove docs matching the query dictionary
Args:
criteria: query dictionary to match
"""
# Update criteria and properties based on aliases
lazy_substitute(criteria, self.reverse_aliases)
self.store.remove_docs(criteria)
def count(self, criteria: Optional[Dict] = None) -> int:
"""
Counts the number of documents matching the query criteria
Args:
criteria: PyMongo filter for documents to count in
"""
criteria = criteria if criteria else {}
lazy_substitute(criteria, self.reverse_aliases)
return self.store.count(criteria)
def query_one(self, criteria=None, properties=None, **kwargs):
if isinstance(properties, list):
properties = {p: 1 for p in properties}
criteria = criteria if criteria else {}
substitute(properties, self.reverse_aliases)
lazy_substitute(criteria, self.reverse_aliases)
d = self.store.query_one(properties=properties, criteria=criteria, **kwargs)
substitute(d, self.aliases)
return d
def query(self, properties=None, criteria=None, **kwargs):
if isinstance(properties, list):
properties = {p: 1 for p in properties}
criteria = criteria if criteria else {}
substitute(properties, self.reverse_aliases)
lazy_substitute(criteria, self.reverse_aliases)
for d in self.store.query(properties, criteria, **kwargs):
substitute(d, self.aliases)
yield d
def groupby(self, keys, criteria=None, properties=None, **kwargs):
# Convert to a list
keys = keys if isinstance(keys, list) else [keys]
# Make the aliasing transformations on keys
keys = [self.aliases[k] if k in self.aliases else k for k in keys]
# Update criteria and properties based on aliases
criteria = criteria if criteria else {}
substitute(properties, self.reverse_aliases)
lazy_substitute(criteria, self.reverse_aliases)
return self.store.groupby(keys=keys, properties=properties, criteria=criteria, **kwargs)
def distinct(self,
field: str,
criteria: Optional[Dict] = None,
all_exist: bool = False) -> List:
"""
Get all distinct values for a field
Args:
field: the field(s) to get distinct values for
criteria: PyMongo filter for documents to search in
"""
criteria = criteria if criteria else {}
lazy_substitute(criteria, self.reverse_aliases)
# substitute forward
return self.store.distinct(self.aliases[field], criteria=criteria)
def groupby(
self,
keys: Union[List[str], str],
criteria: Optional[Dict] = None,
properties: Union[Dict, List, None] = None,
sort: Optional[Dict[str, Union[Sort, int]]] = None,
skip: int = 0,
limit: int = 0,
) -> Iterator[Tuple[Dict, List[Dict]]]:
"""
Simple grouping function that will group documents
by keys.
Args:
keys: fields to group documents
criteria: PyMongo filter for documents to search in
properties: properties to return in grouped documents
sort: Dictionary of sort order for fields. Keys are field names and
values are 1 for ascending or -1 for descending.
skip: number documents to skip
limit: limit on total number of documents returned
Returns:
generator returning tuples of (dict, list of docs)
"""
# Convert to a list
keys = keys if isinstance(keys, list) else [keys]
# Make the aliasing transformations on keys
keys = [self.aliases[k] if k in self.aliases else k for k in keys]
# Update criteria and properties based on aliases
criteria = criteria if criteria else {}
if properties is not None:
if isinstance(properties, list):
properties = {p: 1 for p in properties}
substitute(properties, self.reverse_aliases)
lazy_substitute(criteria, self.reverse_aliases)
return self.store.groupby(keys=keys,
properties=properties,
criteria=criteria,
skip=skip,
limit=limit)
def distinct(self, key, criteria=None, all_exist=True, **kwargs):
if isinstance(key, list):
criteria = criteria if criteria else {}
# Update to ensure keys are there
if all_exist:
criteria.update({k: {"$exists": True} for k in key if k not in criteria})
results = []
for d in self.groupby(key, properties=key, criteria=criteria):
results.append(d["_id"])
return results
else:
criteria = criteria if criteria else {}
lazy_substitute(criteria, self.reverse_aliases)
key = self.aliases[key] if key in self.aliases else key
return self.collection.distinct(key, filter=criteria, **kwargs)
