Function test_subscribe_fetch_wait

src/environmentd/tests/sql.rs:1036–1095

()

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_subscribe_fetch_wait() {
    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();

    while let Some(expect) = next {
        // FETCH with no timeout will wait for at least 1 result.
        let rows = client.query("FETCH c", &[]).unwrap();
        assert_eq!(rows.len(), 1);
        assert_eq!(rows[0].get::<_, i64>(2), *expect);
        next = expected_iter.next();
    }

    // Try again with FETCH ALL. ALL only guarantees that all available rows will
    // be returned, but it's up to the system to decide what is available. This
    // means that we could still get only one row per request, and we won't know
    // how many rows will come back otherwise.
    client
        .batch_execute(
            "COMMIT; BEGIN;
        DECLARE c CURSOR FOR SUBSCRIBE t;",
        )
        .unwrap();
    let mut expected_iter = expected.iter().peekable();
    while expected_iter.peek().is_some() {
        let rows = client.query("FETCH ALL c", &[]).unwrap();
        assert!(rows.len() > 0);
        for row in rows {
            let next = expected_iter.next().unwrap();
            assert_eq!(*next, row.get::<_, i64>(2));
        }
    }

    // Verify that the wait only happens for SUBSCRIBE. A SELECT with 0 rows should not
    // block.
    client.batch_execute("COMMIT").unwrap();
    client.batch_execute("CREATE TABLE empty ()").unwrap();
    client
        .batch_execute(
            "BEGIN;
         DECLARE c CURSOR FOR SELECT * FROM empty;",
        )
        .unwrap();
    let rows = client.query("FETCH c", &[]).unwrap();