foxhunt/data/examples/download_mbp10_data.rs

//! Download MBP-10 (Market By Price, 10 levels) data from Databento
//!
//! This example downloads Level-2 order book data (MBP-10) for ES.FUT
//! to support TLOB neural network training.
//!
//! ## Usage
//!
//! ```bash
//! export DATABENTO_API_KEY="your_api_key_here"
//! cargo run --example download_mbp10_data --release
//! ```
//!
//! ## Data Specifications
//!
//! - **Dataset**: GLBX.MDP3 (CME Globex MDP 3.0)
//! - **Schema**: mbp-10 (Market By Price, 10 levels)
//! - **Symbols**: ES.FUT (E-mini S&P 500 continuous front month)
//! - **Date Range**: 2024-01-02 to 2024-01-10 (7 trading days)
//! - **Encoding**: dbn (Databento Binary Encoding v2)
//! - **Compression**: zstd (Zstandard)
//! - **Expected Size**: ~50-70 GB compressed

use anyhow::{Context, Result};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::env;
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use std::time::Duration;
use tokio::time::sleep;

#[derive(Debug, Serialize)]
struct BatchSubmitRequest {
    dataset: String,
    schema: String,
    symbols: Vec<String>,
    stype_in: String,
    start: String,
    end: String,
    encoding: String,
    compression: String,
}

#[derive(Debug, Deserialize)]
struct BatchSubmitResponse {
    job_id: String,
}

#[derive(Debug, Deserialize)]
struct BatchStatusResponse {
    state: String,
    download_url: Option<String>,
    record_count: Option<u64>,
    size_bytes: Option<u64>,
}

const DATABENTO_API_BASE: &str = "https://hist.databento.com/v0";

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    // Get API key from environment
    let api_key =
        env::var("DATABENTO_API_KEY").context("DATABENTO_API_KEY environment variable not set")?;

    // Create output directory
    let output_dir = PathBuf::from("test_data/mbp10");
    std::fs::create_dir_all(&output_dir).context("Failed to create output directory")?;

    let output_file = output_dir.join("ES.FUT.mbp10.2024-01-02_to_2024-01-10.dbn.zst");

    println!("=== Databento MBP-10 Data Download ===\n");
    println!("Dataset: GLBX.MDP3");
    println!("Schema: mbp-10 (Market By Price, 10 levels)");
    println!("Symbol: ES.FUT");
    println!("Date Range: 2024-01-02 to 2024-01-10 (7 trading days)");
    println!("Output: {}", output_file.display());
    println!("\n{}", "=".repeat(50));

    // Step 1: Submit batch download job
    println!("\n[1/4] Submitting batch download job...");
    let job_id = submit_batch_job(&api_key).await?;
    println!("✓ Job submitted successfully: {}", job_id);

    // Step 2: Poll for job completion
    println!("\n[2/4] Waiting for job to complete...");
    let download_url = poll_job_status(&api_key, &job_id).await?;
    println!("✓ Job completed successfully");

    // Step 3: Download the file
    println!("\n[3/4] Downloading data file...");
    download_file(&api_key, &download_url, &output_file).await?;
    println!("✓ Download completed successfully");

    // Step 4: Validate the file
    println!("\n[4/4] Validating downloaded file...");
    validate_file(&output_file)?;
    println!("✓ File validated successfully");

    println!("\n{}", "=".repeat(50));
    println!(
        "SUCCESS: MBP-10 data downloaded to {}",
        output_file.display()
    );
    println!("\nNext Steps:");
    println!("1. Decompress: zstd -d {}", output_file.display());
    println!("2. Validate schema: cargo run --example validate_dbn_schema");
    println!("3. Implement MBP-10 parser in data/src/providers/databento/");

    Ok(())
}

async fn submit_batch_job(api_key: &str) -> Result<String> {
    let client = Client::new();

    let request = BatchSubmitRequest {
        dataset: "GLBX.MDP3".to_string(),
        schema: "mbp-10".to_string(),
        symbols: vec!["ES.FUT".to_string()],
        stype_in: "continuous".to_string(),
        start: "2024-01-02T00:00:00Z".to_string(),
        end: "2024-01-10T23:59:59Z".to_string(),
        encoding: "dbn".to_string(),
        compression: "zstd".to_string(),
    };

    let url = format!("{}/batch.submit_job", DATABENTO_API_BASE);
    let response = client
        .post(&url)
        .header("Authorization", format!("Bearer {}", api_key))
        .json(&request)
        .send()
        .await
        .context("Failed to submit batch job")?;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        anyhow::bail!("API request failed: {} - {}", status, body);
    }

    let result: BatchSubmitResponse = response
        .json()
        .await
        .context("Failed to parse batch submit response")?;

    Ok(result.job_id)
}

async fn poll_job_status(api_key: &str, job_id: &str) -> Result<String> {
    let client = Client::new();
    let url = format!("{}/batch.status?job_id={}", DATABENTO_API_BASE, job_id);

    loop {
        let response = client
            .get(&url)
            .header("Authorization", format!("Bearer {}", api_key))
            .send()
            .await
            .context("Failed to check job status")?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            anyhow::bail!("API request failed: {} - {}", status, body);
        }

        let status: BatchStatusResponse = response
            .json()
            .await
            .context("Failed to parse batch status response")?;

        match status.state.as_str() {
            "done" => {
                if let Some(url) = status.download_url {
                    let size_mb = status.size_bytes.unwrap_or(0) as f64 / 1_048_576.0;
                    let records = status.record_count.unwrap_or(0);
                    println!("  • Records: {}", records);
                    println!("  • Size: {:.2} MB", size_mb);
                    return Ok(url);
                } else {
                    anyhow::bail!("Job completed but no download URL provided");
                }
            },
            "error" => {
                anyhow::bail!("Job failed with error state");
            },
            state => {
                print!("\r  • Status: {} ... ", state);
                std::io::stdout().flush()?;
                sleep(Duration::from_secs(5)).await;
            },
        }
    }
}

async fn download_file(api_key: &str, download_url: &str, output_path: &PathBuf) -> Result<()> {
    let client = Client::new();

    let response = client
        .get(download_url)
        .header("Authorization", format!("Bearer {}", api_key))
        .send()
        .await
        .context("Failed to start download")?;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        anyhow::bail!("Download failed: {} - {}", status, body);
    }

    let total_size = response.content_length().unwrap_or(0);
    let mut file = File::create(output_path).context("Failed to create output file")?;

    let mut downloaded: u64 = 0;
    let mut stream = response.bytes_stream();

    use futures_util::stream::StreamExt;

    while let Some(chunk) = stream.next().await {
        let chunk = chunk.context("Failed to read chunk")?;
        file.write_all(&chunk)
            .context("Failed to write chunk to file")?;

        downloaded += chunk.len() as u64;

        if total_size > 0 {
            let percent = (downloaded as f64 / total_size as f64) * 100.0;
            print!(
                "\r  • Progress: {:.2}% ({:.2} MB / {:.2} MB)",
                percent,
                downloaded as f64 / 1_048_576.0,
                total_size as f64 / 1_048_576.0
            );
            std::io::stdout().flush()?;
        }
    }

    println!();
    Ok(())
}

fn validate_file(path: &PathBuf) -> Result<()> {
    let metadata = std::fs::metadata(path).context("Failed to read file metadata")?;

    let size_mb = metadata.len() as f64 / 1_048_576.0;
    println!("  • File size: {:.2} MB", size_mb);

    if size_mb < 1.0 {
        anyhow::bail!("File is suspiciously small (< 1 MB), may be corrupted");
    }

    println!("  • File appears valid");
    Ok(())
}