tango/resources/
entities.rs

1//! `GET /api/entities/` — list and stream federal vendor / recipient records.
2
3use crate::client::Client;
4use crate::error::{Error, Result};
5use crate::internal::{apply_pagination, push_opt};
6use crate::pagination::{FetchFn, Page, PageStream};
7use crate::resources::agencies::urlencoding;
8use crate::Record;
9use bon::Builder;
10use std::collections::BTreeMap;
11use std::sync::Arc;
12
13/// Options for [`Client::list_entities`] and [`Client::iterate_entities`].
14///
15/// Mirrors the Go SDK's `ListEntitiesOptions`.
16#[derive(Debug, Clone, Default, Builder, PartialEq, Eq)]
17#[non_exhaustive]
18pub struct ListEntitiesOptions {
19    // ----- Pagination + shape -----
20    /// 1-based page number. Mutually exclusive with [`cursor`](Self::cursor).
21    #[builder(into)]
22    pub page: Option<u32>,
23    /// Page size.
24    #[builder(into)]
25    pub limit: Option<u32>,
26    /// Keyset cursor.
27    #[builder(into)]
28    pub cursor: Option<String>,
29    /// Comma-separated field selector.
30    #[builder(into)]
31    pub shape: Option<String>,
32    /// Collapse nested objects into dot-separated keys.
33    #[builder(default)]
34    pub flat: bool,
35    /// When [`flat`](Self::flat) is also true, flatten list-valued fields.
36    #[builder(default)]
37    pub flat_lists: bool,
38
39    // ----- Resource filters -----
40    /// Free-text search filter.
41    #[builder(into)]
42    pub search: Option<String>,
43    /// CAGE code filter.
44    #[builder(into)]
45    pub cage_code: Option<String>,
46    /// Alias for [`cage_code`](Self::cage_code): both filter the same field.
47    /// The API rejects a request that sets both, so use one or the other.
48    #[builder(into)]
49    pub cage: Option<String>,
50    /// NAICS code filter.
51    #[builder(into)]
52    pub naics: Option<String>,
53    /// Entity name filter.
54    #[builder(into)]
55    pub name: Option<String>,
56    /// PSC code filter.
57    #[builder(into)]
58    pub psc: Option<String>,
59    /// SAM purpose-of-registration code.
60    #[builder(into)]
61    pub purpose_of_registration_code: Option<String>,
62    /// Socioeconomic indicator filter.
63    #[builder(into)]
64    pub socioeconomic: Option<String>,
65    /// US state filter (2-letter code).
66    #[builder(into)]
67    pub state: Option<String>,
68    /// Lower bound for total obligated awards (dollars).
69    #[builder(into)]
70    pub total_awards_obligated_gte: Option<String>,
71    /// Upper bound for total obligated awards (dollars).
72    #[builder(into)]
73    pub total_awards_obligated_lte: Option<String>,
74    /// UEI filter.
75    #[builder(into)]
76    pub uei: Option<String>,
77    /// ZIP code filter.
78    #[builder(into)]
79    pub zip_code: Option<String>,
80
81    /// Escape hatch for filter keys not yet first-classed on this struct.
82    #[builder(default)]
83    pub extra: BTreeMap<String, String>,
84}
85
86impl ListEntitiesOptions {
87    pub(crate) fn to_query(&self) -> Vec<(String, String)> {
88        let mut q = Vec::new();
89        apply_pagination(
90            &mut q,
91            self.page,
92            self.limit,
93            self.cursor.as_deref(),
94            self.shape.as_deref(),
95            self.flat,
96            self.flat_lists,
97        );
98        push_opt(&mut q, "search", self.search.as_deref());
99        push_opt(&mut q, "cage_code", self.cage_code.as_deref());
100        push_opt(&mut q, "cage", self.cage.as_deref());
101        push_opt(&mut q, "naics", self.naics.as_deref());
102        push_opt(&mut q, "name", self.name.as_deref());
103        push_opt(&mut q, "psc", self.psc.as_deref());
104        push_opt(
105            &mut q,
106            "purpose_of_registration_code",
107            self.purpose_of_registration_code.as_deref(),
108        );
109        push_opt(&mut q, "socioeconomic", self.socioeconomic.as_deref());
110        push_opt(&mut q, "state", self.state.as_deref());
111        push_opt(
112            &mut q,
113            "total_awards_obligated_gte",
114            self.total_awards_obligated_gte.as_deref(),
115        );
116        push_opt(
117            &mut q,
118            "total_awards_obligated_lte",
119            self.total_awards_obligated_lte.as_deref(),
120        );
121        push_opt(&mut q, "uei", self.uei.as_deref());
122        push_opt(&mut q, "zip_code", self.zip_code.as_deref());
123
124        for (k, v) in &self.extra {
125            if !v.is_empty() {
126                q.push((k.clone(), v.clone()));
127            }
128        }
129        q
130    }
131}
132
133/// Options for [`Client::get_entity`].
134#[derive(Debug, Clone, Default, Builder, PartialEq, Eq)]
135#[non_exhaustive]
136pub struct GetEntityOptions {
137    /// Shape selector. When empty, the server returns its comprehensive default.
138    #[builder(into)]
139    pub shape: Option<String>,
140    /// Flatten nested objects into dot-separated keys.
141    #[builder(default)]
142    pub flat: bool,
143    /// When `flat=true`, also flatten list-valued nested fields.
144    #[builder(default)]
145    pub flat_lists: bool,
146}
147
148impl GetEntityOptions {
149    pub(crate) fn to_query(&self) -> Vec<(String, String)> {
150        let mut q = Vec::new();
151        push_opt(&mut q, "shape", self.shape.as_deref());
152        if self.flat {
153            q.push(("flat".into(), "true".into()));
154        }
155        if self.flat_lists {
156            q.push(("flat_lists".into(), "true".into()));
157        }
158        q
159    }
160}
161
162impl Client {
163    /// `GET /api/entities/` — one page of entity records.
164    pub async fn list_entities(&self, opts: ListEntitiesOptions) -> Result<Page<Record>> {
165        let q = opts.to_query();
166        let bytes = self.get_bytes("/api/entities/", &q).await?;
167        Page::decode(&bytes)
168    }
169
170    /// `GET /api/entities/{uei}/` — fetch a single entity by UEI or CAGE code.
171    pub async fn get_entity(&self, uei: &str, opts: Option<GetEntityOptions>) -> Result<Record> {
172        if uei.is_empty() {
173            return Err(Error::Validation {
174                message: "get_entity: uei is required".into(),
175                response: None,
176            });
177        }
178        let q = opts.unwrap_or_default().to_query();
179        let path = format!("/api/entities/{}/", urlencoding(uei));
180        self.get_json::<Record>(&path, &q).await
181    }
182
183    /// Stream every entity record matching `opts`.
184    pub fn iterate_entities(&self, opts: ListEntitiesOptions) -> PageStream<Record> {
185        let opts = Arc::new(opts);
186        let fetch: FetchFn<Record> = Box::new(move |client, page, cursor| {
187            let mut next = (*opts).clone();
188            next.page = page;
189            next.cursor = cursor;
190            Box::pin(async move { client.list_entities(next).await })
191        });
192        PageStream::new(self.clone(), fetch)
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199
200    fn get_q(q: &[(String, String)], k: &str) -> Option<String> {
201        q.iter().find(|(kk, _)| kk == k).map(|(_, v)| v.clone())
202    }
203
204    #[test]
205    fn list_entities_all_filters_emit() {
206        let opts = ListEntitiesOptions::builder()
207            .search("Acme")
208            .cage_code("1ABC5")
209            .cage("1ABC5")
210            .naics("541512")
211            .name("Acme Corp")
212            .psc("D302")
213            .purpose_of_registration_code("Z1")
214            .socioeconomic("A5")
215            .state("VA")
216            .total_awards_obligated_gte("100000")
217            .total_awards_obligated_lte("999999")
218            .uei("UEI123456789")
219            .zip_code("22201")
220            .build();
221        let q = opts.to_query();
222        assert_eq!(get_q(&q, "search").as_deref(), Some("Acme"));
223        assert_eq!(get_q(&q, "cage_code").as_deref(), Some("1ABC5"));
224        assert_eq!(get_q(&q, "cage").as_deref(), Some("1ABC5"));
225        assert_eq!(get_q(&q, "naics").as_deref(), Some("541512"));
226        assert_eq!(get_q(&q, "name").as_deref(), Some("Acme Corp"));
227        assert_eq!(get_q(&q, "psc").as_deref(), Some("D302"));
228        assert_eq!(
229            get_q(&q, "purpose_of_registration_code").as_deref(),
230            Some("Z1")
231        );
232        assert_eq!(get_q(&q, "socioeconomic").as_deref(), Some("A5"));
233        assert_eq!(get_q(&q, "state").as_deref(), Some("VA"));
234        assert_eq!(
235            get_q(&q, "total_awards_obligated_gte").as_deref(),
236            Some("100000")
237        );
238        assert_eq!(
239            get_q(&q, "total_awards_obligated_lte").as_deref(),
240            Some("999999")
241        );
242        assert_eq!(get_q(&q, "uei").as_deref(), Some("UEI123456789"));
243        assert_eq!(get_q(&q, "zip_code").as_deref(), Some("22201"));
244    }
245
246    #[test]
247    fn list_entities_zero_value_omitted() {
248        let opts = ListEntitiesOptions::builder().build();
249        let q = opts.to_query();
250        assert!(q.is_empty(), "expected empty query, got {q:?}");
251    }
252
253    #[test]
254    fn list_entities_pagination_applied() {
255        let opts = ListEntitiesOptions::builder()
256            .page(2u32)
257            .limit(25u32)
258            .build();
259        let q = opts.to_query();
260        assert_eq!(get_q(&q, "page").as_deref(), Some("2"));
261        assert_eq!(get_q(&q, "limit").as_deref(), Some("25"));
262    }
263
264    #[test]
265    fn get_entity_opts_emits_shape_and_flat() {
266        let opts = GetEntityOptions::builder()
267            .shape(crate::SHAPE_ENTITIES_MINIMAL)
268            .flat(true)
269            .flat_lists(true)
270            .build();
271        let q = opts.to_query();
272        assert!(q.contains(&("shape".into(), crate::SHAPE_ENTITIES_MINIMAL.into())));
273        assert!(q.contains(&("flat".into(), "true".into())));
274        assert!(q.contains(&("flat_lists".into(), "true".into())));
275    }
276
277    #[tokio::test]
278    async fn get_entity_empty_uei_returns_validation() {
279        let client = Client::builder().api_key("x").build().expect("build");
280        let err = client.get_entity("", None).await.expect_err("must error");
281        match err {
282            Error::Validation { message, .. } => assert!(message.contains("uei")),
283            other => panic!("expected Validation, got {other:?}"),
284        }
285    }
286}
tango/resources/entities.rs

tango/resources/
entities.rs
