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

python 
import heapq
 
 
class MovingMedian:

Running median with two heaps

heaps streaming invariants
Advanced 8 steps
typescript 
interface RunningStats {
  count: number;
  total: number;
  average: number;

Streaming running averages in TypeScript

generators streaming incremental-aggregation
Intermediate 6 steps
rust 
use axum::{
    extract::{FromRequestParts, Path, Query},
    http::{request::Parts, StatusCode},
    response::{IntoResponse, Redirect},

Signed download links as an Axum extractor

hmac custom-extractor authentication
Advanced 9 steps
rust 
use axum::{
    extract::{FromRef, FromRequestParts},
    http::{header, request::Parts, StatusCode},
    RequestPartsExt,

How a JWT extractor works in Axum

jwt authentication extractors
Intermediate 8 steps
rust 
use std::collections::HashMap;
 
#[derive(Clone, Copy, PartialEq)]
enum Color {

Detecting cycles with three-color DFS in Rust

graph-algorithms cycle-detection depth-first-search
Intermediate 9 steps
rust 
#[derive(Deserialize)]
pub struct CreateArticle {
    title: String,
    body: String,

Building a create endpoint in Axum

extractors json-deserialization sqlx
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