def __init__(self, query=None, explain=None, columns=None, sort=None,
width=None, filter=None, stream_results=False,
complete_results=True, **kwargs):
super(EntityFind, self).__init__(**kwargs)
if query:
self.query = entity_query.Query(query)
else:
self.query = entity_query.Query(self.search)
if columns is not None:
self.columns = columns
if sort is not None:
self.sort = sort
if width is not None:
self.width = width
if filter is not None:
self.display_filter = entity_query.Query(filter)
if stream_results is not None:
if sort:
raise ValueError(
"Can't stream results and sort at the same time.")
self.stream_results = stream_results
if complete_results is not None:
self.complete_results = complete_results
self.explain = explain
def find_by_component(self, component, complete=True):
"""Finds all entities that have the component.
Arguments:
complete: If True, will attempt to collect the component.
"""
query = entity_query.Query(expression.ComponentLiteral(component))
if complete:
self.collect_for(query)
return list(self.lookup_tables["components"].lookup(component))
def stream(self, query, handler, query_params=None):
query = entity_query.Query(query, params=query_params)
seen = set()
def _deduplicator(entity):
if entity in seen:
return
seen.add(entity)
handler(entity)
self.collect_for(query, result_stream_handler=_deduplicator)
for entity in self.find(query, complete=False):
_deduplicator(entity)
def get_referencing_entities(self, key, complete=True):
"""Finds entities that reference this entity by its identity.
If other entities have an attribute that stores the identity of this
entity then this function will find those entities.
For example, calling this on a process, one can find all the handles
owned by the process by calling
process.get_referencing_entities("Handle/process").
Arguments:
key: The property path to the attribute on the other entities.
As usual, form is Component/attribute.
"""
# Automatically ask the entity manager to add indexing for
# identity-based attributes. This is a good heuristic for optimal
# performance.
self.manager.add_attribute_lookup(key)
return self.manager.find(query.Query("%s is {}" % key,
params=[self.identity]),
complete=complete)
def find(self, query, complete=True, validate=True, query_params=None):
"""Runs the query and yields entities that match.
Arguments:
query: Either an instance of the query AST, a query string, or a
dictionary of queries. If a dict is given, a new dict will
be returned with the same keys and values replaced with
results.
complete: If True, will trigger collectors as necessary, to ensure
completness of results.
validate: Will cause the query to be validated first (mostly for
type errors.)
"""
if isinstance(query, dict):
results = {}
for query_name, expr in query.iteritems():
results[query_name] = self.find(expr,
complete=complete,
validate=validate,
query_params=query_params)
return results
if not isinstance(query, entity_query.Query):
query = entity_query.Query(query, params=query_params)
if validate:
query.execute("QueryValidator")
if complete:
self.collect_for(query)
# Try to satisfy the query using available lookup tables.
search = entity_lookup.EntityQuerySearch(query)
return search.search(self.entities, self.lookup_tables)
def assertQueryMatches(self, query, bindings):
m = entity_query.Query(query).execute("QueryMatcher",
"match",
bindings=bindings)
self.assertIsNotNone(m)
def analyze(self, wanted):
"""Finds collectors and indexing suggestions for the query.
Returns a dict of:
- collectors: list of collectors to run
- lookups: list of attributes to consider indexing for
- dependencies: list of SimpleDependency instances to include
- exclusions: list of SimpleDependency instances to exclude
"""
if not isinstance(wanted, entity_query.Query):
wanted = entity_query.Query(wanted)
# We cache by the source and not the query because we want to reanalyze
# queries that are logically equivalent, but expressed differently, in
# order to have the right cursor positions stored for highlighting in
# GUI.
cache_key = wanted.source
analysis = self._cached_query_analyses.get(cache_key, None)
if analysis:
# We want to make a copy exactly one level deep.
analysis_copy = {}
for key, value in analysis.iteritems():
analysis_copy[key] = copy.copy(value)
return analysis_copy
analyzer = wanted.execute("QueryAnalyzer")
include = analyzer.include
exclude = analyzer.exclude
suggested_indices = analyzer.latest_indices
# A collector is a match if any of its promises match any of the
# dependencies of the query.
matched_collectors = []
for collector in self.collectors.itervalues():
for promise, dependency in itertools.product(
collector.promises, include):
if dependency.match(promise):
matched_collectors.append(collector)
break
# A collector is yielded unless each one of its promises matches
# an exclusion from dependencies.
collectors = set()
for collector in matched_collectors:
for promise, exclusion in itertools.product(
collector.promises, exclude):
if not exclusion.match(promise):
collectors.add(collector)
break
else:
# No exclusions.
collectors.add(collector)
# A component is guaranteed if any dependency lists it. It is likely
# if collectors we depend on output it (though not guaranteed).
guaranteed_components = set(analyzer.expected_components)
possible_components = set()
for dependency in include:
component = dependency.component
if component in guaranteed_components:
continue
possible_components.add(dependency.component)
for collector in collectors:
for promise in collector.promises:
component = promise.component
if component in guaranteed_components:
continue
possible_components.add(component)
analysis = dict(collectors=list(collectors),
lookups=suggested_indices,
dependencies=include,
exclusions=exclude,
guaranteed_components=guaranteed_components,
possible_components=possible_components)
self._cached_query_analyses[cache_key] = analysis
return analysis
def _ensure_compile_queries(self):
if self.collect_queries or self.collect_args is None:
return
for arg, source in self.collect_args.iteritems():
self.collect_queries[arg] = entity_query.Query(source)
def __init__(self, query=None, **kwargs):
super(EntityAnalyze, self).__init__(**kwargs)
self.query = entity_query.Query(query)
self.analysis = self.session.entities.analyze(self.query)
def assertQueryRaises(self, query):
q = entity_query.Query(query)
self.assertRaises(validator.ValidationError, q.execute,
"QueryValidator")
def assertQueryPasses(self, query):
q = entity_query.Query(query)
try:
q.execute("QueryValidator")
except validator.ValidationError as e:
self.fail("%s raised exception on validation:\n%s" % (q, e))
def _subquery(self, expr):
return entity_query.Query(expression=expr, source=self.query.source)
def analyze(self, query):
return entity_query.Query(query).execute("QueryAnalyzer")
