Function test_subscribe_fetch_timeout

src/environmentd/tests/sql.rs:939–1032

()

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_subscribe_fetch_timeout() {
    let server = test_util::TestHarness::default().start_blocking();
    let mut client = server.connect(postgres::NoTls).unwrap();

    client.batch_execute("CREATE TABLE t (i INT8)").unwrap();
    client
        .batch_execute("INSERT INTO t VALUES (1), (2), (3);")
        .unwrap();
    client
        .batch_execute(
            "BEGIN;
         DECLARE c CURSOR FOR SUBSCRIBE t;",
        )
        .unwrap();

    let expected: Vec<i64> = vec![1, 2, 3];
    let mut expected_iter = expected.iter();
    let mut next = expected_iter.next();

    // Test 0s timeouts.
    while let Some(expect) = next {
        let rows = client.query("FETCH c WITH (TIMEOUT = '0s')", &[]).unwrap();
        // It is fine for there to be no rows ready yet. Immediately try again because
        // they should be ready soon.
        if rows.len() != 1 {
            continue;
        }
        assert_eq!(rows[0].get::<_, i64>(2), *expect);

        next = expected_iter.next();
    }

    // Test a 1s timeout and make sure we waited for at least that long.
    let before = Instant::now();
    let rows = client.query("FETCH c WITH (TIMEOUT = '1s')", &[]).unwrap();
    let duration = before.elapsed();
    assert_eq!(rows.len(), 0);
    // Make sure we waited at least 1s but also not too long.
    assert!(duration >= Duration::from_secs(1));
    assert!(duration < Duration::from_secs(10));

    // Make a new cursor. Try to fetch more rows from it than exist. Verify that
    // we got all the rows we expect and also waited for at least the timeout
    // duration. Cursor may take a moment to be ready, so do it in a loop.
    client
        .batch_execute(
            "COMMIT; BEGIN;
        DECLARE c CURSOR FOR SUBSCRIBE t",
        )
        .unwrap();
    loop {
        let before = Instant::now();
        let rows = client
            .query("FETCH 4 c WITH (TIMEOUT = '1s')", &[])
            .unwrap();
        let duration = before.elapsed();
        if rows.len() != 0 {
            assert_eq!(rows.len(), expected.len());