Skip to main content

faucet_source_mysql/
stream.rs

1//! MySQL source implementation.
2
3use crate::config::MysqlSourceConfig;
4use async_trait::async_trait;
5use faucet_core::shard::{
6    PkShardBounds, ShardSpec, parse_pk_shard, pk_bounds_query, pk_shards_from_bounds,
7};
8use faucet_core::{FaucetError, Stream, StreamPage};
9use futures::TryStreamExt;
10use serde_json::Value;
11use sqlx::mysql::MySqlPoolOptions;
12use sqlx::{Column, MySqlPool, Row};
13use std::pin::Pin;
14use std::sync::Mutex;
15
16/// A source that executes a SQL query against MySQL and returns rows as JSON.
17pub struct MysqlSource {
18    config: MysqlSourceConfig,
19    pool: MySqlPool,
20    /// Shard applied by the cluster coordinator (Mode B), if any. `None` (or the
21    /// whole-dataset shard) means the full query is streamed. Stored behind a
22    /// `Mutex` so `apply_shard(&self, …)` can record it before streaming.
23    applied_shard: Mutex<Option<PkShardBounds>>,
24}
25
26/// Quote a MySQL identifier with backticks (MySQL's default identifier
27/// quoting — double quotes require the non-default `ANSI_QUOTES` sql_mode).
28/// Embedded backticks are doubled, preventing identifier injection.
29fn quote_ident_mysql(name: &str) -> String {
30    format!("`{}`", name.replace('`', "``"))
31}
32
33impl MysqlSource {
34    /// Create a new MySQL source. Establishes a connection pool.
35    pub async fn new(config: MysqlSourceConfig) -> Result<Self, FaucetError> {
36        faucet_core::validate_batch_size(config.batch_size)?;
37
38        let pool = MySqlPoolOptions::new()
39            .max_connections(config.max_connections)
40            .connect(&config.connection_url)
41            .await
42            .map_err(|e| FaucetError::Config(format!("MySQL connection failed: {e}")))?;
43
44        Ok(Self {
45            config,
46            pool,
47            applied_shard: Mutex::new(None),
48        })
49    }
50
51    /// Apply the currently-set shard (if any) to a resolved query string.
52    fn shard_wrap(&self, query: String) -> String {
53        match &*self.applied_shard.lock().expect("shard mutex poisoned") {
54            Some(bounds) => bounds.wrap(&query, quote_ident_mysql),
55            None => query,
56        }
57    }
58}
59
60/// Convert a MySQL row column value to a `serde_json::Value`.
61///
62/// Attempts common types in order of likelihood. Falls back to `Value::Null`
63/// for unsupported or null columns.
64fn mysql_value_to_json(row: &sqlx::mysql::MySqlRow, col_name: &str) -> Value {
65    // Try JSON first
66    if let Ok(v) = row.try_get::<Value, _>(col_name) {
67        return v;
68    }
69
70    // Try common scalar types
71    if let Ok(v) = row.try_get::<String, _>(col_name) {
72        return Value::String(v);
73    }
74    if let Ok(v) = row.try_get::<i64, _>(col_name) {
75        return Value::Number(v.into());
76    }
77    if let Ok(v) = row.try_get::<i32, _>(col_name) {
78        return Value::Number(v.into());
79    }
80    if let Ok(v) = row.try_get::<i16, _>(col_name) {
81        return Value::Number(v.into());
82    }
83    // UNSIGNED integer columns (#264). sqlx-mysql treats UNSIGNED as a
84    // distinct type from the signed decoders above, so without these arms
85    // UNSIGNED columns fall through to the `bool` probe (TINYINT UNSIGNED ->
86    // bool) or to the `Null` fall-through (larger UNSIGNED -> null), silently
87    // corrupting every unsigned column including UNSIGNED primary keys.
88    //
89    // These are placed *after* the signed probes so a signed column always
90    // matches a signed arm first, and *before* `bool`/`f64`/`f32` so a
91    // `TINYINT UNSIGNED` decodes as a number rather than a bool (MySQL's
92    // boolean is `TINYINT(1)`). `u64` fits `serde_json::Number` exactly, so
93    // BIGINT UNSIGNED values above `i64::MAX` round-trip losslessly.
94    if let Ok(v) = row.try_get::<u64, _>(col_name) {
95        return Value::Number(v.into());
96    }
97    if let Ok(v) = row.try_get::<u32, _>(col_name) {
98        return Value::Number(v.into());
99    }
100    if let Ok(v) = row.try_get::<u16, _>(col_name) {
101        return Value::Number(v.into());
102    }
103    if let Ok(v) = row.try_get::<u8, _>(col_name) {
104        return Value::Number(v.into());
105    }
106    if let Ok(v) = row.try_get::<f64, _>(col_name) {
107        return serde_json::Number::from_f64(v)
108            .map(Value::Number)
109            .unwrap_or(Value::Null);
110    }
111    if let Ok(v) = row.try_get::<f32, _>(col_name) {
112        return serde_json::Number::from_f64(v as f64)
113            .map(Value::Number)
114            .unwrap_or(Value::Null);
115    }
116    if let Ok(v) = row.try_get::<bool, _>(col_name) {
117        return Value::Bool(v);
118    }
119
120    // Richer types that would otherwise silently decode to Null (#78/#43).
121    if let Ok(v) =
122        row.try_get::<sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc>, _>(col_name)
123    {
124        return Value::String(v.to_rfc3339());
125    }
126    if let Ok(v) = row.try_get::<sqlx::types::chrono::NaiveDateTime, _>(col_name) {
127        return Value::String(v.to_string());
128    }
129    if let Ok(v) = row.try_get::<sqlx::types::chrono::NaiveDate, _>(col_name) {
130        return Value::String(v.to_string());
131    }
132    if let Ok(v) = row.try_get::<sqlx::types::chrono::NaiveTime, _>(col_name) {
133        return Value::String(v.to_string());
134    }
135    // DECIMAL → string, preserving exact precision.
136    if let Ok(v) = row.try_get::<sqlx::types::BigDecimal, _>(col_name) {
137        return Value::String(v.to_string());
138    }
139    // BLOB / BINARY → base64.
140    if let Ok(v) = row.try_get::<Vec<u8>, _>(col_name) {
141        use base64::Engine as _;
142        return Value::String(base64::engine::general_purpose::STANDARD.encode(v));
143    }
144
145    Value::Null
146}
147
148/// Build the effective SQL query and ordered context-bind values for a given
149/// parent context. Returns the literal query when there is no context.
150fn resolve_query(
151    config: &MysqlSourceConfig,
152    context: &std::collections::HashMap<String, Value>,
153) -> (String, Vec<Value>) {
154    if context.is_empty() {
155        (config.query.clone(), Vec::new())
156    } else {
157        faucet_core::util::substitute_context_bind_params(&config.query, context, 1, |_| {
158            "?".to_string()
159        })
160    }
161}
162
163/// How a numeric bind value should be bound onto a sqlx query.
164///
165/// Classifying *before* binding keeps the integer/float decision in one pure,
166/// unit-testable place and — critically — binds any integer in
167/// `[i64::MIN, i64::MAX]` as an exact `i64` rather than an `f64`. Binding an
168/// integer above `2^53` as `f64` silently rounds it (audit F38), so a large
169/// 64-bit id threaded into `WHERE id = ?` would compare against the *wrong*
170/// value and return wrong rows.
171#[derive(Debug, Clone, Copy, PartialEq, Eq)]
172enum NumberBind {
173    /// Exact `i64` — covers every integer in `[i64::MIN, i64::MAX]`.
174    I64,
175    /// Value above `i64::MAX`; bind as `u64` (MySQL has native UNSIGNED).
176    U64,
177    /// Genuine floating-point value — bind as `f64`.
178    F64,
179}
180
181/// Classify a JSON number into the bind category to use.
182///
183/// `is_i64()` losslessly covers `[i64::MIN, i64::MAX]` (including the
184/// `(2^53, i64::MAX]` range that `f64` would round); `is_u64()` covers values
185/// above `i64::MAX`; everything else is a real float.
186fn classify_number(n: &serde_json::Number) -> NumberBind {
187    if n.is_i64() {
188        NumberBind::I64
189    } else if n.is_u64() {
190        NumberBind::U64
191    } else {
192        NumberBind::F64
193    }
194}
195
196/// Apply context-derived bind values onto a sqlx query.
197fn bind_params<'q>(
198    mut query: sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>,
199    bind_values: &'q [Value],
200) -> sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments> {
201    for value in bind_values {
202        query = match value {
203            Value::String(s) => query.bind(s.clone()),
204            Value::Number(n) => match classify_number(n) {
205                // `unwrap()` is sound: the classifier proves the predicate.
206                NumberBind::I64 => query.bind(n.as_i64().unwrap()),
207                // MySQL has a native UNSIGNED BIGINT type, so bind the `u64`
208                // directly — values above `i64::MAX` round-trip losslessly.
209                NumberBind::U64 => query.bind(n.as_u64().unwrap()),
210                NumberBind::F64 => query.bind(n.as_f64().unwrap_or(0.0)),
211            },
212            Value::Bool(b) => query.bind(*b),
213            Value::Null => query.bind(None::<String>),
214            _ => query.bind(value.to_string()),
215        };
216    }
217    query
218}
219
220/// One flattened `information_schema.columns` row used by [`discover`].
221///
222/// (table, column, data_type, is_nullable, estimated_rows)
223type CatalogRow = (String, String, String, bool, Option<u64>);
224
225/// In-progress per-table accumulator while grouping catalog rows:
226/// `(table, estimated_rows, columns)`.
227type PendingTable = Option<(String, Option<u64>, Vec<(String, Value)>)>;
228
229/// Group flattened catalog rows (ordered by table name, ordinal position)
230/// into one [`DatasetDescriptor`] per table. Pure — unit-testable without a
231/// live server.
232///
233/// The dataset name is the bare table name: a MySQL connection is scoped to a
234/// single database (named in the connection URL), so the generated `SELECT`
235/// needs no database qualifier.
236fn descriptors_from_catalog(rows: Vec<CatalogRow>) -> Vec<faucet_core::DatasetDescriptor> {
237    let mut out: Vec<faucet_core::DatasetDescriptor> = Vec::new();
238    let mut current: PendingTable = None;
239
240    let flush = |cur: PendingTable, out: &mut Vec<faucet_core::DatasetDescriptor>| {
241        if let Some((table, est, cols)) = cur {
242            let query = format!("SELECT * FROM {}", quote_ident_mysql(&table));
243            let mut d = faucet_core::DatasetDescriptor::new(
244                table,
245                "table",
246                serde_json::json!({ "query": query }),
247            )
248            .with_schema(faucet_core::columns_to_schema(cols));
249            // NULL table_rows (e.g. a view snuck through, or stats missing)
250            // means no estimate.
251            if let Some(n) = est {
252                d = d.with_estimated_rows(n);
253            }
254            out.push(d);
255        }
256    };
257
258    for (table, column, data_type, is_nullable, est) in rows {
259        let same = current.as_ref().is_some_and(|(t, _, _)| *t == table);
260        if !same {
261            flush(current.take(), &mut out);
262            current = Some((table, est, Vec::new()));
263        }
264        let mut fragment = faucet_core::sql_type_to_json_schema(&data_type);
265        if is_nullable {
266            fragment = faucet_core::nullable_type(fragment);
267        }
268        if let Some((_, _, cols)) = current.as_mut() {
269            cols.push((column, fragment));
270        }
271    }
272    flush(current, &mut out);
273    out
274}
275
276/// Convert a single `MySqlRow` into a JSON object whose keys are the row's
277/// column names.
278fn row_to_json(row: &sqlx::mysql::MySqlRow) -> Value {
279    let mut map = serde_json::Map::new();
280    for col in row.columns() {
281        let name = col.name().to_string();
282        let value = mysql_value_to_json(row, &name);
283        map.insert(name, value);
284    }
285    Value::Object(map)
286}
287
288#[async_trait]
289impl faucet_core::Source for MysqlSource {
290    async fn fetch_with_context(
291        &self,
292        context: &std::collections::HashMap<String, serde_json::Value>,
293    ) -> Result<Vec<Value>, FaucetError> {
294        let (query_str, bind_values) = resolve_query(&self.config, context);
295        let query_str = self.shard_wrap(query_str);
296        let query = bind_params(sqlx::query(&query_str), &bind_values);
297
298        let rows = query
299            .fetch_all(&self.pool)
300            .await
301            .map_err(|e| FaucetError::Config(format!("MySQL query failed: {e}")))?;
302
303        let records: Vec<Value> = rows.iter().map(row_to_json).collect();
304        tracing::info!(rows = records.len(), query = %self.config.query, "MySQL source fetch complete");
305        Ok(records)
306    }
307
308    /// Stream rows from the underlying sqlx cursor without buffering the full
309    /// result set. Each emitted [`StreamPage`] holds up to
310    /// [`MysqlSourceConfig::batch_size`] rows.
311    ///
312    /// The trait-level `batch_size` argument is ignored in favour of the
313    /// config field — the config is the user-facing knob the README
314    /// documents, and routing the pipeline-supplied hint through it would
315    /// silently override an explicit config value.
316    ///
317    /// `batch_size = 0` drains the entire cursor into a single page. The
318    /// mysql query source has no incremental-replication mode today, so
319    /// every emitted page carries `bookmark: None`.
320    fn stream_pages<'a>(
321        &'a self,
322        context: &'a std::collections::HashMap<String, Value>,
323        _batch_size: usize,
324    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
325        let batch_size = self.config.batch_size;
326
327        Box::pin(async_stream::try_stream! {
328            let (query_str, bind_values) = resolve_query(&self.config, context);
329            let query_str = self.shard_wrap(query_str);
330            let query = bind_params(sqlx::query(&query_str), &bind_values);
331
332            let mut rows = query.fetch(&self.pool);
333            let chunk = if batch_size == 0 { usize::MAX } else { batch_size };
334            let initial_capacity = if batch_size == 0 { 1024 } else { batch_size };
335            let mut buffer: Vec<Value> = Vec::with_capacity(initial_capacity);
336            let mut total = 0usize;
337
338            while let Some(row) = rows
339                .try_next()
340                .await
341                .map_err(|e| FaucetError::Config(format!("MySQL query failed: {e}")))?
342            {
343                buffer.push(row_to_json(&row));
344                if buffer.len() >= chunk {
345                    let page = std::mem::replace(&mut buffer, Vec::with_capacity(initial_capacity));
346                    total += page.len();
347                    yield StreamPage { records: page, bookmark: None };
348                }
349            }
350            if !buffer.is_empty() {
351                total += buffer.len();
352                yield StreamPage { records: buffer, bookmark: None };
353            }
354
355            tracing::info!(
356                rows = total,
357                batch_size,
358                query = %self.config.query,
359                "MySQL source stream complete",
360            );
361        })
362    }
363
364    fn connector_name(&self) -> &'static str {
365        "mysql"
366    }
367
368    fn config_schema(&self) -> serde_json::Value {
369        serde_json::to_value(faucet_core::schema_for!(MysqlSourceConfig))
370            .expect("schema serialization")
371    }
372
373    fn dataset_uri(&self) -> String {
374        format!(
375            "{}?query={}",
376            faucet_core::redact_uri_credentials(&self.config.connection_url),
377            self.config.query
378        )
379    }
380
381    fn supports_discover(&self) -> bool {
382        true
383    }
384
385    /// Enumerate every base table in the connection's current database, with
386    /// column types from `information_schema.columns` and a row estimate from
387    /// `information_schema.tables.table_rows` (catalog metadata only — no
388    /// data scan).
389    async fn discover(&self) -> Result<Vec<faucet_core::DatasetDescriptor>, FaucetError> {
390        // The lowercase aliases matter: MySQL 8 returns information_schema
391        // result columns as UPPERCASE (`TABLE_NAME`, …) without them. The
392        // CAST(… AS CHAR) wrappers matter too: MySQL 8 reports several
393        // information_schema string columns with a binary collation, which
394        // sqlx surfaces as VARBINARY and refuses to decode as String.
395        let sql = "\
396            SELECT CAST(c.table_name AS CHAR) AS table_name, \
397                   CAST(c.column_name AS CHAR) AS column_name, \
398                   CAST(c.data_type AS CHAR) AS data_type, \
399                   CAST(c.is_nullable AS CHAR) AS is_nullable, \
400                   t.table_rows AS estimated_rows \
401              FROM information_schema.columns c \
402              JOIN information_schema.tables t \
403                ON t.table_schema = c.table_schema AND t.table_name = c.table_name \
404             WHERE t.table_type = 'BASE TABLE' \
405               AND c.table_schema = DATABASE() \
406             ORDER BY c.table_name, c.ordinal_position";
407        let rows = sqlx::query(sql)
408            .fetch_all(&self.pool)
409            .await
410            .map_err(|e| FaucetError::Source(format!("mysql: catalog discovery failed: {e}")))?;
411
412        let catalog: Vec<CatalogRow> = rows
413            .iter()
414            .map(|row| -> Result<CatalogRow, FaucetError> {
415                let decode = |col: &str| -> Result<String, FaucetError> {
416                    row.try_get::<String, _>(col).map_err(|e| {
417                        FaucetError::Source(format!("mysql: catalog decode failed ({col}): {e}"))
418                    })
419                };
420                Ok((
421                    decode("table_name")?,
422                    decode("column_name")?,
423                    decode("data_type")?,
424                    decode("is_nullable")?.eq_ignore_ascii_case("yes"),
425                    // NULL (or an unexpected type) → no estimate.
426                    row.try_get::<u64, _>("estimated_rows").ok(),
427                ))
428            })
429            .collect::<Result<_, _>>()?;
430
431        Ok(descriptors_from_catalog(catalog))
432    }
433
434    /// Shardable when a [`ShardConfig`](crate::config::ShardConfig) is set.
435    fn is_shardable(&self) -> bool {
436        self.config.shard.is_some()
437    }
438
439    /// Enumerate contiguous primary-key range shards by computing the `key`
440    /// column's `MIN`/`MAX` over the (unsharded) base query and splitting that
441    /// range into ~`target` slices. Returns a single whole-dataset shard when no
442    /// `shard` config is set or the result set is empty.
443    async fn enumerate_shards(&self, target: usize) -> Result<Vec<ShardSpec>, FaucetError> {
444        let Some(shard_cfg) = &self.config.shard else {
445            return Ok(vec![ShardSpec::whole()]);
446        };
447
448        let bounds_sql = pk_bounds_query(
449            &self.config.query,
450            &quote_ident_mysql(&shard_cfg.key),
451            "SIGNED",
452        );
453        let row = sqlx::query(&bounds_sql)
454            .fetch_one(&self.pool)
455            .await
456            .map_err(|e| {
457                FaucetError::Source(format!(
458                    "mysql: failed to compute shard bounds for key {:?} \
459                     (it must be an integer-typed column): {e}",
460                    shard_cfg.key
461                ))
462            })?;
463
464        let lo: Option<i64> = row
465            .try_get("lo")
466            .map_err(|e| FaucetError::Source(format!("mysql: shard bounds decode failed: {e}")))?;
467        let hi: Option<i64> = row
468            .try_get("hi")
469            .map_err(|e| FaucetError::Source(format!("mysql: shard bounds decode failed: {e}")))?;
470        Ok(pk_shards_from_bounds(&shard_cfg.key, lo, hi, target))
471    }
472
473    /// Narrow this source to a single PK-range shard. The whole-dataset shard
474    /// clears any applied range (streams the full query).
475    async fn apply_shard(&self, shard: &ShardSpec) -> Result<(), FaucetError> {
476        *self.applied_shard.lock().expect("shard mutex poisoned") = parse_pk_shard(shard, "mysql")?;
477        Ok(())
478    }
479}
480
481#[cfg(test)]
482mod tests {
483    use super::*;
484    use faucet_core::shard::plan_pk_shards;
485
486    #[tokio::test]
487    async fn new_rejects_out_of_range_batch_size() {
488        let mut config = MysqlSourceConfig::new("mysql://localhost/test", "SELECT 1");
489        config.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
490        match MysqlSource::new(config).await {
491            Err(faucet_core::FaucetError::Config(m)) => {
492                assert!(m.contains("batch_size"), "got: {m}")
493            }
494            _ => panic!("expected a batch_size Config error"),
495        }
496    }
497
498    // dataset_uri is a pure-config method; the source requires a live DB to
499    // construct so we verify the credential-stripping logic directly.
500    #[test]
501    fn dataset_uri_strips_credentials() {
502        let redacted = faucet_core::redact_uri_credentials("mysql://u:p@h:3306/db");
503        let uri = format!("{}?query={}", redacted, "SELECT 1");
504        assert_eq!(uri, "mysql://h:3306/db?query=SELECT 1");
505    }
506
507    // ── F38: numeric bind classification (precision-safe) ───────────────────
508
509    fn num(v: serde_json::Value) -> serde_json::Number {
510        match v {
511            serde_json::Value::Number(n) => n,
512            _ => panic!("not a number"),
513        }
514    }
515
516    #[test]
517    fn classify_small_int_is_i64() {
518        assert_eq!(
519            classify_number(&num(serde_json::json!(42))),
520            NumberBind::I64
521        );
522        assert_eq!(
523            classify_number(&num(serde_json::json!(-7))),
524            NumberBind::I64
525        );
526        assert_eq!(classify_number(&num(serde_json::json!(0))), NumberBind::I64);
527    }
528
529    #[test]
530    fn classify_above_2_pow_53_stays_i64_not_f64() {
531        // 2^53 + 1 must NOT be bound as f64 (which would round it). It is a
532        // valid i64, so it must classify as I64.
533        let v = 9_007_199_254_740_993i64; // 2^53 + 1
534        assert_eq!(classify_number(&num(serde_json::json!(v))), NumberBind::I64);
535    }
536
537    #[test]
538    fn classify_i64_boundaries_are_i64() {
539        assert_eq!(
540            classify_number(&num(serde_json::json!(i64::MAX))),
541            NumberBind::I64
542        );
543        assert_eq!(
544            classify_number(&num(serde_json::json!(i64::MIN))),
545            NumberBind::I64
546        );
547    }
548
549    #[test]
550    fn classify_above_i64_max_is_u64() {
551        let v: u64 = i64::MAX as u64 + 1;
552        assert_eq!(classify_number(&num(serde_json::json!(v))), NumberBind::U64);
553        assert_eq!(
554            classify_number(&num(serde_json::json!(u64::MAX))),
555            NumberBind::U64
556        );
557    }
558
559    #[test]
560    fn classify_float_is_f64() {
561        assert_eq!(
562            classify_number(&num(serde_json::json!(3.5))),
563            NumberBind::F64
564        );
565    }
566
567    // ── PK-range sharding (Mode B, #262) ─────────────────────────────────────
568
569    #[test]
570    fn quote_ident_mysql_backticks_and_escapes() {
571        assert_eq!(quote_ident_mysql("id"), "`id`");
572        // Embedded backticks are doubled — identifier injection is inert.
573        assert_eq!(quote_ident_mysql("we`ird"), "`we``ird`");
574    }
575
576    #[test]
577    fn shard_wrap_uses_backtick_quoting() {
578        let spec = faucet_core::shard::ShardSpec::new(
579            "1",
580            serde_json::json!({"key": "id", "lo": 100, "hi": 200, "lo_unbounded": false, "hi_unbounded": false}),
581        );
582        let bounds = PkShardBounds::from_spec(&spec).unwrap();
583        let sql = bounds.wrap("SELECT * FROM t", quote_ident_mysql);
584        assert!(sql.contains("(SELECT * FROM t) AS _faucet_shard"), "{sql}");
585        assert!(sql.contains("`id` >= 100"), "backtick-quoted key: {sql}");
586        assert!(sql.contains("`id` < 200"), "half-open upper bound: {sql}");
587    }
588
589    #[test]
590    fn last_shard_wrap_covers_null_keys() {
591        let shards = plan_pk_shards("id", 0, 99, 3);
592        let last = PkShardBounds::from_spec(shards.last().unwrap()).unwrap();
593        let sql = last.wrap("SELECT * FROM t", quote_ident_mysql);
594        assert!(
595            sql.contains("`id` IS NULL"),
596            "last shard must match NULL keys: {sql}"
597        );
598    }
599
600    /// Build a source over a lazy pool (no server needed) so the shard glue —
601    /// `apply_shard`, `shard_wrap`, and `enumerate_shards`' non-I/O branches —
602    /// is testable without Docker.
603    fn lazy_source(config: MysqlSourceConfig) -> MysqlSource {
604        let pool = MySqlPoolOptions::new()
605            // Fail fast at first checkout — these tests never reach a server.
606            .acquire_timeout(std::time::Duration::from_millis(200))
607            .connect_lazy(&config.connection_url)
608            .expect("lazy pool");
609        MysqlSource {
610            config,
611            pool,
612            applied_shard: Mutex::new(None),
613        }
614    }
615
616    #[tokio::test]
617    async fn apply_shard_then_shard_wrap_narrows_query() {
618        use faucet_core::Source as _;
619        let mut config = MysqlSourceConfig::new("mysql://root@127.0.0.1:1/db", "SELECT * FROM t");
620        config.shard = Some(crate::config::ShardConfig { key: "id".into() });
621        let source = lazy_source(config);
622        assert!(source.is_shardable());
623
624        // No shard applied / whole shard applied → query passes through.
625        assert_eq!(source.shard_wrap("SELECT 1".into()), "SELECT 1");
626        source
627            .apply_shard(&faucet_core::ShardSpec::whole())
628            .await
629            .unwrap();
630        assert_eq!(source.shard_wrap("SELECT 1".into()), "SELECT 1");
631
632        // A real shard narrows with backtick quoting.
633        let spec = &plan_pk_shards("id", 0, 99, 2)[0];
634        source.apply_shard(spec).await.unwrap();
635        let wrapped = source.shard_wrap("SELECT * FROM t".into());
636        assert!(wrapped.contains("`id`"), "got: {wrapped}");
637        assert!(wrapped.contains("_faucet_shard"), "got: {wrapped}");
638
639        // Malformed descriptor is rejected.
640        let bad = faucet_core::ShardSpec::new("0", serde_json::json!({ "key": "id" }));
641        assert!(source.apply_shard(&bad).await.is_err());
642    }
643
644    // ── discover: pure catalog-row grouping (#211) ───────────────────────────
645
646    #[test]
647    fn descriptors_group_catalog_rows_per_table() {
648        let rows: Vec<CatalogRow> = vec![
649            (
650                "orders".to_string(),
651                "id".to_string(),
652                "bigint".to_string(),
653                false,
654                Some(120u64),
655            ),
656            (
657                "orders".to_string(),
658                "note".to_string(),
659                "varchar".to_string(),
660                true,
661                Some(120u64),
662            ),
663            (
664                "users".to_string(),
665                "total".to_string(),
666                "decimal".to_string(),
667                false,
668                None,
669            ),
670        ];
671        let ds = descriptors_from_catalog(rows);
672        assert_eq!(ds.len(), 2, "rows group into one descriptor per table");
673
674        assert_eq!(ds[0].name, "orders", "bare table name — no db qualifier");
675        assert_eq!(ds[0].kind, "table");
676        assert_eq!(ds[0].estimated_rows, Some(120));
677        assert_eq!(ds[0].config_patch["query"], "SELECT * FROM `orders`");
678        let schema = ds[0].schema.as_ref().unwrap();
679        assert_eq!(schema["type"], "object");
680        assert_eq!(schema["properties"]["id"]["type"], "integer");
681        assert_eq!(
682            schema["properties"]["note"]["type"],
683            serde_json::json!(["string", "null"]),
684            "nullable column"
685        );
686
687        assert_eq!(ds[1].name, "users");
688        assert_eq!(ds[1].estimated_rows, None, "NULL table_rows = no estimate");
689        assert_eq!(
690            ds[1].schema.as_ref().unwrap()["properties"]["total"]["type"],
691            "number"
692        );
693    }
694
695    #[test]
696    fn descriptors_quote_hostile_identifiers() {
697        let rows: Vec<CatalogRow> = vec![(
698            "we`ird".to_string(),
699            "id".to_string(),
700            "int".to_string(),
701            false,
702            None,
703        )];
704        let ds = descriptors_from_catalog(rows);
705        assert_eq!(
706            ds[0].config_patch["query"], "SELECT * FROM `we``ird`",
707            "embedded backticks are doubled"
708        );
709    }
710
711    #[test]
712    fn descriptors_empty_catalog_is_empty() {
713        assert!(descriptors_from_catalog(Vec::new()).is_empty());
714    }
715
716    #[tokio::test]
717    async fn source_advertises_discover() {
718        use faucet_core::Source as _;
719        let source = lazy_source(MysqlSourceConfig::new(
720            "mysql://root@127.0.0.1:1/db",
721            "SELECT 1",
722        ));
723        assert!(source.supports_discover());
724        // Against an unreachable server the catalog query surfaces the typed
725        // discovery error (exercises the error path without Docker).
726        let err = source.discover().await.unwrap_err();
727        assert!(
728            err.to_string().contains("catalog discovery failed"),
729            "typed error: {err}"
730        );
731    }
732
733    #[tokio::test]
734    async fn enumerate_shards_without_config_is_whole_and_with_config_needs_db() {
735        use faucet_core::Source as _;
736        // No `shard` config → single whole shard, no I/O.
737        let plain = lazy_source(MysqlSourceConfig::new(
738            "mysql://root@127.0.0.1:1/db",
739            "SELECT 1",
740        ));
741        assert!(!plain.is_shardable());
742        let shards = plain.enumerate_shards(4).await.unwrap();
743        assert_eq!(shards.len(), 1);
744        assert!(shards[0].is_whole());
745
746        // With config, enumeration must reach the (unreachable) server → the
747        // bounds-probe error path surfaces as FaucetError::Source.
748        let mut config = MysqlSourceConfig::new("mysql://root@127.0.0.1:1/db", "SELECT 1");
749        config.shard = Some(crate::config::ShardConfig { key: "id".into() });
750        let sharded = lazy_source(config);
751        let err = sharded.enumerate_shards(4).await.unwrap_err();
752        assert!(
753            err.to_string().contains("shard bounds"),
754            "expected bounds-probe error, got: {err}"
755        );
756    }
757}