quickhouse 0.3.3__tar.gz → 0.3.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (34) hide show
  1. {quickhouse-0.3.3 → quickhouse-0.3.4}/Cargo.lock +2 -2
  2. {quickhouse-0.3.3 → quickhouse-0.3.4}/Cargo.toml +1 -1
  3. {quickhouse-0.3.3 → quickhouse-0.3.4}/PKG-INFO +19 -1
  4. {quickhouse-0.3.3 → quickhouse-0.3.4}/README.md +18 -0
  5. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/config.rs +34 -0
  6. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/ddl.rs +1 -0
  7. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/sink/bigquery.rs +1 -0
  8. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/source/postgres.rs +15 -7
  9. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/sync.rs +134 -2
  10. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/transform.rs +137 -1
  11. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/types.rs +43 -2
  12. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-py/src/lib.rs +3 -0
  13. {quickhouse-0.3.3 → quickhouse-0.3.4}/pyproject.toml +1 -1
  14. {quickhouse-0.3.3 → quickhouse-0.3.4}/python/quickhouse/_quickhouse.pyi +33 -0
  15. {quickhouse-0.3.3 → quickhouse-0.3.4}/LICENSE +0 -0
  16. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/Cargo.toml +0 -0
  17. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/archive.rs +0 -0
  18. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/decimal.rs +0 -0
  19. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/decode.rs +0 -0
  20. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/decode_bigquery.rs +0 -0
  21. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/decode_mysql.rs +0 -0
  22. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/error.rs +0 -0
  23. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/lib.rs +0 -0
  24. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/memory.rs +0 -0
  25. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/sink/bigquery_proto.rs +0 -0
  26. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/sink/clickhouse.rs +0 -0
  27. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/sink/mod.rs +0 -0
  28. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/source/bigquery.rs +0 -0
  29. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/source/mod.rs +0 -0
  30. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-core/src/source/mysql.rs +0 -0
  31. {quickhouse-0.3.3 → quickhouse-0.3.4}/crates/quickhouse-py/Cargo.toml +0 -0
  32. {quickhouse-0.3.3 → quickhouse-0.3.4}/python/quickhouse/__init__.py +0 -0
  33. {quickhouse-0.3.3 → quickhouse-0.3.4}/python/quickhouse/progress.py +0 -0
  34. {quickhouse-0.3.3 → quickhouse-0.3.4}/python/quickhouse/py.typed +0 -0
@@ -2634,7 +2634,7 @@ dependencies = [
2634
2634
 
2635
2635
  [[package]]
2636
2636
  name = "quickhouse-core"
2637
- version = "0.3.3"
2637
+ version = "0.3.4"
2638
2638
  dependencies = [
2639
2639
  "anyhow",
2640
2640
  "arrow",
@@ -2670,7 +2670,7 @@ dependencies = [
2670
2670
 
2671
2671
  [[package]]
2672
2672
  name = "quickhouse-py"
2673
- version = "0.3.3"
2673
+ version = "0.3.4"
2674
2674
  dependencies = [
2675
2675
  "pyo3",
2676
2676
  "quickhouse-core",
@@ -3,7 +3,7 @@ resolver = "2"
3
3
  members = ["crates/quickhouse-core", "crates/quickhouse-py"]
4
4
 
5
5
  [workspace.package]
6
- version = "0.3.3"
6
+ version = "0.3.4"
7
7
  edition = "2021"
8
8
  license = "MIT"
9
9
  authors = ["quickhouse contributors"]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: quickhouse
3
- Version: 0.3.3
3
+ Version: 0.3.4
4
4
  Classifier: Development Status :: 3 - Alpha
5
5
  Classifier: Intended Audience :: Developers
6
6
  Classifier: License :: OSI Approved :: MIT License
@@ -72,6 +72,16 @@ print(result) # rows_read, rows_written, bytes_written, duration_secs, new_wat
72
72
  legacy quirks like MySQL zero-dates or out-of-range timestamps are coerced to
73
73
  `NULL` with a warning instead of aborting the run.
74
74
 
75
+ - **It's gentle on a small production database.** Set `read_max_rows_per_sec`
76
+ and quickhouse paces the read to that aggregate rate across all partitions;
77
+ because `COPY`/streaming results only produce as fast as the client consumes,
78
+ the source scan itself backs off — you're throttling the database's work, not
79
+ just your own. Combine it with `parallelism=1` (one connection), incremental
80
+ mode (only new rows), and a `statement_timeout_secs`, point it at a read
81
+ replica, and a bulk export stops competing with your app. The Postgres
82
+ connection also shows up as `application_name = 'quickhouse'` in
83
+ `pg_stat_activity`, so a DBA can see and kill it.
84
+
75
85
  - **There's nothing to stand up.** `pip install quickhouse` and you're done —
76
86
  no JVM, no Spark cluster, no separate service. It's an ordinary Python
77
87
  dependency that runs wherever your jobs already run: cron, Airflow, Dagster,
@@ -218,6 +228,7 @@ errors are surfaced verbatim rather than wrapped in something generic.
218
228
  | `parallelism` | Concurrent read streams |
219
229
  | `batch_rows` / `batch_bytes` | Per-batch size knobs (rows, or estimated bytes) |
220
230
  | `max_memory_bytes` | Hard ceiling on total in-flight memory; decoding blocks when hit (default 512 MiB, `0` = unbounded) |
231
+ | `read_max_rows_per_sec` | Cap the aggregate source read rate to be gentle on a small DB; `None` = unlimited (default). Postgres/MySQL only |
221
232
  | `type_overrides` | Force a destination column type, e.g. `{"qty": "Decimal(18, 3)"}` |
222
233
  | `rename`, `include`, `exclude` | Column renames and allow/deny lists |
223
234
  | `on_progress` | Progress callback |
@@ -240,6 +251,13 @@ knowing:
240
251
  front, rather than silently falling back to `Float64`.
241
252
  - **`TIME`** columns transfer as canonical text into a `String` column
242
253
  (ClickHouse has no time-of-day type).
254
+ - **MySQL `DATETIME`/`TIMESTAMP`** map to a UTC-aware timestamp (BigQuery
255
+ `TIMESTAMP`, ClickHouse `DateTime64(6, 'UTC')`) — the wall-clock value is read
256
+ as UTC, matching how a `TIMESTAMP` column expects it. To land a column as a
257
+ naive BigQuery `DATETIME` instead, opt out per-column with
258
+ `type_overrides={"col": "DATETIME"}` — that flips the actual encoding, not
259
+ just the declared type. (PostgreSQL keeps the distinction natively:
260
+ `timestamptz` → UTC-aware, `timestamp` → naive.)
243
261
  - **Out-of-range dates** (and MySQL zero-dates like `0000-00-00`) coerce to
244
262
  `NULL` with a warning rather than failing the transfer.
245
263
  - **Nullable** source columns stay nullable in the destination.
@@ -40,6 +40,16 @@ print(result) # rows_read, rows_written, bytes_written, duration_secs, new_wat
40
40
  legacy quirks like MySQL zero-dates or out-of-range timestamps are coerced to
41
41
  `NULL` with a warning instead of aborting the run.
42
42
 
43
+ - **It's gentle on a small production database.** Set `read_max_rows_per_sec`
44
+ and quickhouse paces the read to that aggregate rate across all partitions;
45
+ because `COPY`/streaming results only produce as fast as the client consumes,
46
+ the source scan itself backs off — you're throttling the database's work, not
47
+ just your own. Combine it with `parallelism=1` (one connection), incremental
48
+ mode (only new rows), and a `statement_timeout_secs`, point it at a read
49
+ replica, and a bulk export stops competing with your app. The Postgres
50
+ connection also shows up as `application_name = 'quickhouse'` in
51
+ `pg_stat_activity`, so a DBA can see and kill it.
52
+
43
53
  - **There's nothing to stand up.** `pip install quickhouse` and you're done —
44
54
  no JVM, no Spark cluster, no separate service. It's an ordinary Python
45
55
  dependency that runs wherever your jobs already run: cron, Airflow, Dagster,
@@ -186,6 +196,7 @@ errors are surfaced verbatim rather than wrapped in something generic.
186
196
  | `parallelism` | Concurrent read streams |
187
197
  | `batch_rows` / `batch_bytes` | Per-batch size knobs (rows, or estimated bytes) |
188
198
  | `max_memory_bytes` | Hard ceiling on total in-flight memory; decoding blocks when hit (default 512 MiB, `0` = unbounded) |
199
+ | `read_max_rows_per_sec` | Cap the aggregate source read rate to be gentle on a small DB; `None` = unlimited (default). Postgres/MySQL only |
189
200
  | `type_overrides` | Force a destination column type, e.g. `{"qty": "Decimal(18, 3)"}` |
190
201
  | `rename`, `include`, `exclude` | Column renames and allow/deny lists |
191
202
  | `on_progress` | Progress callback |
@@ -208,6 +219,13 @@ knowing:
208
219
  front, rather than silently falling back to `Float64`.
209
220
  - **`TIME`** columns transfer as canonical text into a `String` column
210
221
  (ClickHouse has no time-of-day type).
222
+ - **MySQL `DATETIME`/`TIMESTAMP`** map to a UTC-aware timestamp (BigQuery
223
+ `TIMESTAMP`, ClickHouse `DateTime64(6, 'UTC')`) — the wall-clock value is read
224
+ as UTC, matching how a `TIMESTAMP` column expects it. To land a column as a
225
+ naive BigQuery `DATETIME` instead, opt out per-column with
226
+ `type_overrides={"col": "DATETIME"}` — that flips the actual encoding, not
227
+ just the declared type. (PostgreSQL keeps the distinction natively:
228
+ `timestamptz` → UTC-aware, `timestamp` → naive.)
211
229
  - **Out-of-range dates** (and MySQL zero-dates like `0000-00-00`) coerce to
212
230
  `NULL` with a warning rather than failing the transfer.
213
231
  - **Nullable** source columns stay nullable in the destination.
@@ -256,6 +256,17 @@ pub struct TransferConfig {
256
256
  /// Column used to split the table into parallel range partitions.
257
257
  /// Defaults to the first `key` column, else the sync falls back to a single stream.
258
258
  pub partition_column: Option<String>,
259
+ /// Optional cap on how many source rows are pulled **per second**, summed
260
+ /// across all parallel partitions (a global limiter, not per-connection).
261
+ /// Deliberately paces the read so a small/production database isn't
262
+ /// hammered by a bulk export: after each batch is read, the reader sleeps
263
+ /// long enough to hold the aggregate rate at this ceiling, which — because
264
+ /// `COPY TO STDOUT` streams only as fast as the client consumes — makes the
265
+ /// server-side scan itself back off (TCP backpressure), not just the
266
+ /// client. `None` (default) reads as fast as possible. Applies to the
267
+ /// PostgreSQL and MySQL sources; ignored for a BigQuery source (its read
268
+ /// path is a managed, separately-metered API).
269
+ pub read_max_rows_per_sec: Option<u64>,
259
270
 
260
271
  // ---- transforms ----
261
272
  /// Per-column destination type overrides (column name -> the
@@ -330,6 +341,11 @@ impl TransferConfig {
330
341
  "max_memory_bytes must be 0 (unbounded) or >= 65536",
331
342
  ));
332
343
  }
344
+ if self.read_max_rows_per_sec == Some(0) {
345
+ return Err(EtlError::config(
346
+ "read_max_rows_per_sec must be None (unlimited) or >= 1",
347
+ ));
348
+ }
333
349
  Ok(())
334
350
  }
335
351
  }
@@ -367,6 +383,7 @@ mod tests {
367
383
  batch_bytes: 0,
368
384
  max_memory_bytes: 0,
369
385
  partition_column: None,
386
+ read_max_rows_per_sec: None,
370
387
  type_overrides: HashMap::new(),
371
388
  rename: HashMap::new(),
372
389
  include: vec![],
@@ -423,4 +440,21 @@ mod tests {
423
440
  c.lookback_seconds = 60;
424
441
  assert!(c.validate().is_ok());
425
442
  }
443
+
444
+ #[test]
445
+ fn validate_rejects_zero_read_rate() {
446
+ let mut c = cfg(SyncMode::Full, None);
447
+ c.read_max_rows_per_sec = Some(0);
448
+ let err = c.validate().unwrap_err().to_string();
449
+ assert!(err.contains("read_max_rows_per_sec"), "got: {err}");
450
+ }
451
+
452
+ #[test]
453
+ fn validate_accepts_none_or_positive_read_rate() {
454
+ let mut c = cfg(SyncMode::Full, None);
455
+ c.read_max_rows_per_sec = None;
456
+ assert!(c.validate().is_ok());
457
+ c.read_max_rows_per_sec = Some(10_000);
458
+ assert!(c.validate().is_ok());
459
+ }
426
460
  }
@@ -147,6 +147,7 @@ mod tests {
147
147
  batch_bytes: 0,
148
148
  max_memory_bytes: 0,
149
149
  partition_column: None,
150
+ read_max_rows_per_sec: None,
150
151
  type_overrides: HashMap::new(),
151
152
  rename: HashMap::new(),
152
153
  include: vec![],
@@ -945,6 +945,7 @@ mod tests {
945
945
  batch_bytes: 0,
946
946
  max_memory_bytes: 0,
947
947
  partition_column: None,
948
+ read_max_rows_per_sec: None,
948
949
  type_overrides: HashMap::new(),
949
950
  rename: HashMap::new(),
950
951
  include: vec![],
@@ -92,14 +92,22 @@ impl PgSource {
92
92
  tracing::warn!("postgres connection error: {e}");
93
93
  }
94
94
  });
95
+ // Session setup, run as one round-trip. `application_name` makes this
96
+ // export visible (and killable) in `pg_stat_activity` — exactly what a
97
+ // DBA wants when watching load from a bulk read. When a statement
98
+ // timeout is configured we also cap `idle_in_transaction_session_timeout`
99
+ // at the same value: defense-in-depth so a stalled connection can never
100
+ // sit inside an open transaction holding an MVCC snapshot (which would
101
+ // block vacuum and bloat the source). The bulk COPY itself runs in
102
+ // autocommit, so statement_timeout is the primary guard; this covers the
103
+ // edge/future transactional paths.
104
+ let mut setup = String::from("SET application_name = 'quickhouse'");
95
105
  if self.statement_timeout_secs > 0 {
96
- client
97
- .batch_execute(&format!(
98
- "SET statement_timeout = {}",
99
- self.statement_timeout_secs * 1000
100
- ))
101
- .await?;
106
+ let ms = self.statement_timeout_secs * 1000;
107
+ setup.push_str(&format!("; SET statement_timeout = {ms}"));
108
+ setup.push_str(&format!("; SET idle_in_transaction_session_timeout = {ms}"));
102
109
  }
110
+ client.batch_execute(&setup).await?;
103
111
  Ok(client)
104
112
  }
105
113
 
@@ -132,7 +140,7 @@ impl PgSource {
132
140
  column: c.name().to_string(),
133
141
  type_name: c.type_().name().to_string(),
134
142
  })?;
135
- let nullable = !not_null.contains(&c.name().to_string());
143
+ let nullable = !not_null.contains(c.name());
136
144
  cols.push(ColumnType {
137
145
  name: c.name().to_string(),
138
146
  type_id: pg_oid,
@@ -5,8 +5,8 @@
5
5
  //! source-agnostic.
6
6
 
7
7
  use std::sync::atomic::{AtomicU64, Ordering};
8
- use std::sync::Arc;
9
- use std::time::Instant;
8
+ use std::sync::{Arc, Mutex};
9
+ use std::time::{Duration, Instant};
10
10
 
11
11
  use arrow_array::RecordBatch;
12
12
  use arrow_schema::SchemaRef;
@@ -54,6 +54,60 @@ struct Counters {
54
54
  bytes_written: AtomicU64,
55
55
  }
56
56
 
57
+ /// Global source-read rate limiter, shared across every parallel partition so
58
+ /// the ceiling is an *aggregate* rows/sec (not per-connection). Uses a
59
+ /// virtual-scheduling (GCRA-style) clock: each `acquire(n)` reserves `n / rate`
60
+ /// seconds of read time by pushing a shared `next_available` instant forward,
61
+ /// and the caller sleeps until its reserved slot. Because reads are gated
62
+ /// *after* each batch, and `COPY TO STDOUT` (and MySQL's streaming result) only
63
+ /// produce as fast as the client consumes, pausing here applies TCP
64
+ /// backpressure that slows the server-side scan itself — the whole point of the
65
+ /// knob. No burst credit accumulates while idle (`next_available` is clamped
66
+ /// forward to `now`, never left in the past), which is the safe choice for
67
+ /// "be gentle to the source".
68
+ struct ReadThrottle {
69
+ rate_per_sec: f64,
70
+ next_available: Mutex<Instant>,
71
+ }
72
+
73
+ impl ReadThrottle {
74
+ fn new(rows_per_sec: u64) -> Self {
75
+ Self {
76
+ rate_per_sec: rows_per_sec as f64,
77
+ next_available: Mutex::new(Instant::now()),
78
+ }
79
+ }
80
+
81
+ /// Reserve capacity for `rows` just-read rows and return how long the
82
+ /// caller must wait before pulling more. Pure of any `.await` (so the
83
+ /// `std::sync::Mutex` is never held across a suspension point, and the math
84
+ /// is unit-testable without sleeping); `acquire` wraps it with the sleep.
85
+ fn reserve(&self, rows: u64) -> Duration {
86
+ // A zero rate can't happen via the public API (`validate` rejects
87
+ // `Some(0)` before any throttle is built), but guard anyway so a Rust
88
+ // caller constructing `TransferConfig` directly gets a safe no-op
89
+ // instead of an `inf` `Duration` panic (`rows / 0.0`) in the read loop.
90
+ // `rate_per_sec` is a `u64 as f64`, so it's always finite and >= 0.
91
+ if rows == 0 || self.rate_per_sec <= 0.0 {
92
+ return Duration::ZERO;
93
+ }
94
+ let mut next = self.next_available.lock().unwrap();
95
+ let now = Instant::now();
96
+ // Idle time grants no credit: never start earlier than `now`.
97
+ let start = (*next).max(now);
98
+ let cost = Duration::from_secs_f64(rows as f64 / self.rate_per_sec);
99
+ *next = start + cost;
100
+ start.saturating_duration_since(now)
101
+ }
102
+
103
+ async fn acquire(&self, rows: u64) {
104
+ let wait = self.reserve(rows);
105
+ if !wait.is_zero() {
106
+ tokio::time::sleep(wait).await;
107
+ }
108
+ }
109
+ }
110
+
57
111
  /// Shared handle for uploading finished batches. Cloned cheaply into each
58
112
  /// partition task (all fields are `Arc`/`Copy`), so every partition spawns
59
113
  /// sends against the *same* memory budget and counters.
@@ -69,6 +123,9 @@ struct SendCtx {
69
123
  /// `None` otherwise (including always for BigQuery). See
70
124
  /// `ArchiveRunInfo::writer_for`.
71
125
  archive: Option<Arc<ArchiveRunInfo>>,
126
+ /// `Some` when `read_max_rows_per_sec` is set: a single limiter shared by
127
+ /// all partition tasks, so the cap is an aggregate across the whole read.
128
+ throttle: Option<Arc<ReadThrottle>>,
72
129
  }
73
130
 
74
131
  impl SendCtx {
@@ -354,6 +411,11 @@ async fn run_transfer_impl(
354
411
 
355
412
  // --- Fan out partitions with bounded concurrency. ---
356
413
  let counters = Arc::new(Counters::default());
414
+ // One limiter shared across every partition, so `read_max_rows_per_sec`
415
+ // caps the *aggregate* read rate regardless of `parallelism`.
416
+ let throttle = cfg
417
+ .read_max_rows_per_sec
418
+ .map(|r| Arc::new(ReadThrottle::new(r)));
357
419
  let cfg = Arc::new(cfg);
358
420
  let extra_filter = Arc::new(extra_filter);
359
421
  let target_table = Arc::new(target_table);
@@ -365,6 +427,7 @@ async fn run_transfer_impl(
365
427
  progress: progress.clone(),
366
428
  started,
367
429
  archive: archive_info.clone(),
430
+ throttle,
368
431
  };
369
432
 
370
433
  let mut results = futures::stream::iter(partitions.into_iter().map(|part| {
@@ -549,6 +612,10 @@ async fn run_transfer_bigquery(
549
612
  progress: progress.clone(),
550
613
  started,
551
614
  archive: archive_info,
615
+ // read_max_rows_per_sec is a Postgres/MySQL knob; the BigQuery
616
+ // Storage Read API is a separately-metered managed service, so this
617
+ // path never throttles.
618
+ throttle: None,
552
619
  };
553
620
 
554
621
  let mut iter = source
@@ -767,10 +834,16 @@ async fn transfer_partition_postgres(
767
834
  let chunk = chunk?;
768
835
  let batches = decoder.feed(&chunk)?;
769
836
  for batch in batches {
837
+ let rows = batch.num_rows() as u64;
770
838
  if let Some(w) = archive_writer.as_mut() {
771
839
  w.write(&batch).await?;
772
840
  }
773
841
  ctx.spawn_upload(&mut sends, schema.clone(), batch).await;
842
+ // Pace the read: pausing before pulling the next chunk applies TCP
843
+ // backpressure to the COPY stream, slowing the server-side scan.
844
+ if let Some(t) = &ctx.throttle {
845
+ t.acquire(rows).await;
846
+ }
774
847
  }
775
848
  reap(&mut sends, false).await?; // surface any upload error promptly
776
849
  }
@@ -853,11 +926,17 @@ async fn transfer_partition_mysql(
853
926
  while let Some(row) = stream.next().await {
854
927
  let row = row.map_err(|e| EtlError::other(format!("mysql row error: {e}")))?;
855
928
  if let Some(batch) = batcher.append_row(row)? {
929
+ let rows = batch.num_rows() as u64;
856
930
  if let Some(w) = archive_writer.as_mut() {
857
931
  w.write(&batch).await?;
858
932
  }
859
933
  ctx.spawn_upload(&mut sends, schema.clone(), batch).await;
860
934
  reap(&mut sends, false).await?; // surface any upload error promptly
935
+ // Pace the read: pausing before fetching more rows applies
936
+ // backpressure to the streaming result set, slowing the scan.
937
+ if let Some(t) = &ctx.throttle {
938
+ t.acquire(rows).await;
939
+ }
861
940
  }
862
941
  }
863
942
  if let Some(batch) = batcher.finish()? {
@@ -1374,6 +1453,59 @@ mod tests {
1374
1453
  assert_ne!(staging_name("orders", "1"), staging_name("orders", "2"));
1375
1454
  }
1376
1455
 
1456
+ #[test]
1457
+ fn throttle_reserve_zero_is_immediate() {
1458
+ let t = ReadThrottle::new(1000);
1459
+ assert!(t.reserve(0).is_zero());
1460
+ }
1461
+
1462
+ #[test]
1463
+ fn throttle_zero_rate_is_safe_noop_not_panic() {
1464
+ // `validate` prevents this via the public API, but a direct Rust
1465
+ // caller must not trip an `inf` Duration panic (rows / 0.0).
1466
+ let t = ReadThrottle::new(0);
1467
+ assert!(t.reserve(1_000_000).is_zero());
1468
+ }
1469
+
1470
+ #[test]
1471
+ fn throttle_first_reserve_does_not_wait() {
1472
+ // The limiter starts unfilled: the first read is never delayed
1473
+ // (`next_available` begins at construction time, already in the past by
1474
+ // the time we reserve, so `start` clamps to `now` and the wait is 0).
1475
+ let t = ReadThrottle::new(1000);
1476
+ assert!(t.reserve(500).is_zero());
1477
+ }
1478
+
1479
+ #[test]
1480
+ fn throttle_paces_subsequent_reads_by_rate() {
1481
+ // At 1000 rows/s, reserving 1000 rows pushes the schedule ~1s forward,
1482
+ // so the very next reservation must wait close to a full second. Bounds
1483
+ // are wide (>0.5s, <1s) so only genuine pacing — not scheduling jitter
1484
+ // in the microseconds between two calls — can satisfy them.
1485
+ let t = ReadThrottle::new(1000);
1486
+ assert!(t.reserve(1000).is_zero());
1487
+ let wait = t.reserve(1000);
1488
+ assert!(
1489
+ wait > Duration::from_millis(500) && wait < Duration::from_secs(1),
1490
+ "expected ~1s pacing wait, got {wait:?}"
1491
+ );
1492
+ }
1493
+
1494
+ #[test]
1495
+ fn throttle_is_shared_across_handles() {
1496
+ // Two clones of the same Arc share one schedule, so the cap is an
1497
+ // aggregate across all partitions rather than per-connection: after one
1498
+ // handle consumes a full second of budget, the other must wait.
1499
+ let t = Arc::new(ReadThrottle::new(1000));
1500
+ let t2 = t.clone();
1501
+ assert!(t.reserve(1000).is_zero());
1502
+ let wait = t2.reserve(1000);
1503
+ assert!(
1504
+ wait > Duration::from_millis(500),
1505
+ "a second handle should be paced by the first's reservation, got {wait:?}"
1506
+ );
1507
+ }
1508
+
1377
1509
  fn col(name: &str) -> ColumnType {
1378
1510
  col_typed(name, DataType::Int64)
1379
1511
  }
@@ -13,15 +13,60 @@
13
13
  //! `decimal.rs`.
14
14
 
15
15
  use std::collections::HashSet;
16
+ use std::sync::Arc;
16
17
 
17
18
  use arrow_array::types::{validate_decimal_precision_and_scale, Decimal128Type};
18
- use arrow_schema::DataType;
19
+ use arrow_schema::{DataType, TimeUnit};
19
20
 
20
21
  use crate::config::TransferConfig;
21
22
  use crate::decimal::parse_decimal_override;
22
23
  use crate::error::{EtlError, Result};
23
24
  use crate::types::{may_coerce_to_null, ColumnType};
24
25
 
26
+ /// Classify a destination datetime type name into the Arrow timezone it
27
+ /// implies: `Some(Some(tz))` = tz-aware (BigQuery `TIMESTAMP` / ClickHouse
28
+ /// `DateTime64(P, 'tz')`), `Some(None)` = naive (BigQuery `DATETIME` /
29
+ /// ClickHouse `DateTime64(P)`), and `None` = not a datetime type at all (so
30
+ /// the source Arrow type is left untouched).
31
+ ///
32
+ /// This is what makes a `type_overrides` entry flip a datetime column's
33
+ /// *encoding*, not just its DDL string. MySQL datetimes now default to
34
+ /// tz-aware UTC (see `types::map_mysql_type`), so `type_overrides={col:
35
+ /// "DATETIME"}` is the per-column opt-out back to naive — and because the Arrow
36
+ /// tz flag drives the BigQuery Storage Write proto encoding (Int64 micros vs
37
+ /// civil string) and the destination column type, flipping it here fixes the
38
+ /// whole chain, which a DDL-only relabel could not (it would leave a
39
+ /// TIMESTAMP-encoded payload aimed at a DATETIME column, or vice-versa).
40
+ ///
41
+ /// Recognizes BigQuery names (`TIMESTAMP`/`DATETIME`) and ClickHouse names
42
+ /// (`DateTime`/`DateTime64(...)`), case-insensitively.
43
+ fn datetime_override_tz(dest_type: &str) -> Option<Option<Arc<str>>> {
44
+ let trimmed = dest_type.trim();
45
+ let upper = trimmed.to_ascii_uppercase();
46
+ // BigQuery scalar names are exact.
47
+ if upper == "TIMESTAMP" {
48
+ return Some(Some(Arc::from("UTC")));
49
+ }
50
+ if upper == "DATETIME" {
51
+ return Some(None);
52
+ }
53
+ // ClickHouse `DateTime` / `DateTime64(P[, 'tz']])`: tz-aware iff a quoted
54
+ // timezone argument is present.
55
+ if upper.starts_with("DATETIME") {
56
+ return Some(extract_quoted_tz(trimmed));
57
+ }
58
+ None
59
+ }
60
+
61
+ /// Extract a single-quoted timezone argument (e.g. `UTC` from
62
+ /// `DateTime64(6, 'UTC')`), or `None` when the type carries no timezone.
63
+ fn extract_quoted_tz(s: &str) -> Option<Arc<str>> {
64
+ let start = s.find('\'')?;
65
+ let rest = &s[start + 1..];
66
+ let end = rest.find('\'')?;
67
+ Some(Arc::from(&rest[..end]))
68
+ }
69
+
25
70
  #[derive(Debug)]
26
71
  pub struct SelectPlan {
27
72
  /// Source column names to read, in order (drives the COPY SELECT list).
@@ -133,6 +178,20 @@ pub fn plan(source: &[ColumnType], cfg: &TransferConfig) -> Result<SelectPlan> {
133
178
  }
134
179
  None => c.arrow.clone(),
135
180
  }
181
+ } else if matches!(c.arrow, DataType::Timestamp(_, _)) {
182
+ // A datetime column's tz-awareness follows its (possibly
183
+ // overridden) destination type — this is the per-column opt-out
184
+ // that lets `type_overrides={col: "DATETIME"}` render a
185
+ // (UTC-by-default) MySQL datetime as a naive BigQuery DATETIME, and
186
+ // the reverse. With no override, `ch_inner` is the source's own
187
+ // default mapping string, so this reproduces the source tz exactly
188
+ // (no behavior change). A non-datetime override (e.g. "String")
189
+ // returns `None` and leaves the Arrow timestamp untouched, matching
190
+ // the pre-existing DDL-only-relabel behavior for such overrides.
191
+ match datetime_override_tz(&ch_inner) {
192
+ Some(tz) => DataType::Timestamp(TimeUnit::Microsecond, tz),
193
+ None => c.arrow.clone(),
194
+ }
136
195
  } else {
137
196
  c.arrow.clone()
138
197
  };
@@ -219,6 +278,7 @@ mod tests {
219
278
  batch_bytes: 0,
220
279
  max_memory_bytes: 0,
221
280
  partition_column: None,
281
+ read_max_rows_per_sec: None,
222
282
  type_overrides: HashMap::new(),
223
283
  rename: HashMap::new(),
224
284
  include: vec![],
@@ -427,4 +487,80 @@ mod tests {
427
487
  assert!(err.contains("amount"), "must name the column: {err}");
428
488
  assert!(err.contains("Decimal256"), "must mention the Decimal256 follow-up: {err}");
429
489
  }
490
+
491
+ fn ts(tz: Option<&str>) -> DataType {
492
+ DataType::Timestamp(TimeUnit::Microsecond, tz.map(Arc::from))
493
+ }
494
+
495
+ /// The classifier that drives the per-column datetime opt-out: BigQuery and
496
+ /// ClickHouse type names, tz-aware vs naive, and non-datetime strings.
497
+ #[test]
498
+ fn datetime_override_tz_classifies_type_names() {
499
+ // tz-aware
500
+ assert_eq!(datetime_override_tz("TIMESTAMP"), Some(Some(Arc::from("UTC"))));
501
+ assert_eq!(datetime_override_tz("timestamp"), Some(Some(Arc::from("UTC"))));
502
+ assert_eq!(datetime_override_tz("DateTime64(6, 'UTC')"), Some(Some(Arc::from("UTC"))));
503
+ assert_eq!(datetime_override_tz("DateTime('Asia/Jakarta')"), Some(Some(Arc::from("Asia/Jakarta"))));
504
+ // naive
505
+ assert_eq!(datetime_override_tz("DATETIME"), Some(None));
506
+ assert_eq!(datetime_override_tz("DateTime64(6)"), Some(None));
507
+ assert_eq!(datetime_override_tz("DateTime"), Some(None));
508
+ // not a datetime type at all -> leave the Arrow type alone
509
+ assert_eq!(datetime_override_tz("String"), None);
510
+ assert_eq!(datetime_override_tz("Int64"), None);
511
+ assert_eq!(datetime_override_tz("Decimal(18, 2)"), None);
512
+ }
513
+
514
+ /// A `type_overrides={col: "DATETIME"}` on a (UTC-by-default) MySQL-style
515
+ /// timestamp column flips its Arrow tz to naive — the per-column opt-out
516
+ /// that makes it land as a BigQuery DATETIME. This drives the Storage Write
517
+ /// proto encoding, not just the DDL string, so it's the only thing that
518
+ /// actually works (a DDL-only relabel would leave a mismatched payload).
519
+ #[test]
520
+ fn datetime_override_flips_utc_timestamp_to_naive() {
521
+ let src = vec![typed_col("event_at", ts(Some("UTC")), true)];
522
+ let mut cfg = cfg();
523
+ cfg.type_overrides = HashMap::from([("event_at".to_string(), "DATETIME".to_string())]);
524
+ let p = plan(&src, &cfg).unwrap();
525
+ assert_eq!(p.dest_columns[0].arrow, ts(None), "override must flip tz-aware -> naive");
526
+ assert_eq!(p.dest_columns[0].clickhouse_inner, "DATETIME");
527
+ }
528
+
529
+ /// The reverse: an override to `TIMESTAMP` promotes a naive timestamp
530
+ /// column to tz-aware UTC (BigQuery TIMESTAMP).
531
+ #[test]
532
+ fn datetime_override_promotes_naive_timestamp_to_utc() {
533
+ let src = vec![typed_col("event_at", ts(None), true)];
534
+ let mut cfg = cfg();
535
+ cfg.type_overrides = HashMap::from([("event_at".to_string(), "TIMESTAMP".to_string())]);
536
+ let p = plan(&src, &cfg).unwrap();
537
+ assert_eq!(p.dest_columns[0].arrow, ts(Some("UTC")), "override must flip naive -> tz-aware UTC");
538
+ }
539
+
540
+ /// With no override, a timestamp column keeps its source tz exactly — the
541
+ /// classifier reproduces the default from the source's own `clickhouse_inner`
542
+ /// mapping string, so this is a no-op for existing behavior.
543
+ #[test]
544
+ fn timestamp_without_override_keeps_source_tz() {
545
+ let mut utc_col = typed_col("a", ts(Some("UTC")), true);
546
+ utc_col.clickhouse_inner = "DateTime64(6, 'UTC')".into();
547
+ let mut naive_col = typed_col("b", ts(None), true);
548
+ naive_col.clickhouse_inner = "DateTime64(6)".into();
549
+ let p = plan(&[utc_col, naive_col], &cfg()).unwrap();
550
+ assert_eq!(p.dest_columns[0].arrow, ts(Some("UTC")), "UTC source stays UTC");
551
+ assert_eq!(p.dest_columns[1].arrow, ts(None), "naive source stays naive");
552
+ }
553
+
554
+ /// A non-datetime override (e.g. storing a timestamp as text) leaves the
555
+ /// Arrow timestamp untouched — preserving the pre-existing DDL-only-relabel
556
+ /// behavior for such overrides (ClickHouse coerces at insert time).
557
+ #[test]
558
+ fn non_datetime_override_leaves_timestamp_arrow_untouched() {
559
+ let src = vec![typed_col("event_at", ts(Some("UTC")), true)];
560
+ let mut cfg = cfg();
561
+ cfg.type_overrides = HashMap::from([("event_at".to_string(), "String".to_string())]);
562
+ let p = plan(&src, &cfg).unwrap();
563
+ assert_eq!(p.dest_columns[0].arrow, ts(Some("UTC")), "arrow type unchanged by a non-datetime override");
564
+ assert_eq!(p.dest_columns[0].clickhouse_inner, "String", "DDL string still honored");
565
+ }
430
566
  }
@@ -335,12 +335,29 @@ pub mod mysql {
335
335
  MyType::MYSQL_TYPE_DATE | MyType::MYSQL_TYPE_NEWDATE => {
336
336
  (DataType::Date32, "Date32".to_string())
337
337
  }
338
+ // MySQL DATETIME/TIMESTAMP default to a tz-aware UTC mapping
339
+ // (Arrow `Some("UTC")` -> BigQuery `TIMESTAMP` / ClickHouse
340
+ // `DateTime64(6, 'UTC')`), NOT the tz-naive arm. Rationale: the
341
+ // decoder already interprets the wall-clock value as UTC
342
+ // (`dt.and_utc().timestamp_micros()`, decode_mysql.rs) and MySQL
343
+ // has no wire type that would ever reach a tz-aware arm otherwise,
344
+ // so a MySQL datetime could previously only ever land as BigQuery
345
+ // DATETIME — making it impossible to sync into an existing BigQuery
346
+ // `TIMESTAMP` column (fails the staging->dest MERGE with "Value of
347
+ // type DATETIME cannot be assigned to <col>, which has type
348
+ // TIMESTAMP"). This matches the legacy pandas->to_gbq semantics
349
+ // (naive wall-clock stored into BQ TIMESTAMP as UTC — same instant)
350
+ // and the Postgres `TIMESTAMPTZ` mapping. A column that genuinely
351
+ // wants naive (BigQuery DATETIME / `DateTime64(6)`) opts out
352
+ // per-column via `type_overrides={col: "DATETIME"}` — see
353
+ // `transform::datetime_override_tz`, which flips the Arrow tz flag
354
+ // (and thus the Storage Write proto encoding), not just the DDL.
338
355
  MyType::MYSQL_TYPE_DATETIME
339
356
  | MyType::MYSQL_TYPE_DATETIME2
340
357
  | MyType::MYSQL_TYPE_TIMESTAMP
341
358
  | MyType::MYSQL_TYPE_TIMESTAMP2 => (
342
- DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None),
343
- "DateTime64(6)".to_string(),
359
+ DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("UTC".into())),
360
+ "DateTime64(6, 'UTC')".to_string(),
344
361
  ),
345
362
  // TIME -> text into a ClickHouse String column. MySQL TIME can be
346
363
  // negative and exceed 24h (range +/-838:59:59), which no time-of-day
@@ -389,6 +406,30 @@ mod tests {
389
406
  assert!(map_oid(oid::TIMESTAMPTZ).unwrap().1.contains("UTC"));
390
407
  }
391
408
 
409
+ /// MySQL DATETIME/TIMESTAMP default to a tz-aware UTC mapping so they can
410
+ /// land in a BigQuery `TIMESTAMP` column (see the map arm's comment). The
411
+ /// decoder computes the same UTC epoch micros regardless of this flag, so
412
+ /// this only changes the destination type, not the stored instant.
413
+ #[test]
414
+ fn mysql_datetime_defaults_to_utc_aware() {
415
+ use super::mysql::map_mysql_type;
416
+ use mysql_async::consts::ColumnType as MyType;
417
+ for ty in [
418
+ MyType::MYSQL_TYPE_DATETIME,
419
+ MyType::MYSQL_TYPE_DATETIME2,
420
+ MyType::MYSQL_TYPE_TIMESTAMP,
421
+ MyType::MYSQL_TYPE_TIMESTAMP2,
422
+ ] {
423
+ let (arrow, ch) = map_mysql_type(ty, false, false).unwrap();
424
+ assert_eq!(
425
+ arrow,
426
+ DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("UTC".into())),
427
+ "{ty:?} must map to a tz-aware (UTC) Arrow timestamp"
428
+ );
429
+ assert_eq!(ch, "DateTime64(6, 'UTC')", "{ty:?} ClickHouse type");
430
+ }
431
+ }
432
+
392
433
  #[test]
393
434
  fn nullable_wrapping() {
394
435
  let c = ColumnType {
@@ -479,6 +479,7 @@ fn parse_parquet_compression(c: &str) -> PyResult<core::ParquetCompression> {
479
479
  batch_bytes=4_194_304,
480
480
  max_memory_bytes=536_870_912,
481
481
  partition_column=None,
482
+ read_max_rows_per_sec=None,
482
483
  type_overrides=None,
483
484
  rename=None,
484
485
  include=None,
@@ -507,6 +508,7 @@ fn sync(
507
508
  batch_bytes: usize,
508
509
  max_memory_bytes: usize,
509
510
  partition_column: Option<String>,
511
+ read_max_rows_per_sec: Option<u64>,
510
512
  type_overrides: Option<HashMap<String, String>>,
511
513
  rename: Option<HashMap<String, String>>,
512
514
  include: Option<Vec<String>>,
@@ -534,6 +536,7 @@ fn sync(
534
536
  batch_bytes,
535
537
  max_memory_bytes,
536
538
  partition_column,
539
+ read_max_rows_per_sec,
537
540
  type_overrides: type_overrides.unwrap_or_default(),
538
541
  rename: rename.unwrap_or_default(),
539
542
  include: include.unwrap_or_default(),
@@ -4,7 +4,7 @@ build-backend = "maturin"
4
4
 
5
5
  [project]
6
6
  name = "quickhouse"
7
- version = "0.3.3"
7
+ version = "0.3.4"
8
8
  description = "Fast PostgreSQL/MySQL/BigQuery -> ClickHouse ETL with a Rust engine (parallel, bounded-memory, Arrow-based)."
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.9"
@@ -232,6 +232,7 @@ def sync(
232
232
  batch_bytes: int = 4_194_304,
233
233
  max_memory_bytes: int = 536_870_912,
234
234
  partition_column: Optional[str] = None,
235
+ read_max_rows_per_sec: Optional[int] = None,
235
236
  type_overrides: Optional[Mapping[str, str]] = None,
236
237
  rename: Optional[Mapping[str, str]] = None,
237
238
  include: Optional[Sequence[str]] = None,
@@ -291,6 +292,38 @@ def sync(
291
292
  with concurrent uploads and blocks (backpressure) when this ceiling is
292
293
  reached, so peak RSS stays bounded regardless of ``parallelism`` or row
293
294
  width. Default 512 MiB; ``0`` disables the ceiling (unbounded).
295
+
296
+ Being gentle to a small source database:
297
+
298
+ - ``read_max_rows_per_sec`` caps how many source rows are pulled per
299
+ second, summed across *all* parallel partitions (a global limiter, not
300
+ per-connection). After each batch is read the reader pauses to hold the
301
+ aggregate rate at this ceiling; because ``COPY``/streaming results only
302
+ produce as fast as the client consumes, that pause pushes back on the
303
+ server-side scan itself, so the source does proportionally less work —
304
+ not just quickhouse. ``None`` (default) reads as fast as possible.
305
+ Applies to PostgreSQL and MySQL sources; ignored for a BigQuery source
306
+ (its read path is a separately-metered managed API). For the lightest
307
+ possible footprint on a small instance, combine a modest
308
+ ``read_max_rows_per_sec`` with ``parallelism=1`` (one connection, one
309
+ scan), ``mode="incremental"`` (reads only new rows, not the whole
310
+ table), and a ``statement_timeout_secs`` on the ``Postgres``/``MySQL``
311
+ connection. The Postgres connection also reports itself as
312
+ ``application_name = 'quickhouse'`` so the export is visible (and
313
+ killable) in ``pg_stat_activity``.
314
+
315
+ Datetime/timezone handling:
316
+
317
+ - MySQL ``DATETIME``/``TIMESTAMP`` map to a UTC-aware timestamp — BigQuery
318
+ ``TIMESTAMP``, ClickHouse ``DateTime64(6, 'UTC')`` — reading the
319
+ wall-clock value as UTC (the same instant the legacy pandas/``to_gbq``
320
+ path stored, and what an existing BigQuery ``TIMESTAMP`` column expects).
321
+ To land a column as a naive BigQuery ``DATETIME`` (or ClickHouse
322
+ ``DateTime64(6)``) instead, opt out per-column with
323
+ ``type_overrides={"col": "DATETIME"}``; this flips the actual wire
324
+ encoding, not just the declared destination type, so it works on the
325
+ Storage Write path too. PostgreSQL keeps the distinction from the source
326
+ type: ``timestamptz`` → UTC-aware, ``timestamp`` → naive.
294
327
  """
295
328
  ...
296
329
 
File without changes