1use crate::config::{GraphqlAuth, GraphqlPagination, GraphqlStreamConfig};
4use async_trait::async_trait;
5use base64::Engine as _;
6use faucet_core::util::{self, DEFAULT_ERROR_BODY_MAX_LEN};
7use faucet_core::{AuthSpec, Credential, FaucetError, SharedAuthProvider, Stream, StreamPage};
8use jsonpath_rust::JsonPath;
9use reqwest::Client;
10use serde_json::{Value, json};
11use std::collections::HashMap;
12use std::pin::Pin;
13use std::time::Duration;
14
15const RETRY_MAX_ATTEMPTS: u32 = 3;
17const RETRY_BASE_BACKOFF: Duration = Duration::from_millis(500);
19
20pub struct GraphqlStream {
22 config: GraphqlStreamConfig,
23 client: Client,
24 auth_provider: Option<SharedAuthProvider>,
28 retry_policy: faucet_core::RetryPolicy,
32}
33
34fn credential_to_auth(cred: Credential) -> GraphqlAuth {
37 match cred {
38 Credential::Bearer(token) => GraphqlAuth::Bearer { token },
39 Credential::Token(token) => GraphqlAuth::Custom {
40 headers: HashMap::from([("Authorization".into(), token)]),
41 },
42 Credential::Header { name, value } => GraphqlAuth::Custom {
43 headers: HashMap::from([(name, value)]),
44 },
45 Credential::Basic { username, password } => GraphqlAuth::Custom {
46 headers: HashMap::from([(
47 "Authorization".into(),
48 format!(
49 "Basic {}",
50 base64::engine::general_purpose::STANDARD
51 .encode(format!("{username}:{password}"))
52 ),
53 )]),
54 },
55 }
56}
57
58impl GraphqlStream {
59 pub fn new(config: GraphqlStreamConfig) -> Self {
61 Self {
62 config,
63 client: Client::new(),
64 auth_provider: None,
65 retry_policy: faucet_core::RetryPolicy {
69 max_attempts: RETRY_MAX_ATTEMPTS + 1,
70 backoff: faucet_core::BackoffKind::Exponential,
71 base: RETRY_BASE_BACKOFF,
72 max: Duration::from_secs(60),
73 jitter: true,
74 retry_on: faucet_core::RetryClassSet::default(),
75 },
76 }
77 }
78
79 pub fn with_retry_policy(mut self, policy: faucet_core::RetryPolicy) -> Self {
84 self.retry_policy = policy;
85 self
86 }
87
88 pub fn with_auth_provider(mut self, provider: SharedAuthProvider) -> Self {
94 self.auth_provider = Some(provider);
95 self
96 }
97
98 pub async fn fetch_all(&self) -> Result<Vec<Value>, FaucetError> {
100 self.fetch_all_with_context(&std::collections::HashMap::new())
101 .await
102 }
103
104 async fn fetch_all_with_context(
106 &self,
107 context: &std::collections::HashMap<String, Value>,
108 ) -> Result<Vec<Value>, FaucetError> {
109 let mut all_records = Vec::new();
110 let mut cursor: Option<String> = None;
111 let mut pages_fetched = 0usize;
112 let mut warned_unresolved_has_next = false;
113
114 loop {
115 if let Some(max) = self.config.max_pages
116 && pages_fetched >= max
117 {
118 tracing::warn!("max pages ({max}) reached");
119 break;
120 }
121
122 let body = self.execute_query(&cursor, context).await?;
123 let records = self.extract_records(&body)?;
124 all_records.extend(records);
125 pages_fetched += 1;
126
127 match &self.config.pagination {
129 Some(pag) => {
130 let (step, unresolved) = decide_next_page(&body, pag, cursor.as_deref());
131 if unresolved && !warned_unresolved_has_next {
132 tracing::warn!(
133 path = %pag.has_next_page_path,
134 "GraphQL has_next_page path did not resolve to a boolean; \
135 deferring to cursor presence to decide pagination"
136 );
137 warned_unresolved_has_next = true;
138 }
139 match step {
140 PageStep::Stop => break,
141 PageStep::StopLoop => {
142 tracing::warn!("cursor loop detected, stopping pagination");
143 break;
144 }
145 PageStep::Advance(next) => cursor = Some(next),
146 }
147 }
148 None => break,
149 }
150 }
151
152 tracing::info!(
153 records = all_records.len(),
154 pages = pages_fetched,
155 "GraphQL fetch complete"
156 );
157 Ok(all_records)
158 }
159
160 async fn execute_query(
162 &self,
163 cursor: &Option<String>,
164 context: &std::collections::HashMap<String, Value>,
165 ) -> Result<Value, FaucetError> {
166 let mut variables = self.config.variables.clone();
167
168 if !context.is_empty()
170 && let Value::Object(ref mut map) = variables
171 {
172 for (key, value) in context {
173 map.insert(key.clone(), value.clone());
174 }
175 }
176
177 if let (Some(pag), Some(cursor_val)) = (&self.config.pagination, cursor)
179 && let Value::Object(ref mut map) = variables
180 {
181 map.insert(pag.cursor_variable.clone(), json!(cursor_val));
182 }
183 if let Some(pag) = &self.config.pagination
187 && self.config.batch_size != 0
188 && let Value::Object(map) = &mut variables
189 {
190 map.insert(
191 pag.page_size_variable.clone(),
192 json!(self.config.batch_size),
193 );
194 }
195
196 let payload = json!({
197 "query": self.config.query,
198 "variables": variables,
199 });
200
201 let mut req = self
202 .client
203 .post(&self.config.endpoint)
204 .headers(self.config.headers.clone())
205 .json(&payload);
206
207 let effective_auth: GraphqlAuth = if let Some(provider) = &self.auth_provider {
211 credential_to_auth(provider.credential().await?)
212 } else {
213 match &self.config.auth {
214 AuthSpec::Inline(a) => a.clone(),
215 AuthSpec::Reference(r) => {
216 return Err(FaucetError::Auth(format!(
217 "auth references provider '{}' but no provider was supplied; \
218 set one via the CLI `auth:` catalog or `with_auth_provider`",
219 r.name
220 )));
221 }
222 }
223 };
224
225 match effective_auth {
227 GraphqlAuth::None => {}
228 GraphqlAuth::Bearer { token } => {
229 req = req.bearer_auth(token);
230 }
231 GraphqlAuth::Custom { headers } => {
232 let mut hm = reqwest::header::HeaderMap::new();
233 for (name, value) in &headers {
234 let n =
235 reqwest::header::HeaderName::from_bytes(name.as_bytes()).map_err(|e| {
236 FaucetError::Auth(format!("invalid custom header name {name:?}: {e}"))
237 })?;
238 let v = reqwest::header::HeaderValue::from_str(value).map_err(|e| {
239 FaucetError::Auth(format!("invalid custom header value for {name:?}: {e}"))
240 })?;
241 hm.insert(n, v);
242 }
243 req = req.headers(hm);
244 }
245 }
246
247 let body: Value = faucet_core::execute_with_policy(&self.retry_policy, None, || {
252 let attempt = req.try_clone();
253 async move {
254 let req = attempt.ok_or_else(|| {
255 FaucetError::Source("graphql: request is not cloneable for retry".into())
256 })?;
257 let resp = req.send().await.map_err(FaucetError::Http)?;
258 let resp = util::check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
259 resp.json().await.map_err(FaucetError::Http)
260 }
261 })
262 .await?;
263
264 if let Some(errors) = body.get("errors")
266 && let Some(arr) = errors.as_array()
267 && !arr.is_empty()
268 {
269 let msg = arr
270 .iter()
271 .filter_map(|e| e.get("message").and_then(|m| m.as_str()))
272 .collect::<Vec<_>>()
273 .join("; ");
274 let lower = msg.to_lowercase();
280 if self.config.batch_size == 0
281 && let Some(pag) = &self.config.pagination
282 {
283 let var_name = pag.page_size_variable.to_lowercase();
284 if lower.contains(&var_name)
285 && (lower.contains("non-null")
286 || lower.contains("non null")
287 || lower.contains("must not be null")
288 || lower.contains("cannot be null")
289 || lower.contains("required"))
290 {
291 return Err(FaucetError::Config(format!(
292 "batch_size = 0 requires the upstream to accept a null {}: argument \
293 (GraphQL errors: {msg})",
294 pag.page_size_variable
295 )));
296 }
297 }
298 return Err(FaucetError::HttpStatus {
299 status: 200,
300 url: self.config.endpoint.clone(),
301 body: format!("GraphQL errors: {msg}"),
302 });
303 }
304
305 Ok(body)
306 }
307
308 fn extract_records(&self, body: &Value) -> Result<Vec<Value>, FaucetError> {
310 match &self.config.records_path {
311 Some(path) => util::extract_records(body, Some(path)),
312 None => {
313 match body.get("data") {
319 Some(Value::Null) | None => Ok(Vec::new()),
320 Some(data) => Ok(vec![data.clone()]),
321 }
322 }
323 }
324 }
325
326 fn stream_pages_inner(
339 &self,
340 context: &std::collections::HashMap<String, Value>,
341 ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + '_>> {
342 let owned_context: std::collections::HashMap<String, Value> = context.clone();
344
345 Box::pin(async_stream::try_stream! {
346 let mut cursor: Option<String> = None;
347 let mut pages_fetched = 0usize;
348 let mut warned_unresolved_has_next = false;
349 let running_max: Option<Value> = None;
354 let mut bookmark_emitted = false;
355
356 loop {
357 if let Some(max) = self.config.max_pages
358 && pages_fetched >= max
359 {
360 tracing::warn!("max pages ({max}) reached");
361 break;
362 }
363
364 let body = self.execute_query(&cursor, &owned_context).await?;
365 let records = self.extract_records(&body)?;
366 pages_fetched += 1;
367
368 let has_next = match &self.config.pagination {
371 Some(pag) => {
372 let (step, unresolved) =
373 decide_next_page(&body, pag, cursor.as_deref());
374 if unresolved && !warned_unresolved_has_next {
375 tracing::warn!(
376 path = %pag.has_next_page_path,
377 "GraphQL has_next_page path did not resolve to a boolean; \
378 deferring to cursor presence to decide pagination"
379 );
380 warned_unresolved_has_next = true;
381 }
382 match step {
383 PageStep::Stop => false,
384 PageStep::StopLoop => {
385 tracing::warn!("cursor loop detected, stopping pagination");
386 false
387 }
388 PageStep::Advance(next) => {
389 cursor = Some(next);
390 true
391 }
392 }
393 }
394 None => false,
395 };
396
397 if has_next {
398 yield StreamPage { records, bookmark: None };
400 } else {
401 bookmark_emitted = running_max.is_some();
404 yield StreamPage {
405 records,
406 bookmark: running_max.clone(),
407 };
408 break;
409 }
410 }
411
412 if !bookmark_emitted && running_max.is_some() {
419 yield StreamPage {
420 records: Vec::new(),
421 bookmark: running_max,
422 };
423 }
424
425 tracing::info!(
426 pages = pages_fetched,
427 batch_size = self.config.batch_size,
428 "GraphQL source stream complete",
429 );
430 })
431 }
432}
433
434#[async_trait]
435impl faucet_core::Source for GraphqlStream {
436 async fn fetch_with_context(
437 &self,
438 context: &std::collections::HashMap<String, serde_json::Value>,
439 ) -> Result<Vec<Value>, FaucetError> {
440 self.fetch_all_with_context(context).await
441 }
442
443 fn stream_pages<'a>(
450 &'a self,
451 context: &'a std::collections::HashMap<String, Value>,
452 _batch_size: usize,
453 ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
454 self.stream_pages_inner(context)
455 }
456
457 fn config_schema(&self) -> serde_json::Value {
458 serde_json::to_value(faucet_core::schema_for!(GraphqlStreamConfig))
459 .expect("schema serialization")
460 }
461
462 fn dataset_uri(&self) -> String {
463 faucet_core::redact_uri_credentials(&self.config.endpoint)
464 }
465}
466
467fn extract_string(body: &Value, path: &str) -> Option<String> {
468 let results = body.query(path).ok()?;
469 match results.first()? {
470 Value::String(s) => Some(s.clone()),
471 _ => None,
472 }
473}
474
475fn extract_bool(body: &Value, path: &str) -> Option<bool> {
476 let results = body.query(path).ok()?;
477 results.first()?.as_bool()
478}
479
480#[derive(Debug, PartialEq)]
482enum PageStep {
483 Stop,
485 StopLoop,
488 Advance(String),
490}
491
492fn decide_next_page(
501 body: &Value,
502 pag: &GraphqlPagination,
503 prev_cursor: Option<&str>,
504) -> (PageStep, bool) {
505 let (stop, unresolved) = match extract_bool(body, &pag.has_next_page_path) {
506 Some(false) => (true, false),
507 Some(true) => (false, false),
508 None => (false, true),
510 };
511 if stop {
512 return (PageStep::Stop, unresolved);
513 }
514 match extract_string(body, &pag.cursor_path) {
515 None => (PageStep::Stop, unresolved),
516 Some(next) if Some(next.as_str()) == prev_cursor => (PageStep::StopLoop, unresolved),
517 Some(next) => (PageStep::Advance(next), unresolved),
518 }
519}
520
521#[cfg(test)]
522mod tests {
523 use super::*;
524
525 #[test]
526 fn extract_string_from_json() {
527 let body = json!({"data": {"users": {"pageInfo": {"endCursor": "abc123"}}}});
528 assert_eq!(
529 extract_string(&body, "$.data.users.pageInfo.endCursor"),
530 Some("abc123".into())
531 );
532 }
533
534 #[test]
535 fn extract_bool_from_json() {
536 let body = json!({"data": {"users": {"pageInfo": {"hasNextPage": true}}}});
537 assert_eq!(
538 extract_bool(&body, "$.data.users.pageInfo.hasNextPage"),
539 Some(true)
540 );
541 }
542
543 fn pageinfo_pagination() -> GraphqlPagination {
544 GraphqlPagination {
545 has_next_page_path: "$.data.users.pageInfo.hasNextPage".into(),
546 cursor_path: "$.data.users.pageInfo.endCursor".into(),
547 ..GraphqlPagination::default()
548 }
549 }
550
551 #[test]
552 fn decide_next_page_advances_when_has_next_true() {
553 let body =
554 json!({"data": {"users": {"pageInfo": {"hasNextPage": true, "endCursor": "c2"}}}});
555 let (step, unresolved) = decide_next_page(&body, &pageinfo_pagination(), Some("c1"));
556 assert_eq!(step, PageStep::Advance("c2".into()));
557 assert!(!unresolved);
558 }
559
560 #[test]
561 fn decide_next_page_stops_when_has_next_false() {
562 let body =
563 json!({"data": {"users": {"pageInfo": {"hasNextPage": false, "endCursor": "c2"}}}});
564 let (step, unresolved) = decide_next_page(&body, &pageinfo_pagination(), Some("c1"));
565 assert_eq!(step, PageStep::Stop);
566 assert!(!unresolved);
567 }
568
569 #[test]
570 fn decide_next_page_detects_cursor_loop() {
571 let body =
572 json!({"data": {"users": {"pageInfo": {"hasNextPage": true, "endCursor": "c1"}}}});
573 let (step, _) = decide_next_page(&body, &pageinfo_pagination(), Some("c1"));
574 assert_eq!(step, PageStep::StopLoop);
575 }
576
577 #[test]
578 fn decide_next_page_defers_to_cursor_when_has_next_unresolved() {
579 let body = json!({"data": {"users": {"pageInfo": {"endCursor": "c2"}}}}); let (step, unresolved) = decide_next_page(&body, &pageinfo_pagination(), Some("c1"));
584 assert_eq!(
585 step,
586 PageStep::Advance("c2".into()),
587 "unresolved has-next must defer to cursor presence, not stop"
588 );
589 assert!(unresolved, "the caller is told to warn once");
590
591 let body_no_cursor = json!({"data": {"users": {"pageInfo": {}}}});
593 let (step, unresolved) =
594 decide_next_page(&body_no_cursor, &pageinfo_pagination(), Some("c1"));
595 assert_eq!(step, PageStep::Stop);
596 assert!(unresolved);
597 }
598
599 #[test]
600 fn extract_records_with_path() {
601 let config =
602 GraphqlStreamConfig::new("https://api.example.com/graphql", "query { users { id } }")
603 .records_path("$.data.users[*]");
604 let stream = GraphqlStream::new(config);
605 let body = json!({"data": {"users": [{"id": 1}, {"id": 2}]}});
606 let records = stream.extract_records(&body).unwrap();
607 assert_eq!(records.len(), 2);
608 assert_eq!(records[0]["id"], 1);
609 }
610
611 #[test]
612 fn extract_records_without_path_returns_data() {
613 let config =
614 GraphqlStreamConfig::new("https://api.example.com/graphql", "query { user { id } }");
615 let stream = GraphqlStream::new(config);
616 let body = json!({"data": {"user": {"id": 1}}});
617 let records = stream.extract_records(&body).unwrap();
618 assert_eq!(records.len(), 1);
619 assert_eq!(records[0]["user"]["id"], 1);
620 }
621
622 #[test]
623 fn extract_records_without_path_null_data_yields_empty() {
624 let config =
628 GraphqlStreamConfig::new("https://api.example.com/graphql", "query { user { id } }");
629 let stream = GraphqlStream::new(config);
630 let body = json!({ "data": null });
631 let records = stream.extract_records(&body).unwrap();
632 assert!(
633 records.is_empty(),
634 "expected empty Vec for null `data`, got {records:?}"
635 );
636 }
637
638 #[test]
639 fn extract_records_without_path_absent_data_yields_empty() {
640 let config =
643 GraphqlStreamConfig::new("https://api.example.com/graphql", "query { user { id } }");
644 let stream = GraphqlStream::new(config);
645 let body = json!({ "extensions": { "foo": 1 } });
646 let records = stream.extract_records(&body).unwrap();
647 assert!(
648 records.is_empty(),
649 "expected empty Vec when `data` is absent, got {records:?}"
650 );
651 }
652
653 #[test]
654 fn dataset_uri_returns_endpoint() {
655 use faucet_core::Source;
656 let stream = GraphqlStream::new(GraphqlStreamConfig::new(
657 "https://api.example.com/graphql",
658 "query { id }",
659 ));
660 assert_eq!(stream.dataset_uri(), "https://api.example.com/graphql");
661 }
662
663 #[test]
664 fn dataset_uri_redacts_credentials() {
665 use faucet_core::Source;
666 let stream = GraphqlStream::new(GraphqlStreamConfig::new(
667 "https://user:pw@api.example.com/graphql",
668 "query { id }",
669 ));
670 assert_eq!(stream.dataset_uri(), "https://api.example.com/graphql");
671 }
672
673 #[test]
674 fn default_retry_policy_reproduces_legacy_constants() {
675 let stream = GraphqlStream::new(GraphqlStreamConfig::new(
676 "https://api.example.com/graphql",
677 "query { id }",
678 ));
679 assert_eq!(stream.retry_policy.max_attempts, RETRY_MAX_ATTEMPTS + 1);
680 assert_eq!(stream.retry_policy.base, RETRY_BASE_BACKOFF);
681 }
682
683 #[test]
684 fn with_retry_policy_overrides_the_default() {
685 let policy = faucet_core::RetryPolicy {
686 max_attempts: 9,
687 base: Duration::from_secs(7),
688 ..faucet_core::RetryPolicy::default()
689 };
690 let stream = GraphqlStream::new(GraphqlStreamConfig::new(
691 "https://api.example.com/graphql",
692 "query { id }",
693 ))
694 .with_retry_policy(policy);
695 assert_eq!(stream.retry_policy.max_attempts, 9);
696 assert_eq!(stream.retry_policy.base, Duration::from_secs(7));
697 }
698}