diff --git a/usaspending_api/awards/tests/unit/test_location_filter_geocode.py b/usaspending_api/awards/tests/unit/test_location_filter_geocode.py index b26364ab47..783e2ccfb7 100644 --- a/usaspending_api/awards/tests/unit/test_location_filter_geocode.py +++ b/usaspending_api/awards/tests/unit/test_location_filter_geocode.py @@ -2,14 +2,15 @@ from usaspending_api.awards.v2.filters.location_filter_geocode import ( create_nested_object, + geocode_filter_locations, get_fields_list, location_error_handling, validate_location_keys, ) from usaspending_api.common.exceptions import InvalidParameterException from usaspending_api.common.helpers.api_helper import ( - INCOMPATIBLE_DISTRICT_LOCATION_PARAMETERS, DUPLICATE_DISTRICT_LOCATION_PARAMETERS, + INCOMPATIBLE_DISTRICT_LOCATION_PARAMETERS, ) @@ -101,3 +102,24 @@ def test_get_fields_list(): assert get_fields_list("county_code", "01") == ["1", "01", "1.0"] assert get_fields_list("feet", "01") == ["01"] assert get_fields_list("congressional_code", "abc") == ["abc"] + + +def test_geocode_filter_locations_for_both_district_original_and_current_across_elements(): + """ + Cross-element same-state current and original must OR not overwrite + """ + + values = [ + {"country": "USA", "state": "VA", "district_original": "08"}, + {"country": "USA", "state": "VA", "district_current": "11"} + ] + # validate_location_keys only checks for moth district types within a single element + assert validate_location_keys(values) is None + nested = create_nested_object(values) + assert nested["USA"]["VA"]["district_original"] == ["08"] + assert nested["USA"]["VA"]["district_current"] == ["11"] + + q = geocode_filter_locations("test", values) + q_repr = repr(q) + assert "test_congressional_code_current__in" in q_repr + assert "test_congressional_code__in" in q_repr diff --git a/usaspending_api/awards/v2/filters/location_filter_geocode.py b/usaspending_api/awards/v2/filters/location_filter_geocode.py index 76b9450458..3afaa430bf 100644 --- a/usaspending_api/awards/v2/filters/location_filter_geocode.py +++ b/usaspending_api/awards/v2/filters/location_filter_geocode.py @@ -1,3 +1,5 @@ +from typing import Any + from django.db.models import Q from usaspending_api.common.exceptions import InvalidParameterException @@ -10,57 +12,72 @@ ALL_FOREIGN_COUNTRIES = "FOREIGN" -def geocode_filter_locations(scope: str, values: list) -> Q: +def geocode_filter_locations(scope: str, values: list[dict[str, Any]]) -> Q: """ Function filter querysets on location table scope- place of performance or recipient location mappings values- array of location requests returns queryset """ - or_queryset = Q() - - # creates a dictionary with all of the locations organized by country - # Counties and congressional districts are nested under state codes nested_values = create_nested_object(values) - # In this for-loop a django Q filter object is created from the python dict + or_queryset = Q() for country, state_zip in nested_values.items(): - country_qs = None - if country != ALL_FOREIGN_COUNTRIES: - country_qs = Q(**{f"{scope}_country_code__exact": country}) - state_qs = Q() - - for state_zip_key, location_values in state_zip.items(): - - if state_zip_key == "city": - state_inner_qs = Q(**{f"{scope}_city_name__in": location_values}) - elif state_zip_key == "zip": - state_inner_qs = Q(**{f"{scope}_zip5__in": location_values}) - else: - state_inner_qs = Q(**{f"{scope}_state_code__exact": state_zip_key.upper()}) - county_qs = Q() - district_qs = Q() - city_qs = Q() - - if location_values["county"]: - county_qs = Q(**{f"{scope}_county_code__in": location_values["county"]}) - if location_values["district_current"]: - district_qs = Q(**{f"{scope}_congressional_code_current__in": location_values["district_current"]}) - if location_values["district_original"]: - district_qs = Q(**{f"{scope}_congressional_code__in": location_values["district_original"]}) - if location_values["city"]: - city_qs = Q(**{f"{scope}_city_name__in": location_values["city"]}) - state_inner_qs &= county_qs | district_qs | city_qs - - state_qs |= state_inner_qs - if country_qs: - or_queryset |= country_qs & state_qs - else: - or_queryset |= state_qs + country_qs = _build_country_query(scope, country) + state_qs = _build_state_queries(scope, state_zip) + + or_queryset |= (country_qs & state_qs) if country_qs else state_qs + return or_queryset -def validate_location_keys(values): +def _build_country_query(scope: str, country: str) -> Q | None: + """Build country-level query filter""" + if country == ALL_FOREIGN_COUNTRIES: + return None + return Q(**{f"{scope}_country_code__exact": country}) + + +def _build_state_queries(scope: str, state_zip: dict[str, Any]) -> Q: + """Build state-level query filters""" + state_qs = Q() + + for state_zip_key, location_values in state_zip.items(): + if state_zip_key == "city": + state_inner_qs = Q(**{f"{scope}_city_name__in": location_values}) + elif state_zip_key == "zip": + state_inner_qs = Q(**{f"{scope}_zip5__in": location_values}) + else: + state_inner_qs = _build_state_location_query(scope, state_zip_key, location_values) + + state_qs |= state_inner_qs + + return state_qs + + +def _build_state_location_query(scope: str, state_code: str, location_values: dict[str, Any]) -> Q: + """Build query for state with nested county/district/city filters""" + state_qs = Q(**{f"{scope}_state_code__exact": state_code.upper()}) + + # Build sub-filters + sub_filters = Q() + + if location_values.get("county"): + sub_filters |= Q(**{f"{scope}_county_code__in": location_values["county"]}) + + if location_values.get("district_current"): + sub_filters |= Q(**{f"{scope}_congressional_code_current__in": location_values["district_current"]}) + + if location_values.get("district_original"): + sub_filters |= Q(**{f"{scope}_congressional_code__in": location_values["district_original"]}) + + if location_values.get("city"): + sub_filters |= Q(**{f"{scope}_city_name__in": location_values["city"]}) + + return state_qs & sub_filters + + +def validate_location_keys(values: list[dict[str, Any]]) -> None: """Validate that the keys provided are sufficient and match properly.""" for v in values: state = v.get("state") @@ -78,70 +95,107 @@ def validate_location_keys(values): location_error_handling(v.keys()) -def create_nested_object(values): +def create_nested_object(values: list[dict[str, Any]]) -> dict[str, Any]: """Makes sure keys provided are valid""" validate_location_keys(values) - nested_locations = {} + nested_locations: dict[str, Any] = {} for v in values: upper_case_dict_values(v) - city = v.get("city") - country = v.get("country") - county = v.get("county") - district_original = v.get("district_original") - district_current = v.get("district_current") - state = v.get("state") - zip = v.get("zip") - # First level in location filtering in country - # All location requests must have a country otherwise there will be a key error - if nested_locations.get(country) is None: - nested_locations[country] = {} - - # Initialize the list - if zip and not nested_locations[country].get("zip"): - nested_locations[country]["zip"] = [] - - if city and not nested_locations[country].get("city"): - nested_locations[country]["city"] = [] - - # Second level of filtering is zip and state - # Requests must have a country+zip or country+state combination - if zip: - # Appending zips so we don't overwrite - nested_locations[country]["zip"].append(zip) - - # If we have a state, add it to the list - if state and nested_locations[country].get(state) is None: - nested_locations[country][state] = { - "county": [], - "district_original": [], - "district_current": [], - "city": [], - } - - # Based on previous checks, there will always be a state if either of these exist - if county: - nested_locations[country][state]["county"].extend(get_fields_list("county", county)) - - if district_current: - nested_locations[country][state]["district_current"].extend( - get_fields_list("district_current", district_current) - ) - - if district_original: - nested_locations[country][state]["district_original"].extend( - get_fields_list("district_original", district_original) - ) - - if city and state: - nested_locations[country][state]["city"].append(city) - elif city: - nested_locations[country]["city"].append(city) + location_data = _extract_location_data(v) + _process_location(nested_locations, location_data) return nested_locations -def location_error_handling(fields): +def _extract_location_data(location_dict: dict[str, Any]) -> dict[str, Any]: + """Extract location fields from the input dictionary""" + return { + "city": location_dict.get("city"), + "country": location_dict.get("country"), + "county": location_dict.get("county"), + "district_original": location_dict.get("district_original"), + "district_current": location_dict.get("district_current"), + "state": location_dict.get("state"), + "zip_code": location_dict.get("zip"), + } + + +def _process_location(nested_locations: dict[str, Any], location_data: dict[str, Any]) -> None: + """Process a single location and add it to the nested structure""" + country = location_data["country"] + + # Initialize country if needed + if country not in nested_locations: + nested_locations[country] = {} + + # Process zip codes + _process_zip_code(nested_locations[country], location_data["zip_code"]) + + # Process city (country-level) + _process_country_level_city(nested_locations[country], location_data["city"]) + + # Process state-level data + if location_data["state"]: + _process_state_data(nested_locations[country], location_data) + elif location_data["city"]: + # City without state goes to country level + nested_locations[country]["city"].append(location_data["city"]) + + +def _process_zip_code(country_data: dict[str, Any], zip_code: str | None) -> None: + """Process zip code for a country""" + if not zip_code: + return + + if "zip" not in country_data: + country_data["zip"] = [] + + country_data["zip"].append(zip_code) + + +def _process_country_level_city(country_data: dict[str, Any], city: str | None) -> None: + """Initialize city list at country level if needed""" + if city and "city" not in country_data: + country_data["city"] = [] + + +def _process_state_data(country_data: dict[str, Any], location_data: dict[str, Any]) -> None: + """Process state-level location data""" + state = location_data["state"] + + # Initialize state structure if needed + if state not in country_data: + country_data[state] = { + "county": [], + "district_original": [], + "district_current": [], + "city": [], + } + + state_data = country_data[state] + + # Add county data + if location_data["county"]: + state_data["county"].extend(get_fields_list("county", location_data["county"])) + + # Add district data + if location_data["district_current"]: + state_data["district_current"].extend( + get_fields_list("district_current", location_data["district_current"]) + ) + + if location_data["district_original"]: + state_data["district_original"].extend( + get_fields_list("district_original", location_data["district_original"]) + ) + + # Add city data + if location_data["city"]: + state_data["city"].append(location_data["city"]) + + +def location_error_handling(fields: Any) -> None: """Raise the relevant error for location keys.""" # Request must have country, and can only have 3 fields, and must have state if there is county if "country" not in fields: @@ -151,7 +205,7 @@ def location_error_handling(fields): raise InvalidParameterException("Invalid filter: Missing necessary location field: state.") -def get_fields_list(scope, field_value): +def get_fields_list(scope: str, field_value: str) -> list[str]: """List of values to search for; `field_value`, plus possibly variants on it""" if scope in ["congressional_code", "county_code"]: try: diff --git a/usaspending_api/awards/v2/filters/sub_award.py b/usaspending_api/awards/v2/filters/sub_award.py index 09032aeb00..b959634e85 100644 --- a/usaspending_api/awards/v2/filters/sub_award.py +++ b/usaspending_api/awards/v2/filters/sub_award.py @@ -1,7 +1,8 @@ import itertools import logging +from typing import Any -from django.db.models import Exists, OuterRef, Q +from django.db.models import Exists, OuterRef, Q, QuerySet from usaspending_api.awards.models import TransactionNormalized from usaspending_api.awards.models.financial_accounts_by_awards import FinancialAccountsByAwards @@ -20,7 +21,7 @@ logger = logging.getLogger(__name__) -def subaward_download(filters): +def subaward_download(filters: dict[str, Any]) -> QuerySet: """Used by the Custom download""" return subaward_filter(filters, for_downloads=True) @@ -32,11 +33,22 @@ def geocode_filter_subaward_locations(scope: str, values: list) -> Q: values- array of location requests returns queryset """ + location_mappings = _get_location_mappings(scope) + nested_values = create_nested_object(values) + or_queryset = Q() + for country, state_zip in nested_values.items(): + country_qs = _build_country_query(scope, country, location_mappings) + state_qs = _build_state_queries(scope, state_zip, location_mappings) + + or_queryset |= (country_qs & state_qs) if country_qs else state_qs + + return or_queryset - # Yes, these are mostly the same, but congressional_code is different - # and I'd rather have them all laid out here versus burying a extra couple lines for congressional_code - location_mappings = { + +def _get_location_mappings(scope: str) -> dict[str, str]: + """Extract location field mappings for the given scope""" + all_mappings = { "country_code": {"sub_legal_entity": "country_code", "sub_place_of_perform": "country_co"}, "zip5": {"sub_legal_entity": "zip5", "sub_place_of_perform": "zip5"}, "city_name": {"sub_legal_entity": "city_name", "sub_place_of_perform": "city_name"}, @@ -45,337 +57,334 @@ def geocode_filter_subaward_locations(scope: str, values: list) -> Q: "congressional_code": {"sub_legal_entity": "congressional", "sub_place_of_perform": "congressio"}, "current_congressional_code": { "sub_legal_entity": "sub_legal_entity_congressional_current", - # Due to the rigidness of how we map values to columns - # it's required that the column start with sub_place_of_perform - # however, when current congressional codes were implemented - # the column name chosen did not match this pattern. - # That's why we have the full column name in the value "sub_place_of_perform": "sub_place_of_performance_congressional_current", }, } - location_mappings = {location_type: field_dict[scope] for location_type, field_dict in location_mappings.items()} + return {location_type: field_dict[scope] for location_type, field_dict in all_mappings.items()} - # creates a dictionary with all of the locations organized by country - # Counties and congressional districts are nested under state codes - nested_values = create_nested_object(values) - # In this for-loop a django Q filter object is created from the python dict - for country, state_zip in nested_values.items(): - country_qs = None - if country != ALL_FOREIGN_COUNTRIES: - country_qs = Q(**{f"{scope}_{location_mappings['country_code']}__exact": country}) - state_qs = Q() - - for state_zip_key, location_values in state_zip.items(): - if state_zip_key == "city": - state_inner_qs = Q(**{f"{scope}_{location_mappings['city_name']}__in": location_values}) - elif state_zip_key == "zip": - state_inner_qs = Q(**{f"{scope}_{location_mappings['zip5']}__in": location_values}) - else: - state_inner_qs = Q(**{f"{scope}_{location_mappings['state_code']}__exact": state_zip_key.upper()}) - county_qs = Q() - district_qs = Q() - city_qs = Q() - - if location_values["county"]: - county_qs = Q(**{f"{scope}_{location_mappings['county_code']}__in": location_values["county"]}) - if location_values["district_current"]: - district_qs = Q( - **{ - f"{location_mappings['current_congressional_code']}__in": location_values[ - "district_current" - ] - } - ) - if location_values["district_original"]: - district_qs = Q( - **{ - f"{scope}_{location_mappings['congressional_code']}__in": location_values[ - "district_original" - ] - } - ) - if location_values["city"]: - city_qs = Q(**{f"{scope}_{location_mappings['city_name']}__in": location_values["city"]}) - state_inner_qs &= county_qs | district_qs | city_qs - - state_qs |= state_inner_qs - if country_qs: - or_queryset |= country_qs & state_qs +def _build_country_query(scope: str, country: str, location_mappings: dict[str, str]) -> Q | None: + """Build country-level query filter""" + if country == ALL_FOREIGN_COUNTRIES: + return None + return Q(**{f"{scope}_{location_mappings['country_code']}__exact": country}) + + +def _build_state_queries(scope: str, state_zip: dict, location_mappings: dict[str, str]) -> Q: + """Build state-level query filters""" + state_qs = Q() + + for state_zip_key, location_values in state_zip.items(): + if state_zip_key == "city": + state_inner_qs = Q(**{f"{scope}_{location_mappings['city_name']}__in": location_values}) + elif state_zip_key == "zip": + state_inner_qs = Q(**{f"{scope}_{location_mappings['zip5']}__in": location_values}) else: - or_queryset |= state_qs - return or_queryset + state_inner_qs = _build_state_location_query(scope, state_zip_key, location_values, location_mappings) + + state_qs |= state_inner_qs + + return state_qs + + +def _build_state_location_query( + scope: str, state_code: str, location_values: dict, location_mappings: dict[str, str] +) -> Q: + """Build query for state with nested county/district/city filters""" + state_qs = Q(**{f"{scope}_{location_mappings['state_code']}__exact": state_code.upper()}) + + # Build sub-filters + sub_filters = Q() + + if location_values.get("county"): + sub_filters |= Q(**{f"{scope}_{location_mappings['county_code']}__in": location_values["county"]}) + + if location_values.get("district_current"): + sub_filters |= Q( + **{f"{location_mappings['current_congressional_code']}__in": location_values["district_current"]} + ) + + if location_values.get("district_original"): + sub_filters |= Q( + **{f"{scope}_{location_mappings['congressional_code']}__in": location_values["district_original"]} + ) + + if location_values.get("city"): + sub_filters |= Q(**{f"{scope}_{location_mappings['city_name']}__in": location_values["city"]}) + + return state_qs & sub_filters # TODO: Performance when multiple false values are initially provided -def subaward_filter(filters, for_downloads=False): +def subaward_filter(filters: dict[str, Any], for_downloads: bool = False) -> QuerySet: queryset = SubawardSearch.objects.all() recipient_scope_q = Q(sub_legal_entity_country_code="USA") | Q(sub_legal_entity_country_name="UNITED STATES") pop_scope_q = Q(sub_place_of_perform_country_co="USA") | Q(sub_place_of_perform_country_name="UNITED STATES") + # Define valid filter keys + valid_keys = [ + "keywords", "description", "transaction_keyword_search", "time_period", + "award_type_codes", "prime_and_sub_award_types", "agencies", "legal_entities", + "recipient_search_text", "recipient_scope", "recipient_locations", + "recipient_type_names", "place_of_performance_scope", "place_of_performance_locations", + "award_amounts", "award_ids", "program_numbers", "naics_codes", + PSCCodes.underscore_name, "contract_pricing_type_codes", + "set_aside_type_codes", "extent_competed_type_codes", + TasCodes.underscore_name, TreasuryAccounts.underscore_name, + "def_codes", "program_activities", + ] + + # Create filter handlers mapping + filter_handlers = { + "keywords": _handle_keywords_filter, + "description": _handle_description_filter, + "transaction_keyword_search": _handle_transaction_keyword_filter, + "time_period": lambda qs, val: _handle_time_period_filter(qs, val, for_downloads), + "award_type_codes": _handle_award_type_codes_filter, + "prime_and_sub_award_types": _handle_prime_and_sub_award_types_filter, + "agencies": _handle_agencies_filter, + "legal_entities": _handle_legal_entities_filter, + "recipient_search_text": _handle_recipient_search_text_filter, + "recipient_scope": lambda qs, val: _handle_scope_filter(qs, val, recipient_scope_q, "recipient_scope"), + "recipient_locations": lambda qs, val: _handle_locations_filter(qs, val, "sub_legal_entity"), + "recipient_type_names": _handle_recipient_type_names_filter, + "place_of_performance_scope": lambda qs, val: _handle_scope_filter(qs, val, pop_scope_q, + "place_of_performance_scope"), + "place_of_performance_locations": lambda qs, val: _handle_locations_filter(qs, val, "sub_place_of_perform"), + "award_amounts": lambda qs, val: _handle_award_amounts_filter(qs, val, filters), + "award_ids": _handle_award_ids_filter, + PSCCodes.underscore_name: _handle_psc_codes_filter, + "contract_pricing_type_codes": _handle_contract_pricing_filter, + "program_numbers": _handle_program_numbers_filter, + "set_aside_type_codes": lambda qs, val: _handle_set_aside_extent_filter(qs, val, "set_aside_type_codes", + "type_set_aside"), + "extent_competed_type_codes": lambda qs, val: _handle_set_aside_extent_filter(qs, val, + "extent_competed_type_codes", + "extent_competed"), + TasCodes.underscore_name: lambda qs, val: _handle_tas_codes_filter(qs, val, filters), + "def_codes": _handle_def_codes_filter, + "program_activities": _handle_program_activities_filter, + } + for key, value in filters.items(): if value is None: - raise InvalidParameterException("Invalid filter: " + key + " has null as its value.") - - key_list = [ - "keywords", - "description", - "transaction_keyword_search", - "time_period", - "award_type_codes", - "prime_and_sub_award_types", - "agencies", - "legal_entities", - "recipient_search_text", - "recipient_scope", - "recipient_locations", - "recipient_type_names", - "place_of_performance_scope", - "place_of_performance_locations", - "award_amounts", - "award_ids", - "program_numbers", - "naics_codes", - PSCCodes.underscore_name, - "contract_pricing_type_codes", - "set_aside_type_codes", - "extent_competed_type_codes", - TasCodes.underscore_name, - TreasuryAccounts.underscore_name, - "def_codes", - "program_activities", - ] - - if key not in key_list: - raise InvalidParameterException("Invalid filter: " + key + " does not exist.") - - if key == "keywords": - - def keyword_parse(keyword): - # keyword_ts_vector & award_ts_vector are Postgres TS_vectors. - # keyword_ts_vector = recipient_name + psc_description + subaward_description - # award_ts_vector = piid + fain + uri + subaward_number - filter_obj = Q(keyword_ts_vector=keyword) | Q(award_ts_vector=keyword) - # Commenting out until NAICS is associated with subawards in GSDM 1.3.1 - # if keyword.isnumeric(): - # filter_obj |= Q(naics_code__contains=keyword) - if len(keyword) == 4 and PSC.objects.filter(code__iexact=keyword).exists(): - filter_obj |= Q(product_or_service_code__iexact=keyword) - - return filter_obj - - filter_obj = Q() - for keyword in value: - filter_obj |= keyword_parse(keyword) - - # Search for DUNS - potential_duns = list(filter((lambda x: len(x) == 9), value)) - if len(potential_duns) > 0: - filter_obj |= Q(sub_awardee_or_recipient_uniqu__in=potential_duns) | Q( - sub_ultimate_parent_unique_ide__in=potential_duns - ) - - # Search for UEI - potential_ueis = list(filter((lambda x: len(x) == 12), value)) - potential_ueis = [uei.upper() for uei in potential_ueis] - if len(potential_ueis) > 0: - filter_obj |= Q(sub_awardee_or_recipient_uei__in=potential_ueis) | Q( - sub_ultimate_parent_uei__in=potential_ueis - ) - - queryset = queryset.filter(filter_obj) - - elif key == "description": - queryset = queryset.filter(subaward_description__icontains=value) - - elif key == "transaction_keyword_search": - keyword = value - transaction_ids = elasticsearch_helper.get_download_ids(keyword=keyword, field="transaction_id") - # flatten IDs - transaction_ids = list(itertools.chain.from_iterable(transaction_ids)) - logger.info("Found {} transactions based on keyword: {}".format(len(transaction_ids), keyword)) - transaction_ids = [str(transaction_id) for transaction_id in transaction_ids] - queryset = queryset.filter(latest_transaction__isnull=False) - - # Prepare a SQL snippet to include in the predicate for searching an array of transaction IDs - # TODO: Now that SubawardSearch has an FK field to TransactionSearch, we don't need the extra (raw SQL) - # Look to add a Django filter that does the same as below - sql_fragment = ( - '"subaward_search"."latest_transaction_id" = ANY(\'{{{}}}\'::int[])' # int[] -> int array type - ) - queryset = queryset.extra(where=[sql_fragment.format(",".join(transaction_ids))]) - - elif key == "time_period": - min_date = API_SEARCH_MIN_DATE - if for_downloads: - min_date = API_MIN_DATE - queryset &= combine_date_range_queryset(value, SubawardSearch, min_date, API_MAX_DATE, is_subaward=True) - - elif key == "award_type_codes": - queryset = queryset.filter(prime_award_type__in=value) - - elif key == "prime_and_sub_award_types": - award_types = value.get("sub_awards") - if award_types: - queryset = queryset.filter(prime_award_group__in=award_types) - - elif key == "agencies": - # TODO: Make function to match agencies in award filter throwing dupe error - funding_toptier = Q() - funding_subtier = Q() - awarding_toptier = Q() - awarding_subtier = Q() - for v in value: - type = v["type"] - tier = v["tier"] - name = v["name"] - if type == "funding": - if tier == "toptier": - funding_toptier |= Q(funding_toptier_agency_name=name) - elif tier == "subtier": - if "toptier_name" in v: - funding_subtier |= Q(funding_subtier_agency_name=name) & Q( - funding_toptier_agency_name=v["toptier_name"] - ) - else: - funding_subtier |= Q(funding_subtier_agency_name=name) - - elif type == "awarding": - if tier == "toptier": - awarding_toptier |= Q(awarding_toptier_agency_name=name) - elif tier == "subtier": - if "toptier_name" in v: - awarding_subtier |= Q(awarding_subtier_agency_name=name) & Q( - awarding_toptier_agency_name=v["toptier_name"] - ) - else: - awarding_subtier |= Q(awarding_subtier_agency_name=name) - - awarding_queryfilter = Q() - funding_queryfilter = Q() - - # Since these are Q filters, no DB hits for boolean checks - if funding_toptier: - funding_queryfilter |= funding_toptier - if funding_subtier: - funding_queryfilter |= funding_subtier - if awarding_toptier: - awarding_queryfilter |= awarding_toptier - if awarding_subtier: - awarding_queryfilter |= awarding_subtier - - queryset = queryset.filter(funding_queryfilter & awarding_queryfilter) - - elif key == "legal_entities": - # This filter key has effectively become obsolete by recipient_search_text - msg = 'API request included "{}" key. No filtering will occur with provided value "{}"' - logger.info(msg.format(key, value)) - - elif key == "recipient_search_text": - - def recip_string_parse(recipient_string): - upper_recipient_string = recipient_string.upper() - - # recipient_name_ts_vector is a postgres TS_Vector - filter_obj = Q(recipient_name_ts_vector=upper_recipient_string) - if len(upper_recipient_string) == 9 and upper_recipient_string[:5].isnumeric(): - filter_obj |= Q(sub_awardee_or_recipient_uniqu=upper_recipient_string) - elif len(upper_recipient_string) == 12: - filter_obj |= Q(sub_awardee_or_recipient_uei=upper_recipient_string) - return filter_obj - - filter_obj = Q() - for recipient in value: - filter_obj |= recip_string_parse(recipient) - queryset = queryset.filter(filter_obj) - - elif key == "recipient_scope": - if value == "domestic": - queryset = queryset.filter(recipient_scope_q) - elif value == "foreign": - queryset = queryset.exclude(recipient_scope_q) - else: - raise InvalidParameterException("Invalid filter: recipient_scope type is invalid.") - - elif key == "recipient_locations": - queryset = queryset.filter(geocode_filter_subaward_locations("sub_legal_entity", value)) - - elif key == "recipient_type_names": - if len(value) != 0: - queryset = queryset.filter(business_categories__overlap=value) - - elif key == "place_of_performance_scope": - if value == "domestic": - queryset = queryset.filter(pop_scope_q) - elif value == "foreign": - queryset = queryset.exclude(pop_scope_q) - else: - raise InvalidParameterException("Invalid filter: place_of_performance_scope is invalid.") - - elif key == "place_of_performance_locations": - queryset = queryset.filter(geocode_filter_subaward_locations("sub_place_of_perform", value)) - - elif key == "award_amounts": - queryset &= total_obligation_queryset(value, SubawardSearch, filters, is_subaward=True) - - elif key == "award_ids": - queryset = build_award_ids_filter(queryset, value, ("piid", "fain")) - - elif key == PSCCodes.underscore_name: - q = PSCCodes.build_tas_codes_filter(value) - queryset = queryset.filter(q) if q else queryset - - # add "naics_codes" (column naics) after NAICS are mapped to subawards - elif key in ("contract_pricing_type_codes"): - if len(value) != 0: - queryset &= SubawardSearch.objects.filter(type_of_contract_pricing__in=value) - - elif key == "program_numbers": - if len(value) != 0: - queryset = queryset.filter( - Exists( - TransactionNormalized.objects.filter( - award_id=OuterRef("award_id"), - assistance_data__cfda_number__in=value, - ) - ) - ) - - elif key in ("set_aside_type_codes", "extent_competed_type_codes"): - or_queryset = Q() - filter_to_col = {"set_aside_type_codes": "type_set_aside", "extent_competed_type_codes": "extent_competed"} - in_query = [v for v in value] - for v in value: - or_queryset |= Q(**{"{}__exact".format(filter_to_col[key]): in_query}) - queryset = queryset.filter(or_queryset) - - # Because these two filters OR with each other, we need to know about the presence of both filters to know what to do - # This filter was picked arbitrarily to be the one that checks for the other - elif key == TasCodes.underscore_name: - q = TasCodes.build_tas_codes_filter(queryset, value) - if TreasuryAccounts.underscore_name in filters.keys(): - q |= TreasuryAccounts.build_tas_codes_filter(queryset, filters[TreasuryAccounts.underscore_name]) - queryset = queryset.filter(q) - - elif key == TreasuryAccounts.underscore_name and TasCodes.underscore_name not in filters.keys(): + raise InvalidParameterException(f"Invalid filter: {key} has null as its value.") + + if key not in valid_keys: + raise InvalidParameterException(f"Invalid filter: {key} does not exist.") + + # Handle TreasuryAccounts special case + if key == TreasuryAccounts.underscore_name and TasCodes.underscore_name not in filters: queryset = queryset.filter(TreasuryAccounts.build_tas_codes_filter(queryset, value)) + continue + + # Skip TreasuryAccounts if TasCodes is present (handled in TasCodes handler) + if key == TreasuryAccounts.underscore_name: + continue + + # Apply filter handler + handler = filter_handlers.get(key) + if handler: + queryset = handler(queryset, value) + + return queryset + + +# Filter handler functions +def _handle_keywords_filter(queryset: QuerySet, value: list[str]) -> QuerySet: + def keyword_parse(keyword: str) -> Q: + filter_obj = Q(keyword_ts_vector=keyword) | Q(award_ts_vector=keyword) + if len(keyword) == 4 and PSC.objects.filter(code__iexact=keyword).exists(): + filter_obj |= Q(product_or_service_code__iexact=keyword) + return filter_obj + + filter_obj = Q() + for keyword in value: + filter_obj |= keyword_parse(keyword) + + # Search for DUNS + potential_duns = [x for x in value if len(x) == 9] + if potential_duns: + filter_obj |= Q(sub_awardee_or_recipient_uniqu__in=potential_duns) | Q( + sub_ultimate_parent_unique_ide__in=potential_duns + ) + + # Search for UEI + potential_ueis = [uei.upper() for uei in value if len(uei) == 12] + if potential_ueis: + filter_obj |= Q(sub_awardee_or_recipient_uei__in=potential_ueis) | Q( + sub_ultimate_parent_uei__in=potential_ueis + ) + + return queryset.filter(filter_obj) + + +def _handle_description_filter(queryset: QuerySet, value: str) -> QuerySet: + return queryset.filter(subaward_description__icontains=value) + + +def _handle_transaction_keyword_filter(queryset: QuerySet, value: str) -> QuerySet: + transaction_ids = elasticsearch_helper.get_download_ids(keyword=value, field="transaction_id") + transaction_ids = list(itertools.chain.from_iterable(transaction_ids)) + logger.info(f"Found {len(transaction_ids)} transactions based on keyword: {value}") + transaction_ids = [str(tid) for tid in transaction_ids] + + queryset = queryset.filter(latest_transaction__isnull=False) + sql_fragment = '"subaward_search"."latest_transaction_id" = ANY(\'{{{}}}\'::int[])' + return queryset.extra(where=[sql_fragment.format(",".join(transaction_ids))]) + + +def _handle_time_period_filter(queryset: QuerySet, value: list[dict[str, Any]], for_downloads: bool) -> QuerySet: + min_date = API_MIN_DATE if for_downloads else API_SEARCH_MIN_DATE + return queryset & combine_date_range_queryset(value, SubawardSearch, min_date, API_MAX_DATE, is_subaward=True) + + +def _handle_award_type_codes_filter(queryset: QuerySet, value: list[str]) -> QuerySet: + return queryset.filter(prime_award_type__in=value) - elif key == "def_codes": - queryset = queryset.filter(DefCodes.build_def_codes_filter(value)) - - elif key == "program_activities": - query_filter_predicates = [Q(program_activity_id__isnull=False)] - award_ids_filtered_by_program_activities = [] - for program_activity in value: - if "name" in program_activity: - query_filter_predicates.append( - Q(program_activity__program_activity_name=program_activity["name"].upper()) - ) - if "code" in program_activity: - query_filter_predicates.append(Q(program_activity__program_activity_code=program_activity["code"])) - filter_ = FinancialAccountsByAwards.objects.filter(*query_filter_predicates) - award_ids_filtered_by_program_activities.extend(list(filter_.values_list("award_id", flat=True))) - - queryset &= SubawardSearch.objects.filter(award_id__in=award_ids_filtered_by_program_activities) + +def _handle_prime_and_sub_award_types_filter(queryset: QuerySet, value: dict[str, Any]) -> QuerySet: + award_types = value.get("sub_awards") + return queryset.filter(prime_award_group__in=award_types) if award_types else queryset + + +def _handle_agencies_filter(queryset: QuerySet, value: list[dict[str, str]]) -> QuerySet: + funding_toptier = Q() + funding_subtier = Q() + awarding_toptier = Q() + awarding_subtier = Q() + + for v in value: + agency_type = v["type"] + tier = v["tier"] + name = v["name"] + + if agency_type == "funding": + if tier == "toptier": + funding_toptier |= Q(funding_toptier_agency_name=name) + elif tier == "subtier": + base_q = Q(funding_subtier_agency_name=name) + funding_subtier |= base_q & Q( + funding_toptier_agency_name=v["toptier_name"]) if "toptier_name" in v else base_q + elif agency_type == "awarding": + if tier == "toptier": + awarding_toptier |= Q(awarding_toptier_agency_name=name) + elif tier == "subtier": + base_q = Q(awarding_subtier_agency_name=name) + awarding_subtier |= base_q & Q( + awarding_toptier_agency_name=v["toptier_name"]) if "toptier_name" in v else base_q + + funding_filter = funding_toptier | funding_subtier + awarding_filter = awarding_toptier | awarding_subtier + + return queryset.filter(funding_filter & awarding_filter) + + +def _handle_legal_entities_filter(queryset: QuerySet, value: Any) -> QuerySet: + logger.info(f'API request included "legal_entities" key. No filtering will occur with provided value "{value}"') return queryset + + +def _handle_recipient_search_text_filter(queryset: QuerySet, value: list[str]) -> QuerySet: + def recip_string_parse(recipient_string: str) -> Q: + upper_recipient_string = recipient_string.upper() + filter_obj = Q(recipient_name_ts_vector=upper_recipient_string) + + if len(upper_recipient_string) == 9 and upper_recipient_string[:5].isnumeric(): + filter_obj |= Q(sub_awardee_or_recipient_uniqu=upper_recipient_string) + elif len(upper_recipient_string) == 12: + filter_obj |= Q(sub_awardee_or_recipient_uei=upper_recipient_string) + + return filter_obj + + filter_obj = Q() + for recipient in value: + filter_obj |= recip_string_parse(recipient) + + return queryset.filter(filter_obj) + + +def _handle_scope_filter(queryset: QuerySet, value: str, scope_q: Q, filter_name: str) -> QuerySet: + if value == "domestic": + return queryset.filter(scope_q) + if value == "foreign": + return queryset.exclude(scope_q) + raise InvalidParameterException(f"Invalid filter: {filter_name} type is invalid.") + + +def _handle_locations_filter(queryset: QuerySet, value: list[dict[str, Any]], prefix: str) -> QuerySet: + return queryset.filter(geocode_filter_subaward_locations(prefix, value)) + + +def _handle_recipient_type_names_filter(queryset: QuerySet, value: list[str]) -> QuerySet: + return queryset.filter(business_categories__overlap=value) if value else queryset + + +def _handle_award_amounts_filter(queryset: QuerySet, value: dict[str, Any], filters: dict[str, Any]) -> QuerySet: + return queryset & total_obligation_queryset(value, SubawardSearch, filters, is_subaward=True) + + +def _handle_award_ids_filter(queryset: QuerySet, value: list[str]) -> QuerySet: + return build_award_ids_filter(queryset, value, ("piid", "fain")) + + +def _handle_psc_codes_filter(queryset: QuerySet, value: list[str]) -> QuerySet: + q = PSCCodes.build_tas_codes_filter(value) + return queryset.filter(q) if q else queryset + + +def _handle_contract_pricing_filter(queryset: QuerySet, value: list[str]) -> QuerySet: + return queryset & SubawardSearch.objects.filter(type_of_contract_pricing__in=value) if value else queryset + + +def _handle_program_numbers_filter(queryset: QuerySet, value: list[str]) -> QuerySet: + if not value: + return queryset + + return queryset.filter( + Exists( + TransactionNormalized.objects.filter( + award_id=OuterRef("award_id"), + assistance_data__cfda_number__in=value, + ) + ) + ) + + +def _handle_set_aside_extent_filter(queryset: QuerySet, value: list[str], key: str, column: str) -> QuerySet: + or_queryset = Q() + for item in value: + or_queryset |= Q(**{f"{column}__exact": item}) + return queryset.filter(or_queryset) + + +def _handle_tas_codes_filter(queryset: QuerySet, value: list[dict[str, Any]], filters: dict[str, Any]) -> QuerySet: + q = TasCodes.build_tas_codes_filter(queryset, value) + if TreasuryAccounts.underscore_name in filters: + q |= TreasuryAccounts.build_tas_codes_filter(queryset, filters[TreasuryAccounts.underscore_name]) + return queryset.filter(q) + + +def _handle_def_codes_filter(queryset: QuerySet, value: list[str]) -> QuerySet: + return queryset.filter(DefCodes.build_def_codes_filter(value)) + + +def _handle_program_activities_filter(queryset: QuerySet, value: list[dict[str, str]]) -> QuerySet: + query_filter_predicates = [Q(program_activity_id__isnull=False)] + award_ids_filtered = [] + + for program_activity in value: + if "name" in program_activity: + query_filter_predicates.append( + Q(program_activity__program_activity_name=program_activity["name"].upper()) + ) + if "code" in program_activity: + query_filter_predicates.append( + Q(program_activity__program_activity_code=program_activity["code"]) + ) + + filter_ = FinancialAccountsByAwards.objects.filter(*query_filter_predicates) + award_ids_filtered.extend(list(filter_.values_list("award_id", flat=True))) + + return queryset & SubawardSearch.objects.filter(award_id__in=award_ids_filtered)