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_common_redshift::data_api::{self, DataColumn, Effect};
15use faucet_common_redshift::{DataApiClient, RedshiftPool};
16use faucet_core::replication::cursor::{CursorBookmark, CursorTracker};
17use faucet_core::util::{quote_ident, substitute_context_bind_params};
18use faucet_core::{FaucetError, Source, Stream, StreamPage};
19use serde_json::Value;
20
21use crate::config::{BOOKMARK_TOKEN, RedshiftReplication, RedshiftSourceConfig};
22use crate::convert::{bind_params, row_to_json};
23
24/// A source that executes a SQL query against Amazon Redshift and returns rows
25/// as JSON objects.
26pub struct RedshiftSource {
27    config: RedshiftSourceConfig,
28    backend: Backend,
29    /// Bookmark applied via [`Source::apply_start_bookmark`] (incremental only).
30    start_bookmark: Mutex<Option<CursorBookmark>>,
31    /// Backend-call context (#880): the query and every page of rows run
32    /// through it.
33    calls: faucet_core::CallSlot,
34}
35
36/// How statements reach Redshift.
37enum Backend {
38    /// The PostgreSQL wire protocol (`password`, `iam`).
39    Wire(RedshiftPool),
40    /// The Redshift Data API over HTTPS (`redshift_data_api`).
41    DataApi(Box<DataApiClient>),
42}
43
44/// A backend ready to run a query: the wire pool with fresh credentials, or
45/// the Data API client.
46enum Conn<'a> {
47    Wire(sqlx::PgPool),
48    DataApi(&'a DataApiClient),
49}
50
51/// Where the rows of a running query come from.
52enum Cursor<'a> {
53    Wire(futures::stream::BoxStream<'a, Result<sqlx::postgres::PgRow, sqlx::Error>>),
54    DataApi {
55        api: &'a DataApiClient,
56        id: String,
57        token: Option<String>,
58        columns: Vec<DataColumn>,
59        done: bool,
60    },
61}
62
63/// Incremental-read context (column + starting position).
64struct IncrementalCtx {
65    column: String,
66    start: CursorBookmark,
67}
68
69impl RedshiftSource {
70    /// Create a new source. Validates config and builds a lazily-connected pool
71    /// (no I/O — connectivity is verified on first query or via
72    /// [`Source::check`]).
73    pub fn new(config: RedshiftSourceConfig) -> Result<Self, FaucetError> {
74        config.validate()?;
75        let backend = match DataApiClient::from_connection(&config.connection, FaucetError::Source)?
76        {
77            Some(api) => Backend::DataApi(Box::new(api)),
78            None => Backend::Wire(faucet_common_redshift::connect_lazy(
79                &config.connection,
80                config.max_connections,
81            )?),
82        };
83        Ok(Self {
84            config,
85            backend,
86            start_bookmark: Mutex::new(None),
87            calls: faucet_core::CallSlot::new(faucet_core::CallContext::new().without_timeout()),
88        })
89    }
90
91    /// The effective incremental start bookmark (persisted bookmark, else the
92    /// configured `initial_value`), or `None` for full replication.
93    fn incremental_ctx(&self) -> Option<IncrementalCtx> {
94        match &self.config.replication {
95            RedshiftReplication::Full => None,
96            RedshiftReplication::Incremental {
97                column,
98                initial_value,
99            } => {
100                let start = self
101                    .start_bookmark
102                    .lock()
103                    .expect("start_bookmark mutex poisoned")
104                    .clone()
105                    .unwrap_or_else(|| CursorBookmark::initial(initial_value.clone()));
106                Some(IncrementalCtx {
107                    column: column.clone(),
108                    start,
109                })
110            }
111        }
112    }
113
114    /// Build the effective SQL and ordered positional bind values for a parent
115    /// context + optional incremental cursor, whose bookmark binds as
116    /// `CAST($n AS <cast>)` (the cursor column's type).
117    ///
118    /// Bind order is: static [`config.params`](RedshiftSourceConfig::params),
119    /// then parent-context `{key}` values, then the `${bookmark}` value (when
120    /// the query contains that token), then the bookmark of the outer
121    /// `column >= …` filter.
122    fn resolve_query(
123        &self,
124        context: &HashMap<String, Value>,
125        incr: Option<&IncrementalCtx>,
126        cast: &str,
127    ) -> (String, Vec<Value>) {
128        let mut binds = self.config.params.clone();
129        let mut sql = self.config.query.clone();
130
131        if !context.is_empty() {
132            let (rewritten, ctx_values) =
133                substitute_context_bind_params(&sql, context, binds.len() + 1, |i| format!("${i}"));
134            sql = rewritten;
135            binds.extend(ctx_values);
136        }
137
138        let Some(ctx) = incr else {
139            return (sql, binds);
140        };
141        let cast = bookmark_cast(cast, &ctx.start.value);
142        if sql.contains(BOOKMARK_TOKEN) {
143            let marker = format!("CAST(${} AS {cast})", binds.len() + 1);
144            sql = sql.replace(BOOKMARK_TOKEN, &marker);
145            binds.push(ctx.start.value.clone());
146        }
147        binds.push(ctx.start.value.clone());
148        let c = quote_ident(&ctx.column);
149        let inner = sql.trim_end_matches(|ch: char| ch.is_whitespace() || ch == ';');
150        (
151            format!(
152                "SELECT * FROM ({inner}\n) AS q WHERE q.{c} >= CAST(${n} AS {cast}) ORDER BY q.{c}",
153                n = binds.len()
154            ),
155            binds,
156        )
157    }
158
159    /// The cursor column's type, as the server describes the query's output.
160    async fn cursor_cast(
161        &self,
162        context: &HashMap<String, Value>,
163        column: &str,
164    ) -> Result<String, FaucetError> {
165        let described = self.describe(context).await.map_err(|e| match e {
166            FaucetError::Source(m) => FaucetError::Source(format!(
167                "redshift: could not describe the incremental query: {m}"
168            )),
169            other => other,
170        })?;
171        let name = described
172            .iter()
173            .find(|(c, _, _)| c == column)
174            .map(|(_, t, _)| t.clone())
175            .ok_or_else(|| {
176                FaucetError::Config(format!(
177                    "redshift: incremental replication `column` {column:?} is not a column of \
178                     the query's output"
179                ))
180            })?;
181        cast_type(&name)
182    }
183
184    /// The query's output columns as `(name, type name, nullable)`: a
185    /// describe on the wire, a `LIMIT 0` run of the query through the Data
186    /// API (which has no describe).
187    async fn describe(
188        &self,
189        context: &HashMap<String, Value>,
190    ) -> Result<Vec<(String, String, Option<bool>)>, FaucetError> {
191        use sqlx::Executor;
192        let (probe, binds) = self.resolve_query(context, None, "");
193        let probe = probe.replace(BOOKMARK_TOKEN, "NULL");
194        let calls = self.calls.get();
195        match &self.backend {
196            Backend::Wire(pool) => {
197                let pool = pool.get(&calls).await?;
198                let described = faucet_common_redshift::calls::once(&calls, "describe", async {
199                    (&pool)
200                        .describe(&probe)
201                        .await
202                        .map_err(|e| FaucetError::Source(e.to_string()))
203                })
204                .await?;
205                Ok(described
206                    .columns()
207                    .iter()
208                    .enumerate()
209                    .map(|(i, c)| {
210                        (
211                            sqlx::Column::name(c).to_string(),
212                            sqlx::TypeInfo::name(sqlx::Column::type_info(c)).to_string(),
213                            described.nullable(i),
214                        )
215                    })
216                    .collect())
217            }
218            Backend::DataApi(api) => {
219                let inner = probe.trim_end_matches(|ch: char| ch.is_whitespace() || ch == ';');
220                let sql = format!("SELECT * FROM ({inner}\n) AS q LIMIT 0");
221                let (columns, _) = api
222                    .query(&calls, &data_api::named_parameters(&sql, &binds))
223                    .await?;
224                Ok(columns
225                    .into_iter()
226                    .map(|c| (c.name, c.type_name, c.nullable))
227                    .collect())
228            }
229        }
230    }
231
232    /// Start `sql` and return where its rows come from.
233    async fn open<'a>(
234        calls: &faucet_core::io::CallContext,
235        conn: &'a Conn<'a>,
236        sql: &'a str,
237        binds: &'a [Value],
238    ) -> Result<Cursor<'a>, FaucetError> {
239        use futures::StreamExt;
240        match conn {
241            Conn::DataApi(api) => {
242                let done = api
243                    .execute(calls, Effect::Read, &data_api::named_parameters(sql, binds))
244                    .await?;
245                Ok(Cursor::DataApi {
246                    api,
247                    id: done.id,
248                    token: None,
249                    columns: Vec::new(),
250                    done: !done.has_result_set,
251                })
252            }
253            Conn::Wire(pool) => {
254                let query = bind_params(sqlx::query(sql), binds)?;
255                calls.record("query");
256                Ok(Cursor::Wire(query.fetch(pool).boxed()))
257            }
258        }
259    }
260
261    /// Up to `max` more rows of `cursor` as JSON objects, and whether it is
262    /// exhausted. A Data API page is read whole.
263    async fn next_chunk(
264        calls: &faucet_core::io::CallContext,
265        cursor: &mut Cursor<'_>,
266        max: usize,
267    ) -> Result<(Vec<Value>, bool), FaucetError> {
268        match cursor {
269            Cursor::Wire(rows) => {
270                let reads = calls.clone().with_roundtrip_metering(false);
271                let fetched =
272                    faucet_common_redshift::calls::next_rows(&reads, "fetch", rows, max, |e| {
273                        FaucetError::Source(format!("redshift query failed: {e}"))
274                    })
275                    .await?;
276                let exhausted = fetched.len() < max;
277                let rows = fetched.iter().map(row_to_json).collect::<Result<_, _>>()?;
278                Ok((rows, exhausted))
279            }
280            Cursor::DataApi {
281                api,
282                id,
283                token,
284                columns,
285                done,
286            } => {
287                if *done {
288                    return Ok((Vec::new(), true));
289                }
290                let page = api.result_page(calls, id, token.as_deref()).await?;
291                if columns.is_empty() {
292                    *columns = page.columns;
293                }
294                let rows = page
295                    .rows
296                    .iter()
297                    .map(|f| data_api::row_to_json(columns, f).map(Value::Object))
298                    .collect::<Result<_, _>>()?;
299                *token = page.next_token;
300                *done = token.is_none();
301                Ok((rows, *done))
302            }
303        }
304    }
305}
306
307/// The cast a bookmark binds through. The server describes a column's type
308/// without its precision and scale, and a bare `NUMERIC` / `DECIMAL` on
309/// Redshift is `DECIMAL(18,0)`, which would round a fractional bookmark and
310/// skip the rows just above it. A decimal cursor is cast at the widest
311/// precision and the bookmark's own scale instead, which is exact (the
312/// bookmark is rendered at the column's scale).
313fn bookmark_cast(cast: &str, bookmark: &Value) -> String {
314    let base = cast.trim().to_ascii_uppercase();
315    if !matches!(base.as_str(), "NUMERIC" | "DECIMAL") {
316        return cast.to_string();
317    }
318    let text = match bookmark {
319        Value::String(s) => s.trim().to_string(),
320        Value::Number(n) => n.to_string(),
321        _ => String::new(),
322    };
323    let mantissa = text.split(['e', 'E']).next().unwrap_or("");
324    let scale = mantissa
325        .split_once('.')
326        .map_or(0, |(_, frac)| {
327            frac.chars().filter(char::is_ascii_digit).count()
328        })
329        .min(37);
330    format!("NUMERIC(38,{scale})")
331}
332
333/// A server-reported type name usable in `CAST(… AS <name>)`.
334fn cast_type(name: &str) -> Result<String, FaucetError> {
335    let ok = !name.is_empty()
336        && name
337            .chars()
338            .all(|c| c.is_ascii_alphanumeric() || matches!(c, '_' | ' ' | '"' | '[' | ']'));
339    if ok {
340        Ok(name.to_string())
341    } else {
342        Err(FaucetError::Config(format!(
343            "redshift: the incremental cursor column has type {name:?}, which cannot be \
344             compared to a bookmark; cast it in the query"
345        )))
346    }
347}
348
349/// Derive a stable state key from the host, database, and query.
350fn default_state_key(config: &RedshiftSourceConfig) -> String {
351    // Fixed FNV-1a: the key is durable and `DefaultHasher` is not stable across Rust releases.
352    let material = format!(
353        "{}\u{1f}{}\u{1f}{}\u{1f}{}",
354        config.connection.endpoint_label(),
355        config.connection.port,
356        config.connection.database,
357        config.query
358    );
359    format!(
360        "redshift:{:016x}",
361        faucet_core::shard::shard_hash(&material)
362    )
363}
364
365/// The logical type of a described column as this source renders it (#955).
366fn declared_type(type_name: &str) -> faucet_core::LogicalType {
367    use faucet_core::LogicalType as L;
368    match type_name.to_ascii_uppercase().as_str() {
369        "TIMETZ" | "INTERVAL" | "OID" => L::String,
370        _ => L::from_sql_type(type_name, None, None).unwrap_or(L::String),
371    }
372}
373
374#[async_trait]
375impl Source for RedshiftSource {
376    fn connector_name(&self) -> &'static str {
377        "redshift"
378    }
379
380    fn capabilities(&self) -> faucet_core::Capabilities {
381        faucet_core::Capabilities::of_source(self)
382            .to_builder()
383            .declared_columns(true)
384            .build()
385    }
386
387    /// The query's columns from a describe of it (#955). The server reports
388    /// no decimal precision, so a `NUMERIC` is declared unbounded.
389    async fn declared_columns(
390        &self,
391        context: &HashMap<String, Value>,
392    ) -> Result<Option<faucet_core::DeclaredColumns>, FaucetError> {
393        let Ok(described) = self.describe(context).await else {
394            return Ok(None);
395        };
396        Ok(Some(faucet_core::DeclaredColumns::new(
397            described.into_iter().map(|(name, ty, nullable)| {
398                faucet_core::DeclaredColumn::new(name, declared_type(&ty), nullable != Some(false))
399            }),
400        )))
401    }
402
403    fn set_call_context(&self, ctx: faucet_core::io::CallContext) {
404        self.calls.install(ctx);
405    }
406
407    fn set_roundtrip_recorder(
408        &self,
409        recorder: std::sync::Arc<faucet_core::observability::RoundtripRecorder>,
410    ) {
411        self.calls.install_recorder(recorder);
412    }
413
414    fn config_schema(&self) -> Value {
415        serde_json::to_value(faucet_core::schema_for!(RedshiftSourceConfig))
416            .expect("schema serialization")
417    }
418
419    fn dataset_uri(&self) -> String {
420        let conn = &self.config.connection;
421        if conn.uses_data_api() {
422            return format!(
423                "redshift://{}/{}?query={}",
424                conn.endpoint_label(),
425                conn.database,
426                self.config.query
427            );
428        }
429        format!(
430            "redshift://{}:{}/{}?query={}",
431            self.config.connection.host,
432            self.config.connection.port,
433            self.config.connection.database,
434            self.config.query
435        )
436    }
437
438    fn state_key(&self) -> Option<String> {
439        match &self.config.replication {
440            RedshiftReplication::Full => None,
441            RedshiftReplication::Incremental { .. } => Some(
442                self.config
443                    .state_key
444                    .clone()
445                    .unwrap_or_else(|| default_state_key(&self.config)),
446            ),
447        }
448    }
449
450    async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
451        let start = CursorBookmark::from_state(bookmark)?;
452        *self
453            .start_bookmark
454            .lock()
455            .expect("start_bookmark mutex poisoned") = Some(start);
456        Ok(())
457    }
458
459    async fn fetch_with_context(
460        &self,
461        context: &HashMap<String, Value>,
462    ) -> Result<Vec<Value>, FaucetError> {
463        Ok(self.fetch_with_context_incremental(context).await?.0)
464    }
465
466    async fn fetch_with_context_incremental(
467        &self,
468        context: &HashMap<String, Value>,
469    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
470        use futures::StreamExt;
471        let mut records = Vec::new();
472        let mut bookmark = None;
473        let mut s = self.stream_pages(context, self.config.batch_size);
474        while let Some(page) = s.next().await {
475            let page = page?;
476            records.extend(page.records);
477            bookmark = page.bookmark.or(bookmark);
478        }
479        Ok((records, bookmark))
480    }
481
482    /// Stream rows from the `sqlx` cursor without buffering the full result set.
483    /// Each emitted [`StreamPage`] holds up to
484    /// [`RedshiftSourceConfig::batch_size`] rows.
485    ///
486    /// For incremental replication rows arrive in cursor order, the rows
487    /// already emitted at the bookmark value are dropped, and every page
488    /// carries the bookmark after its last row (the final page carries one even
489    /// when empty).
490    fn stream_pages<'a>(
491        &'a self,
492        context: &'a HashMap<String, Value>,
493        _batch_size: usize,
494    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
495        let batch_size = self.config.batch_size;
496
497        Box::pin(async_stream::try_stream! {
498            let incr = self.incremental_ctx();
499            let cast = match &incr {
500                Some(ic) => self.cursor_cast(context, &ic.column).await?,
501                None => String::new(),
502            };
503            let (query_str, binds) = self.resolve_query(context, incr.as_ref(), &cast);
504            let mut tracker = incr
505                .as_ref()
506                .map(|ic| CursorTracker::new(ic.column.clone(), ic.start.clone()));
507            let calls = self.calls.get();
508            let conn = match &self.backend {
509                Backend::Wire(p) => Conn::Wire(p.get(&calls).await?),
510                Backend::DataApi(api) => Conn::DataApi(api),
511            };
512            let mut cursor = Self::open(&calls, &conn, &query_str, &binds).await?;
513            let chunk = if batch_size == 0 { usize::MAX } else { batch_size };
514            let cap = if batch_size == 0 { 1024 } else { batch_size };
515            let mut buffer: Vec<Value> = Vec::with_capacity(cap);
516            let mut total = 0usize;
517
518            loop {
519            let (fetched, exhausted) = Self::next_chunk(&calls, &mut cursor, chunk).await?;
520            for obj in fetched {
521                if tracker.as_mut().is_none_or(|t| t.admit(&obj)) {
522                    buffer.push(obj);
523                }
524                if buffer.len() >= chunk {
525                    let page = std::mem::replace(&mut buffer, Vec::with_capacity(cap));
526                    total += page.len();
527                    let bookmark = tracker.as_ref().map(|t| t.bookmark().to_value());
528                    yield StreamPage { records: page, bookmark };
529                }
530            }
531            if exhausted {
532                break;
533            }
534            }
535
536            if !buffer.is_empty() || tracker.is_some() {
537                total += buffer.len();
538                let bookmark = tracker.as_ref().map(|t| t.bookmark().to_value());
539                yield StreamPage { records: buffer, bookmark };
540            }
541            if let Some(missing) = tracker.as_ref().map(CursorTracker::missing).filter(|n| *n > 0) {
542                tracing::warn!(
543                    missing,
544                    "Redshift incremental read: rows without a cursor value were emitted \
545                     without moving the bookmark"
546                );
547            }
548
549            tracing::info!(
550                rows = total,
551                batch_size,
552                query = %self.config.query,
553                "Redshift source stream complete",
554            );
555        })
556    }
557
558    /// Preflight probe for `faucet doctor`: acquire a connection and run
559    /// `SELECT 1` (non-scanning — never executes the configured query).
560    async fn check(
561        &self,
562        ctx: &faucet_core::check::CheckContext,
563    ) -> Result<faucet_core::check::CheckReport, FaucetError> {
564        use faucet_core::check::{CheckReport, Probe};
565
566        let started = std::time::Instant::now();
567        let calls = self.calls.get();
568        let ping = async {
569            match &self.backend {
570                Backend::Wire(pool) => {
571                    let pool = pool.get(&calls).await?;
572                    faucet_common_redshift::calls::once(&calls, "query", async {
573                        sqlx::query("SELECT 1")
574                            .execute(&pool)
575                            .await
576                            .map(|_| ())
577                            .map_err(|e| FaucetError::Source(e.to_string()))
578                    })
579                    .await
580                }
581                Backend::DataApi(api) => api
582                    .query(&calls, &data_api::Statement::new("SELECT 1"))
583                    .await
584                    .map(|_| ()),
585            }
586        };
587        let probe = match tokio::time::timeout(ctx.timeout, ping).await {
588            Ok(Ok(_)) => Probe::pass("auth", started.elapsed()),
589            Ok(Err(e)) => Probe::fail_hint(
590                "auth",
591                started.elapsed(),
592                e.to_string(),
593                "check host/port/database/user/credentials and that the cluster is reachable",
594            ),
595            Err(_) => Probe::fail_hint(
596                "auth",
597                started.elapsed(),
598                "timed out",
599                "check host/port/database/user/credentials and that the cluster is reachable",
600            ),
601        };
602        Ok(CheckReport::single(probe))
603    }
604}
605
606#[cfg(test)]
607mod tests;
faucet_source_redshift/stream.rs

faucet_source_redshift/
stream.rs
