Coverage for src/usaspending/queries/funding_search.py: 100%

1"""Funding search query builder for USASpending data.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any, Dict, TYPE_CHECKING 

6 

7from ..exceptions import ValidationError 

8from ..models.funding import Funding 

9from .query_builder import QueryBuilder 

10from ..logging_config import USASpendingLogger 

11 

12if TYPE_CHECKING: 

13 from ..client import USASpending 

14 

15logger = USASpendingLogger.get_logger(__name__) 

16 

17 

18class FundingSearch(QueryBuilder["Funding"]): 

19 """ 

20 Builds and executes a funding search query, allowing for retrieval 

21 of federal account funding data for a specific award. 

22 """ 

23 

24 # Map user-friendly sort field names to API field names 

25 SORT_FIELD_MAP = { 

26 "account_title": "account_title", 

27 "awarding_agency": "awarding_agency_name", 

28 "disaster_code": "disaster_emergency_fund_code", 

29 "federal_account": "federal_account", 

30 "funding_agency": "funding_agency_name", 

31 "gross_outlay": "gross_outlay_amount", 

32 "object_class": "object_class", 

33 "program_activity": "program_activity", 

34 "reporting_date": "reporting_fiscal_date", 

35 "fiscal_date": "reporting_fiscal_date", 

36 "obligated_amount": "transaction_obligated_amount", 

37 "obligation": "transaction_obligated_amount", 

38 } 

39 

40 def __init__(self, client: "USASpending"): 

41 """ 

42 Initializes the FundingSearch query builder. 

43 

44 Args: 

45 client: The USASpending client instance. 

46 """ 

47 super().__init__(client) 

48 self._award_id: str = None 

49 self._sort_field: str = "reporting_fiscal_date" 

50 self._sort_order: str = "desc" 

51 

52 @property 

53 def _endpoint(self) -> str: 

54 """The API endpoint for this query.""" 

55 return "/awards/funding/" 

56 

57 def _clone(self) -> FundingSearch: 

58 """Creates an immutable copy of the query builder.""" 

59 clone = super()._clone() 

60 clone._award_id = self._award_id 

61 clone._sort_field = self._sort_field 

62 clone._sort_order = self._sort_order 

63 return clone 

64 

65 def _build_payload(self, page: int) -> Dict[str, Any]: 

66 """Constructs the final API request payload.""" 

67 if not self._award_id: 

68 raise ValidationError( 

69 "An award_id is required. Use the .for_award() method." 

70 ) 

71 

72 payload = { 

73 "award_id": self._award_id, 

74 "limit": self._get_effective_page_size(), 

75 "page": page, 

76 "sort": self._sort_field, 

77 "order": self._sort_order, 

78 } 

79 

80 return payload 

81 

82 def _transform_result(self, result: Dict[str, Any]) -> Funding: 

83 """Transforms a single API result item into a Funding model.""" 

84 return Funding(result) 

85 

86 def count(self) -> int: 

87 """ 

88 Counts the number of funding records for the award. 

89 

90 Since the funding endpoint doesn't provide a count API, 

91 we need to iterate through all pages to get the count. 

92 """ 

93 logger.debug(f"{self.__class__.__name__}.count() called") 

94 

95 if not self._award_id: 

96 raise ValidationError( 

97 "An award_id is required. Use the .for_award() method." 

98 ) 

99 

100 # Iterate through all results to count 

101 count = 0 

102 for _ in self: 

103 count += 1 

104 

105 logger.info( 

106 f"{self.__class__.__name__}.count() = {count} funding records " 

107 f"for award {self._award_id}" 

108 ) 

109 return count 

110 

111 # ========================================================================== 

112 # Filter Methods 

113 # ========================================================================== 

114 

115 def for_award(self, award_id: str) -> FundingSearch: 

116 """ 

117 Filter funding records for a specific award. 

118 

119 Args: 

120 award_id: The unique award identifier. 

121 

122 Returns: 

123 A new FundingSearch instance with the award filter applied. 

124 """ 

125 if not award_id: 

126 raise ValidationError("award_id cannot be empty") 

127 

128 clone = self._clone() 

129 clone._award_id = str(award_id).strip() 

130 return clone 

131 

132 def order_by(self, field: str, direction: str = "desc") -> FundingSearch: 

133 """ 

134 Set the sort order for results. 

135 

136 Args: 

137 field: The field to sort by. Can be a user-friendly name or API field name. 

138 User-friendly names include: 

139 - 'account_title', 'awarding_agency', 'disaster_code' 

140 - 'federal_account', 'funding_agency', 'gross_outlay' 

141 - 'object_class', 'program_activity', 'reporting_date' 

142 - 'fiscal_date', 'obligated_amount', 'obligation' 

143 direction: Sort direction - 'asc' or 'desc' (default: 'desc') 

144 

145 Returns: 

146 A new FundingSearch instance with the sort configuration applied. 

147 """ 

148 # Validate direction 

149 if direction not in ["asc", "desc"]: 

150 raise ValidationError( 

151 f"Invalid sort direction: {direction}. Must be 'asc' or 'desc'." 

152 ) 

153 

154 # Map user-friendly field names to API field names 

155 api_field = self.SORT_FIELD_MAP.get(field.lower(), field) 

156 

157 # Validate that the field is supported by the API 

158 valid_api_fields = [ 

159 "account_title", 

160 "awarding_agency_name", 

161 "disaster_emergency_fund_code", 

162 "federal_account", 

163 "funding_agency_name", 

164 "gross_outlay_amount", 

165 "object_class", 

166 "program_activity", 

167 "reporting_fiscal_date", 

168 "transaction_obligated_amount", 

169 ] 

170 

171 if api_field not in valid_api_fields: 

172 raise ValidationError( 

173 f"Invalid sort field: {field}. " 

174 f"Valid fields are: {', '.join(self.SORT_FIELD_MAP.keys())}" 

175 ) 

176 

177 clone = self._clone() 

178 clone._sort_field = api_field 

179 clone._sort_order = direction 

180 return clone