def get_sum_and_count_aggregation_results(keyword):
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(
{"keyword_search": [es_minimal_sanitize(keyword)]})
search = TransactionSearch().filter(filter_query)
search.aggs.bucket("prime_awards_obligation_amount",
{"sum": {
"field": "transaction_amount"
}})
search.aggs.bucket("prime_awards_count",
{"value_count": {
"field": "transaction_id"
}})
response = search.handle_execute()
if response is not None:
try:
results = {}
results["prime_awards_count"] = response["aggregations"][
"prime_awards_count"]["value"]
results["prime_awards_obligation_amount"] = round(
response["aggregations"]["prime_awards_obligation_amount"]
["value"], 2)
return results
except KeyError:
logger.exception("Unexpected Response")
else:
return None
def perform_elasticsearch_search(self, loans=False) -> Response:
filters = {f"nested_{key}": val for key, val in self.filters.items() if key != "award_type_codes"}
if self.filters.get("award_type_codes") is not None:
filters["award_type_codes"] = self.filters["award_type_codes"]
# Need to update the value of "query" to have the fields to search on
query = filters.pop("nested_query", None)
if query:
filters["nested_query"] = {"text": query, "fields": self.query_fields}
# Ensure that only non-zero values are taken into consideration
filters["nested_nonzero_fields"] = list(self.nested_nonzero_fields.values())
self.filter_query = QueryWithFilters.generate_accounts_elasticsearch_query(filters)
# using a set value here as doing an extra ES query is detrimental to performance
# And the dimensions on which group-by aggregations are performed so far
# (agency, TAS, object_class) all have cardinality less than this number
# If the data increases to a point where there are more results than this, it should be changed
self.bucket_count = 1000
messages = []
if self.pagination.sort_key in ("id", "code"):
messages.append(
(
f"Notice! API Request to sort on '{self.pagination.sort_key}' field isn't fully implemented."
" Results were actually sorted using 'description' field."
)
)
response = self.query_elasticsearch(loans)
response["page_metadata"] = get_pagination_metadata(
len(response["results"]), self.pagination.limit, self.pagination.page
)
response["results"] = response["results"][self.pagination.lower_limit : self.pagination.upper_limit]
if messages:
response["messages"] = messages
return Response(response)
def perform_search(self, validated_payload: dict, original_filters: dict) -> dict:
self.filters = validated_payload.get("filters", {})
self.subawards = validated_payload["subawards"]
self.pagination = self._get_pagination(validated_payload)
if self.subawards:
base_queryset = subaward_filter(self.filters)
self.obligation_column = "amount"
results = self.query_django_for_subawards(base_queryset)
else:
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(self.filters)
results = self.query_elasticsearch_for_prime_awards(filter_query)
page_metadata = get_simple_pagination_metadata(len(results), self.pagination.limit, self.pagination.page)
response = {
"category": self.category.name,
"limit": self.pagination.limit,
"page_metadata": page_metadata,
"results": results[: self.pagination.limit],
"messages": self._get_messages(original_filters),
}
return response
def query_elasticsearch(self, time_periods: list) -> list:
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(
self.filters)
search = TransactionSearch().filter(filter_query)
self.apply_elasticsearch_aggregations(search)
response = search.handle_execute()
return self.build_elasticsearch_result(response.aggs, time_periods)
def perform_search(self, validated_payload: dict, original_filters: dict) -> dict:
self.filters = validated_payload.get("filters", {})
self.elasticsearch = validated_payload.get("elasticsearch")
self.subawards = validated_payload["subawards"]
self.pagination = self._get_pagination(validated_payload)
if self.subawards:
base_queryset = subaward_filter(self.filters)
self.obligation_column = "amount"
results = self.query_django(base_queryset)
elif self.elasticsearch:
logger.info(
f"Using experimental Elasticsearch functionality for 'spending_by_category/{self.category.name}'"
)
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(self.filters)
results = self.query_elasticsearch(filter_query)
else:
base_queryset = spending_by_category_view_queryset(self.category.name, self.filters)
self.obligation_column = "generated_pragmatic_obligation"
results = self.query_django(base_queryset)
page_metadata = get_simple_pagination_metadata(len(results), self.pagination.limit, self.pagination.page)
response = {
"category": self.category.name,
"limit": self.pagination.limit,
"page_metadata": page_metadata,
"results": results[: self.pagination.limit],
"messages": self._get_messages(original_filters),
}
return response
def perform_search(self, validated_payload: dict,
original_filters: dict) -> dict:
self.filters = validated_payload.get("filters", {})
self.subawards = validated_payload["subawards"]
self.pagination = self._get_pagination(validated_payload)
if self.subawards:
if self.category.name in self.sub_awards_not_implemented:
self._raise_not_implemented()
base_queryset = subaward_filter(self.filters)
self.obligation_column = "amount"
results = self.query_django(base_queryset)
else:
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(
self.filters)
results = self.query_elasticsearch(filter_query)
# else:
# base_queryset = spending_by_category_view_queryset(self.category.name, self.filters)
# self.obligation_column = "generated_pragmatic_obligation"
# results = self.query_django(base_queryset)
page_metadata = get_simple_pagination_metadata(len(results),
self.pagination.limit,
self.pagination.page)
response = {
"category": self.category.name,
"limit": self.pagination.limit,
"page_metadata": page_metadata,
"results": results[:self.pagination.limit],
"messages": self._get_messages(original_filters),
}
return response
def get_total_results(keyword):
group_by_agg_key_values = {
"filters": {
category: {
"terms": {
"type": types
}
}
for category, types in INDEX_ALIASES_TO_AWARD_TYPES.items()
}
}
aggs = A("filters", **group_by_agg_key_values)
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(
{"keyword_search": [es_minimal_sanitize(keyword)]})
search = TransactionSearch().filter(filter_query)
search.aggs.bucket("types", aggs)
response = search.handle_execute()
if response is not None:
try:
return response["aggregations"]["types"]["buckets"]
except KeyError:
logger.error("Unexpected Response")
else:
logger.error("No Response")
return None
def query_elasticsearch(self, filters) -> list:
filter_query = QueryWithFilters.generate_awards_elasticsearch_query(filters)
s = AwardSearch().filter(filter_query)
s.aggs.bucket(
"types",
"filters",
filters={category: Q("terms", type=types) for category, types in all_award_types_mappings.items()},
)
results = s.handle_execute()
contracts = results.aggregations.types.buckets.contracts.doc_count
idvs = results.aggregations.types.buckets.idvs.doc_count
grants = results.aggregations.types.buckets.grants.doc_count
direct_payments = results.aggregations.types.buckets.direct_payments.doc_count
loans = results.aggregations.types.buckets.loans.doc_count
other = results.aggregations.types.buckets.other_financial_assistance.doc_count
response = {
"contracts": contracts,
"direct_payments": direct_payments,
"grants": grants,
"idvs": idvs,
"loans": loans,
"other": other,
}
return response
def get_sum_aggregation_results(keyword, field="transaction_amount"):
group_by_agg_key_values = {"field": field}
aggs = A("sum", **group_by_agg_key_values)
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(
{"keywords": es_minimal_sanitize(keyword)})
search = TransactionSearch().filter(filter_query)
search.aggs.bucket("transaction_sum", aggs)
response = search.handle_execute()
if response:
return response["aggregations"]
else:
return None
def post(self, request: Request) -> Response:
filter_query = QueryWithFilters.generate_awards_elasticsearch_query(
self.filters)
# Ensure that only non-zero values are taken into consideration
non_zero_columns = list(
ElasticsearchLoansPaginationMixin.sum_column_mapping.values())
non_zero_queries = []
for field in non_zero_columns:
non_zero_queries.append(ES_Q("range", **{field: {"gt": 0}}))
non_zero_queries.append(ES_Q("range", **{field: {"lt": 0}}))
filter_query.must.append(
ES_Q("bool", should=non_zero_queries, minimum_should_match=1))
bucket_count = get_number_of_unique_terms_for_awards(
filter_query, "cfda_number.hash")
return Response({"count": bucket_count})
def query_elasticsearch(self) -> list:
filter_query = QueryWithFilters.generate_awards_elasticsearch_query(
self.filters)
sort_field = self.get_elastic_sort_by_fields()
sorts = [{
field: self.pagination["sort_order"]
} for field in sort_field]
search = ((AwardSearch().filter(filter_query).sort(*sorts).extra(
search_after=[self.last_value, self.last_id]
)[0:self.pagination["limit"]]) if self.last_value and self.last_id else
(AwardSearch().filter(filter_query).sort(
*sorts)[((self.pagination["page"] - 1) *
self.pagination["limit"]):(
((self.pagination["page"] - 1) *
self.pagination["limit"]) +
self.pagination["limit"])]))
response = search.handle_execute()
return response
def query_elasticsearch(self) -> list:
filter_query = QueryWithFilters.generate_awards_elasticsearch_query(
self.filters)
sort_field = self.get_elastic_sort_by_fields()
sorts = [{
field: self.pagination["sort_order"]
} for field in sort_field]
record_num = (self.pagination["page"] - 1) * self.pagination["limit"]
# random page jumping was removed due to performance concerns
if (self.last_record_sort_value is None and self.last_record_unique_id
is not None) or (self.last_record_sort_value is not None
and self.last_record_unique_id is None):
# malformed request
raise Exception(
"Using search_after functionality in Elasticsearch requires both last_record_sort_value and last_record_unique_id."
)
if record_num >= settings.ES_AWARDS_MAX_RESULT_WINDOW and (
self.last_record_unique_id is None
and self.last_record_sort_value is None):
raise UnprocessableEntityException(
"Page #{page} with limit {limit} is over the maximum result limit {es_limit}. Please provide the 'last_record_sort_value' and 'last_record_unique_id' to paginate sequentially."
.format(
page=self.pagination["page"],
limit=self.pagination["limit"],
es_limit=settings.ES_AWARDS_MAX_RESULT_WINDOW,
))
# Search_after values are provided in the API request - use search after
if self.last_record_sort_value is not None and self.last_record_unique_id is not None:
search = (
AwardSearch().filter(filter_query).sort(*sorts).extra(
search_after=[
self.last_record_sort_value, self.last_record_unique_id
])[:self.pagination["limit"] +
1] # add extra result to check for next page
)
# no values, within result window, use regular elasticsearch
else:
search = AwardSearch().filter(filter_query).sort(
*sorts)[record_num:record_num + self.pagination["limit"]]
response = search.handle_execute()
return response
def post(self, request: Request) -> Response:
# Need to update the value of "query" to have the fields to search on
query = self.filters.pop("query", None)
if query:
self.filters["query"] = {
"text": query,
"fields": self.query_fields
}
self.filter_query = QueryWithFilters.generate_awards_elasticsearch_query(
self.filters)
# Ensure that only non-zero values are taken into consideration
# TODO: Refactor to use new NonzeroFields filter in QueryWithFilters
non_zero_queries = []
for field in self.sum_column_mapping.values():
non_zero_queries.append(ES_Q("range", **{field: {"gt": 0}}))
non_zero_queries.append(ES_Q("range", **{field: {"lt": 0}}))
self.filter_query.must.append(
ES_Q("bool", should=non_zero_queries, minimum_should_match=1))
self.bucket_count = get_number_of_unique_terms_for_awards(
self.filter_query, f"{self.agg_key.replace('.keyword', '')}.hash")
messages = []
if self.pagination.sort_key in ("id", "code"):
messages.append((
f"Notice! API Request to sort on '{self.pagination.sort_key}' field isn't fully implemented."
" Results were actually sorted using 'description' field."))
if self.bucket_count > 10000 and self.agg_key == settings.ES_ROUTING_FIELD:
self.bucket_count = 10000
messages.append((
"Notice! API Request is capped at 10,000 results. Either download to view all results or"
" filter using the 'query' attribute."))
response = self.query_elasticsearch()
response["page_metadata"] = get_pagination_metadata(
self.bucket_count, self.pagination.limit, self.pagination.page)
if messages:
response["messages"] = messages
return Response(response)
def post(self, request):
"""Returns boolean of whether a download request is greater than the max limit. """
models = [{
"name": "subawards",
"key": "subawards",
"type": "boolean",
"default": False
}]
models.extend(copy.deepcopy(AWARD_FILTER))
self.original_filters = request.data.get("filters")
json_request = TinyShield(models).block(request.data)
# If no filters in request return empty object to return all transactions
filters = json_request.get("filters", {})
if json_request["subawards"]:
total_count = subaward_filter(filters).count()
else:
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(
filters)
search = TransactionSearch().filter(filter_query)
total_count = search.handle_count()
if total_count is None:
total_count = 0
result = {
"calculated_transaction_count":
total_count,
"maximum_transaction_limit":
settings.MAX_DOWNLOAD_LIMIT,
"transaction_rows_gt_limit":
total_count > settings.MAX_DOWNLOAD_LIMIT,
"messages": [
get_generic_filters_message(
self.original_filters.keys(),
[elem["name"] for elem in AWARD_FILTER])
],
}
return Response(result)
def get_download_ids(keyword, field, size=10000):
"""
returns a generator that
yields list of transaction ids in chunksize SIZE
Note: this only works for fields in ES of integer type.
"""
n_iter = DOWNLOAD_QUERY_SIZE // size
results = get_total_results(keyword)
if results is None:
logger.error(
"Error retrieving total results. Max number of attempts reached")
return
total = sum(results[category]["doc_count"]
for category in INDEX_ALIASES_TO_AWARD_TYPES.keys())
required_iter = (total // size) + 1
n_iter = min(max(1, required_iter), n_iter)
for i in range(n_iter):
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(
{"keyword_search": [es_minimal_sanitize(keyword)]})
search = TransactionSearch().filter(filter_query)
group_by_agg_key_values = {
"field": field,
"include": {
"partition": i,
"num_partitions": n_iter
},
"size": size,
"shard_size": size,
}
aggs = A("terms", **group_by_agg_key_values)
search.aggs.bucket("results", aggs)
response = search.handle_execute()
if response is None:
raise Exception("Breaking generator, unable to reach cluster")
results = []
for result in response["aggregations"]["results"]["buckets"]:
results.append(result["key"])
yield results
def post(self, request):
models = [{
"name": "fields",
"key": "fields",
"type": "array",
"array_type": "text",
"text_type": "search",
"optional": False,
}]
models.extend(copy.deepcopy(AWARD_FILTER))
models.extend(copy.deepcopy(PAGINATION))
for m in models:
if m["name"] in ("keywords", "award_type_codes", "sort"):
m["optional"] = False
validated_payload = TinyShield(models).block(request.data)
record_num = (validated_payload["page"] -
1) * validated_payload["limit"]
if record_num >= settings.ES_TRANSACTIONS_MAX_RESULT_WINDOW:
raise UnprocessableEntityException(
"Page #{page} of size {limit} is over the maximum result limit ({es_limit}). Consider using custom data downloads to obtain large data sets."
.format(
page=validated_payload["page"],
limit=validated_payload["limit"],
es_limit=settings.ES_TRANSACTIONS_MAX_RESULT_WINDOW,
))
if validated_payload["sort"] not in validated_payload["fields"]:
raise InvalidParameterException(
"Sort value not found in fields: {}".format(
validated_payload["sort"]))
if "filters" in validated_payload and "no intersection" in validated_payload[
"filters"]["award_type_codes"]:
# "Special case": there will never be results when the website provides this value
return Response({
"limit": validated_payload["limit"],
"results": [],
"page_metadata": {
"page": validated_payload["page"],
"next": None,
"previous": None,
"hasNext": False,
"hasPrevious": False,
},
})
sorts = {
TRANSACTIONS_LOOKUP[validated_payload["sort"]]:
validated_payload["order"]
}
lower_limit = (validated_payload["page"] -
1) * validated_payload["limit"]
upper_limit = (
validated_payload["page"]) * validated_payload["limit"] + 1
validated_payload["filters"]["keyword_search"] = [
es_minimal_sanitize(x)
for x in validated_payload["filters"]["keywords"]
]
validated_payload["filters"].pop("keywords")
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(
validated_payload["filters"])
search = TransactionSearch().filter(filter_query).sort(
sorts)[lower_limit:upper_limit]
response = search.handle_execute()
return Response(
self.build_elasticsearch_result(validated_payload, response))
def post(self, request: Request) -> Response:
models = [
{
"name": "subawards",
"key": "subawards",
"type": "boolean",
"default": False
},
{
"name": "scope",
"key": "scope",
"type": "enum",
"optional": False,
"enum_values": ["place_of_performance", "recipient_location"],
},
{
"name": "geo_layer",
"key": "geo_layer",
"type": "enum",
"optional": False,
"enum_values": ["state", "county", "district"],
},
{
"name": "geo_layer_filters",
"key": "geo_layer_filters",
"type": "array",
"array_type": "text",
"text_type": "search",
},
]
models.extend(copy.deepcopy(AWARD_FILTER))
models.extend(copy.deepcopy(PAGINATION))
original_filters = request.data.get("filters")
json_request = TinyShield(models).block(request.data)
agg_key_dict = {
"county": "county_agg_key",
"district": "congressional_agg_key",
"state": "state_agg_key",
}
location_dict = {
"county": "county_code",
"district": "congressional_code",
"state": "state_code"
}
model_dict = {
"place_of_performance": "pop",
"recipient_location": "recipient_location",
# 'subawards_place_of_performance': 'pop',
# 'subawards_recipient_location': 'recipient_location'
}
self.scope_field_name = model_dict[json_request["scope"]]
self.agg_key = f"{self.scope_field_name}_{agg_key_dict[json_request['geo_layer']]}"
self.filters = json_request.get("filters")
self.geo_layer = GeoLayer(json_request["geo_layer"])
self.geo_layer_filters = json_request.get("geo_layer_filters")
self.loc_field_name = location_dict[self.geo_layer.value]
self.loc_lookup = f"{self.scope_field_name}_{self.loc_field_name}"
self.subawards = json_request["subawards"]
if self.subawards:
# We do not use matviews for Subaward filtering, just the Subaward download filters
self.model_name = SubawardView
self.queryset = subaward_filter(self.filters)
self.obligation_column = "amount"
result = self.query_django()
elif is_experimental_elasticsearch_api(request):
if self.scope_field_name == "pop":
scope_filter_name = "place_of_performance_scope"
else:
scope_filter_name = "recipient_scope"
# Only search for values within USA, but don't overwrite a user's search
if scope_filter_name not in self.filters:
self.filters[scope_filter_name] = "domestic"
self.obligation_column = "generated_pragmatic_obligation"
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(
self.filters)
result = self.query_elasticsearch(filter_query)
else:
self.queryset, self.model_name = spending_by_geography(
self.filters)
self.obligation_column = "generated_pragmatic_obligation"
result = self.query_django()
return Response({
"scope":
json_request["scope"],
"geo_layer":
self.geo_layer.value,
"results":
result,
"messages":
get_generic_filters_message(
original_filters.keys(),
[elem["name"] for elem in AWARD_FILTER]),
})
def post(self, request: Request) -> Response:
models = [
{
"key":
"geo_layer",
"name":
"geo_layer",
"type":
"enum",
"enum_values":
sorted([geo_layer.value for geo_layer in list(GeoLayer)]),
"text_type":
"search",
"allow_nulls":
False,
"optional":
False,
},
{
"key": "geo_layer_filters",
"name": "geo_layer_filters",
"type": "array",
"array_type": "text",
"text_type": "search",
},
{
"key": "spending_type",
"name": "spending_type",
"type": "enum",
"enum_values": ["obligation", "outlay", "face_value_of_loan"],
"allow_nulls": False,
"optional": False,
},
{
"name": "scope",
"key": "scope",
"type": "enum",
"optional": True,
"enum_values": ["place_of_performance", "recipient_location"],
"default": "recipient_location",
},
]
# NOTE: filter object in request handled in base class: see self.filters
json_request = TinyShield(models).block(request.data)
agg_key_dict = {
"county": "county_agg_key",
"district": "congressional_agg_key",
"state": "state_agg_key",
}
scope_dict = {
"place_of_performance": "pop",
"recipient_location": "recipient_location"
}
location_dict = {
"county": "county_code",
"district": "congressional_code",
"state": "state_code"
}
self.geo_layer = GeoLayer(json_request["geo_layer"])
scope_field_name = scope_dict[json_request["scope"]]
loc_field_name = location_dict[self.geo_layer.value]
self.agg_key = f"{scope_field_name}_{agg_key_dict[json_request['geo_layer']]}"
self.geo_layer_filters = json_request.get("geo_layer_filters")
self.spending_type = json_request.get("spending_type")
self.loc_lookup = f"{scope_field_name}_{loc_field_name}"
# Set which field will be the aggregation amount
if self.spending_type == "obligation":
self.metric_field = "total_covid_obligation"
elif self.spending_type == "outlay":
self.metric_field = "total_covid_outlay"
elif self.spending_type == "face_value_of_loan":
self.metric_field = "total_loan_value"
else:
raise UnprocessableEntityException(
f"Unrecognized value '{self.spending_type}' for field "
f"'spending_type'")
filter_query = QueryWithFilters.generate_awards_elasticsearch_query(
self.filters)
result = self.query_elasticsearch(filter_query)
return Response({
"geo_layer": self.geo_layer.value,
"spending_type": self.spending_type,
"scope": json_request["scope"],
"results": result,
})
def obtain_recipient_totals(recipient_id, children=False, year="latest"):
"""Extract the total amount and transaction count for the recipient_hash given the time frame
Args:
recipient_id: string of hash(duns, name)-[recipient-level]
children: whether or not to group by children
year: the year the totals/counts are based on
Returns:
list of dictionaries representing hashes and their totals/counts
"""
filters = reshape_filters(recipient_id=recipient_id, year=year)
filter_query = QueryWithFilters.generate_transactions_elasticsearch_query(
filters)
search = TransactionSearch().filter(filter_query)
if children:
group_by_field = "recipient_agg_key"
elif recipient_id[-2:] == "-P":
group_by_field = "parent_recipient_hash"
else:
group_by_field = "recipient_hash"
bucket_count = get_number_of_unique_terms_for_transactions(
filter_query, f"{group_by_field}.hash")
if bucket_count == 0:
return []
# Not setting the shard_size since the number of child recipients under a
# parent recipient will not exceed 10k
group_by_recipient = A("terms", field=group_by_field, size=bucket_count)
sum_obligation = get_scaled_sum_aggregations(
"generated_pragmatic_obligation")["sum_field"]
filter_loans = A("filter", terms={"type": list(loan_type_mapping.keys())})
sum_face_value_loan = get_scaled_sum_aggregations(
"face_value_loan_guarantee")["sum_field"]
search.aggs.bucket("group_by_recipient", group_by_recipient)
search.aggs["group_by_recipient"].metric("sum_obligation", sum_obligation)
search.aggs["group_by_recipient"].bucket("filter_loans", filter_loans)
search.aggs["group_by_recipient"]["filter_loans"].metric(
"sum_face_value_loan", sum_face_value_loan)
response = search.handle_execute()
response_as_dict = response.aggs.to_dict()
recipient_info_buckets = response_as_dict.get("group_by_recipient",
{}).get("buckets", [])
result_list = []
for bucket in recipient_info_buckets:
result = {}
if children:
recipient_info = json.loads(bucket.get("key"))
hash_with_level = recipient_info.get("hash_with_level") or None
result = {
"recipient_hash":
hash_with_level[:-2] if hash_with_level else None,
"recipient_unique_id": recipient_info.get("unique_id"),
"recipient_name": recipient_info.get("name"),
}
loan_info = bucket.get("filter_loans", {})
result.update({
"total_obligation_amount":
int(bucket.get("sum_obligation", {"value": 0})["value"]) /
Decimal("100"),
"total_obligation_count":
bucket.get("doc_count", 0),
"total_face_value_loan_amount":
int(loan_info.get("sum_face_value_loan", {"value": 0})["value"]) /
Decimal("100"),
"total_face_value_loan_count":
loan_info.get("doc_count", 0),
})
result_list.append(result)
return result_list