Function test_subscribe_progress

src/environmentd/tests/sql.rs:685–797

()

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_subscribe_progress() {
    let server = test_util::TestHarness::default().start_blocking();

    for has_initial_data in [false, true] {
        for has_index in [false, true] {
            for has_snapshot in [false, true] {
                let mut client_writes = server.connect(postgres::NoTls).unwrap();
                let mut client_reads = server.connect(postgres::NoTls).unwrap();

                info!(
                    msg = "Running test",
                    has_initial_data, has_index, has_snapshot
                );

                client_writes
                    .batch_execute("CREATE TABLE t1 (data text)")
                    .unwrap();
                if has_index {
                    client_writes
                        .batch_execute("CREATE INDEX i1 on t1(data)")
                        .unwrap();
                }
                if has_initial_data {
                    client_writes
                        .batch_execute("INSERT INTO t1 VALUES ('snapdata')")
                        .unwrap();
                }
                client_reads
                    .batch_execute(&format!(
                        "COMMIT; BEGIN;
                DECLARE c1 CURSOR FOR SUBSCRIBE t1 WITH (PROGRESS, SNAPSHOT = {})",
                        has_snapshot
                    ))
                    .unwrap();

                // Asserts that the next data message is `data`. Ignores any progress
                // messages that occur first.
                //
                // Returns the timestamp at which that data arrived.
                fn await_data(
                    client: &mut postgres::Client,
                    last_seen_ts: &mut u64,
                    data: &str,
                ) -> u64 {
                    // We have to try several times. It might be that the FETCH gets a
                    // progress statements rather than data. We retry until we get the batch
                    // that has the data.
                    debug!("awaiting data: {data}");
                    loop {
                        let rows = client.query("FETCH 1 c1", &[]).unwrap();
                        debug!(row = ?rows.first());
                        let data_row = match rows.first() {
                            Some(row) if row.try_get::<_, String>("data").is_ok() => row,
                            _ => continue, // retry
                        };
                        assert_eq!(data_row.get::<_, bool>("mz_progressed"), false);
                        assert_eq!(data_row.get::<_, i64>("mz_diff"), 1);
                        assert_eq!(data_row.get::<_, String>("data"), data);