rust 68 lines · 8 steps

Streaming a DB migration with Axum

An Axum handler paginates through a table in a background task and streams each row to the client as NDJSON.

Explained by highlit
Axum Advanced streaming keyset-pagination backpressure async-channels ndjson
1use axum::{
2 body::Body,
3 extract::State,
4 http::{header, StatusCode},
5 response::{IntoResponse, Response},
6};
7use futures::stream::StreamExt;
8use serde_json::json;
9use sqlx::PgPool;
10use tokio::sync::mpsc;
11use tokio_stream::wrappers::ReceiverStream;
12 
13#[derive(sqlx::FromRow, serde::Serialize)]
14struct LegacyUser {
15 id: i64,
16 email: String,
17 display_name: Option<String>,
18}
19 
20pub async fn migrate_users(State(pool): State<PgPool>) -> Response {
21 let (tx, rx) = mpsc::channel::<Result<String, std::io::Error>>(16);
22 
23 tokio::spawn(async move {
24 let mut last_id: i64 = 0;
25 loop {
26 let batch = sqlx::query_as::<_, LegacyUser>(
27 "SELECT id, email, display_name FROM legacy_users \
28 WHERE id > $1 ORDER BY id ASC LIMIT $2",
29 )
30 .bind(last_id)
31 .bind(500_i64)
32 .fetch_all(&pool)
33 .await;
34 
35 let rows = match batch {
36 Ok(rows) if rows.is_empty() => break,
37 Ok(rows) => rows,
38 Err(e) => {
39 let _ = tx
40 .send(Ok(json!({ "error": e.to_string() }).to_string() + "\n"))
41 .await;
42 return;
43 }
44 };
45 
46 last_id = rows.last().map(|u| u.id).unwrap_or(last_id);
47 
48 for user in &rows {
49 let line = serde_json::to_string(user).unwrap() + "\n";
50 if tx.send(Ok(line)).await.is_err() {
51 return;
52 }
53 }
54 }
55 });
56 
57 let stream = ReceiverStream::new(rx).map(|res| res.map(Body::from));
58 let body = Body::from_stream(stream.map(|r| r.map(|_| unreachable!())).boxed().into_inner().map(|_| unreachable!()));
59 let _ = body;
60 
61 Response::builder()
62 .status(StatusCode::OK)
63 .header(header::CONTENT_TYPE, "application/x-ndjson")
64 .header(header::TRANSFER_ENCODING, "chunked")
65 .body(Body::from_stream(ReceiverStream::new_from(rx)))
66 .unwrap()
67 .into_response()
68}
01 / 01
STEP 01

Walkthrough

Space play ←→ step click any line
Three takeaways
  1. 1A bounded mpsc channel lets a producer and an HTTP response stream cooperate with natural backpressure.
  2. 2Keyset pagination (WHERE id > last_id) scans large tables efficiently without OFFSET's cost.
  3. 3Spawning the query loop keeps the handler responsive while rows are pushed to the client incrementally.

Related explainers

rust 
use serde::Deserialize;
 
#[derive(Debug, Deserialize)]
#[serde(untagged)]

Parsing flexible JSON shapes with serde

deserialization enums json
Intermediate 6 steps
rust 
use std::cmp::{Ordering, Reverse};
use std::collections::BinaryHeap;
use std::fs::File;
use std::io::{self, BufRead, BufReader, BufWriter, Lines, Write};

K-way merge of sorted logs in Rust

binary-heap k-way-merge streaming-io
Intermediate 8 steps
rust 
use axum::{
    extract::{Path, State},
    response::sse::{Event, KeepAlive, Sse},
};

Streaming import progress with SSE in Axum

server-sent-events streams watch-channel
Advanced 7 steps
rust 
use std::f64::consts::PI;
 
#[derive(Debug, Clone, Copy)]
pub struct LatLng {

Building geographic bounding boxes in Rust

geospatial value-types option
Intermediate 7 steps
python 
import random
from typing import Iterator, List
 
 

How reservoir sampling picks k items

reservoir-sampling streaming randomness
Intermediate 5 steps
rust 
use chrono::{Duration, NaiveDate};
 
#[derive(Debug)]
pub struct DateRange {

Parsing and iterating date ranges in Rust

error-handling iterators parsing
Intermediate 7 steps

Share this explainer

Here's the card — post it anywhere.

Streaming a DB migration with Axum — share card
Made with highlit — turn any snippet into a walkthrough like this in about a minute.
Explain your code