Function test_github_12546

src/environmentd/tests/sql.rs:1588–1635

()

#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 2))]
#[allow(clippy::disallowed_methods)]
async fn test_github_12546() {
    let server = test_util::TestHarness::default()
        .with_propagate_crashes(false)
        .unsafe_mode()
        .start()
        .await;
    server
        .enable_feature_flags(&["unsafe_enable_unsafe_functions"])
        .await;

    let (client, conn_task) = server.connect().with_handle().await.unwrap();

    client
        .batch_execute("CREATE TABLE test(a text);")
        .await
        .unwrap();
    client
        .batch_execute("INSERT INTO test VALUES ('a');")
        .await
        .unwrap();

    let query = client.query("SELECT mz_unsafe.mz_panic(a) FROM test", &[]);
    let timeout = tokio::time::timeout(Duration::from_secs(2), query);
    // We expect the timeout to trigger because the query should be crashing the
    // compute instance.
    assert_eq!(
        timeout.await.unwrap_err().to_string(),
        "deadline has elapsed"
    );

    // Aborting the connection should cause its pending queries to be cancelled,
    // allowing the compute instances to stop crashing while trying to execute
    // them.
    //
    // We need to await `conn_task` to actually deliver the `abort`.
    conn_task.abort_and_wait().await;

    // Make a new connection to verify the compute instance can now start.
    let client = server.connect().await.unwrap();
    assert_eq!(
        client
            .query_one("SELECT count(*) FROM test", &[])
            .await
            .unwrap()
            .get::<_, i64>(0),
        1,
    );
}

/// Regression test for database-issues#3721.
#[mz_ore::test]