1use 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
16pub struct MysqlSource {
18 config: MysqlSourceConfig,
19 pool: MySqlPool,
20 applied_shard: Mutex<Option<PkShardBounds>>,
24}
25
26fn quote_ident_mysql(name: &str) -> String {
30 format!("`{}`", name.replace('`', "``"))
31}
32
33impl MysqlSource {
34 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 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
60fn mysql_value_to_json(row: &sqlx::mysql::MySqlRow, col_name: &str) -> Value {
65 if let Ok(v) = row.try_get::<Value, _>(col_name) {
67 return v;
68 }
69
70 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 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 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 if let Ok(v) = row.try_get::<sqlx::types::BigDecimal, _>(col_name) {
137 return Value::String(v.to_string());
138 }
139 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
148fn 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
172enum NumberBind {
173 I64,
175 U64,
177 F64,
179}
180
181fn 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
196fn 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 NumberBind::I64 => query.bind(n.as_i64().unwrap()),
207 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
220type CatalogRow = (String, String, String, bool, Option<u64>);
224
225type PendingTable = Option<(String, Option<u64>, Vec<(String, Value)>)>;
228
229fn 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 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
276fn 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 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 async fn discover(&self) -> Result<Vec<faucet_core::DatasetDescriptor>, FaucetError> {
390 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 row.try_get::<u64, _>("estimated_rows").ok(),
427 ))
428 })
429 .collect::<Result<_, _>>()?;
430
431 Ok(descriptors_from_catalog(catalog))
432 }
433
434 fn is_shardable(&self) -> bool {
436 self.config.shard.is_some()
437 }
438
439 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 "e_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 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 #[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 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 let v = 9_007_199_254_740_993i64; 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 #[test]
570 fn quote_ident_mysql_backticks_and_escapes() {
571 assert_eq!(quote_ident_mysql("id"), "`id`");
572 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 fn lazy_source(config: MysqlSourceConfig) -> MysqlSource {
604 let pool = MySqlPoolOptions::new()
605 .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 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 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 let bad = faucet_core::ShardSpec::new("0", serde_json::json!({ "key": "id" }));
641 assert!(source.apply_shard(&bad).await.is_err());
642 }
643
644 #[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 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 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 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}