faucet-source-bigquery 1.0.0

//! Integration tests for the BigQuery source.
//!
//! Drives [`BigQuerySource`] against a wiremock server mocking the BigQuery
//! REST API. The same OAuth-token + service-account dance the sink tests
//! use lets us run the real `gcp_bigquery_client` client end-to-end without
//! reaching out to GCP — we just point both the auth and `v2` base URLs at
//! the local mock.

use std::collections::HashMap;

use faucet_core::Source;
use faucet_source_bigquery::{BigQueryCredentials, BigQuerySource, BigQuerySourceConfig};
use futures::StreamExt;
use gcp_bigquery_client::client_builder::ClientBuilder;
use serde::Serialize;
use serde_json::{Value, json};
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};

const PROJECT_ID: &str = "test-project";
const AUTH_TOKEN_PATH: &str = "/:o/oauth2/token";
const AUTH_SCOPE_BASE: &str = "/auth/bigquery";

#[derive(Serialize)]
struct FakeToken {
    access_token: &'static str,
    token_type: &'static str,
    expires_in: u32,
}

fn fake_token() -> FakeToken {
    FakeToken {
        access_token: "fake-token",
        token_type: "bearer",
        expires_in: 9_999_999,
    }
}

fn dummy_service_account_json(oauth_server: &str) -> Value {
    let token_uri = format!("{oauth_server}{AUTH_TOKEN_PATH}");
    json!({
        "type": "service_account",
        "project_id": "dummy",
        "private_key_id": "dummy",
        "private_key": "-----BEGIN PRIVATE KEY-----\nMIIEvwIBADANBgkqhkiG9w0BAQEFAASCBKkwggSlAgEAAoIBAQDNk6cKkWP/4NMu\nWb3s24YHfM639IXzPtTev06PUVVQnyHmT1bZgQ/XB6BvIRaReqAqnQd61PAGtX3e\n8XocTw+u/ZfiPJOf+jrXMkRBpiBh9mbyEIqBy8BC20OmsUc+O/YYh/qRccvRfPI7\n3XMabQ8eFWhI6z/t35oRpvEVFJnSIgyV4JR/L/cjtoKnxaFwjBzEnxPiwtdy4olU\nKO/1maklXexvlO7onC7CNmPAjuEZKzdMLzFszikCDnoKJC8k6+2GZh0/JDMAcAF4\nwxlKNQ89MpHVRXZ566uKZg0MqZqkq5RXPn6u7yvNHwZ0oahHT+8ixPPrAEjuPEKM\nUPzVRz71AgMBAAECggEAfdbVWLW5Befkvam3hea2+5xdmeN3n3elrJhkiXxbAhf3\nE1kbq9bCEHmdrokNnI34vz0SWBFCwIiWfUNJ4UxQKGkZcSZto270V8hwWdNMXUsM\npz6S2nMTxJkdp0s7dhAUS93o9uE2x4x5Z0XecJ2ztFGcXY6Lupu2XvnW93V9109h\nkY3uICLdbovJq7wS/fO/AL97QStfEVRWW2agIXGvoQG5jOwfPh86GZZRYP9b8VNw\ntkAUJe4qpzNbWs9AItXOzL+50/wsFkD/iWMGWFuU8DY5ZwsL434N+uzFlaD13wtZ\n63D+tNAxCSRBfZGQbd7WxJVFfZe/2vgjykKWsdyNAQKBgQDnEBgSI836HGSRk0Ub\nDwiEtdfh2TosV+z6xtyU7j/NwjugTOJEGj1VO/TMlZCEfpkYPLZt3ek2LdNL66n8\nDyxwzTT5Q3D/D0n5yE3mmxy13Qyya6qBYvqqyeWNwyotGM7hNNOix1v9lEMtH5Rd\nUT0gkThvJhtrV663bcAWCALmtQKBgQDjw2rYlMUp2TUIa2/E7904WOnSEG85d+nc\norhzthX8EWmPgw1Bbfo6NzH4HhebTw03j3NjZdW2a8TG/uEmZFWhK4eDvkx+rxAa\n6EwamS6cmQ4+vdep2Ac4QCSaTZj02YjHb06Be3gptvpFaFrotH2jnpXxggdiv8ul\n6x+ooCffQQKBgQCR3ykzGoOI6K/c75prELyR+7MEk/0TzZaAY1cSdq61GXBHLQKT\nd/VMgAN1vN51pu7DzGBnT/dRCvEgNvEjffjSZdqRmrAVdfN/y6LSeQ5RCfJgGXSV\nJoWVmMxhCNrxiX3h01Xgp/c9SYJ3VD54AzeR/dwg32/j/oEAsDraLciXGQKBgQDF\nMNc8k/DvfmJv27R06Ma6liA6AoiJVMxgfXD8nVUDW3/tBCVh1HmkFU1p54PArvxe\nchAQqoYQ3dUMBHeh6ZRJaYp2ATfxJlfnM99P1/eHFOxEXdBt996oUMBf53bZ5cyJ\n/lAVwnQSiZy8otCyUDHGivJ+mXkTgcIq8BoEwERFAQKBgQDmImBaFqoMSVihqHIf\nDa4WZqwM7ODqOx0JnBKrKO8UOc51J5e1vpwP/qRpNhUipoILvIWJzu4efZY7GN5C\nImF9sN3PP6Sy044fkVPyw4SYEisxbvp9tfw8Xmpj/pbmugkB2ut6lz5frmEBoJSN\n3osZlZTgx+pM3sO6ITV6U4ID2Q==\n-----END PRIVATE KEY-----\n",
        "client_email": "dummy@developer.gserviceaccount.com",
        "client_id": "dummy",
        "auth_uri": format!("{oauth_server}/o/oauth2/auth"),
        "token_uri": token_uri,
        "auth_provider_x509_cert_url": format!("{oauth_server}/oauth2/v1/certs"),
        "client_x509_cert_url": format!("{oauth_server}/robot/v1/metadata/x509/dummy"),
    })
}

async fn mount_token(server: &MockServer) {
    Mock::given(method("POST"))
        .and(path(AUTH_TOKEN_PATH))
        .respond_with(ResponseTemplate::new(200).set_body_json(fake_token()))
        .mount(server)
        .await;
}

async fn build_source(
    server: &MockServer,
    config: BigQuerySourceConfig,
) -> (BigQuerySource, tempfile::NamedTempFile) {
    let sa_json = dummy_service_account_json(&server.uri());
    let sa_file = tempfile::NamedTempFile::new().unwrap();
    std::fs::write(
        sa_file.path(),
        serde_json::to_string_pretty(&sa_json).unwrap(),
    )
    .unwrap();

    let client = ClientBuilder::new()
        .with_auth_base_url(format!("{}{AUTH_SCOPE_BASE}", server.uri()))
        .with_v2_base_url(server.uri())
        .build_from_service_account_key_file(sa_file.path().to_str().unwrap())
        .await
        .expect("build bigquery client against mock");

    (BigQuerySource::from_parts(config, client), sa_file)
}

fn schema_two_cols() -> Value {
    json!({
        "fields": [
            {"name": "id", "type": "INTEGER"},
            {"name": "name", "type": "STRING"}
        ]
    })
}

fn rows(start: i64, count: i64) -> Vec<Value> {
    (start..start + count)
        .map(|i| json!({"f": [{"v": i.to_string()}, {"v": format!("name-{i}")}]}))
        .collect()
}

fn default_config() -> BigQuerySourceConfig {
    BigQuerySourceConfig::new(
        PROJECT_ID,
        BigQueryCredentials::ApplicationDefault,
        "SELECT id, name FROM events",
    )
    .with_batch_size(10)
}

#[tokio::test]
async fn fetch_all_returns_rows_when_first_response_is_complete() {
    let server = MockServer::start().await;
    mount_token(&server).await;
    Mock::given(method("POST"))
        .and(path(format!("/projects/{PROJECT_ID}/queries")))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "jobComplete": true,
            "schema": schema_two_cols(),
            "jobReference": {"projectId": PROJECT_ID, "jobId": "job-1"},
            "rows": rows(0, 3),
        })))
        .mount(&server)
        .await;

    let (src, _f) = build_source(&server, default_config()).await;
    let rows = src.fetch_all().await.unwrap();
    assert_eq!(rows.len(), 3);
    assert_eq!(rows[0], json!({"id": 0, "name": "name-0"}));
    assert_eq!(rows[2], json!({"id": 2, "name": "name-2"}));
}

#[tokio::test]
async fn fetch_all_paginates_via_page_token() {
    let server = MockServer::start().await;
    mount_token(&server).await;

    // jobs.query: first page (rows 0..5) + pageToken.
    Mock::given(method("POST"))
        .and(path(format!("/projects/{PROJECT_ID}/queries")))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "jobComplete": true,
            "schema": schema_two_cols(),
            "jobReference": {"projectId": PROJECT_ID, "jobId": "job-page"},
            "rows": rows(0, 5),
            "pageToken": "tok-1",
        })))
        .mount(&server)
        .await;

    // getQueryResults: second page (rows 5..10) — final page (no further token).
    Mock::given(method("GET"))
        .and(path(format!("/projects/{PROJECT_ID}/queries/job-page")))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "jobComplete": true,
            "schema": schema_two_cols(),
            "rows": rows(5, 5),
        })))
        .mount(&server)
        .await;

    let (src, _f) = build_source(&server, default_config()).await;
    let rows = src.fetch_all().await.unwrap();
    assert_eq!(rows.len(), 10);
    assert_eq!(rows[0]["id"], 0);
    assert_eq!(rows[9]["id"], 9);
}

#[tokio::test]
async fn stream_pages_chunks_by_batch_size() {
    let server = MockServer::start().await;
    mount_token(&server).await;
    Mock::given(method("POST"))
        .and(path(format!("/projects/{PROJECT_ID}/queries")))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "jobComplete": true,
            "schema": schema_two_cols(),
            "jobReference": {"projectId": PROJECT_ID, "jobId": "job-stream"},
            "rows": rows(0, 13),
        })))
        .mount(&server)
        .await;

    // batch_size = 5 → 13 rows yield 3 pages of [5, 5, 3].
    let (src, _f) = build_source(&server, default_config().with_batch_size(5)).await;
    let ctx = HashMap::new();
    let mut sizes = Vec::new();
    let mut pages = src.stream_pages(&ctx, 0);
    while let Some(page) = pages.next().await {
        let page = page.unwrap();
        assert!(page.bookmark.is_none());
        sizes.push(page.records.len());
    }
    assert_eq!(sizes, vec![5, 5, 3]);
}

#[tokio::test]
async fn stream_pages_batch_size_zero_emits_one_page() {
    let server = MockServer::start().await;
    mount_token(&server).await;
    Mock::given(method("POST"))
        .and(path(format!("/projects/{PROJECT_ID}/queries")))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "jobComplete": true,
            "schema": schema_two_cols(),
            "jobReference": {"projectId": PROJECT_ID, "jobId": "job-flat"},
            "rows": rows(0, 7),
        })))
        .mount(&server)
        .await;

    let (src, _f) = build_source(&server, default_config().with_batch_size(0)).await;
    let ctx = HashMap::new();
    let pages: Vec<_> = src.stream_pages(&ctx, 0).collect().await;
    assert_eq!(pages.len(), 1);
    assert_eq!(pages[0].as_ref().unwrap().records.len(), 7);
}

#[tokio::test]
async fn polls_until_job_complete_when_first_response_is_incomplete() {
    let server = MockServer::start().await;
    mount_token(&server).await;

    // jobs.query returns jobComplete=false (statement still running).
    Mock::given(method("POST"))
        .and(path(format!("/projects/{PROJECT_ID}/queries")))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "jobComplete": false,
            "jobReference": {"projectId": PROJECT_ID, "jobId": "job-poll"},
        })))
        .mount(&server)
        .await;

    // First poll: still not complete.
    Mock::given(method("GET"))
        .and(path(format!("/projects/{PROJECT_ID}/queries/job-poll")))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "jobComplete": false,
        })))
        .up_to_n_times(1)
        .mount(&server)
        .await;
    // Subsequent poll: complete with rows.
    Mock::given(method("GET"))
        .and(path(format!("/projects/{PROJECT_ID}/queries/job-poll")))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "jobComplete": true,
            "schema": schema_two_cols(),
            "rows": rows(0, 2),
        })))
        .mount(&server)
        .await;

    let (src, _f) = build_source(&server, default_config()).await;
    let rows = src.fetch_all().await.unwrap();
    assert_eq!(rows.len(), 2);
    assert_eq!(rows[0]["id"], 0);
}

#[tokio::test]
async fn submits_positional_bindings_for_params_and_context() {
    let server = MockServer::start().await;
    mount_token(&server).await;
    Mock::given(method("POST"))
        .and(path(format!("/projects/{PROJECT_ID}/queries")))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "jobComplete": true,
            "schema": schema_two_cols(),
            "jobReference": {"projectId": PROJECT_ID, "jobId": "job-binds"},
            "rows": rows(0, 1),
        })))
        .mount(&server)
        .await;

    let cfg = default_config()
        .with_params(vec![json!("us-east")])
        .with_location("EU");
    let mut cfg = cfg;
    cfg.query = "SELECT * FROM events WHERE region = ? AND id > {parent.min_id}".into();
    let (src, _f) = build_source(&server, cfg).await;

    let mut ctx = HashMap::new();
    ctx.insert("parent.min_id".to_string(), json!(100));
    src.fetch_with_context(&ctx).await.unwrap();

    let reqs = server.received_requests().await.unwrap();
    let query_req = reqs
        .iter()
        .find(|r| r.url.path().ends_with("/queries"))
        .expect("captured jobs.query request");
    let body: Value = serde_json::from_slice(&query_req.body).unwrap();
    assert_eq!(
        body["query"],
        "SELECT * FROM events WHERE region = ? AND id > ?"
    );
    assert_eq!(body["parameterMode"], "POSITIONAL");
    let params = body["queryParameters"].as_array().unwrap();
    assert_eq!(params.len(), 2);
    assert_eq!(params[0]["parameterValue"]["value"], "us-east");
    assert_eq!(params[1]["parameterValue"]["value"], "100");
    assert_eq!(body["location"], "EU");
}