faucet_source_redshift/
stream.rs

1//! Amazon Redshift query source implementation.
2//!
3//! Redshift is PostgreSQL wire-compatible, so the source streams rows through
4//! `sqlx`'s Postgres cursor, re-framing them into `batch_size`-sized pages. It
5//! supports full and incremental replication: an incremental read wraps the
6//! query in a `>=` cursor filter ordered by the cursor and drops the rows it
7//! already emitted at the bookmark value (`CursorTracker`).
8
9use std::collections::HashMap;
10use std::pin::Pin;
11use std::sync::Mutex;
12
13use async_trait::async_trait;
14use faucet_core::replication::cursor::{CursorBookmark, CursorTracker};
15use faucet_core::util::{quote_ident, substitute_context_bind_params};
16use faucet_core::{FaucetError, Source, Stream, StreamPage};
17use serde_json::Value;
18use sqlx::PgPool;
19
20use crate::config::{BOOKMARK_TOKEN, RedshiftReplication, RedshiftSourceConfig};
21use crate::convert::{bind_params, row_to_json};
22
23/// A source that executes a SQL query against Amazon Redshift and returns rows
24/// as JSON objects.
25pub struct RedshiftSource {
26    config: RedshiftSourceConfig,
27    pool: PgPool,
28    /// Bookmark applied via [`Source::apply_start_bookmark`] (incremental only).
29    start_bookmark: Mutex<Option<CursorBookmark>>,
30    /// Backend-call context (#880): the query and every page of rows run
31    /// through it.
32    calls: faucet_core::CallSlot,
33}
34
35/// Incremental-read context (column + starting position).
36struct IncrementalCtx {
37    column: String,
38    start: CursorBookmark,
39}
40
41impl RedshiftSource {
42    /// Create a new source. Validates config and builds a lazily-connected pool
43    /// (no I/O — connectivity is verified on first query or via
44    /// [`Source::check`]).
45    pub fn new(config: RedshiftSourceConfig) -> Result<Self, FaucetError> {
46        config.validate()?;
47        let pool =
48            faucet_common_redshift::build_pool_lazy(&config.connection, config.max_connections)?;
49        Ok(Self {
50            config,
51            pool,
52            start_bookmark: Mutex::new(None),
53            calls: faucet_core::CallSlot::new(faucet_core::CallContext::new().without_timeout()),
54        })
55    }
56
57    /// The effective incremental start bookmark (persisted bookmark, else the
58    /// configured `initial_value`), or `None` for full replication.
59    fn incremental_ctx(&self) -> Option<IncrementalCtx> {
60        match &self.config.replication {
61            RedshiftReplication::Full => None,
62            RedshiftReplication::Incremental {
63                column,
64                initial_value,
65            } => {
66                let start = self
67                    .start_bookmark
68                    .lock()
69                    .expect("start_bookmark mutex poisoned")
70                    .clone()
71                    .unwrap_or_else(|| CursorBookmark::initial(initial_value.clone()));
72                Some(IncrementalCtx {
73                    column: column.clone(),
74                    start,
75                })
76            }
77        }
78    }
79
80    /// Build the effective SQL and ordered positional bind values for a parent
81    /// context + optional incremental cursor, whose bookmark binds as
82    /// `CAST($n AS <cast>)` (the cursor column's type).
83    ///
84    /// Bind order is: static [`config.params`](RedshiftSourceConfig::params),
85    /// then parent-context `{key}` values, then the `${bookmark}` value (when
86    /// the query contains that token), then the bookmark of the outer
87    /// `column >= …` filter.
88    fn resolve_query(
89        &self,
90        context: &HashMap<String, Value>,
91        incr: Option<&IncrementalCtx>,
92        cast: &str,
93    ) -> (String, Vec<Value>) {
94        let mut binds = self.config.params.clone();
95        let mut sql = self.config.query.clone();
96
97        if !context.is_empty() {
98            let (rewritten, ctx_values) =
99                substitute_context_bind_params(&sql, context, binds.len() + 1, |i| format!("${i}"));
100            sql = rewritten;
101            binds.extend(ctx_values);
102        }
103
104        let Some(ctx) = incr else {
105            return (sql, binds);
106        };
107        if sql.contains(BOOKMARK_TOKEN) {
108            let marker = format!("CAST(${} AS {cast})", binds.len() + 1);
109            sql = sql.replace(BOOKMARK_TOKEN, &marker);
110            binds.push(ctx.start.value.clone());
111        }
112        binds.push(ctx.start.value.clone());
113        let c = quote_ident(&ctx.column);
114        let inner = sql.trim_end_matches(|ch: char| ch.is_whitespace() || ch == ';');
115        (
116            format!(
117                "SELECT * FROM ({inner}\n) AS q WHERE q.{c} >= CAST(${n} AS {cast}) ORDER BY q.{c}",
118                n = binds.len()
119            ),
120            binds,
121        )
122    }
123
124    /// The cursor column's type, as the server describes the query's output.
125    async fn cursor_cast(
126        &self,
127        context: &HashMap<String, Value>,
128        column: &str,
129    ) -> Result<String, FaucetError> {
130        use sqlx::Executor;
131        let (probe, _) = self.resolve_query(context, None, "");
132        let probe = probe.replace(BOOKMARK_TOKEN, "NULL");
133        let calls = self.calls.get();
134        let described = faucet_common_redshift::calls::once(&calls, "describe", async {
135            (&self.pool).describe(&probe).await.map_err(|e| {
136                FaucetError::Source(format!(
137                    "redshift: could not describe the incremental query: {e}"
138                ))
139            })
140        })
141        .await?;
142        let name = described
143            .columns()
144            .iter()
145            .find(|c| sqlx::Column::name(*c) == column)
146            .map(|c| sqlx::TypeInfo::name(sqlx::Column::type_info(c)).to_string())
147            .ok_or_else(|| {
148                FaucetError::Config(format!(
149                    "redshift: incremental replication `column` {column:?} is not a column of \
150                     the query's output"
151                ))
152            })?;
153        cast_type(&name)
154    }
155}
156
157/// A server-reported type name usable in `CAST(… AS <name>)`.
158fn cast_type(name: &str) -> Result<String, FaucetError> {
159    let ok = !name.is_empty()
160        && name
161            .chars()
162            .all(|c| c.is_ascii_alphanumeric() || matches!(c, '_' | ' ' | '"' | '[' | ']'));
163    if ok {
164        Ok(name.to_string())
165    } else {
166        Err(FaucetError::Config(format!(
167            "redshift: the incremental cursor column has type {name:?}, which cannot be \
168             compared to a bookmark; cast it in the query"
169        )))
170    }
171}
172
173/// Derive a stable state key from the host, database, and query.
174fn default_state_key(config: &RedshiftSourceConfig) -> String {
175    // Fixed FNV-1a: the key is durable and `DefaultHasher` is not stable across Rust releases.
176    let material = format!(
177        "{}\u{1f}{}\u{1f}{}\u{1f}{}",
178        config.connection.host, config.connection.port, config.connection.database, config.query
179    );
180    format!(
181        "redshift:{:016x}",
182        faucet_core::shard::shard_hash(&material)
183    )
184}
185
186#[async_trait]
187impl Source for RedshiftSource {
188    fn connector_name(&self) -> &'static str {
189        "redshift"
190    }
191
192    fn set_call_context(&self, ctx: faucet_core::io::CallContext) {
193        self.calls.install(ctx);
194    }
195
196    fn set_roundtrip_recorder(
197        &self,
198        recorder: std::sync::Arc<faucet_core::observability::RoundtripRecorder>,
199    ) {
200        self.calls.install_recorder(recorder);
201    }
202
203    fn config_schema(&self) -> Value {
204        serde_json::to_value(faucet_core::schema_for!(RedshiftSourceConfig))
205            .expect("schema serialization")
206    }
207
208    fn dataset_uri(&self) -> String {
209        format!(
210            "redshift://{}:{}/{}?query={}",
211            self.config.connection.host,
212            self.config.connection.port,
213            self.config.connection.database,
214            self.config.query
215        )
216    }
217
218    fn state_key(&self) -> Option<String> {
219        match &self.config.replication {
220            RedshiftReplication::Full => None,
221            RedshiftReplication::Incremental { .. } => Some(
222                self.config
223                    .state_key
224                    .clone()
225                    .unwrap_or_else(|| default_state_key(&self.config)),
226            ),
227        }
228    }
229
230    async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
231        let start = CursorBookmark::from_state(bookmark)?;
232        *self
233            .start_bookmark
234            .lock()
235            .expect("start_bookmark mutex poisoned") = Some(start);
236        Ok(())
237    }
238
239    async fn fetch_with_context(
240        &self,
241        context: &HashMap<String, Value>,
242    ) -> Result<Vec<Value>, FaucetError> {
243        Ok(self.fetch_with_context_incremental(context).await?.0)
244    }
245
246    async fn fetch_with_context_incremental(
247        &self,
248        context: &HashMap<String, Value>,
249    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
250        use futures::StreamExt;
251        let mut records = Vec::new();
252        let mut bookmark = None;
253        let mut s = self.stream_pages(context, self.config.batch_size);
254        while let Some(page) = s.next().await {
255            let page = page?;
256            records.extend(page.records);
257            bookmark = page.bookmark.or(bookmark);
258        }
259        Ok((records, bookmark))
260    }
261
262    /// Stream rows from the `sqlx` cursor without buffering the full result set.
263    /// Each emitted [`StreamPage`] holds up to
264    /// [`RedshiftSourceConfig::batch_size`] rows.
265    ///
266    /// For incremental replication rows arrive in cursor order, the rows
267    /// already emitted at the bookmark value are dropped, and every page
268    /// carries the bookmark after its last row (the final page carries one even
269    /// when empty).
270    fn stream_pages<'a>(
271        &'a self,
272        context: &'a HashMap<String, Value>,
273        _batch_size: usize,
274    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
275        let batch_size = self.config.batch_size;
276
277        Box::pin(async_stream::try_stream! {
278            let incr = self.incremental_ctx();
279            let cast = match &incr {
280                Some(ic) => self.cursor_cast(context, &ic.column).await?,
281                None => String::new(),
282            };
283            let (query_str, binds) = self.resolve_query(context, incr.as_ref(), &cast);
284            let mut tracker = incr
285                .as_ref()
286                .map(|ic| CursorTracker::new(ic.column.clone(), ic.start.clone()));
287            let query = bind_params(sqlx::query(&query_str), &binds)?;
288
289            let mut rows = query.fetch(&self.pool);
290            let calls = self.calls.get();
291            calls.record("query");
292            let reads = calls.with_roundtrip_metering(false);
293            let chunk = if batch_size == 0 { usize::MAX } else { batch_size };
294            let cap = if batch_size == 0 { 1024 } else { batch_size };
295            let mut buffer: Vec<Value> = Vec::with_capacity(cap);
296            let mut total = 0usize;
297
298            loop {
299            let fetched = faucet_common_redshift::calls::next_rows(
300                &reads,
301                "fetch",
302                &mut rows,
303                chunk,
304                |e| FaucetError::Source(format!("redshift query failed: {e}")),
305            )
306            .await?;
307            let exhausted = fetched.len() < chunk;
308            for row in fetched {
309                let obj = row_to_json(&row)?;
310                if tracker.as_mut().is_none_or(|t| t.admit(&obj)) {
311                    buffer.push(obj);
312                }
313                if buffer.len() >= chunk {
314                    let page = std::mem::replace(&mut buffer, Vec::with_capacity(cap));
315                    total += page.len();
316                    let bookmark = tracker.as_ref().map(|t| t.bookmark().to_value());
317                    yield StreamPage { records: page, bookmark };
318                }
319            }
320            if exhausted {
321                break;
322            }
323            }
324
325            if !buffer.is_empty() || tracker.is_some() {
326                total += buffer.len();
327                let bookmark = tracker.as_ref().map(|t| t.bookmark().to_value());
328                yield StreamPage { records: buffer, bookmark };
329            }
330            if let Some(missing) = tracker.as_ref().map(CursorTracker::missing).filter(|n| *n > 0) {
331                tracing::warn!(
332                    missing,
333                    "Redshift incremental read: rows without a cursor value were emitted \
334                     without moving the bookmark"
335                );
336            }
337
338            tracing::info!(
339                rows = total,
340                batch_size,
341                query = %self.config.query,
342                "Redshift source stream complete",
343            );
344        })
345    }
346
347    /// Preflight probe for `faucet doctor`: acquire a connection and run
348    /// `SELECT 1` (non-scanning — never executes the configured query).
349    async fn check(
350        &self,
351        ctx: &faucet_core::check::CheckContext,
352    ) -> Result<faucet_core::check::CheckReport, FaucetError> {
353        use faucet_core::check::{CheckReport, Probe};
354
355        let started = std::time::Instant::now();
356        let calls = self.calls.get();
357        let ping = faucet_common_redshift::calls::once(&calls, "query", async {
358            sqlx::query("SELECT 1")
359                .execute(&self.pool)
360                .await
361                .map_err(|e| FaucetError::Source(e.to_string()))
362        });
363        let probe = match tokio::time::timeout(ctx.timeout, ping).await {
364            Ok(Ok(_)) => Probe::pass("auth", started.elapsed()),
365            Ok(Err(e)) => Probe::fail_hint(
366                "auth",
367                started.elapsed(),
368                e.to_string(),
369                "check host/port/database/user/credentials and that the cluster is reachable",
370            ),
371            Err(_) => Probe::fail_hint(
372                "auth",
373                started.elapsed(),
374                "timed out",
375                "check host/port/database/user/credentials and that the cluster is reachable",
376            ),
377        };
378        Ok(CheckReport::single(probe))
379    }
380}
381
382#[cfg(test)]
383mod tests {
384    use super::*;
385    use crate::config::RedshiftReplication;
386    use faucet_common_redshift::RedshiftConnection;
387    use serde_json::json;
388
389    fn base_config() -> RedshiftSourceConfig {
390        RedshiftSourceConfig {
391            connection: RedshiftConnection::new("host", "db", "user", "pw"),
392            query: "SELECT * FROM t".into(),
393            params: Vec::new(),
394            max_connections: 10,
395            batch_size: 1000,
396            replication: RedshiftReplication::Full,
397            state_key: None,
398        }
399    }
400
401    fn source(c: RedshiftSourceConfig) -> RedshiftSource {
402        RedshiftSource::new(c).unwrap()
403    }
404
405    #[test]
406    fn new_rejects_invalid_batch_size() {
407        let mut c = base_config();
408        c.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
409        match RedshiftSource::new(c) {
410            Err(FaucetError::Config(m)) => assert!(m.contains("batch_size"), "got: {m}"),
411            _ => panic!("expected a batch_size Config error"),
412        }
413    }
414
415    #[test]
416    fn new_surfaces_unsupported_credentials() {
417        let mut c = base_config();
418        c.connection.credentials = faucet_common_redshift::RedshiftCredentials::Iam {
419            region: None,
420            cluster_identifier: None,
421            db_user: None,
422        };
423        assert!(matches!(
424            RedshiftSource::new(c),
425            Err(FaucetError::Config(_))
426        ));
427    }
428
429    #[tokio::test]
430    async fn connector_name_is_redshift() {
431        assert_eq!(source(base_config()).connector_name(), "redshift");
432    }
433
434    #[tokio::test]
435    async fn dataset_uri_has_host_db_and_query() {
436        let s = source(base_config());
437        assert_eq!(
438            s.dataset_uri(),
439            "redshift://host:5439/db?query=SELECT * FROM t"
440        );
441    }
442
443    #[tokio::test]
444    async fn config_schema_reports_required_fields() {
445        let s = source(base_config());
446        let schema = s.config_schema();
447        assert!(schema["properties"]["query"].is_object());
448        let required = schema["required"].as_array().expect("required array");
449        assert!(required.iter().any(|v| v == "query"));
450    }
451
452    #[tokio::test]
453    async fn full_mode_has_no_state_key() {
454        assert!(source(base_config()).state_key().is_none());
455    }
456
457    #[tokio::test]
458    async fn incremental_state_key_derived_and_stable() {
459        let mut c = base_config();
460        c.replication = RedshiftReplication::Incremental {
461            column: "ts".into(),
462            initial_value: json!("2026-01-01"),
463        };
464        let k1 = source(c.clone()).state_key().unwrap();
465        let k2 = source(c).state_key().unwrap();
466        assert_eq!(k1, k2);
467        assert!(k1.starts_with("redshift:"));
468        let c = base_config();
469        let material = format!(
470            "{}\u{1f}{}\u{1f}{}\u{1f}{}",
471            c.connection.host, c.connection.port, c.connection.database, c.query
472        );
473        assert_eq!(
474            k1,
475            format!(
476                "redshift:{:016x}",
477                faucet_core::shard::shard_hash(&material)
478            )
479        );
480    }
481
482    #[tokio::test]
483    async fn explicit_state_key_wins_and_bookmark_overrides_initial() {
484        let mut c = base_config();
485        c.state_key = Some("my-key".into());
486        c.replication = RedshiftReplication::Incremental {
487            column: "ts".into(),
488            initial_value: json!("2026-01-01"),
489        };
490        let s = source(c);
491        assert_eq!(s.state_key().as_deref(), Some("my-key"));
492        s.apply_start_bookmark(json!("2026-06-01")).await.unwrap();
493        assert_eq!(
494            s.incremental_ctx().unwrap().start.value,
495            json!("2026-06-01")
496        );
497        assert!(
498            s.apply_start_bookmark(json!({"value": null}))
499                .await
500                .is_err()
501        );
502    }
503
504    #[tokio::test]
505    async fn resolve_query_no_context_no_incremental_is_verbatim() {
506        let mut c = base_config();
507        c.params = vec![json!(7)];
508        let s = source(c);
509        let (sql, binds) = s.resolve_query(&HashMap::new(), None, "");
510        assert_eq!(sql, "SELECT * FROM t");
511        assert_eq!(binds, vec![json!(7)]);
512    }
513
514    #[tokio::test]
515    async fn resolve_query_substitutes_context_positionally() {
516        let mut c = base_config();
517        c.query = "SELECT * FROM t WHERE id = {parent.id}".into();
518        let s = source(c);
519        let mut ctx = HashMap::new();
520        ctx.insert("parent.id".to_string(), json!(42));
521        let (sql, binds) = s.resolve_query(&ctx, None, "");
522        assert_eq!(sql, "SELECT * FROM t WHERE id = $1");
523        assert_eq!(binds, vec![json!(42)]);
524    }
525
526    #[tokio::test]
527    async fn resolve_query_pushes_down_bookmark_after_params_and_context() {
528        let mut c = base_config();
529        c.params = vec![json!(1)];
530        c.query = "SELECT * FROM t WHERE tenant = {p.t} AND ts >= ${bookmark};".into();
531        c.replication = RedshiftReplication::Incremental {
532            column: "ts".into(),
533            initial_value: json!("2026-01-01"),
534        };
535        let s = source(c);
536        let incr = s.incremental_ctx();
537        let mut ctx = HashMap::new();
538        ctx.insert("p.t".to_string(), json!("acme"));
539        let (sql, binds) = s.resolve_query(&ctx, incr.as_ref(), "TIMESTAMP");
540        // $1 = static param, $2 = context value, $3 = bookmark, $4 = outer filter.
541        assert_eq!(
542            sql,
543            "SELECT * FROM (SELECT * FROM t WHERE tenant = $2 AND ts >= CAST($3 AS TIMESTAMP)\n) \
544             AS q WHERE q.\"ts\" >= CAST($4 AS TIMESTAMP) ORDER BY q.\"ts\""
545        );
546        assert_eq!(
547            binds,
548            vec![
549                json!(1),
550                json!("acme"),
551                json!("2026-01-01"),
552                json!("2026-01-01")
553            ]
554        );
555    }
556
557    #[test]
558    fn cast_type_accepts_type_names_only() {
559        assert_eq!(cast_type("TIMESTAMPTZ").unwrap(), "TIMESTAMPTZ");
560        assert_eq!(cast_type("character varying").unwrap(), "character varying");
561        assert!(cast_type("INT4); DROP TABLE t; --").is_err());
562        assert!(cast_type("").is_err());
563    }
564}
faucet_source_redshift/stream.rs

faucet_source_redshift/
stream.rs
