server.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! Integration tests for Materialize server.

#![recursion_limit = "256"]

use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::fmt::Write;
use std::io::Write as _;
use std::net::{IpAddr, Ipv4Addr};
use std::path::PathBuf;
use std::process::{Command, Stdio};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use std::{iter, thread};

use anyhow::bail;
use chrono::{DateTime, Utc};
use flate2::Compression;
use flate2::write::GzEncoder;
use futures::FutureExt;
use http::Request;
use itertools::Itertools;
use jsonwebtoken::{DecodingKey, EncodingKey};
use mz_auth::password::Password;
use mz_environmentd::test_util::{self, Ca, KAFKA_ADDRS, PostgresErrorExt, make_pg_tls};
use mz_environmentd::{WebSocketAuth, WebSocketResponse};
use mz_frontegg_auth::{
    Authenticator as FronteggAuthentication, AuthenticatorConfig as FronteggConfig,
    DEFAULT_REFRESH_DROP_FACTOR, DEFAULT_REFRESH_DROP_LRU_CACHE_SIZE,
};
use mz_frontegg_mock::{FronteggMockServer, models::ApiToken, models::UserConfig};
use mz_ore::cast::CastFrom;
use mz_ore::cast::CastLossy;
use mz_ore::cast::TryCastFrom;
use mz_ore::collections::CollectionExt;
use mz_ore::error::ErrorExt;
use mz_ore::metrics::MetricsRegistry;
use mz_ore::now::{NowFn, SYSTEM_TIME, to_datetime};
use mz_ore::retry::Retry;
use mz_ore::{assert_contains, task::RuntimeExt};
use mz_ore::{assert_err, assert_none, assert_ok, task};
use mz_pgrepr::UInt8;
use mz_repr::UNKNOWN_COLUMN_NAME;
use mz_sql::session::user::{ANALYTICS_USER, HTTP_DEFAULT_USER, SYSTEM_USER};
use mz_sql_parser::ast::display::AstDisplay;
use openssl::ssl::{SslConnectorBuilder, SslVerifyMode};
use openssl::x509::X509;
use postgres::config::SslMode;
use postgres_array::Array;
use rand::RngCore;
use rdkafka::ClientConfig;
use rdkafka::admin::{AdminClient, AdminOptions, NewTopic, TopicReplication};
use rdkafka_sys::RDKafkaErrorCode;
use regex::Regex;
use reqwest::blocking::Client;
use reqwest::header::{CONTENT_ENCODING, CONTENT_TYPE};
use reqwest::{StatusCode, Url};
use serde::{Deserialize, Serialize};
use tempfile::TempDir;
use tokio::sync::oneshot;
use tokio_postgres::error::SqlState;
use tracing::info;
use tungstenite::error::ProtocolError;
use tungstenite::{Error, Message, Utf8Bytes};
use uuid::Uuid;

// Allow the use of banned rdkafka methods, because we are just in tests.
#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_persistence() {
    let data_dir = tempfile::tempdir().unwrap();
    let harness = test_util::TestHarness::default()
        .data_directory(data_dir.path())
        .unsafe_mode();

    {
        let server = harness.clone().start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();

        server.runtime().block_on(async {
            let admin: AdminClient<_> = ClientConfig::new()
                .set("bootstrap.servers", &*KAFKA_ADDRS)
                .create()
                .expect("Admin client creation failed");
            let new_topic = NewTopic::new("foo", 1, TopicReplication::Fixed(1));
            let topic_results = admin
                .create_topics([&new_topic], &AdminOptions::new())
                .await
                .expect("topic creation failed");
            match topic_results[0] {
                Ok(_) | Err((_, RDKafkaErrorCode::TopicAlreadyExists)) => {}
                Err((ref err, _)) => panic!("failed to ensure topic: {err}"),
            }
        });

        client
            .batch_execute(&format!(
                "CREATE CONNECTION kafka_conn TO KAFKA (BROKER '{}', SECURITY PROTOCOL PLAINTEXT)",
                &*KAFKA_ADDRS,
            ))
            .unwrap();
        client
            .batch_execute(
                "CREATE SOURCE src FROM KAFKA CONNECTION kafka_conn (TOPIC 'foo') FORMAT BYTES",
            )
            .unwrap();
        client
            .batch_execute("CREATE VIEW constant AS SELECT 1")
            .unwrap();
        client.batch_execute(
            "CREATE VIEW mat (a, a_data, c, c_data) AS SELECT 'a', data, 'c' AS c, data FROM src",
        ).unwrap();
        client.batch_execute("CREATE DEFAULT INDEX ON mat").unwrap();
        client.batch_execute("CREATE DATABASE d").unwrap();
        client.batch_execute("CREATE SCHEMA d.s").unwrap();
        client
            .batch_execute("CREATE VIEW d.s.v AS SELECT 1")
            .unwrap();
    }

    let server = harness.start_blocking();
    let mut client = server.connect(postgres::NoTls).unwrap();
    assert_eq!(
        client
            .query("SHOW VIEWS", &[])
            .unwrap()
            .into_iter()
            .map(|row| row.get(0))
            .collect::<Vec<String>>(),
        &["constant", "mat"]
    );
    assert_eq!(
        client
            .query_one("SHOW INDEXES ON mat", &[])
            .unwrap()
            .get::<_, Vec<String>>("key"),
        &["a", "a_data", "c", "c_data"],
    );
    assert_eq!(
        client
            .query("SHOW VIEWS FROM d.s", &[])
            .unwrap()
            .into_iter()
            .map(|row| row.get(0))
            .collect::<Vec<String>>(),
        &["v"]
    );

    // Test that catalog recovery correctly populates `mz_objects`.
    assert_eq!(
        client
            .query(
                "SELECT id FROM mz_objects WHERE id LIKE 'u%' ORDER BY 1",
                &[]
            )
            .unwrap()
            .into_iter()
            .map(|row| row.get(0))
            .collect::<Vec<String>>(),
        vec!["u1", "u2", "u3", "u4", "u5", "u6", "u7"]
    );
}

/// A wrapper around `TestServerWithRuntime` that runs statement logging checks when dropped.
///
/// This guard ensures that all statements have finished executing (have non-NULL `finished_at`
/// and `finished_status` in `mz_internal.mz_recent_activity_log`) before the test completes.
struct TestServerWithStatementLoggingChecks {
    server: test_util::TestServerWithRuntime,
}

impl TestServerWithStatementLoggingChecks {
    /// Connect to the __internal__ SQL port of the running `environmentd` server.
    pub fn connect_internal<T>(&self, tls: T) -> Result<postgres::Client, anyhow::Error>
    where
        T: postgres::tls::MakeTlsConnect<postgres::Socket> + Send + 'static,
        T::TlsConnect: Send,
        T::Stream: Send,
        <T::TlsConnect as postgres::tls::TlsConnect<postgres::Socket>>::Future: Send,
    {
        self.server.connect_internal(tls)
    }

    /// Returns the metrics registry for the test server.
    pub fn metrics_registry(&self) -> &MetricsRegistry {
        self.server.metrics_registry()
    }
}

/// Helper to get statement logging record counts from the metrics registry.
/// Returns (sampled_true_count, sampled_false_count).
#[allow(clippy::disallowed_methods)]
fn get_statement_logging_record_counts(
    server: &TestServerWithStatementLoggingChecks,
) -> (u64, u64) {
    let metrics = server.metrics_registry().gather();
    let record_count_metric = metrics
        .into_iter()
        .find(|m| m.name() == "mz_statement_logging_record_count")
        .expect("mz_statement_logging_record_count metric should exist");

    let metric_entries = record_count_metric.get_metric();
    let sampled_true = metric_entries
        .iter()
        .find(|m| {
            m.get_label()
                .iter()
                .any(|l| l.name() == "sample" && l.value() == "true")
        })
        .map(|m| u64::cast_lossy(m.get_counter().value()))
        .unwrap_or(0);
    let sampled_false = metric_entries
        .iter()
        .find(|m| {
            m.get_label()
                .iter()
                .any(|l| l.name() == "sample" && l.value() == "false")
        })
        .map(|m| u64::cast_lossy(m.get_counter().value()))
        .unwrap_or(0);

    (sampled_true, sampled_false)
}

impl Drop for TestServerWithStatementLoggingChecks {
    #[allow(clippy::disallowed_methods)]
    fn drop(&mut self) {
        // Don't run checks if we're already panicking, as this could mask the original error.
        if std::thread::panicking() {
            return;
        }

        let mut mz_client = self
            .server
            .connect_internal(postgres::NoTls)
            .expect("Failed to connect to internal SQL port for statement logging check");

        // Disable RBAC checks so we can query mz_internal tables.
        // (We don't need to restore this afterwards, since no more tests run in the same system.)
        mz_client
            .batch_execute("ALTER SYSTEM SET enable_rbac_checks = false")
            .expect("Failed to disable RBAC checks");

        // The statement log has a 5-second buffer flush interval, so allow sufficient time.
        Retry::default()
            .max_duration(Duration::from_secs(30))
            .retry(|_| {
                let result = mz_client.query_one(
                    "SELECT count(*)
                     FROM mz_internal.mz_recent_activity_log
                     WHERE
                       (finished_at IS NULL OR finished_status IS NULL)
                       AND sql NOT LIKE '%__FILTER-OUT-THIS-QUERY__%'
                       AND finished_status != 'aborted'",
                    &[],
                );

                match result {
                    Ok(row) => {
                        let count: i64 = row.get(0);
                        if count == 0 {
                            Ok(())
                        } else {
                            Err(format!("{} statements have not finished", count))
                        }
                    }
                    Err(e) => Err(format!("Query failed: {}", e)),
                }
            })
            .expect("All statements should have finished executing");
    }
}

fn setup_statement_logging_core(
    max_sample_rate: f64,
    sample_rate: f64,
    target_data_rate: &str,
    test_harness: test_util::TestHarness,
) -> (TestServerWithStatementLoggingChecks, postgres::Client) {
    let server = test_harness
        .with_system_parameter_default(
            "statement_logging_max_sample_rate".to_string(),
            max_sample_rate.to_string(),
        )
        .with_system_parameter_default(
            "statement_logging_default_sample_rate".to_string(),
            sample_rate.to_string(),
        )
        .with_system_parameter_default(
            "statement_logging_max_data_credit".to_string(),
            "".to_string(),
        )
        .with_system_parameter_default(
            "statement_logging_target_data_rate".to_string(),
            target_data_rate.to_string(),
        )
        .with_system_parameter_default(
            "statement_logging_use_reproducible_rng".to_string(),
            "true".to_string(),
        )
        .start_blocking();
    let client = server.connect(postgres::NoTls).unwrap();
    let server = TestServerWithStatementLoggingChecks { server };
    (server, client)
}

fn setup_statement_logging(
    max_sample_rate: f64,
    sample_rate: f64,
    target_data_rate: &str,
) -> (TestServerWithStatementLoggingChecks, postgres::Client) {
    setup_statement_logging_core(
        max_sample_rate,
        sample_rate,
        target_data_rate,
        test_util::TestHarness::default(),
    )
}

// Test that we log various kinds of statement whose execution terminates in the coordinator.
#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_statement_logging_immediate() {
    let (server, mut client) = setup_statement_logging(1.0, 1.0, "");

    let mut mz_client = server.connect_internal(postgres::NoTls).unwrap();
    mz_client
        .batch_execute("ALTER SYSTEM SET enable_statement_lifecycle_logging = false")
        .unwrap();
    mz_client
        .batch_execute("ALTER SYSTEM SET statement_logging_max_sample_rate = 1")
        .unwrap();
    mz_client
        .batch_execute("ALTER SYSTEM SET statement_logging_default_sample_rate = 1")
        .unwrap();
    mz_client
        .batch_execute("ALTER SYSTEM SET enable_load_generator_counter = true")
        .unwrap();

    let successful_immediates: &[&str] = &[
        "CREATE VIEW v AS SELECT 1;",
        "CREATE DEFAULT INDEX i ON v;",
        "CREATE TABLE t (x bigint);",
        "INSERT INTO t VALUES (1), (2), (3)",
        "UPDATE t SET x=x+1",
        "DELETE FROM t;",
        "CREATE SECRET s AS 'hunter2';",
        "DROP SECRET s;",
        "",
        "CREATE SOURCE s FROM LOAD GENERATOR COUNTER",
        "PREPARE foo AS SELECT * FROM t",
        "EXECUTE foo",
        "BEGIN",
        "DECLARE c CURSOR FOR SELECT * FROM t",
        "FETCH FORWARD ALL FROM c",
        "COMMIT",
        "BEGIN",
        "ROLLBACK",
        "SET application_name='my_application'",
        "SHOW ALL",
        "SHOW application_name",
    ];
    let constants: &[&str] = &["1", "2", "3", "hunter2", "my_application"];

    for &statement in successful_immediates {
        client.execute(statement, &[]).unwrap();

        // Enforce a small delay to avoid duplicate `began_at` times, which would make the ordering
        // of logged statements non-deterministic when we retrieve them below.
        thread::sleep(Duration::from_millis(10));
    }

    let mut client = server.connect_internal(postgres::NoTls).unwrap();
    let seh_query = "
        SELECT
            mseh.sample_rate,
            mseh.began_at,
            mseh.finished_at,
            mseh.finished_status,
            mst.sql,
            mpsh.prepared_at,
            mst.redacted_sql
        FROM mz_internal.mz_statement_execution_history AS mseh
        LEFT JOIN
            mz_internal.mz_prepared_statement_history AS mpsh
            ON mseh.prepared_statement_id = mpsh.id
        JOIN
            (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst
            ON mpsh.sql_hash = mst.sql_hash
        WHERE
            mst.sql !~~ '%mz_statement_execution_history%' AND
            mseh.finished_at IS NOT NULL
        ORDER BY mseh.began_at";

    // Statement logging happens async, retry until we get the expected number of logged
    // statements.
    let mut sl = Vec::new();
    for _ in 0..10 {
        thread::sleep(Duration::from_secs(1));
        sl = client.query(seh_query, &[]).unwrap();
        if sl.len() >= successful_immediates.len() {
            break;
        }
    }

    assert_eq!(sl.len(), successful_immediates.len());

    #[derive(Debug)]
    struct Record {
        sample_rate: f64,
        began_at: DateTime<Utc>,
        finished_at: DateTime<Utc>,
        finished_status: String,
        sql: String,
        prepared_at: DateTime<Utc>,
        redacted_sql: String,
    }
    for (r, stmt) in std::iter::zip(sl.iter(), successful_immediates) {
        let r = Record {
            sample_rate: r.get(0),
            began_at: r.get(1),
            finished_at: r.get(2),
            finished_status: r.get(3),
            sql: r.get(4),
            prepared_at: r.get(5),
            redacted_sql: r.get(6),
        };
        assert_eq!(r.sample_rate, 1.0);

        let expected_sql = if r.sql.contains("SECRET")
            || r.sql.contains("INSERT")
            || r.sql.contains("UPDATE")
            || r.sql.contains("EXECUTE")
        {
            mz_sql::parse::parse(&r.sql)
                .unwrap()
                .into_element()
                .ast
                .to_ast_string_redacted()
        } else {
            stmt.chars().filter(|&ch| ch != ';').collect::<String>()
        };
        assert_eq!(r.sql, expected_sql);
        assert_eq!(r.finished_status, "success");
        assert!(r.prepared_at <= r.began_at);
        assert!(r.began_at <= r.finished_at);
        // NB[btv] -- It would be a bit nicer if we could separately mock
        // both the start and end time, but the `NowFn` mechanism doesn't
        // appear to give us any way to do that. Instead, let's just check
        // that none of these statements took longer than 5s wall-clock time.
        assert!(r.finished_at - r.began_at <= chrono::Duration::try_seconds(5).unwrap());
        if !r.sql.is_empty() {
            let expected_redacted = mz_sql::parse::parse(&r.sql)
                .unwrap()
                .into_element()
                .ast
                .to_ast_string_redacted();
            assert_eq!(r.redacted_sql, expected_redacted);
            for constant in constants {
                assert!(!r.redacted_sql.contains(constant));
            }
        }
    }
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_statement_logging_basic() {
    let (server, mut client) = setup_statement_logging(1.0, 1.0, "");
    client.execute("SELECT 1", &[]).unwrap();
    // We test that queries of this view execute on a cluster.
    // If we ever change the threshold for constant folding such that
    // this gets to run on environmentd, change this query.
    client
        .execute(
            "CREATE VIEW v AS SELECT * FROM generate_series(1, 10001)",
            &[],
        )
        .unwrap();
    client.execute("SELECT * FROM v", &[]).unwrap();
    client.execute("CREATE DEFAULT INDEX i ON v", &[]).unwrap();
    client.execute("SELECT * FROM v", &[]).unwrap();
    let _ = client.execute("SELECT 1/0", &[]);
    client.execute("CREATE TABLE t (x int)", &[]).unwrap();
    client.execute("SELECT * FROM t", &[]).unwrap();

    #[derive(Debug)]
    struct Record {
        sample_rate: f64,
        began_at: DateTime<Utc>,
        finished_at: DateTime<Utc>,
        finished_status: String,
        error_message: Option<String>,
        prepared_at: DateTime<Utc>,
        execution_strategy: Option<String>,
        result_size: Option<i64>,
        rows_returned: Option<i64>,
        execution_timestamp: Option<u64>,
    }

    let mut client = server.connect_internal(postgres::NoTls).unwrap();

    let result = Retry::default()
        .max_duration(Duration::from_secs(30))
        .retry(|_| {
            let sl_results = client
                .query(
                    "SELECT
    mseh.sample_rate,
    mseh.began_at,
    mseh.finished_at,
    mseh.finished_status,
    mseh.error_message,
    mpsh.prepared_at,
    mseh.execution_strategy,
    mseh.result_size,
    mseh.rows_returned,
    mseh.execution_timestamp
FROM
    mz_internal.mz_statement_execution_history AS mseh
        LEFT JOIN
            mz_internal.mz_prepared_statement_history AS mpsh
            ON mseh.prepared_statement_id = mpsh.id
        JOIN
            (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) AS mst
            ON mpsh.sql_hash = mst.sql_hash
WHERE (mst.sql ~~ 'SELECT%'
AND mst.sql !~~ '%unique string to prevent this query showing up in results after retries%'
AND mst.sql !~~ '%pg_catalog.pg_type%' --this gets executed behind the scenes by tokio-postgres
OR mst.sql ~~ 'CREATE TABLE%')
AND mseh.finished_at IS NOT NULL
ORDER BY mseh.began_at",
                    &[],
                )
                .unwrap();

            if sl_results.len() == 6 {
                Ok(sl_results)
            } else {
                Err(sl_results.len())
            }
        });
    let sl_results = match result {
        Ok(rows) => rows
            .into_iter()
            .map(|r| Record {
                sample_rate: r.get(0),
                began_at: r.get(1),
                finished_at: r.get(2),
                finished_status: r.get(3),
                error_message: r.get(4),
                prepared_at: r.get(5),
                execution_strategy: r.get(6),
                result_size: r.get(7),
                rows_returned: r.get(8),
                execution_timestamp: r.get::<_, Option<UInt8>>(9).map(|UInt8(val)| val),
            })
            .collect::<Vec<_>>(),
        Err(rows) => {
            panic!("number of results never became correct: {rows}");
        }
    };
    // The two queries on generate_series(1,10001) execute at the maximum timestamp
    assert_eq!(
        sl_results
            .iter()
            .filter(|r| r.execution_timestamp == Some(u64::MAX))
            .count(),
        2
    );
    // The two queries that can be satisfied by envd (SELECT 1 and SELECT 1/0) have no execution timestamp
    assert_eq!(
        sl_results
            .iter()
            .filter(|r| r.execution_timestamp.is_none())
            .count(),
        2
    );
    // All other queries have an execution timestamp, in particular, including `CREATE TABLE`.
    assert_eq!(sl_results.len(), 6);
    for r in &sl_results {
        assert_eq!(r.sample_rate, 1.0);
        assert!(r.prepared_at <= r.began_at);
        assert!(r.began_at <= r.finished_at);
        // It would be nice to be able to control
        // execution timestamp via a `NowFn`, but
        // that is hard to get right and interferes with our logic
        // about when to flush to persist. So instead, just check that they're sane.
        if let Some(ts) = r.execution_timestamp {
            if ts != u64::MAX {
                let ts = to_datetime(ts);
                assert!((ts - r.prepared_at).abs() < chrono::Duration::try_seconds(5).unwrap())
            }
        }
    }
    assert!(sl_results[0].result_size.unwrap_or(0) > 0);
    assert_eq!(sl_results[0].rows_returned, Some(1));
    assert_eq!(sl_results[0].finished_status, "success");
    assert_eq!(
        sl_results[0].execution_strategy.as_ref().unwrap(),
        "constant"
    );
    assert!(sl_results[1].result_size.unwrap_or(0) > 0);
    assert_eq!(sl_results[1].rows_returned, Some(10001));
    assert_eq!(sl_results[1].finished_status, "success");
    assert_eq!(
        sl_results[1].execution_strategy.as_ref().unwrap(),
        "standard"
    );
    assert!(sl_results[2].result_size.unwrap_or(0) > 0);
    assert_eq!(sl_results[2].rows_returned, Some(10001));
    assert_eq!(sl_results[2].finished_status, "success");
    assert_eq!(
        sl_results[2].execution_strategy.as_ref().unwrap(),
        "fast-path"
    );
    assert_eq!(sl_results[3].finished_status, "error");
    assert!(
        sl_results[3]
            .error_message
            .as_ref()
            .unwrap()
            .contains("division by zero")
    );
    assert_none!(sl_results[3].result_size);
    assert_none!(sl_results[3].rows_returned);

    // Verify metrics show all statements were sampled (100% sample rate means no unsampled).
    let (sampled_true, sampled_false) = get_statement_logging_record_counts(&server);
    assert!(
        sampled_true > 0,
        "some statements should be sampled with 100% rate"
    );
    assert_eq!(
        sampled_false, 0,
        "no statements should be unsampled with 100% rate"
    );

    // Verify statement_logging_actual_bytes metric is being tracked.
    // With 100% sample rate, actual_bytes should equal unsampled_bytes.
    let metrics = server.metrics_registry().gather();
    let actual_bytes = metrics
        .iter()
        .find(|m| m.name() == "mz_statement_logging_actual_bytes")
        .expect("mz_statement_logging_actual_bytes metric should exist")
        .get_metric()[0]
        .get_counter()
        .value();
    let unsampled_bytes = metrics
        .iter()
        .find(|m| m.name() == "mz_statement_logging_unsampled_bytes")
        .expect("mz_statement_logging_unsampled_bytes metric should exist")
        .get_metric()[0]
        .get_counter()
        .value();
    assert!(
        actual_bytes > 0.0,
        "actual_bytes should be > 0 with 100% sample rate"
    );
    assert_eq!(
        actual_bytes, unsampled_bytes,
        "with 100% sample rate, actual_bytes should equal unsampled_bytes"
    );
}

#[allow(clippy::disallowed_methods)]
fn run_throttling_test(use_prepared_statement: bool) {
    // The `target_data_rate` should be
    // - high enough so that the `SELECT 1` queries get throttled (even with high CPU load due to
    //   other tests running in parallel),
    // - but low enough that the `SELECT 2` query after the sleep doesn't get throttled.
    let (server, mut client) = setup_statement_logging(1.0, 1.0, "200");
    thread::sleep(Duration::from_secs(2));

    if use_prepared_statement {
        let statement = client.prepare("SELECT 1").unwrap();
        for _ in 0..100 {
            client.execute(&statement, &[]).unwrap();
        }
    } else {
        for _ in 0..100 {
            client.execute("SELECT 1", &[]).unwrap();
        }
    }

    thread::sleep(Duration::from_secs(4));
    client.execute("SELECT 2", &[]).unwrap();
    let mut client = server.connect_internal(postgres::NoTls).unwrap();
    let logs = Retry::default()
        .max_duration(Duration::from_secs(60))
        .retry(|_| {
            let sl_results = client
                .query(
                    "SELECT
    sql,
    throttled_count
FROM mz_internal.mz_statement_execution_history mseh
JOIN mz_internal.mz_prepared_statement_history mpsh
ON mseh.prepared_statement_id = mpsh.id
JOIN (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst
ON mpsh.sql_hash = mst.sql_hash
WHERE sql IN ('SELECT 1', 'SELECT 2')",
                    &[],
                )
                .unwrap();

            if sl_results.iter().any(|stmt| {
                let sql: String = stmt.get(0);
                sql == "SELECT 2"
            }) {
                Ok(sl_results)
            } else {
                Err(())
            }
        })
        .expect("Never saw last statement (`SELECT 2`)");
    let throttled_count = logs
        .iter()
        .map(|log| {
            let UInt8(throttled_count) = log.get(1);
            throttled_count
        })
        .sum::<u64>();
    assert!(
        throttled_count > 0,
        "at least some statements should have been throttled"
    );

    assert_eq!(logs.len() + usize::cast_from(throttled_count), 101);
}

#[mz_ore::test]
fn test_statement_logging_throttling() {
    run_throttling_test(false);
}

#[mz_ore::test]
fn test_statement_logging_prepared_statement_throttling() {
    run_throttling_test(true);
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_statement_logging_subscribes() {
    let (server, mut client) = setup_statement_logging(1.0, 1.0, "");
    let cancel_token = client.cancel_token();

    // This should finish
    client
        .execute(
            "SUBSCRIBE TO (SELECT * FROM generate_series(1, 10001))",
            &[],
        )
        .unwrap();

    let handle = thread::spawn(move || {
        client.execute("CREATE TABLE t (x int)", &[]).unwrap();
        // This should not finish until it's canceled.
        let _ = client.execute("SUBSCRIBE TO (SELECT * FROM t)", &[]);
    });

    while !handle.is_finished() {
        thread::sleep(Duration::from_secs(1));
        cancel_token.cancel_query(postgres::NoTls).unwrap();
    }
    handle.join().unwrap();

    let mut client = server.connect_internal(postgres::NoTls).unwrap();
    let seh_query = "
        SELECT
            mseh.sample_rate,
            mseh.began_at,
            mseh.finished_at,
            mseh.finished_status,
            mpsh.prepared_at,
            mseh.execution_strategy
        FROM mz_internal.mz_statement_execution_history AS mseh
        LEFT JOIN
            mz_internal.mz_prepared_statement_history AS mpsh
            ON mseh.prepared_statement_id = mpsh.id
        JOIN
            mz_internal.mz_sql_text AS mst
            ON mpsh.sql_hash = mst.sql_hash
        WHERE
            mst.sql ~~ 'SUBSCRIBE%' AND
            mseh.finished_at IS NOT NULL
        ORDER BY mseh.began_at";

    // Statement logging happens async, retry until we get the expected number of logged
    // statements.
    let mut sl = Vec::new();
    for _ in 0..10 {
        thread::sleep(Duration::from_secs(1));
        sl = client.query(seh_query, &[]).unwrap();
        if sl.len() >= 2 {
            break;
        }
    }

    assert_eq!(sl.len(), 2);

    struct Record {
        sample_rate: f64,
        began_at: DateTime<Utc>,
        finished_at: DateTime<Utc>,
        finished_status: String,
        prepared_at: DateTime<Utc>,
        execution_strategy: Option<String>,
    }

    let sl_subscribes = sl
        .into_iter()
        .map(|r| Record {
            sample_rate: r.get(0),
            began_at: r.get(1),
            finished_at: r.get(2),
            finished_status: r.get(3),
            prepared_at: r.get(4),
            execution_strategy: r.get(5),
        })
        .collect::<Vec<_>>();
    for r in &sl_subscribes {
        assert_eq!(r.sample_rate, 1.0);
        assert!(r.prepared_at <= r.began_at);
        assert!(r.began_at <= r.finished_at);
        assert_none!(r.execution_strategy);
    }
    assert_eq!(sl_subscribes[0].finished_status, "success");
    assert_eq!(sl_subscribes[1].finished_status, "canceled");
}

/// Test that we are sampling approximately 50% of statements.
/// Relies on two assumptions:
/// (1) that the effective sampling rate for the session is 50%,
/// (2) that we are using the deterministic testing RNG.
#[allow(clippy::disallowed_methods)]
fn test_statement_logging_sampling_inner(
    server: TestServerWithStatementLoggingChecks,
    mut client: postgres::Client,
) {
    for i in 0..50 {
        client.execute(&format!("SELECT {i}"), &[]).unwrap();

        // Enforce a small delay to avoid duplicate `began_at` times, which would make the ordering
        // of logged statements non-deterministic when we retrieve them below.
        thread::sleep(Duration::from_millis(10));
    }

    // 23 randomly sampled out of 50 with 50% sampling. Seems legit!
    let expected_sqls = [
        2, 4, 5, 6, 9, 15, 17, 18, 19, 20, 21, 23, 24, 25, 31, 32, 33, 36, 37, 42, 46,
    ]
    .into_iter()
    .map(|i| format!("SELECT {i}"))
    .collect::<Vec<_>>();

    let mut internal_client = server.connect_internal(postgres::NoTls).unwrap();
    let seh_query = "
        SELECT mst.sql
        FROM mz_internal.mz_statement_execution_history AS mseh
        JOIN
            mz_internal.mz_prepared_statement_history AS mpsh
            ON mseh.prepared_statement_id = mpsh.id
        JOIN
            mz_internal.mz_sql_text AS mst
            ON mpsh.sql_hash = mst.sql_hash
        WHERE mst.sql ~~ 'SELECT%' AND mst.sql !~~ '%mz_statement_execution_history%'
        ORDER BY mseh.began_at ASC";

    // Statement logging happens async, retry until we get the expected number of logged
    // statements.
    let mut sl = Vec::new();
    for _ in 0..10 {
        thread::sleep(Duration::from_secs(1));
        sl = internal_client.query(seh_query, &[]).unwrap();
        if sl.len() >= expected_sqls.len() {
            break;
        }
    }

    let sqls: Vec<String> = sl.into_iter().map(|r| r.get(0)).collect();
    assert_eq!(sqls, expected_sqls);

    // Verify the statement_logging_record_count metric correctly tracks sampled vs unsampled.
    // With 50% sampling and deterministic RNG, exactly 21 of 50 statements should be sampled.
    let (sampled_true, sampled_false) = get_statement_logging_record_counts(&server);
    assert_eq!(
        sampled_true, 21,
        "expected 21 statements to be sampled with 50% rate and deterministic RNG"
    );
    assert_eq!(
        sampled_false, 29,
        "expected 29 statements to not be sampled with 50% rate and deterministic RNG"
    );
}

#[mz_ore::test]
fn test_statement_logging_sampling() {
    let (server, client) = setup_statement_logging(1.0, 0.5, "");
    test_statement_logging_sampling_inner(server, client);
}

/// Test that we are not allowed to set `statement_logging_sample_rate`
/// arbitrarily high, but that it is constrained by `statement_logging_max_sample_rate`.
#[mz_ore::test]
fn test_statement_logging_sampling_constrained() {
    let (server, client) = setup_statement_logging(0.5, 1.0, "");
    test_statement_logging_sampling_inner(server, client);
}

/// Test that the `mz_statement_logging_unsampled_bytes` metric tracks the total bytes
/// of SQL text that would have been logged if statement logging were fully enabled.
/// We set `sample_rate=0.0` so no statements are actually sampled/logged, but the
/// unsampled_bytes metric still gets incremented for every executed statement.
#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_statement_logging_unsampled_metrics() {
    // Use sample_rate=0.0 so statements are not sampled, but unsampled_bytes metric is still tracked.
    let (server, mut client) = setup_statement_logging(1.0, 0.0, "");

    let batch_queries = [
        "SELECT 'Hello, world!';SELECT 1;;",
        "SELECT 'Hello, world again!'",
    ];
    let batch_total: usize = batch_queries
        .iter()
        .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count())
        .sum();
    let single_queries = ["SELECT 'foo'", "SELECT 'bar';;;"];
    let single_total: usize = single_queries
        .iter()
        .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count())
        .sum();
    let prepared_queries = ["SELECT 'baz';;;", "SELECT 'quux';"];
    let prepared_total: usize = prepared_queries
        .iter()
        .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count())
        .sum();

    let named_prepared_inner = "SELECT 42";
    let named_prepared_outer = format!("PREPARE p AS {named_prepared_inner};EXECUTE p;");
    let named_prepared_outer_len = named_prepared_outer
        .as_bytes()
        .iter()
        .filter(|&&ch| ch != b';')
        .count();

    for q in batch_queries {
        client.batch_execute(q).unwrap();
    }

    for q in single_queries {
        client.execute(q, &[]).unwrap();
    }

    for q in prepared_queries {
        let s = client.prepare(q).unwrap();
        client.execute(&s, &[]).unwrap();
    }

    client.batch_execute(&named_prepared_outer).unwrap();

    // This should NOT be logged, since we never actually execute it.
    client.prepare("SELECT 'Hello, not counted!'").unwrap();

    let expected_total = batch_total + single_total + prepared_total + named_prepared_outer_len;
    let metric_value = server
        .metrics_registry()
        .gather()
        .into_iter()
        .find(|m| m.name() == "mz_statement_logging_unsampled_bytes")
        .unwrap()
        .take_metric()[0]
        .get_counter()
        .value();
    let metric_value = usize::cast_from(u64::try_cast_from(metric_value).unwrap());
    assert_eq!(expected_total, metric_value);

    // Also verify that statement_logging_record_count shows all statements as not sampled
    // (since we're using 0% sample rate).
    let (sampled_true, _sampled_false) = get_statement_logging_record_counts(&server);
    assert_eq!(
        sampled_true, 0,
        "no statements should be sampled with 0% sample rate"
    );
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_enable_internal_statement_logging() {
    let (server, mut client) = setup_statement_logging_core(
        1.0,
        1.0,
        "",
        test_util::TestHarness::default().with_system_parameter_default(
            "enable_internal_statement_logging".to_string(),
            "true".to_string(),
        ),
    );

    client.execute("SELECT 1", &[]).unwrap();

    let mut client = server.connect_internal(postgres::NoTls).unwrap();
    let num_mz_system_statements = Retry::default()
        .max_duration(Duration::from_secs(30))
        .retry(|_| {
            let sl_results = client
                .query(
                    "SELECT
    count(*)
FROM mz_internal.mz_prepared_statement_history mpsh
JOIN mz_internal.mz_session_history USING (session_id)
WHERE authenticated_user='mz_system'",
                    &[],
                )
                .unwrap();

            let count: i64 = sl_results[0].get(0);

            if count > 0 { Ok(count) } else { Err(()) }
        })
        .expect("at least some statements from mz_system should have been logged");

    assert!(
        num_mz_system_statements > 0,
        "statements executed by mz_system should have been logged"
    );
}

// Test the POST and WS server endpoints.
#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_http_sql() {
    // Datadriven directives for WebSocket are "ws-text" and "ws-binary" to send
    // text or binary websocket messages that are the input. Output is
    // everything until and including the next ReadyForQuery message. An
    // optional "rows=N" argument can be given in the directive to produce
    // datadriven output after N rows. Any directive with rows=N should be the
    // final directive in a file, since it leaves the websocket in a
    // mid-statement state.
    // A "fixtimestamp=true" argument can be given to replace timestamps with "<TIMESTAMP>".
    // A "fixid=true" argument can be given to replace IDs with "<ID>".
    //
    // Datadriven directive for HTTP POST is "http". Input and output are the
    // documented JSON formats.

    let fixtimestamp_replacements = [(Regex::new(r#"\d{13}(\.0)?"#).unwrap(), "<TIMESTAMP>")];
    let fixid_replacements = [
        (
            Regex::new(r#"\\"(User|System|Transient)\\": \d+"#).unwrap(),
            "<ID>",
        ),
        (Regex::new(r#"\b[ust]\d+\b"#).unwrap(), "<ID>"),
        (Regex::new(r#"\\n[ust]\d+\b"#).unwrap(), "\\n<ID>"),
    ];

    datadriven::walk("tests/testdata/http", |f| {
        let server = test_util::TestHarness::default().start_blocking();

        // Grant all privileges to default http user.
        // TODO(jkosh44) The HTTP endpoint has a special default user while the WS endpoint just
        //  uses the "materialize" role. This is probably wrong. They should either both user
        //  materialize, both use the same special default user, each have their own special
        //  default user.
        {
            let mut super_user = server
                .pg_config_internal()
                .user(&SYSTEM_USER.name)
                .connect(postgres::NoTls)
                .unwrap();
            super_user
                .batch_execute(&format!("CREATE ROLE {}", &HTTP_DEFAULT_USER.name))
                .unwrap();
            super_user
                .batch_execute(&format!(
                    "GRANT ALL PRIVILEGES ON SYSTEM TO {}",
                    &HTTP_DEFAULT_USER.name
                ))
                .unwrap();
            super_user
                .batch_execute(&format!(
                    "GRANT ALL PRIVILEGES ON CLUSTER quickstart TO {}",
                    &HTTP_DEFAULT_USER.name
                ))
                .unwrap();
            super_user
                .batch_execute(&format!(
                    "GRANT ALL PRIVILEGES ON DATABASE materialize TO {}",
                    &HTTP_DEFAULT_USER.name
                ))
                .unwrap();
            super_user
                .batch_execute(&format!(
                    "GRANT ALL PRIVILEGES ON SCHEMA materialize.public TO {}",
                    &HTTP_DEFAULT_USER.name
                ))
                .unwrap();
        }

        let ws_url = server.ws_addr();
        let http_url = Url::parse(&format!("http://{}/api/sql", server.http_local_addr())).unwrap();
        let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap();
        let ws_init = test_util::auth_with_ws(&mut ws, BTreeMap::default()).unwrap();

        // Verify ws_init contains roughly what we expect. This varies (rng secret and version
        // numbers), so easier to test here instead of in the ws file.
        assert!(
            ws_init
                .iter()
                .filter(|m| matches!(m, WebSocketResponse::ParameterStatus(_)))
                .count()
                > 1
        );
        assert!(matches!(
            ws_init.last(),
            Some(WebSocketResponse::BackendKeyData(_))
        ));

        f.run(|tc| {
            let msg = match tc.directive.as_str() {
                "ws-text" => Message::Text(tc.input.clone().into()),
                "ws-binary" => Message::Binary(tc.input.as_bytes().to_vec().into()),
                "http" => {
                    let json: serde_json::Value = serde_json::from_str(&tc.input).unwrap();
                    let res = Client::new()
                        .post(http_url.clone())
                        .json(&json)
                        .send()
                        .unwrap();
                    return format!("{}\n{}\n", res.status(), res.text().unwrap());
                }
                _ => panic!("unknown directive {}", tc.directive),
            };
            let mut rows = tc
                .args
                .get("rows")
                .map(|rows| rows.get(0).map(|row| row.parse::<usize>().unwrap()))
                .flatten();

            let mut replacements = Vec::new();
            if tc.args.contains_key("fixtimestamp") {
                replacements.extend_from_slice(&fixtimestamp_replacements);
            }
            if tc.args.contains_key("fixid") {
                replacements.extend_from_slice(&fixid_replacements);
            }

            ws.send(msg).unwrap();
            let mut responses = String::new();
            loop {
                let resp = ws.read().unwrap();
                match resp {
                    Message::Text(mut msg) => {
                        for (re, replace) in &replacements {
                            msg = Utf8Bytes::from(re.replace_all(&msg, *replace).into_owned());
                        }
                        let msg: WebSocketResponse = serde_json::from_str(&msg).unwrap();
                        write!(&mut responses, "{}\n", serde_json::to_string(&msg).unwrap())
                            .unwrap();
                        match msg {
                            WebSocketResponse::ReadyForQuery(_) => return responses,
                            WebSocketResponse::Row(_) => {
                                if let Some(rows) = rows.as_mut() {
                                    *rows -= 1;
                                    if *rows == 0 {
                                        return responses;
                                    }
                                }
                            }
                            _ => {}
                        }
                    }
                    Message::Ping(_) => continue,
                    _ => panic!("unexpected response: {:?}", resp),
                }
            }
        });
    });
}

// Test that the server properly handles cancellation requests.
#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_cancel_long_running_query() {
    let server = test_util::TestHarness::default()
        .unsafe_mode()
        .start_blocking();

    let mut client = server.connect(postgres::NoTls).unwrap();
    let cancel_token = client.cancel_token();

    client.batch_execute("CREATE TABLE t (i INT)").unwrap();
    let (shutdown_tx, shutdown_rx) = std::sync::mpsc::channel();

    let handle = thread::spawn(move || {
        // Repeatedly attempt to abort the query because we're not sure exactly
        // when the SELECT will arrive.
        loop {
            thread::sleep(Duration::from_secs(1));
            match shutdown_rx.try_recv() {
                Ok(()) => return,
                Err(std::sync::mpsc::TryRecvError::Empty) => {
                    let _ = cancel_token.cancel_query(postgres::NoTls);
                }
                _ => panic!("unexpected"),
            }
        }
    });

    match client.simple_query("SELECT * FROM t AS OF 18446744073709551615") {
        Err(e) if e.code() == Some(&postgres::error::SqlState::QUERY_CANCELED) => {}
        Err(e) => panic!("expected error SqlState::QUERY_CANCELED, but got {:?}", e),
        Ok(_) => panic!("expected error SqlState::QUERY_CANCELED, but query succeeded"),
    }

    // Wait for the cancellation thread to stop.
    shutdown_tx.send(()).unwrap();
    handle.join().unwrap();

    client
        .simple_query("SELECT 1")
        .expect("simple query succeeds after cancellation");
}

#[allow(clippy::disallowed_methods)]
fn test_cancellation_cancels_dataflows(query: &str) {
    // Query that returns how many dataflows are currently installed.
    // Ignores introspection subscribe dataflows.
    // Ignores storage operators, whose IDs are offset by STORAGE_ID_OFFSET (1 << 48).
    const DATAFLOW_QUERY: &str = " \
        SELECT count(*) \
        FROM mz_introspection.mz_dataflows \
        WHERE name NOT LIKE '%introspection-subscribe%' \
        AND id < 281474976710656";

    let server = test_util::TestHarness::default()
        .unsafe_mode()
        .start_blocking();

    let mut client1 = server.connect(postgres::NoTls).unwrap();
    let mut client2 = server.connect(postgres::NoTls).unwrap();
    let cancel_token = client1.cancel_token();

    client1.batch_execute("CREATE TABLE t (i INT)").unwrap();
    // No dataflows expected at startup.
    assert_eq!(
        client1
            .query_one(DATAFLOW_QUERY, &[])
            .unwrap()
            .get::<_, i64>(0),
        0
    );

    thread::spawn(move || {
        // Wait until we see the expected dataflow.
        Retry::default()
            .retry(|_state| {
                let count: i64 = client2
                    .query_one(DATAFLOW_QUERY, &[])
                    .map_err(|_| ())
                    .unwrap()
                    .get(0);
                if count == 0 { Err(()) } else { Ok(()) }
            })
            .unwrap();
        cancel_token.cancel_query(postgres::NoTls).unwrap();
    });

    match client1.simple_query(query) {
        Err(e) if e.code() == Some(&postgres::error::SqlState::QUERY_CANCELED) => {}
        Err(e) => panic!("expected error SqlState::QUERY_CANCELED, but got {:?}", e),
        Ok(_) => panic!("expected error SqlState::QUERY_CANCELED, but query succeeded"),
    }
    // Expect the dataflows to shut down.
    Retry::default()
        .retry(|_state| {
            let count: i64 = client1
                .query_one(DATAFLOW_QUERY, &[])
                .map_err(|_| ())
                .unwrap()
                .get(0);
            if count == 0 { Ok(()) } else { Err(()) }
        })
        .unwrap();
}

// Test that dataflow uninstalls cancelled peeks.
#[mz_ore::test]
fn test_cancel_dataflow_removal() {
    test_cancellation_cancels_dataflows("SELECT * FROM t AS OF 9223372036854775807");
}

#[mz_ore::test]
fn test_cancel_long_select() {
    test_cancellation_cancels_dataflows(
        "WITH MUTUALLY RECURSIVE flip(x INTEGER) AS (VALUES(1) EXCEPT ALL SELECT * FROM flip) SELECT * FROM flip;",
    );
}

#[mz_ore::test]
fn test_cancel_insert_select() {
    test_cancellation_cancels_dataflows(
        "INSERT INTO t WITH MUTUALLY RECURSIVE flip(x INTEGER) AS (VALUES(1) EXCEPT ALL SELECT * FROM flip) SELECT * FROM flip;",
    );
}

#[allow(clippy::disallowed_methods)]
fn test_closing_connection_cancels_dataflows(query: String) {
    // Query that returns how many dataflows are currently installed.
    // Ignores introspection subscribe dataflows.
    // Ignores storage operators, whose IDs are offset by STORAGE_ID_OFFSET (1 << 48).
    const DATAFLOW_QUERY: &str = " \
        SELECT count(*) \
        FROM mz_introspection.mz_dataflows \
        WHERE name NOT LIKE '%introspection-subscribe%' \
        AND id < 281474976710656";

    let server = test_util::TestHarness::default()
        .unsafe_mode()
        .start_blocking();

    let mut cmd = Command::new("psql");
    let cmd = cmd
        .args([
            // Ignore .psqlrc so that local execution of testdrive isn't
            // affected by it.
            "--no-psqlrc",
            &format!(
                "postgres://{}:{}/materialize",
                Ipv4Addr::LOCALHOST,
                server.sql_local_addr().port()
            ),
        ])
        .stdin(Stdio::piped());
    tracing::info!("spawning: {cmd:#?}");
    let mut child = cmd.spawn().expect("failed to spawn psql");
    let mut stdin = child.stdin.take().expect("failed to open stdin");
    thread::spawn(move || {
        use std::io::Write;
        stdin
            .write_all("SET STATEMENT_TIMEOUT = \"120s\";".as_bytes())
            .unwrap();
        stdin.write_all(query.as_bytes()).unwrap();
    });

    let spawned_psql = Instant::now();

    let mut client = server.connect(postgres::NoTls).unwrap();

    // Wait until we see the expected dataflow.
    Retry::default()
        .retry(|_state| {
            if spawned_psql.elapsed() > Duration::from_secs(30) {
                panic!("waited too long for psql to send the query");
            }
            let count: i64 = client
                .query_one(DATAFLOW_QUERY, &[])
                .map_err(|_| ())
                .unwrap()
                .get(0);
            if count == 0 { Err(()) } else { Ok(()) }
        })
        .unwrap();

    let started = Instant::now();
    if let Some(wait_status) = child.try_wait().expect("wait shouldn't error") {
        panic!("child should still be running, it exitted with {wait_status}");
    }
    child.kill().expect("killing psql child");

    // Expect the dataflows to shut down.
    Retry::default()
        .retry(|_state| {
            if started.elapsed() > Duration::from_secs(30) {
                // this has to be less than statement timeout
                panic!("waited too long for dataflow cancellation");
            }
            let count: i64 = client
                .query_one(DATAFLOW_QUERY, &[])
                .map_err(|_| ())
                .unwrap()
                .get(0);
            if count == 0 { Ok(()) } else { Err(()) }
        })
        .unwrap();
    info!(
        "Took {:#?} until dataflows were cancelled",
        started.elapsed()
    );
}

#[mz_ore::test]
fn test_closing_connection_for_long_select() {
    test_closing_connection_cancels_dataflows("WITH MUTUALLY RECURSIVE flip(x INTEGER) AS (VALUES(1) EXCEPT ALL SELECT * FROM flip) SELECT * FROM flip;".to_string())
}

#[mz_ore::test]
fn test_closing_connection_for_insert_select() {
    test_closing_connection_cancels_dataflows("CREATE TABLE t1 (a int); INSERT INTO t1 WITH MUTUALLY RECURSIVE flip(x INTEGER) AS (VALUES(1) EXCEPT ALL SELECT * FROM flip) SELECT * FROM flip;".to_string())
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_storage_usage_collection_interval() {
    /// Waits for the next storage collection to occur, then returns the
    /// timestamp at which the collection occurred. The timestamp of the last
    /// collection must be provided
    fn wait_for_next_collection(
        client: &mut postgres::Client,
        last_timestamp: DateTime<Utc>,
    ) -> DateTime<Utc> {
        info!("waiting for next storage usage collection");
        let ts = Retry::default()
            .max_duration(Duration::from_secs(10))
            .retry(|_| {
                let row = client.query_one(
                    "SELECT max(collection_timestamp)
                    FROM mz_internal.mz_storage_usage_by_shard",
                    &[],
                )?;
                // mz_storage_usage_by_shard may not be populated yet, which would result in a NULL ts.
                let ts = row.try_get::<_, DateTime<Utc>>("max")?;
                if ts <= last_timestamp {
                    bail!("next collection has not yet occurred")
                }
                Ok(ts)
            })
            .unwrap();
        info!(%ts, "detected storage usage collection");
        ts
    }

    fn get_shard_id(client: &mut postgres::Client, name: &str) -> String {
        let row = Retry::default()
            .max_duration(Duration::from_secs(10))
            .retry(|_| {
                client.query_one(
                    "SELECT shard_id
                     FROM mz_internal.mz_storage_shards s
                     JOIN mz_objects o ON o.id = s.object_id
                     WHERE o.name = $1",
                    &[&name],
                )
            })
            .unwrap();
        row.get("shard_id")
    }

    fn get_storage_usage(
        client: &mut postgres::Client,
        shard_id: &str,
        collection_timestamp: DateTime<Utc>,
    ) -> u64 {
        let row = Retry::default()
            .max_duration(Duration::from_secs(10))
            .retry(|_| {
                client.query_one(
                    "SELECT coalesce(sum(size_bytes), 0)::uint8 AS size
                     FROM mz_internal.mz_storage_usage_by_shard
                     WHERE shard_id = $1 AND collection_timestamp = $2",
                    &[&shard_id, &collection_timestamp],
                )
            })
            .unwrap();
        row.get::<_, UInt8>("size").0
    }

    let server = test_util::TestHarness::default()
        .with_storage_usage_collection_interval(Duration::from_secs(1))
        .start_blocking();
    let mut client = server.connect(postgres::NoTls).unwrap();

    // Wait for the initial storage usage collection to occur.
    let mut timestamp = wait_for_next_collection(&mut client, DateTime::<Utc>::MIN_UTC);

    // Create a table with no data.
    client
        .batch_execute("CREATE TABLE usage_test (a int)")
        .unwrap();
    let shard_id = get_shard_id(&mut client, "usage_test");
    info!(%shard_id, "created table");

    // Test that the storage usage for the table was zero before it was
    // created.
    let pre_create_storage_usage = get_storage_usage(&mut client, &shard_id, timestamp);
    info!(%pre_create_storage_usage);
    assert_eq!(pre_create_storage_usage, 0);

    // Test that the storage usage for the table is nonzero after it is created
    // (there is some overhead even for empty tables). We wait multiple storage
    // collection intervals (here and below) because the next storage collection
    // may have been concurrent with the previous operation.
    let post_create_storage_usage = Retry::default().max_tries(10).retry(|_| {
        timestamp = wait_for_next_collection(&mut client, timestamp);
        let post_create_storage_usage = get_storage_usage(&mut client, &shard_id, timestamp);
        info!(%post_create_storage_usage);
        if post_create_storage_usage > 0 {
            Ok(post_create_storage_usage)
        } else {
            Err(format!("expected non-zero post create storage usage, found: {post_create_storage_usage}"))
        }
    }).unwrap();

    // Insert some data into the table.
    for _ in 0..3 {
        client
            .batch_execute("INSERT INTO usage_test VALUES (1)")
            .unwrap();
    }

    // Test that the storage usage for the table is larger than it was before.
    Retry::default().max_tries(10).retry(|_| {
        timestamp = wait_for_next_collection(&mut client, timestamp);
        let after_insert_storage_usage = get_storage_usage(&mut client, &shard_id, timestamp);
        info!(%after_insert_storage_usage);
        if after_insert_storage_usage > post_create_storage_usage {
            Ok(())
        } else {
            Err(format!("expected insert storage usage, {after_insert_storage_usage}, to be higher than pre insert storage usage of {post_create_storage_usage}"))
        }
    }).unwrap();

    // Drop the table.
    client.batch_execute("DROP TABLE usage_test").unwrap();

    // Test that the storage usage is reported as zero.
    Retry::default()
        .max_tries(10)
        .retry(|_| {
            timestamp = wait_for_next_collection(&mut client, timestamp);
            let after_drop_storage_usage = get_storage_usage(&mut client, &shard_id, timestamp);
            info!(%after_drop_storage_usage);
            if after_drop_storage_usage == 0 {
                Ok(())
            } else {
                Err(format!(
                    "expected zero storage usage after drop, found {after_drop_storage_usage}"
                ))
            }
        })
        .unwrap();
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_storage_usage_updates_between_restarts() {
    let data_dir = tempfile::tempdir().unwrap();
    let storage_usage_collection_interval = Duration::from_secs(3);
    let harness = test_util::TestHarness::default()
        .with_storage_usage_collection_interval(storage_usage_collection_interval)
        .data_directory(data_dir.path());

    // Wait for initial storage usage collection.
    let initial_timestamp: f64 = {
        let server = harness.clone().start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();
        // Retry because it may take some time for the initial snapshot to be taken.
        Retry::default().max_duration(Duration::from_secs(60)).retry(|_| {
            client
                    .query_one(
                        "SELECT EXTRACT(EPOCH FROM MAX(collection_timestamp))::float8 FROM mz_catalog.mz_storage_usage;",
                        &[],
                    )
                    .map_err(|e| e.to_string()).unwrap()
                    .try_get::<_, f64>(0)
                    .map_err(|e| e.to_string())
        }).unwrap()
    };

    std::thread::sleep(storage_usage_collection_interval);

    // Another storage usage collection should be scheduled immediately.
    {
        let server = harness.start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();

        // Retry until storage usage is updated.
        Retry::default().max_duration(Duration::from_secs(60)).retry(|_| {
            let updated_timestamp = client
                .query_one("SELECT EXTRACT(EPOCH FROM MAX(collection_timestamp))::float8 FROM mz_catalog.mz_storage_usage;", &[])
                .map_err(|e| e.to_string()).unwrap()
                .try_get::<_, f64>(0)
                .map_err(|e| e.to_string()).unwrap();

            if updated_timestamp > initial_timestamp {
                Ok(())
            } else {
                Err(format!("updated storage collection timestamp {updated_timestamp} is not greater than initial timestamp {initial_timestamp}"))
            }
        }).unwrap();
    }
}

// Test that all rows for a single collection use the same timestamp.
#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_storage_usage_collection_interval_timestamps() {
    let storage_interval_s = 2;
    let server = test_util::TestHarness::default()
        .with_storage_usage_collection_interval(Duration::from_secs(storage_interval_s))
        .start_blocking();
    let mut client = server.connect(postgres::NoTls).unwrap();

    // Retry because it may take some time for the initial snapshot to be taken.
    let rows = Retry::default()
        .max_duration(Duration::from_secs(10))
        .retry(|_| {
            let rows = client
                .query(
                    "SELECT EXTRACT(EPOCH FROM collection_timestamp), ARRAY_AGG(object_id) \
                FROM mz_catalog.mz_storage_usage \
                GROUP BY collection_timestamp \
                ORDER BY collection_timestamp ASC;",
                    &[],
                )
                .map_err(|e| e.to_string())
                .unwrap();

            if rows.is_empty() {
                Err("expected some timestamp, instead found None".to_string())
            } else {
                Ok(rows)
            }
        })
        .unwrap();

    // The timestamp is selected after the storage usage is collected, so depending on how long
    // each collection takes, the timestamps can be arbitrarily close together or far apart. If
    // there are multiple timestamps, we ensure that they each contain the full set of object
    // IDs.
    let object_ids_by_timestamp: Vec<BTreeSet<_>> = rows
        .into_iter()
        .map(|row| row.get::<_, Array<String>>(1))
        .map(|array| array.into_iter().collect())
        .collect();
    let mut prev = None;
    for object_ids in object_ids_by_timestamp {
        match prev {
            None => {
                prev = Some(object_ids);
            }
            Some(ref prev_object_ids) => {
                assert_eq!(
                    prev_object_ids, &object_ids,
                    "found different storage collection timestamps, that contained non-matching \
                        object IDs. {prev_object_ids:?} and {object_ids:?}"
                );
            }
        }
    }
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_old_storage_usage_records_are_reaped_on_restart() {
    let now = Arc::new(Mutex::new(0));
    let now_fn = {
        let timestamp = Arc::clone(&now);
        NowFn::from(move || *timestamp.lock().expect("lock poisoned"))
    };
    let data_dir = tempfile::tempdir().unwrap();
    let collection_interval = Duration::from_millis(100);
    let retention_period = Duration::from_millis(200);
    let harness = test_util::TestHarness::default()
        .with_now(now_fn)
        .with_storage_usage_collection_interval(collection_interval)
        .with_storage_usage_retention_period(retention_period)
        .data_directory(data_dir.path());

    let initial_timestamp = {
        let server = harness.clone().start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();

        // Create a table with no data, which should have some overhead and therefore some storage usage
        client
            .batch_execute("CREATE TABLE usage_test (a int)")
            .unwrap();

        *now.lock().expect("lock poisoned") +=
            u64::try_from(collection_interval.as_millis()).expect("known to fit") + 1;

        // Wait for initial storage usage collection, to be sure records are present.
        let initial_timestamp = Retry::default().max_duration(Duration::from_secs(5)).retry(|_| {
                client
                    .query_one(
                        "SELECT (EXTRACT(EPOCH FROM MAX(collection_timestamp)) * 1000)::integer FROM mz_internal.mz_storage_usage_by_shard;",
                        &[],
                    )
                    .map_err(|e| e.to_string())?
                    .try_get::<_, i32>(0)
                    .map_err(|e| e.to_string())
            }).expect("Could not fetch initial timestamp");

        let initial_storage_usage_records = client
            .query_one(
                "SELECT COUNT(*)::integer AS number
                     FROM mz_internal.mz_storage_usage_by_shard",
                &[],
            )
            .unwrap()
            .try_get::<_, i32>(0)
            .expect("Could not get initial count of records");

        info!(%initial_timestamp, %initial_storage_usage_records);
        assert!(
            initial_storage_usage_records >= 1,
            "No initial server usage records!"
        );

        initial_timestamp
    };

    // Push time forward, start a new server, and assert that the previous storage records have been reaped
    *now.lock().expect("lock poisoned") = u64::try_from(initial_timestamp)
        .expect("negative timestamps are impossible")
        + u64::try_from(retention_period.as_millis()).expect("known to fit")
        // Add a second to account for any rounding errors.
        + 1000;

    {
        let server = harness.start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();

        *now.lock().expect("lock poisoned") +=
            u64::try_from(collection_interval.as_millis()).expect("known to fit") + 1000;

        let subsequent_initial_timestamp = Retry::default()
            .max_duration(Duration::from_secs(30))
            .retry(|_| {
                let query = format!(
                    "SELECT ts FROM (\
                     SELECT (EXTRACT(EPOCH FROM \
                     MIN(collection_timestamp)) * 1000)\
                     ::integer as ts FROM \
                     mz_internal.mz_storage_usage_by_shard\
                     ) WHERE ts > {initial_timestamp};"
                );
                client
                    .query_one(&query, &[])
                    .map_err(|e| e.to_string())?
                    .try_get::<_, i32>(0)
                    .map_err(|e| e.to_string())
            })
            .expect("Could not fetch initial timestamp");

        info!(%subsequent_initial_timestamp);
        assert!(
            subsequent_initial_timestamp > initial_timestamp,
            "Records were not reaped!"
        );
    };
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_storage_usage_records_are_not_cleared_on_restart() {
    let data_dir = tempfile::tempdir().unwrap();
    let collection_interval = Duration::from_secs(1);
    // Intentionally set retention period extremely high so that records aren't reaped.
    let retention_period = Duration::from_secs(u64::MAX);
    let harness = test_util::TestHarness::default()
        .with_storage_usage_collection_interval(collection_interval)
        .with_storage_usage_retention_period(retention_period)
        .data_directory(data_dir.path());

    let (initial_timestamp, initial_storage_usage_records) = {
        let server = harness.clone().start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();
        // Create a table with no data, which should have some overhead and therefore some storage usage.
        client
            .batch_execute("CREATE TABLE usage_test (a int)")
            .unwrap();

        // Wait for initial storage usage collection, to be sure records are present.
        let initial_timestamp = Retry::default().max_duration(Duration::from_secs(5)).retry(|_| {
            client
                .query_one(
                    "SELECT (EXTRACT(EPOCH FROM MIN(collection_timestamp)) * 1000)::int8 FROM mz_internal.mz_storage_usage_by_shard;",
                    &[],
                )
                .map_err(|e| e.to_string()).unwrap()
                .try_get::<_, i64>(0)
                .map_err(|e| e.to_string())
        }).expect("Could not fetch initial timestamp");

        let initial_storage_usage_records = client
            .query_one(
                "SELECT COUNT(*)::integer AS number
                     FROM mz_internal.mz_storage_usage_by_shard",
                &[],
            )
            .unwrap()
            .try_get::<_, i32>(0)
            .expect("Could not get initial count of records");

        info!(%initial_timestamp, %initial_storage_usage_records);
        assert!(
            initial_storage_usage_records >= 1,
            "No initial server usage records!"
        );

        (initial_timestamp, initial_storage_usage_records)
    };

    {
        let server = harness.start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();

        let subsequent_initial_timestamp =client
            .query_one(
                "SELECT (EXTRACT(EPOCH FROM MIN(collection_timestamp)) * 1000)::int8 FROM mz_internal.mz_storage_usage_by_shard;",
                &[],
            )
            .map_err(|e| e.to_string()).unwrap()
            .try_get::<_, i64>(0)
            .expect("Could not fetch subsequent initial timestamp");

        let subsequent_storage_usage_records = client
            .query_one(
                "SELECT COUNT(*)::integer AS number
                     FROM mz_internal.mz_storage_usage_by_shard",
                &[],
            )
            .unwrap()
            .try_get::<_, i32>(0)
            .expect("Could not get initial count of records");

        info!(%subsequent_initial_timestamp, %subsequent_storage_usage_records);
        assert_eq!(
            subsequent_initial_timestamp, initial_timestamp,
            "storage usage records were cleared"
        );
        assert!(
            subsequent_storage_usage_records >= initial_storage_usage_records,
            "storage usage was cleared!"
        );
    };
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_default_cluster_sizes() {
    let server = test_util::TestHarness::default()
        .with_builtin_system_cluster_replica_size("scale=1,workers=1".to_string())
        .with_builtin_catalog_server_cluster_replica_size("scale=1,workers=1".to_string())
        .with_default_cluster_replica_size("scale=1,workers=2".to_string())
        .start_blocking();
    let mut client = server.connect(postgres::NoTls).unwrap();

    let builtin_size: String = client
        .query(
            "SELECT size FROM (SHOW CLUSTER REPLICAS WHERE cluster LIKE 'mz_%')",
            &[],
        )
        .unwrap()
        .get(0)
        .unwrap()
        .get(0);
    assert_eq!(builtin_size, "scale=1,workers=1");

    let builtin_size: String = client
        .query(
            "SELECT size FROM (SHOW CLUSTER REPLICAS WHERE cluster = 'quickstart')",
            &[],
        )
        .unwrap()
        .get(0)
        .unwrap()
        .get(0);
    assert_eq!(builtin_size, "scale=1,workers=2");
}

#[mz_ore::test]
#[ignore] // TODO: Reenable when https://github.com/MaterializeInc/database-issues/issues/6931 is fixed
#[allow(clippy::disallowed_methods)]
fn test_max_request_size() {
    let statement = "SELECT $1::text";
    let statement_size = statement.bytes().count();
    let server = test_util::TestHarness::default().start_blocking();

    // pgwire
    {
        let param_size = mz_pgwire_common::MAX_REQUEST_SIZE - statement_size + 1;
        let param = std::iter::repeat("1").take(param_size).join("");
        let mut client = server.connect(postgres::NoTls).unwrap();

        let err = client.query(statement, &[&param]).unwrap_db_error();
        assert_contains!(err.message(), "request larger than");
        assert_err!(client.is_valid(Duration::from_secs(2)));
    }

    // http
    {
        let param_size = mz_environmentd::http::MAX_REQUEST_SIZE - statement_size + 1;
        let param = std::iter::repeat("1").take(param_size).join("");
        let http_url = Url::parse(&format!("http://{}/api/sql", server.http_local_addr())).unwrap();
        let json = format!("{{\"queries\":[{{\"query\":\"{statement}\",\"params\":[{param}]}}]}}");
        let json: serde_json::Value = serde_json::from_str(&json).unwrap();
        let res = Client::new().post(http_url).json(&json).send().unwrap();
        assert_eq!(res.status(), StatusCode::PAYLOAD_TOO_LARGE);
    }

    // ws
    {
        let param_size = mz_environmentd::http::MAX_REQUEST_SIZE - statement_size + 1;
        let param = std::iter::repeat("1").take(param_size).join("");
        let ws_url = server.ws_addr();
        let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap();
        test_util::auth_with_ws(&mut ws, BTreeMap::default()).unwrap();
        let json =
            format!("{{\"queries\":[{{\"query\":\"{statement}\",\"params\":[\"{param}\"]}}]}}");
        let json: serde_json::Value = serde_json::from_str(&json).unwrap();
        ws.send(Message::Text(json.to_string().into())).unwrap();

        // The specific error isn't forwarded to the client, the connection is just closed.
        let err = ws.read().unwrap_err();
        assert!(matches!(
            err,
            Error::Protocol(ProtocolError::ResetWithoutClosingHandshake)
        ));
    }
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn test_max_statement_batch_size() {
    let statement = "SELECT 1;";
    let statement_size = statement.bytes().count();
    let max_statement_size = mz_sql_parser::parser::MAX_STATEMENT_BATCH_SIZE;
    let max_statement_count = max_statement_size / statement_size + 1;
    let statements = iter::repeat(statement).take(max_statement_count).join("");
    let server = test_util::TestHarness::default().start_blocking();

    // pgwire
    {
        let mut client = server.connect(postgres::NoTls).unwrap();

        let err = client
            .batch_execute(&statements)
            .expect_err("statement should be too large")
            .unwrap_db_error();
        assert_eq!(&SqlState::PROGRAM_LIMIT_EXCEEDED, err.code());
        assert!(
            err.message().contains("statement batch size cannot exceed"),
            "error should indicate that the statement was too large: {}",
            err.message()
        );
    }

    // http
    {
        let http_url = Url::parse(&format!("http://{}/api/sql", server.http_local_addr())).unwrap();
        let json = format!("{{\"query\":\"{statements}\"}}");
        let json: serde_json::Value = serde_json::from_str(&json).unwrap();

        let res = Client::new().post(http_url).json(&json).send().unwrap();
        let body: serde_json::Value = res.json().unwrap();
        let is_error = body
            .get("results")
            .unwrap()
            .get(0)
            .unwrap()
            .get("error")
            .is_some();
        assert!(is_error, "statement should result in an error: {body}");
        assert!(
            body.get("results")
                .unwrap()
                .get(0)
                .unwrap()
                .get("error")
                .unwrap()
                .get("message")
                .unwrap()
                .as_str()
                .unwrap()
                .contains("statement batch size cannot exceed"),
            "error should indicate that the statement was too large: {}",
            body
        );
    }

    // ws
    {
        let ws_url = server.ws_addr();
        let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap();
        test_util::auth_with_ws(&mut ws, BTreeMap::default()).unwrap();
        let json = format!("{{\"query\":\"{statements}\"}}");
        let json: serde_json::Value = serde_json::from_str(&json).unwrap();
        ws.send(Message::Text(json.to_string().into())).unwrap();

        // Discard the CommandStarting message
        let _ = ws.read().unwrap();

        let msg = ws.read().unwrap();
        let msg = msg.into_text().expect("response should be text");
        let msg: WebSocketResponse = serde_json::from_str(&msg).unwrap();
        match msg {
            WebSocketResponse::Error(err) => assert!(
                err.message.contains("statement batch size cannot exceed"),
                "error should indicate that the statement was too large: {}",
                err.message,
            ),
            msg => {
                panic!("response should be error: {msg:?}")
            }
        }
    }
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_console_config_endpoint() {
    let server = test_util::TestHarness::default().start_blocking();
    let http_url = Url::parse(&format!(
        "http://{}/api/console/config",
        server.http_local_addr()
    ))
    .unwrap();

    let res = Client::new().get(http_url.clone()).send().unwrap();
    assert_eq!(res.status(), StatusCode::OK);

    // Response should be the default values.
    let body: serde_json::Value = res.json().unwrap();
    assert_eq!(body["oidc_issuer"], "");
    assert_eq!(body["console_oidc_client_id"], "");
    assert_eq!(body["console_oidc_scopes"], "");

    // Setting the dyncfg via the internal SQL port should be reflected.
    let mut internal_client = server.connect_internal(postgres::NoTls).unwrap();
    internal_client
        .batch_execute("ALTER SYSTEM SET oidc_issuer = 'https://my-issuer.com'")
        .unwrap();
    internal_client
        .batch_execute("ALTER SYSTEM SET console_oidc_client_id = 'my-client-id'")
        .unwrap();
    internal_client
        .batch_execute("ALTER SYSTEM SET console_oidc_scopes = 'openid email'")
        .unwrap();

    let res = Client::new().get(http_url).send().unwrap();
    assert_eq!(res.status(), StatusCode::OK);

    let body: serde_json::Value = res.json().unwrap();
    assert_eq!(body["oidc_issuer"], "https://my-issuer.com");
    assert_eq!(body["console_oidc_client_id"], "my-client-id");
    assert_eq!(body["console_oidc_scopes"], "openid email");
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_mz_system_user_admin() {
    let server = test_util::TestHarness::default().start_blocking();
    let mut client = server
        .pg_config_internal()
        .user(&SYSTEM_USER.name)
        .connect(postgres::NoTls)
        .unwrap();
    assert_eq!(
        "on".to_string(),
        client
            .query_one("SHOW is_superuser;", &[])
            .unwrap()
            .get::<_, String>(0)
    );
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
fn test_ws_passes_options() {
    let server = test_util::TestHarness::default().start_blocking();

    // Create our WebSocket.
    let ws_url = server.ws_addr();
    let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap();
    let options = BTreeMap::from([(
        "application_name".to_string(),
        "billion_dollar_idea".to_string(),
    )]);
    test_util::auth_with_ws(&mut ws, options).unwrap();

    // Query to make sure we get back the correct session var, which should be
    // set from the options map we passed with the auth.
    let json = "{\"query\":\"SHOW application_name;\"}";
    let json: serde_json::Value = serde_json::from_str(json).unwrap();
    ws.send(Message::Text(json.to_string().into())).unwrap();

    let mut read_msg = || -> WebSocketResponse {
        let msg = ws.read().unwrap();
        let msg = msg.into_text().expect("response should be text");
        serde_json::from_str(&msg).unwrap()
    };
    let starting = read_msg();
    let columns = read_msg();
    let row_val = read_msg();

    if !matches!(starting, WebSocketResponse::CommandStarting(_)) {
        panic!("wrong message!, {starting:?}");
    };

    if let WebSocketResponse::Rows(rows) = columns {
        let names: Vec<&str> = rows.columns.iter().map(|c| c.name.as_str()).collect();
        assert_eq!(names, ["application_name"]);
    } else {
        panic!("wrong message!, {columns:?}");
    };

    if let WebSocketResponse::Row(row) = row_val {
        let expected = serde_json::Value::String("billion_dollar_idea".to_string());
        assert_eq!(row, [expected]);
    } else {
        panic!("wrong message!, {row_val:?}");
    }
}

#[mz_ore::test]
// Test that statement logging
// doesn't cause a crash with subscribes over web sockets,
// which was previously happening (in staging) due to us
// dropping the `ExecuteContext` on the floor in that case.
fn test_statement_logging_ws_subscribe_no_crash() {
    let (server, _client) = setup_statement_logging(1.0, 1.0, "");

    // Create our WebSocket.
    let ws_url = server.server.ws_addr();
    let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap();
    test_util::auth_with_ws(&mut ws, Default::default()).unwrap();

    let query = "SUBSCRIBE (SELECT 1)";
    let json = format!("{{\"query\":\"{query}\"}}");
    let json: serde_json::Value = serde_json::from_str(&json).unwrap();
    ws.send(Message::Text(json.to_string().into())).unwrap();

    // Give the server time to crash, if it's going to.
    std::thread::sleep(Duration::from_secs(1))
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
fn test_ws_notifies_for_bad_options() {
    let server = test_util::TestHarness::default().start_blocking();

    // Create our WebSocket.
    let ws_url = server.ws_addr();
    let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap();
    let options = BTreeMap::from([("bad_var_name".to_string(), "i_do_not_exist".to_string())]);
    test_util::auth_with_ws(&mut ws, options).unwrap();

    let mut read_msg = || -> WebSocketResponse {
        let msg = ws.read().unwrap();
        let msg = msg.into_text().expect("response should be text");
        serde_json::from_str(&msg).unwrap()
    };
    let notice = read_msg();

    // After startup, we should get a notice that our var name was not set.
    if let WebSocketResponse::Notice(notice) = notice {
        let msg = notice.message();
        assert!(msg.starts_with("startup setting bad_var_name not set"));
    } else {
        panic!("wrong message!, {notice:?}");
    };
}

#[derive(Debug, Deserialize)]
struct HttpResponse<R> {
    results: Vec<R>,
}

#[derive(Debug, Deserialize)]
struct HttpRows {
    rows: Vec<Vec<serde_json::Value>>,
    desc: HttpResponseDesc,
    notices: Vec<Notice>,
}

#[derive(Debug, Deserialize)]
struct HttpResponseDesc {
    columns: Vec<HttpResponseColumn>,
}

#[derive(Debug, Deserialize)]
struct HttpResponseColumn {
    name: String,
}

#[derive(Debug, Deserialize)]
struct Notice {
    message: String,
    #[serde(rename = "severity")]
    _severity: String,
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
fn test_http_options_param() {
    let server = test_util::TestHarness::default().start_blocking();

    #[derive(Debug, Serialize)]
    struct Params {
        options: String,
    }

    let make_request = |params| {
        let http_url = Url::parse(&format!(
            "http://{}/api/sql?{}",
            server.http_local_addr(),
            params
        ))
        .unwrap();

        let json = r#"{ "query": "SHOW application_name;" }"#;
        let json: serde_json::Value = serde_json::from_str(json).unwrap();
        Client::new().post(http_url).json(&json).send().unwrap()
    };

    //
    // Happy path, valid and correctly formatted options dictionary.
    //
    let options = BTreeMap::from([("application_name", "yet_another_client")]);
    let options = serde_json::to_string(&options).unwrap();
    let params = serde_urlencoded::to_string(Params { options }).unwrap();

    let resp = make_request(params);
    assert!(resp.status().is_success());

    let mut result: HttpResponse<HttpRows> = resp.json().unwrap();
    assert_eq!(result.results.len(), 1);

    let mut rows = result.results.pop().unwrap();
    assert!(rows.notices.is_empty());

    let row = rows.rows.pop().unwrap().pop().unwrap();
    let col = rows.desc.columns.pop().unwrap().name;

    assert_eq!(col, "application_name");
    assert_eq!(row, "yet_another_client");

    //
    // Malformed options dictionary.
    //
    let params = "options=not_a_urlencoded_json_object".to_string();
    let resp = make_request(params);
    assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

    //
    // Correctly formed options dictionary, but invalid param.
    //
    let options = BTreeMap::from([("not_a_session_var", "hmmmm")]);
    let options = serde_json::to_string(&options).unwrap();
    let params = serde_urlencoded::to_string(Params { options }).unwrap();

    let resp = make_request(params);
    assert!(resp.status().is_success());

    let result: HttpResponse<HttpRows> = resp.json().unwrap();
    assert_eq!(result.results.len(), 1);
    assert_eq!(result.results[0].notices.len(), 1);

    let notice = &result.results[0].notices[0];
    assert!(
        notice
            .message
            .contains(r#"startup setting not_a_session_var not set"#)
    );
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn test_max_connections_on_all_interfaces() {
    let query = "SELECT 1";
    let server = test_util::TestHarness::default()
        .unsafe_mode()
        .start_blocking();

    let mut mz_client = server
        .pg_config_internal()
        .user(&SYSTEM_USER.name)
        .connect(postgres::NoTls)
        .unwrap();
    mz_client
        .batch_execute("ALTER SYSTEM SET superuser_reserved_connections = 1")
        .unwrap();
    mz_client
        .batch_execute("ALTER SYSTEM SET max_connections = 2")
        .unwrap();

    let client = server.connect(postgres::NoTls).unwrap();

    let ws_url = server.ws_addr();
    let http_url = Url::parse(&format!("http://{}/api/sql", server.http_local_addr())).unwrap();
    let json = format!("{{\"query\":\"{query}\"}}");
    let json: serde_json::Value = serde_json::from_str(&json).unwrap();

    {
        // while postgres client is connected, http connections error out
        let res = Client::new()
            .post(http_url.clone())
            .json(&json)
            .send()
            .unwrap();
        let status = res.status();
        let text = res.text().expect("no body?");
        assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
        assert_contains!(
            text,
            "creating connection would violate max_connections limit (desired: 2, limit: 2, current: 1)"
        );
    }

    {
        // while postgres client is connected, websockets can't auth
        let (mut ws, _resp) = tungstenite::connect(ws_url.clone()).unwrap();
        let err = test_util::auth_with_ws(&mut ws, BTreeMap::default()).unwrap_err();
        assert_contains!(
            err.to_string_with_causes(),
            "creating connection would violate max_connections limit (desired: 2, limit: 2, current: 1)"
        );
    }

    tracing::info!("closing postgres client");
    client.close().unwrap();
    tracing::info!("closed postgres client");

    tracing::info!("waiting for postgres client to close so that the query goes through");
    Retry::default().max_tries(10).retry(|_state| {
            let res = Client::new().post(http_url.clone()).json(&json).send().unwrap();
            let status = res.status();
            if status == StatusCode::INTERNAL_SERVER_ERROR {
                assert_contains!(res.text().expect("expect body"), "creating connection would violate max_connections limit (desired: 2, limit: 2, current: 1)");
                return Err(());
            }
            assert_eq!(status, StatusCode::OK);
            let result: HttpResponse<HttpRows> = res.json().unwrap();
            assert_eq!(result.results.len(), 1);
            assert_eq!(result.results[0].rows, vec![vec!["1"]]);
            Ok(())
        }).unwrap();

    tracing::info!("http query succeeded");
    let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap();
    test_util::auth_with_ws(&mut ws, BTreeMap::default()).unwrap();
    let json = format!("{{\"query\":\"{query}\"}}");
    let json: serde_json::Value = serde_json::from_str(&json).unwrap();
    ws.send(Message::Text(json.to_string().into())).unwrap();

    // The specific error isn't forwarded to the client, the connection is just closed.
    match ws.read() {
        Ok(Message::Text(msg)) => {
            assert_eq!(
                msg,
                r#"{"type":"CommandStarting","payload":{"has_rows":true,"is_streaming":false}}"#
            );
            assert_eq!(
                ws.read().unwrap(),
                Message::Text(format!(
                    r#"{{"type":"Rows","payload":{{"columns":[{{"name":"{UNKNOWN_COLUMN_NAME}","type_oid":23,"type_len":4,"type_mod":-1}}]}}}}"#
                ).into())
            );
            assert_eq!(
                ws.read().unwrap(),
                Message::Text("{\"type\":\"Row\",\"payload\":[\"1\"]}".to_string().into())
            );
            tracing::info!("data: {:?}", ws.read().unwrap());
        }
        Ok(msg) => panic!("unexpected msg: {msg:?}"),
        Err(e) => panic!("{e}"),
    }

    // While the websocket connection is still open, http requests fail
    let res = Client::new().post(http_url).json(&json).send().unwrap();
    tracing::info!("res: {:#?}", res);
    let status = res.status();
    let text = res.text().expect("no body?");
    assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
    assert_contains!(
        text,
        "creating connection would violate max_connections limit (desired: 2, limit: 2, current: 1)"
    );

    // Make sure lowering our max connections below the current limit does not
    // cause a panic.
    mz_client
        .batch_execute("ALTER SYSTEM SET max_connections = 10")
        .unwrap();

    // Open a few connections (kept alive to test connection limits).
    #[allow(clippy::collection_is_never_read)]
    let mut clients = Vec::new();
    for _ in 0..5 {
        let client = server
            .pg_config()
            .connect(postgres::NoTls)
            .expect("success opening client");
        clients.push(client);
    }

    // Lower the connection below the number of open clients.
    mz_client
        .batch_execute("ALTER SYSTEM SET max_connections = 2")
        .unwrap();

    // Opening a new connection should fail.
    let result = server.pg_config().connect(postgres::NoTls);
    let Err(failure) = result else {
        panic!("unexpected success connecting to server");
    };
    assert_contains!(
        failure.to_string_with_causes(),
        "creating connection would violate max_connections limit (desired: 7, limit: 2, current: 6)"
    );
}

// Test max_connections and superuser_reserved_connections.
#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[allow(clippy::disallowed_methods)]
async fn test_max_connections_limits() {
    let ca = Ca::new_root("test ca").unwrap();
    let (server_cert, server_key) = ca
        .request_cert("server", vec![IpAddr::V4(Ipv4Addr::LOCALHOST)])
        .unwrap();
    let metrics_registry = MetricsRegistry::new();

    let tenant_id = Uuid::new_v4();
    let regular_user = UserConfig::generate(tenant_id.clone(), "user@_.com", Vec::new());
    let admin_user = UserConfig::generate(tenant_id.clone(), "admin@_.com", vec!["mzadmin".into()]);
    let users = BTreeMap::from([
        (regular_user.email.clone(), regular_user.clone()),
        (admin_user.email.clone(), admin_user.clone()),
    ]);

    let issuer = "frontegg-mock".to_owned();
    let encoding_key =
        EncodingKey::from_rsa_pem(&ca.pkey.private_key_to_pem_pkcs8().unwrap()).unwrap();
    let decoding_key = DecodingKey::from_rsa_pem(&ca.pkey.public_key_to_pem().unwrap()).unwrap();

    let frontegg_server = FronteggMockServer::start(
        None,
        issuer,
        encoding_key,
        decoding_key,
        users,
        BTreeMap::new(),
        None,
        SYSTEM_TIME.clone(),
        500,
        None,
        None,
    )
    .await
    .unwrap();

    let frontegg_auth = FronteggAuthentication::new(
        FronteggConfig {
            admin_api_token_url: frontegg_server.auth_api_token_url(),
            decoding_key: DecodingKey::from_rsa_pem(&ca.pkey.public_key_to_pem().unwrap()).unwrap(),
            tenant_id: Some(tenant_id),
            now: SYSTEM_TIME.clone(),
            admin_role: "mzadmin".to_string(),
            refresh_drop_lru_size: DEFAULT_REFRESH_DROP_LRU_CACHE_SIZE,
            refresh_drop_factor: DEFAULT_REFRESH_DROP_FACTOR,
        },
        mz_frontegg_auth::Client::default(),
        &metrics_registry,
    );

    let server = test_util::TestHarness::default()
        .with_tls(server_cert, server_key)
        .with_frontegg_auth(&frontegg_auth)
        .with_metrics_registry(metrics_registry)
        .start()
        .await;

    let tls = make_pg_tls(|b| Ok(b.set_verify(SslVerifyMode::NONE)));

    let connect_regular_user = || async {
        server
            .connect()
            .ssl_mode(SslMode::Require)
            .with_tls(tls.clone())
            .user(&regular_user.email)
            .password(&regular_user.frontegg_password())
            .await
    };
    let connect_external_admin = || async {
        server
            .connect()
            .ssl_mode(SslMode::Require)
            .with_tls(tls.clone())
            .user(&admin_user.email)
            .password(&admin_user.frontegg_password())
            .await
    };
    let connect_system_user = || async { server.connect().internal().await };

    let mz_client = connect_system_user().await.unwrap();
    mz_client
        .batch_execute("ALTER SYSTEM SET max_connections TO 1")
        .await
        .unwrap();
    mz_client
        .batch_execute("ALTER SYSTEM SET superuser_reserved_connections = 1")
        .await
        .unwrap();

    // No regular user connections when max_connections = superuser_reserved_connections.
    assert_contains!(
        connect_regular_user()
            .await
            .expect_err("connect should fail")
            .to_string_with_causes(),
        "creating connection would violate max_connections limit"
    );

    // Set this real high to make sure nothing bad happens.
    mz_client
        .batch_execute("ALTER SYSTEM SET superuser_reserved_connections = 100")
        .await
        .unwrap();
    assert_contains!(
        connect_regular_user()
            .await
            .expect_err("connect should fail")
            .to_string_with_causes(),
        "creating connection would violate max_connections limit"
    );

    mz_client
        .batch_execute("ALTER SYSTEM SET superuser_reserved_connections = 0")
        .await
        .unwrap();

    {
        let client1 = connect_regular_user().await.unwrap();
        client1.batch_execute("SELECT 1").await.unwrap();

        assert_contains!(
            connect_regular_user()
                .await
                .expect_err("connect should fail")
                .to_string_with_causes(),
            "creating connection would violate max_connections limit"
        );

        assert_contains!(
            connect_regular_user()
                .await
                .expect_err("connect should fail")
                .to_string_with_causes(),
            "creating connection would violate max_connections limit"
        );

        let mz_client2 = connect_system_user().await.unwrap();
        mz_client2
            .batch_execute("SELECT 1")
            .await
            .expect("super users are still allowed to do queries");

        // External admins can't get around the connection limit.
        assert_contains!(
            connect_external_admin()
                .await
                .expect_err("connect should fail")
                .to_string_with_causes(),
            "creating connection would violate max_connections limit"
        );

        // Bump the limits.
        mz_client
            .batch_execute("ALTER SYSTEM SET max_connections TO 2")
            .await
            .unwrap();
        mz_client
            .batch_execute("ALTER SYSTEM SET superuser_reserved_connections = 1")
            .await
            .unwrap();

        // Regular users still should not connect.
        assert_contains!(
            connect_regular_user()
                .await
                .expect_err("connect should fail")
                .to_string_with_causes(),
            "creating connection would violate max_connections limit"
        );

        // But external admin should succeed.
        connect_external_admin().await.unwrap();
    }
    // after a client disconnects we can connect once the server notices the close
    Retry::default()
        .max_tries(10)
        .retry_async(|_state| async {
            let client = match connect_regular_user().await {
                Err(e) => {
                    assert_contains!(
                        e.to_string_with_causes(),
                        "creating connection would violate max_connections limit"
                    );
                    return Err(());
                }
                Ok(client) => client,
            };
            client.batch_execute("SELECT 1").await.unwrap();
            Ok(())
        })
        .await
        .unwrap();

    mz_client
        .batch_execute("ALTER SYSTEM RESET max_connections")
        .await
        .unwrap();

    // We can create lots of clients now (kept alive to test connection limits).
    #[allow(clippy::collection_is_never_read)]
    let mut clients = Vec::new();
    for _ in 0..10 {
        let client = connect_regular_user().await.unwrap();
        client.batch_execute("SELECT 1").await.unwrap();
        clients.push(client);
    }
}

#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn test_concurrent_id_reuse() {
    let server = test_util::TestHarness::default().start().await;

    {
        let client = server.connect().await.unwrap();
        client
            .batch_execute("CREATE TABLE t (a INT);")
            .await
            .unwrap();
    }

    let http_url = Url::parse(&format!("http://{}/api/sql", server.http_local_addr())).unwrap();
    let select_json = "{\"queries\":[{\"query\":\"SELECT * FROM t;\",\"params\":[]}]}";
    let select_json: serde_json::Value = serde_json::from_str(select_json).unwrap();

    let insert_json = "{\"queries\":[{\"query\":\"INSERT INTO t VALUES (1);\",\"params\":[]}]}";
    let insert_json: serde_json::Value = serde_json::from_str(insert_json).unwrap();

    // The goal here is to start some connection `B`, after another connection `A` has
    // terminated, but while connection `A` still has some asynchronous work in flight. Then
    // connection `A` will terminate it's session after connection `B` has started it's own
    // session. If they use the same connection ID, then `A` will accidentally tear down `B`'s
    // state and `B` will panic at any point it tries to access it's state. If they don't use
    // the same connection ID, then everything will be fine.
    fail::cfg("async_prepare", "return(true)").unwrap();
    for i in 0..50 {
        let http_url = http_url.clone();
        if i % 2 == 0 {
            let fut = reqwest::Client::new()
                .post(http_url)
                .json(&select_json)
                .send();
            let time = tokio::time::sleep(Duration::from_millis(500));
            futures::select! {
                _ = fut.fuse() => {},
                _ = time.fuse() => {},
            }
        } else {
            reqwest::Client::new()
                .post(http_url)
                .json(&insert_json)
                .send()
                .await
                .unwrap();
        }
    }

    let client = server.connect().await.unwrap();
    client.batch_execute("SELECT 1").await.unwrap();
}

#[mz_ore::test]
fn test_internal_console_proxy() {
    let server = test_util::TestHarness::default().start_blocking();

    let res = Client::builder()
        .build()
        .unwrap()
        .get(
            Url::parse(&format!(
                "http://{}/internal-console/",
                server.internal_http_local_addr()
            ))
            .unwrap(),
        )
        .send()
        .unwrap();

    assert_eq!(res.status().is_success(), true);
    assert_contains!(
        res.headers().get(CONTENT_TYPE).unwrap().to_str().unwrap(),
        "text/html"
    );
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_metrics_public_endpoint() {
    let server = test_util::TestHarness::default()
        .with_system_parameter_default(
            "enable_public_metrics_endpoint".to_string(),
            "true".to_string(),
        )
        .start_blocking();

    let cluster_name = "test_cluster_1";
    let label = format!("cluster_name=\"{cluster_name}\"");

    let mut client = server.connect(postgres::NoTls).unwrap();
    client
        .batch_execute(&format!(
            "CREATE CLUSTER {cluster_name} REPLICAS (r1 (SIZE 'scale=2,workers=2'))"
        ))
        .unwrap();

    let url = Url::parse(&format!(
        "http://{}/metrics/public",
        server.http_local_addr()
    ))
    .unwrap();
    let fetch_body = || {
        let res = Client::new().get(url.clone()).send().unwrap();
        assert_eq!(res.status(), StatusCode::OK);
        res.text().unwrap()
    };

    // Retry the scrape until the replica's `http` listener is up;
    // until then `/metrics/public` only returns env's local metrics with no
    // cluster_name labels.
    Retry::default()
        .max_duration(Duration::from_secs(60))
        .retry(|_| {
            if fetch_body().contains(label.as_str()) {
                Ok(())
            } else {
                Err(format!("{label} not yet in /metrics/public"))
            }
        })
        .unwrap();

    client
        .batch_execute(&format!("DROP CLUSTER {cluster_name}"))
        .unwrap();

    assert!(!fetch_body().contains(label.as_str()));
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
fn test_internal_http_auth() {
    let server = test_util::TestHarness::default().start_blocking();
    let json = serde_json::json!({"query": "SELECT current_user;"});
    let url = Url::parse(&format!(
        "http://{}/api/sql",
        server.internal_http_local_addr()
    ))
    .unwrap();

    let res = Client::new()
        .post(url.clone())
        .header("x-materialize-user", "mz_system")
        .json(&json)
        .send()
        .unwrap();

    tracing::info!("response: {res:?}");
    assert_eq!(
        res.status(),
        StatusCode::OK,
        "{:?}",
        res.json::<serde_json::Value>()
    );
    // can be explicitly set to mz_system
    assert!(res.text().unwrap().to_string().contains("mz_system"));

    // Check that mz_system is a superuser
    let json_superuser = serde_json::json!({"query": "SHOW is_superuser;"});
    let res = Client::new()
        .post(url.clone())
        .header("x-materialize-user", "mz_system")
        .json(&json_superuser)
        .send()
        .unwrap();

    assert_eq!(
        res.status(),
        StatusCode::OK,
        "{:?}",
        res.json::<serde_json::Value>()
    );
    assert!(res.text().unwrap().to_string().contains("on"));

    let res = Client::new()
        .post(url.clone())
        .header("x-materialize-user", "mz_support")
        .json(&json)
        .send()
        .unwrap();

    assert_eq!(
        res.status(),
        StatusCode::OK,
        "{:?}",
        res.json::<serde_json::Value>()
    );
    // can be explicitly set to mz_support
    assert!(res.text().unwrap().to_string().contains("mz_support"));

    // Check that mz_support is not a superuser
    let res = Client::new()
        .post(url.clone())
        .header("x-materialize-user", "mz_support")
        .json(&json_superuser)
        .send()
        .unwrap();

    assert_eq!(
        res.status(),
        StatusCode::OK,
        "{:?}",
        res.json::<serde_json::Value>()
    );
    assert!(res.text().unwrap().to_string().contains("off"));

    let res = Client::new()
        .post(url.clone())
        .header("x-materialize-user", "invalid value")
        .json(&json)
        .send()
        .unwrap();

    // invalid header returns an error
    assert_eq!(res.status(), StatusCode::UNAUTHORIZED, "{:?}", res.text());
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
fn test_internal_ws_auth() {
    let server = test_util::TestHarness::default().start_blocking();

    // Create our WebSocket.
    let ws_url = server.internal_ws_addr();
    let make_req = || {
        Request::builder()
            .uri(ws_url.clone())
            .method("GET")
            .header("Host", ws_url.host().unwrap())
            .header("Connection", "Upgrade")
            .header("Upgrade", "websocket")
            .header("Sec-WebSocket-Version", "13")
            .header("Sec-WebSocket-Key", "foobar")
            // Set our user to the mz_support user
            .header("x-materialize-user", "mz_support")
            .body(())
            .unwrap()
    };

    let (mut ws, _resp) = tungstenite::connect(make_req()).unwrap();
    let options = BTreeMap::from([
        (
            "application_name".to_string(),
            "billion_dollar_idea".to_string(),
        ),
        ("welcome_message".to_string(), "off".to_string()),
    ]);
    // We should receive error if sending the standard bearer auth, since that is unexpected
    // for the Internal HTTP API
    assert_err!(test_util::auth_with_ws(&mut ws, options.clone()));

    // Recreate the websocket
    let (mut ws, _resp) = tungstenite::connect(make_req()).unwrap();
    // Auth with OptionsOnly
    test_util::auth_with_ws_impl(
        &mut ws,
        Message::Text(
            serde_json::to_string(&WebSocketAuth::OptionsOnly { options })
                .unwrap()
                .into(),
        ),
    )
    .unwrap();

    // Query to make sure we get back the correct user, which should be
    // set from the headers passed with the websocket request.
    let json = "{\"query\":\"SELECT current_user;\"}";
    let json: serde_json::Value = serde_json::from_str(json).unwrap();
    ws.send(Message::Text(json.to_string().into())).unwrap();

    let mut read_msg = || -> WebSocketResponse {
        let msg = ws.read().unwrap();
        let msg = msg.into_text().expect("response should be text");
        serde_json::from_str(&msg).unwrap()
    };
    let starting = read_msg();
    let columns = read_msg();
    let row_val = read_msg();

    if !matches!(starting, WebSocketResponse::CommandStarting(_)) {
        panic!("wrong message!, {starting:?}");
    };

    if let WebSocketResponse::Rows(rows) = columns {
        let names: Vec<&str> = rows.columns.iter().map(|c| c.name.as_str()).collect();
        assert_eq!(names, ["current_user"]);
    } else {
        panic!("wrong message!, {columns:?}");
    };

    if let WebSocketResponse::Row(row) = row_val {
        let expected = serde_json::Value::String("mz_support".to_string());
        assert_eq!(row, [expected]);
    } else {
        panic!("wrong message!, {row_val:?}");
    }
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn test_leader_promotion_always_using_deploy_generation() {
    let tmpdir = TempDir::new().unwrap();
    let harness = test_util::TestHarness::default()
        .unsafe_mode()
        .data_directory(tmpdir.path())
        .with_deploy_generation(2);
    {
        // propose a deploy generation for the first time
        let server = harness.clone().start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();
        client.simple_query("SELECT 1").unwrap();
    }
    {
        // keep it the same, no need to promote the leader
        let server = harness.start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();
        client.simple_query("SELECT 1").unwrap();

        let http_client = reqwest::blocking::Client::new();

        // check that we're the leader and promotion doesn't do anything
        let status_http_url = Url::parse(&format!(
            "http://{}/api/leader/status",
            server.internal_http_local_addr()
        ))
        .unwrap();
        let res = http_client.get(status_http_url).send().unwrap();
        assert_eq!(res.status(), StatusCode::OK);
        let response = res.text().unwrap();
        assert_eq!(response, r#"{"status":"IsLeader"}"#);

        let promote_http_url = Url::parse(&format!(
            "http://{}/api/leader/promote",
            server.internal_http_local_addr()
        ))
        .unwrap();
        let res = http_client.post(promote_http_url).send().unwrap();
        assert_eq!(res.status(), StatusCode::OK);
        let response = res.text().unwrap();
        assert_eq!(response, r#"{"result":"Success"}"#);
    }
}

#[mz_ore::test(tokio::test(flavor = "multi_thread"))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn test_leader_promotion_mixed_code_version() {
    let tmpdir = TempDir::new().unwrap();
    let this_version = mz_environmentd::BUILD_INFO.semver_version();
    let next_version = semver::Version::new(this_version.major, this_version.minor + 1, 0);
    let harness = test_util::TestHarness::default()
        .unsafe_mode()
        .data_directory(tmpdir.path())
        .with_deploy_generation(1)
        .with_code_version(this_version);

    // Query a server at the current version before we start a second server.
    let server_this = harness.clone().start().await;
    let client_this = server_this.connect().await.unwrap();
    client_this.simple_query("SELECT 1").await.unwrap();

    // Simulate a rolling upgrade and wait for the preflight checks.
    let listeners_next = test_util::Listeners::new(&harness).await.unwrap();
    let internal_http_addr_next = listeners_next.inner.http["internal"].handle.local_addr;
    let config_next = harness
        .clone()
        .with_deploy_generation(2)
        .with_code_version(next_version.clone());
    let _server_next = mz_ore::task::spawn(|| "next version", async move {
        listeners_next.serve(config_next).await.unwrap()
    })
    .abort_on_drop();

    let status_http_url_next = Url::parse(&format!(
        "http://{}/api/leader/status",
        internal_http_addr_next
    ))
    .unwrap();
    // Wait for the new version to be ready to promote.
    Retry::default()
        .retry_async(|_state| async {
            let res = reqwest::Client::new()
                .get(status_http_url_next.clone())
                .send()
                .await
                .unwrap();
            tracing::info!("{} response: {res:?}", next_version);
            assert_eq!(res.status(), StatusCode::OK);
            let response = res.text().await.unwrap();
            tracing::info!("{} response body: {response:?}", next_version);
            assert_ne!(response, r#"{"status":"IsLeader"}"#);
            if response == r#"{"status":"ReadyToPromote"}"# {
                Ok(())
            } else {
                Err(())
            }
        })
        .await
        .unwrap();
    // The next_version preflight checks shouldn't fence out the this_version
    // server.
    client_this.simple_query("SELECT 1").await.unwrap();
}

// Test that websockets observe cancellation.
#[mz_ore::test]
#[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `epoll_wait` on OS `linux`
#[allow(clippy::disallowed_methods)]
fn test_cancel_ws() {
    let server = test_util::TestHarness::default().start_blocking();
    let mut client = server.connect(postgres::NoTls).unwrap();
    client.batch_execute("CREATE TABLE t (i INT);").unwrap();

    // Start a thread to perform the cancel while the SUBSCRIBE is running in this thread.
    let handle = thread::spawn(move || {
        // Wait for the subscription to start.
        let conn_id = Retry::default()
            .retry(|_| {
                let conn_id: String = client
                    .query_one(
                        "SELECT s.connection_id::text FROM mz_internal.mz_subscriptions b JOIN mz_internal.mz_sessions s ON s.id = b.session_id",
                        &[],
                    )?
                    .get(0);
                Ok::<_, postgres::Error>(conn_id)
            })
            .unwrap();

        client
            .query_one(&format!("SELECT pg_cancel_backend({conn_id})"), &[])
            .unwrap();
    });

    let (mut ws, _resp) = tungstenite::connect(server.ws_addr()).unwrap();
    test_util::auth_with_ws(&mut ws, BTreeMap::default()).unwrap();
    let json = r#"{"queries":[{"query":"SUBSCRIBE t"}]}"#;
    let json: serde_json::Value = serde_json::from_str(json).unwrap();
    ws.send(Message::Text(json.to_string().into())).unwrap();

    loop {
        let msg = ws.read().unwrap();
        if let Ok(msg) = msg.into_text() {
            if msg.contains("canceling statement") {
                break;
            }
        }
    }
    handle.join().unwrap();
}

#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn smoketest_webhook_source() {
    let server = test_util::TestHarness::default().start().await;
    let client = server.connect().await.unwrap();

    #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize, Clone)]
    struct WebhookEvent {
        ts: u128,
        name: String,
        attrs: Vec<String>,
    }

    // Create a webhook source.
    client
        .execute(
            "CREATE CLUSTER webhook_cluster REPLICAS (r1 (SIZE 'scale=1,workers=1'));",
            &[],
        )
        .await
        .expect("failed to create cluster");
    client
        .execute(
            "CREATE SOURCE webhook_json IN CLUSTER webhook_cluster FROM WEBHOOK BODY FORMAT JSON",
            &[],
        )
        .await
        .expect("failed to create source");

    let now = || {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .expect("time went backwards")
            .as_nanos()
    };

    // Generate some events.
    const NUM_EVENTS: i64 = 100;
    let events: BTreeSet<_> = (0..NUM_EVENTS)
        .map(|i| WebhookEvent {
            ts: now(),
            name: format!("event_{i}"),
            attrs: (0..i).map(|j| format!("attr_{j}")).collect(),
        })
        .collect();

    let http_client = reqwest::Client::new();
    let webhook_url = Arc::new(format!(
        "http://{}/api/webhook/materialize/public/webhook_json",
        server.http_local_addr()
    ));
    // Send all of our events to our webhook source.
    let mut handles = Vec::with_capacity(events.len());
    for event in &events {
        let webhook_url = Arc::clone(&webhook_url);
        let event = event.clone();
        let http_client_ = http_client.clone();
        let handle = mz_ore::task::spawn(|| "smoketest_webhook_source_event", async move {
            let resp = http_client_
                .post(&*webhook_url)
                .json(&event)
                .send()
                .await
                .expect("failed to POST event");
            assert!(resp.status().is_success());
        });
        handles.push(handle);
    }
    futures::future::join_all(handles).await;

    let total_requests_metric = server
        .metrics_registry
        .gather()
        .into_iter()
        .find(|metric| metric.name() == "mz_http_requests_total")
        .unwrap();
    let total_requests_metric = &total_requests_metric.get_metric()[0];
    assert_eq!(total_requests_metric.get_counter().value(), 100.0);

    let path_label = &total_requests_metric.get_label()[0];
    assert_eq!(
        path_label.value(),
        "/api/webhook/{:database}/{:schema}/{:id}"
    );

    let status_label = &total_requests_metric.get_label()[2];
    assert_eq!(status_label.value(), "200");

    // Wait for the events to be persisted.
    let (client, result) = mz_ore::retry::Retry::default()
        .max_tries(10)
        .retry_async_with_state(client, |_, client| async move {
            let cnt: i64 = client
                .query_one("SELECT COUNT(*) FROM webhook_json", &[])
                .await
                .expect("failed to get count")
                .get(0);

            if cnt != 100 {
                (client, Err(anyhow::anyhow!("not all rows present")))
            } else {
                (client, Ok(()))
            }
        })
        .await;
    result.expect("failed to read events");

    // Read all of our events back.
    let events_roundtrip: BTreeSet<WebhookEvent> = client
        .query("SELECT * FROM webhook_json", &[])
        .await
        .expect("failed to query source")
        .into_iter()
        .map(|row| serde_json::from_value(row.get("body")).expect("failed to deserialize"))
        .collect();

    similar_asserts::assert_eq!(events, events_roundtrip);
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn test_invalid_webhook_body() {
    let server = test_util::TestHarness::default().start_blocking();

    let mut client = server.connect(postgres::NoTls).unwrap();
    let http_client = Client::new();

    // Create a cluster we can install webhook sources on.
    client
        .execute(
            "CREATE CLUSTER webhook_cluster REPLICAS (r1 (SIZE 'scale=1,workers=1'));",
            &[],
        )
        .expect("failed to create cluster");

    // Create a webhook source with a body format of text.
    client
        .execute(
            "CREATE SOURCE webhook_text IN CLUSTER webhook_cluster FROM WEBHOOK BODY FORMAT TEXT",
            &[],
        )
        .expect("failed to create source");
    let webhook_url = format!(
        "http://{}/api/webhook/materialize/public/webhook_text",
        server.http_local_addr()
    );

    // Send non-UTF8 text which will fail to get deserialized.
    let non_utf8 = vec![255, 255, 255, 255];
    assert_err!(std::str::from_utf8(&non_utf8));

    let resp = http_client
        .post(webhook_url)
        .body(non_utf8)
        .send()
        .expect("failed to POST event");
    assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

    // Create a webhook source with a body format of JSON.
    client
        .execute(
            "CREATE SOURCE webhook_json IN CLUSTER webhook_cluster FROM WEBHOOK BODY FORMAT JSON",
            &[],
        )
        .expect("failed to create source");
    let webhook_url = format!(
        "http://{}/api/webhook/materialize/public/webhook_json",
        server.http_local_addr()
    );

    // Send invalid JSON which will fail to get deserialized.
    let resp = http_client
        .post(webhook_url)
        .body("a")
        .send()
        .expect("failed to POST event");
    assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

    // Create a webhook source with a body format of bytes.
    client
        .execute(
            "CREATE SOURCE webhook_bytes IN CLUSTER webhook_cluster FROM WEBHOOK BODY FORMAT BYTES",
            &[],
        )
        .expect("failed to create source");
    let webhook_url = format!(
        "http://{}/api/webhook/materialize/public/webhook_bytes",
        server.http_local_addr()
    );

    // No matter what is in the body, we should always succeed.
    let mut data = [0u8; 128];
    rand::rng().fill_bytes(&mut data);
    println!("Random bytes: {data:?}");
    let resp = http_client
        .post(webhook_url)
        .body(data.to_vec())
        .send()
        .expect("failed to POST event");
    assert!(resp.status().is_success());
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn test_webhook_duplicate_headers() {
    let server = test_util::TestHarness::default().start_blocking();

    let mut client = server.connect(postgres::NoTls).unwrap();
    let http_client = Client::new();

    // Create a webhook source that includes headers.
    client
        .execute(
            "CREATE CLUSTER webhook_cluster REPLICAS (r1 (SIZE 'scale=1,workers=1'));",
            &[],
        )
        .expect("failed to create cluster");
    client
        .execute(
            "CREATE SOURCE webhook_text IN CLUSTER webhook_cluster FROM WEBHOOK BODY FORMAT TEXT INCLUDE HEADERS",
            &[],
        )
        .expect("failed to create source");
    let webhook_url = format!(
        "http://{}/api/webhook/materialize/public/webhook_text",
        server.http_local_addr()
    );

    // Send a request with duplicate headers.
    let resp = http_client
        .post(webhook_url)
        .body("test")
        .header("dupe", "first")
        .header("DUPE", "second")
        .header("dUpE", "third")
        .header("dUpE", "final")
        .send()
        .expect("failed to POST event");
    assert_eq!(resp.status().as_u16(), 401);
}

// Test that websockets observe cancellation and leave the transaction in an idle state.
#[mz_ore::test]
#[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `epoll_wait` on OS `linux`
#[allow(clippy::disallowed_methods)]
fn test_github_20262() {
    let server = test_util::TestHarness::default().start_blocking();
    let mut client = server.connect(postgres::NoTls).unwrap();
    client.batch_execute("CREATE TABLE t (i INT);").unwrap();

    let mut cancel = || {
        // Wait for the subscription to start.
        let conn_id = Retry::default()
            .retry(|_| {
                let conn_id: String = client
                    .query_one(
                        "SELECT s.connection_id::text FROM mz_internal.mz_subscriptions b JOIN mz_internal.mz_sessions s ON s.id = b.session_id",
                        &[],
                    )?
                    .get(0);
                Ok::<_, postgres::Error>(conn_id)
            })
            .unwrap();
        client
            .query_one(&format!("SELECT pg_cancel_backend({conn_id})"), &[])
            .unwrap();
    };

    let subscribe: serde_json::Value =
        serde_json::from_str(r#"{"queries":[{"query":"SUBSCRIBE t"}]}"#).unwrap();
    let subscribe = subscribe.to_string();
    let commit: serde_json::Value =
        serde_json::from_str(r#"{"queries":[{"query":"COMMIT"}]}"#).unwrap();
    let commit = commit.to_string();
    let select: serde_json::Value =
        serde_json::from_str(r#"{"queries":[{"query":"SELECT 1"}]}"#).unwrap();
    let select = select.to_string();

    let (mut ws, _resp) = tungstenite::connect(server.ws_addr()).unwrap();
    test_util::auth_with_ws(&mut ws, BTreeMap::default()).unwrap();
    ws.send(Message::Text(subscribe.into())).unwrap();
    cancel();
    ws.send(Message::Text(commit.into())).unwrap();
    ws.send(Message::Text(select.into())).unwrap();

    let mut expect = VecDeque::from([
        r#"{"type":"CommandStarting","payload":{"has_rows":true,"is_streaming":true}}"#.to_string(),
        r#"{"type":"Rows","payload":{"columns":[{"name":"mz_timestamp","type_oid":1700,"type_len":-1,"type_mod":2555908},{"name":"mz_diff","type_oid":20,"type_len":8,"type_mod":-1},{"name":"i","type_oid":23,"type_len":4,"type_mod":-1}]}}"#.to_string(),
        r#"{"type":"Error","payload":{"message":"canceling statement due to user request","code":"57014"}}"#.to_string(),
        r#"{"type":"ReadyForQuery","payload":"I"}"#.to_string(),
        r#"{"type":"Notice","payload":{"message":"there is no transaction in progress","code":"25P01","severity":"warning"}}"#.to_string(),
        r#"{"type":"CommandStarting","payload":{"has_rows":false,"is_streaming":false}}"#.to_string(),
        r#"{"type":"CommandComplete","payload":"COMMIT"}"#.to_string(),
        r#"{"type":"ReadyForQuery","payload":"I"}"#.to_string(),
        r#"{"type":"CommandStarting","payload":{"has_rows":true,"is_streaming":false}}"#.to_string(),
        format!(r#"{{"type":"Rows","payload":{{"columns":[{{"name":"{UNKNOWN_COLUMN_NAME}","type_oid":23,"type_len":4,"type_mod":-1}}]}}}}"#),
        r#"{"type":"Row","payload":["1"]}"#.to_string(),
        r#"{"type":"CommandComplete","payload":"SELECT 1"}"#.to_string(),
        r#"{"type":"ReadyForQuery","payload":"I"}"#.to_string(),
    ]);
    while !expect.is_empty() {
        if let Message::Text(text) = ws.read().unwrap() {
            let next = expect.pop_front().unwrap();
            assert_eq!(text, next);
        }
    }
}

// Test that the server properly handles cancellation requests of read-then-write queries.
// See database-issues#6134.
#[mz_ore::test]
#[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `epoll_wait` on OS `linux`
#[allow(clippy::disallowed_methods)]
fn test_cancel_read_then_write() {
    let server = test_util::TestHarness::default()
        .unsafe_mode()
        .start_blocking();
    server.enable_feature_flags(&["unsafe_enable_unsafe_functions"]);

    let mut client = server.connect(postgres::NoTls).unwrap();
    client
        .batch_execute("CREATE TABLE foo (a TEXT, ts INT)")
        .unwrap();

    // Lots of races here, so try this whole thing in a loop.
    Retry::default()
        .clamp_backoff(Duration::ZERO)
        .retry(|_state| {
            let mut client1 = server.connect(postgres::NoTls).unwrap();
            let mut client2 = server.connect(postgres::NoTls).unwrap();
            let cancel_token = client2.cancel_token();

            client1.batch_execute("DELETE FROM foo").unwrap();
            client1.batch_execute("SET statement_timeout = '5s'").unwrap();
            client1
                .batch_execute("INSERT INTO foo VALUES ('hello', 10)")
                .unwrap();

            let handle1 = thread::spawn(move || {
                let err =  client1
                    .batch_execute("insert into foo select a, case when mz_unsafe.mz_sleep(ts) > 0 then 0 end as ts from foo")
                    .unwrap_err();
                assert_contains!(
                    err.to_string_with_causes(),
                    "statement timeout"
                );
                client1
            });
            std::thread::sleep(Duration::from_millis(100));
            let handle2 = thread::spawn(move || {
                let err = client2
                .batch_execute("insert into foo values ('blah', 1);")
                .unwrap_err();
                assert_contains!(
                    err.to_string_with_causes(),
                    "canceling statement"
                );
            });
            std::thread::sleep(Duration::from_millis(100));
            cancel_token.cancel_query(postgres::NoTls)?;
            let mut client1 = handle1.join().unwrap();
            handle2.join().unwrap();
            let rows:i64 = client1.query_one ("SELECT count(*) FROM foo", &[]).unwrap().get(0);
            // We ran 3 inserts. First succeeded. Second timedout. Third cancelled.
            if rows !=1 {
                anyhow::bail!("unexpected row count: {rows}");
            }
            Ok::<_, anyhow::Error>(())
        })
        .unwrap();
}

#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[cfg_attr(miri, ignore)] // too slow
async fn test_http_metrics() {
    let server = test_util::TestHarness::default().start().await;

    let http_url = Url::parse(&format!("http://{}/api/sql", server.http_local_addr(),)).unwrap();

    // Handled query (successful)
    let json = r#"{ "query": "SHOW application_name;" }"#;
    let json: serde_json::Value = serde_json::from_str(json).unwrap();
    let response = reqwest::Client::new()
        .post(http_url.clone())
        .json(&json)
        .send()
        .await
        .unwrap();
    assert!(response.status().is_success());

    // Handled query (error)
    let json = r#"{ "query": "invalid sql 123;" }"#;
    let json: serde_json::Value = serde_json::from_str(json).unwrap();
    let response = reqwest::Client::new()
        .post(http_url.clone())
        .json(&json)
        .send()
        .await
        .unwrap();
    let body: serde_json::Value = response.json().await.unwrap();
    let is_error = body
        .get("results")
        .unwrap()
        .get(0)
        .unwrap()
        .get("error")
        .is_some();
    assert!(is_error);

    // Invalid request
    let json = r#"{ "q": "invalid sql 123;" }"#;
    let json: serde_json::Value = serde_json::from_str(json).unwrap();
    let response = reqwest::Client::new()
        .post(http_url)
        .json(&json)
        .send()
        .await
        .unwrap();
    assert!(response.status().is_client_error());

    let metrics = server.metrics_registry.gather();
    let http_metrics: Vec<_> = metrics
        .into_iter()
        .filter(|metric| metric.name().starts_with("mz_http"))
        .collect();

    // Make sure the duration metric exists.
    let duration_count = http_metrics
        .iter()
        .filter(|metric| metric.name() == "mz_http_request_duration_seconds")
        .count();
    assert_eq!(duration_count, 1);
    // Make sure the active count metric exists.
    let active_count = http_metrics
        .iter()
        .filter(|metric| metric.name() == "mz_http_requests_active")
        .count();
    assert_eq!(active_count, 1);

    // Make sure our metrics capture the one successful query and the one failure.
    let mut request_metrics: Vec<_> = http_metrics
        .into_iter()
        .filter(|metric| metric.name() == "mz_http_requests_total")
        .collect();
    assert_eq!(request_metrics.len(), 1);

    let request_metric = request_metrics.pop().unwrap();
    let success_metric = &request_metric.get_metric()[0];
    assert_eq!(success_metric.get_counter().value(), 2.0);
    assert_eq!(success_metric.get_label()[0].value(), "/api/sql");
    assert_eq!(success_metric.get_label()[2].value(), "200");

    let failure_metric = &request_metric.get_metric()[1];
    assert_eq!(failure_metric.get_counter().value(), 1.0);
    assert_eq!(failure_metric.get_label()[0].value(), "/api/sql");
    assert_eq!(failure_metric.get_label()[2].value(), "422");
}

#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 2))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn webhook_concurrent_actions() {
    let server = test_util::TestHarness::default().start().await;
    let client = server.connect().await.unwrap();

    #[derive(Debug, PartialEq, Eq, Deserialize, Serialize)]
    struct WebhookEvent {
        ts: u128,
        name: String,
        thread: usize,
    }

    // Create a webhook source.
    let src_name = "webhook_json";
    client
        .execute(
            "CREATE CLUSTER webhook_cluster_concurrent REPLICAS (r1 (SIZE 'scale=1,workers=1'));",
            &[],
        )
        .await
        .expect("failed to create cluster");
    client
        .execute(
            &format!("CREATE SOURCE {src_name} IN CLUSTER webhook_cluster_concurrent FROM WEBHOOK BODY FORMAT JSON"),
            &[],
        )
        .await
        .expect("failed to create source");

    fn now() -> u128 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .expect("time went backwards")
            .as_nanos()
    }

    // Flag that lets us shutdown our threads.
    let keep_sending = Arc::new(AtomicBool::new(true));
    // Track how many requests were resolved before we dropped the collection.
    let num_requests_before_drop = Arc::new(AtomicUsize::new(0));

    // Spin up tasks that will contiously push data to the webhook.
    let keep_sending_ = Arc::clone(&keep_sending);
    let num_requests_before_drop_ = Arc::clone(&num_requests_before_drop);
    let addr = server.http_local_addr();

    let poster = mz_ore::task::spawn(|| "webhook_concurrent_actions-poster", async move {
        let mut i = 0;

        let http_client = reqwest::Client::new();
        let mut tasks = Vec::with_capacity(500);

        // Keep sending events until we're told to stop.
        while keep_sending_.load(std::sync::atomic::Ordering::Relaxed) {
            let webhook_url = format!(
                "http://{}/api/webhook/materialize/public/webhook_json",
                addr,
            );

            let http_client_ = http_client.clone();
            let num_requests_before_drop__ = Arc::clone(&num_requests_before_drop_);
            let handle =
                mz_ore::task::spawn(|| "webhook_concurrent_actions-appender", async move {
                    // Create an event.
                    let event = WebhookEvent {
                        ts: now(),
                        name: format!("event_{i}"),
                        thread: 1,
                    };

                    // Send all of our events to our webhook source.
                    let resp = http_client_
                        .post(&webhook_url)
                        .json(&event)
                        .send()
                        .await
                        .expect("failed to POST event");
                    num_requests_before_drop__.fetch_add(1, std::sync::atomic::Ordering::Relaxed);

                    resp
                });
            tasks.push(handle);
            i += 1;

            // Wait before sending another request.
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        // Wait on all of the tasks to finish, then return the results.
        futures::future::join_all(tasks).await
    });

    // Let the posting threads run for a bit.
    tokio::time::sleep(Duration::from_secs(5)).await;

    // We should see at least this many successes.
    let num_requests_before_drop =
        num_requests_before_drop.load(std::sync::atomic::Ordering::Relaxed);
    // Drop the source
    client
        .execute(&format!("DROP SOURCE {src_name};"), &[])
        .await
        .expect("failed to drop source");

    // Keep sending for a bit.
    tokio::time::sleep(std::time::Duration::from_secs(2)).await;
    // Stop the threads.
    keep_sending.store(false, std::sync::atomic::Ordering::Relaxed);
    let results = poster.await;

    // Inspect the results.
    let mut results = results.into_iter().collect::<Vec<_>>().into_iter();

    for _ in 0..num_requests_before_drop {
        let response = results.next().expect("element");
        let status = response.status();

        // We expect the following response codes:
        //
        // 1. 200 - Successfully append data.
        // 2. 429 - Rate limited.
        // 3. 404 - Collection not found. Due to the artificial latency we add to appending, we
        //    could send a request, it gets queued, we drop the source, pull all requests from the
        //    queue, find our source is missing. This should cause a 404 Not Found, not a 500.
        //
        assert!(
            status.is_success()
                || status == StatusCode::TOO_MANY_REQUESTS
                || status == StatusCode::NOT_FOUND,
            "response: {response:?}"
        )
    }
    // We should have seen at least 100 requests before dropping the source.
    assert!(num_requests_before_drop > 100);

    // Make sure the source got dropped.
    //
    // Note: this doubles as a check to make sure nothing internally panicked.
    let sources: Vec<String> = client
        .query("SELECT name FROM mz_sources", &[])
        .await
        .expect("success")
        .into_iter()
        .map(|row| row.get(0))
        .collect();
    assert!(!sources.iter().any(|name| name == src_name));

    // Best effort cleanup.
    let _ = client
        .execute("DROP CLUSTER webhook_cluster_concurrent CASCADE", &[])
        .await;
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn webhook_concurrency_limit() {
    let concurrency_limit = 15;
    let server = test_util::TestHarness::default()
        .unsafe_mode()
        .start_blocking();

    // Note: we need enable_unstable_dependencies to use mz_sleep.
    server.enable_feature_flags(&[
        "unsafe_enable_unstable_dependencies",
        "unsafe_enable_unsafe_functions",
    ]);

    // Reduce the webhook concurrency limit;
    let mut mz_client = server
        .pg_config_internal()
        .user(&SYSTEM_USER.name)
        .connect(postgres::NoTls)
        .unwrap();
    mz_client
        .batch_execute(&format!(
            "ALTER SYSTEM SET webhook_concurrent_request_limit = {concurrency_limit}"
        ))
        .unwrap();

    let mut client = server.connect(postgres::NoTls).unwrap();

    // Create a webhook source.
    client
        .execute(
            "CREATE CLUSTER webhook_cluster REPLICAS (r1 (SIZE 'scale=1,workers=1'));",
            &[],
        )
        .expect("failed to create cluster");
    client
        .execute(
            "CREATE SOURCE webhook_text IN CLUSTER webhook_cluster FROM WEBHOOK \
             BODY FORMAT TEXT \
             CHECK ( WITH(BODY) body IS NOT NULL AND mz_unsafe.mz_sleep(5) IS NULL )",
            &[],
        )
        .expect("failed to create source");

    let http_client = reqwest::Client::new();
    let webhook_url = format!(
        "http://{}/api/webhook/materialize/public/webhook_text",
        server.http_local_addr().clone(),
    );
    let mut handles = Vec::with_capacity(concurrency_limit + 5);

    for _ in 0..concurrency_limit + 5 {
        let http_client_ = http_client.clone();
        let webhook_url_ = webhook_url.clone();

        let handle = server.runtime().spawn_named(
            || "webhook-concurrency-limit-test".to_string(),
            async move {
                let resp = http_client_
                    .post(webhook_url_)
                    .body("a")
                    .send()
                    .await
                    .expect("response to succeed");
                resp.status()
            },
        );
        handles.push(handle);
    }
    let results = server
        .runtime()
        .block_on(futures::future::join_all(handles));

    let successes = results
        .iter()
        .filter(|status_code| status_code.is_success())
        .count();
    // Note: if this status code changes, we need to update our docs.
    let rate_limited = results
        .iter()
        .filter(|status_code| **status_code == StatusCode::TOO_MANY_REQUESTS)
        .count();

    // We expect at least the number of requests we allow concurrently to succeed.
    assert!(successes >= concurrency_limit);
    // We send 5 more requests than our limit, but we assert only at least 3 get
    // rate limited to reduce test flakiness.
    assert!(rate_limited >= 3);
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn webhook_too_large_request() {
    let metrics_registry = MetricsRegistry::new();
    let server = test_util::TestHarness::default()
        .with_metrics_registry(metrics_registry.clone())
        .start_blocking();
    let mut client = server.connect(postgres::NoTls).unwrap();

    // Create a webhook source.
    client
        .execute(
            "CREATE CLUSTER webhook_cluster REPLICAS (r1 (SIZE 'scale=1,workers=1'));",
            &[],
        )
        .expect("failed to create cluster");
    client
        .execute(
            "CREATE SOURCE webhook_bytes IN CLUSTER webhook_cluster FROM WEBHOOK BODY FORMAT BYTES",
            &[],
        )
        .expect("failed to create source");

    let http_client = Client::new();
    let webhook_url = format!(
        "http://{}/api/webhook/materialize/public/webhook_bytes",
        server.http_local_addr(),
    );

    // Send an event with a body larger that is exactly our max size.
    let five_mib = usize::cast_from(bytesize::mib(5u64));
    let body = vec![42u8; five_mib];
    let resp = http_client
        .post(&webhook_url)
        .body(body)
        .send()
        .expect("failed to POST event");
    assert!(resp.status().is_success());

    // Send an event that is one larger than our max size.
    let body = vec![42u8; five_mib + 1];
    let resp = http_client
        .post(&webhook_url)
        .body(body)
        .send()
        .expect("failed to POST event");

    // Note: If this changes then we need to update our docs.
    assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE);

    // Ensure we logged a prometheus metric that we responded with a 413.
    let metrics: Vec<_> = metrics_registry.gather().into_iter().collect();
    let payload_too_large = metrics
        .iter()
        .find(|metric_family| metric_family.name() == "mz_http_requests_total")
        .unwrap()
        .get_metric()
        .iter()
        .find(|metric| {
            metric
                .get_label()
                .iter()
                .find(|label_pair| label_pair.value() == "413")
                .is_some()
        })
        .unwrap()
        .get_counter()
        .value();
    assert_eq!(payload_too_large, 1.0);
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn test_webhook_url_notice() {
    let server = test_util::TestHarness::default().start_blocking();
    let (tx, mut rx) = futures::channel::mpsc::unbounded();

    let mut client = server
        .pg_config()
        .notice_callback(move |notice| tx.unbounded_send(notice).expect("send notice"))
        .connect(postgres::NoTls)
        .unwrap();
    let http_client = Client::new();

    // Create a webhook source that includes headers.
    client
        .execute(
            "CREATE CLUSTER webhook_cluster REPLICAS (r1 (SIZE 'scale=1,workers=1'));",
            &[],
        )
        .expect("failed to create cluster");
    client
        .execute(
            "CREATE SOURCE \"webhook-source(with odd/name\" IN CLUSTER webhook_cluster FROM WEBHOOK BODY FORMAT TEXT INCLUDE HEADERS",
            &[],
        )
        .expect("failed to create source");

    let url_notice = rx.try_recv().expect("contains notice");
    // We should only get the one notice.
    assert_err!(rx.try_recv());

    // Print the notice to stderr for future debug-ability.
    eprintln!("notice: {}", url_notice.message());
    let url_in_quotes = url_notice
        .message()
        .strip_prefix("URL to POST data is ")
        .expect("got wrong message");
    // Remove the single quotes.
    let https_url = url_in_quotes.replace('\'', "");
    // Our local test server does not have the certs for https.
    let http_url = https_url.replace("https", "http");

    // Send a request with duplicate headers.
    let resp = http_client
        .post(http_url)
        .body("request_foo")
        .send()
        .expect("failed to POST event");
    assert_eq!(resp.status().as_u16(), 200);

    // Wait for the source to process our request.
    std::thread::sleep(std::time::Duration::from_secs(3));

    let row = client
        .query_one("SELECT body FROM \"webhook-source(with odd/name\"", &[])
        .expect("success");
    let body: String = row.get("body");
    assert_eq!(body, "request_foo");
}

#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 2))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn webhook_concurrent_swap() {
    let server = test_util::TestHarness::default().start().await;
    let mut client = server.connect().await.unwrap();

    // Create our webhook sources.
    let webhook_cluster = "webhook_cluster_concurrent_swap";
    client
        .execute(
            &format!("CREATE CLUSTER {webhook_cluster} REPLICAS (r1 (SIZE 'scale=1,workers=1'));"),
            &[],
        )
        .await
        .expect("failed to create cluster");

    let src_foo = "webhook_foo";
    client
        .execute(
            &format!(
                "CREATE SOURCE {src_foo} IN CLUSTER {webhook_cluster} FROM WEBHOOK BODY FORMAT TEXT"
            ),
            &[],
        )
        .await
        .expect("failed to create source");
    let src_bar = "webhook_bar";
    client
        .execute(
            &format!(
                "CREATE SOURCE {src_bar} IN CLUSTER {webhook_cluster} FROM WEBHOOK BODY FORMAT TEXT"
            ),
            &[],
        )
        .await
        .expect("failed to create source");

    // Spin up tasks that will contiously push data to both webhooks.
    let addr = server.http_local_addr();
    let http_client = reqwest::Client::new();
    let keep_sending = Arc::new(AtomicBool::new(true));

    let create_poster = |name: &'static str,
                         client: reqwest::Client,
                         addr,
                         keep_sending: Arc<AtomicBool>,
                         count: Arc<AtomicUsize>| {
        let url = format!("http://{addr}/api/webhook/materialize/public/{name}");
        mz_ore::task::spawn(|| format!("POST-{name}"), async move {
            let mut tasks = Vec::new();

            while keep_sending.load(std::sync::atomic::Ordering::Relaxed) {
                let http_client = client.clone();
                let webhook_url = url.clone();
                let handle =
                    mz_ore::task::spawn(|| "webhook_concurrent_actions-appender", async move {
                        // Send all of our events to our webhook source.
                        http_client
                            .post(&webhook_url)
                            .body(name)
                            .send()
                            .await
                            .expect("failed to POST event");
                    });
                tasks.push(handle);
                count.fetch_add(1, Ordering::SeqCst);

                // Wait before sending another request.
                tokio::time::sleep(Duration::from_millis(100)).await;
            }

            // Wait for all of the inner sends to finish.
            futures::future::join_all(tasks).await;
        })
    };

    // Spawn tasks that will continuously POST events to our webhook sources.
    let foo_count = Arc::new(AtomicUsize::new(0));
    let foo_poster = create_poster(
        src_foo,
        http_client.clone(),
        addr.clone(),
        Arc::clone(&keep_sending),
        Arc::clone(&foo_count),
    );
    let bar_count = Arc::new(AtomicUsize::new(0));
    let bar_poster = create_poster(
        src_bar,
        http_client.clone(),
        addr.clone(),
        Arc::clone(&keep_sending),
        Arc::clone(&bar_count),
    );

    // Let the posting tasks run for a bit.
    tokio::time::sleep(Duration::from_secs(2)).await;

    // Swap the two sources.
    let txn = client
        .transaction()
        .await
        .expect("failed to start a transaction");

    txn.execute(&format!("ALTER SOURCE {src_foo} RENAME TO temp"), &[])
        .await
        .expect("failed to rename foo");
    txn.execute(&format!("ALTER SOURCE {src_bar} RENAME TO {src_foo}"), &[])
        .await
        .expect("failed to swap bar");
    txn.execute(&format!("ALTER SOURCE temp RENAME TO {src_bar}"), &[])
        .await
        .expect("failed to swap foo");

    txn.commit().await.expect("transaction failed");

    // Now that we've swapped the source, this is the maximum number of original events we should
    // see in each source.
    let max_foo = foo_count.load(Ordering::SeqCst);
    let max_bar = bar_count.load(Ordering::SeqCst);

    // Let the posting tasks run for a bit longer.
    tokio::time::sleep(Duration::from_secs(1)).await;

    // Stop the tasks.
    keep_sending.store(false, Ordering::SeqCst);
    let _ = foo_poster.await;
    let _ = bar_poster.await;

    let foo_row = client
        .query_one(
            &format!("SELECT COUNT(*) FROM {src_bar} WHERE body = '{src_foo}'"),
            &[],
        )
        .await
        .expect("failed to count foos");
    let foo_count: usize = foo_row.get::<_, i64>(0).try_into().expect("invalid count");

    let bar_row = client
        .query_one(
            &format!("SELECT COUNT(*) FROM {src_foo} WHERE body = '{src_bar}'"),
            &[],
        )
        .await
        .expect("failed to count foos");
    let bar_count: usize = bar_row.get::<_, i64>(0).try_into().expect("invalid count");

    assert!(foo_count <= max_foo);
    assert!(bar_count <= max_bar);

    let appender_count = server
        .metrics_registry
        .gather()
        .into_iter()
        .find(|m| m.name() == "mz_webhook_get_appender_count")
        .unwrap()
        .take_metric()[0]
        .get_counter()
        .value();

    // We should only get a webhook appender from the Coordinator 4 times, once for each source
    // when we start posting, and then once again for each source after they are renamed.
    assert_eq!(appender_count, 4.0);

    // Best effort cleanup.
    let _ = client
        .execute(&format!("DROP CLUSTER {webhook_cluster} CASCADE"), &[])
        .await;
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn copy_from() {
    let server = test_util::TestHarness::default().start_blocking();
    let mut client = server.connect(postgres::NoTls).unwrap();

    client
        .execute("CREATE TABLE copy_from_test ( x text )", &[])
        .expect("success");

    let mut writer = client
        .copy_in("COPY copy_from_test FROM STDIN (FORMAT TEXT)")
        .expect("success");
    writer
        .write_all(b"hello\nworld\n")
        .expect("write all to succeed");
    writer.finish().expect("success");

    let rows = client
        .query("SELECT * FROM copy_from_test", &[])
        .expect("success");
    assert_eq!(rows.len(), 2);
}

// Test that a cluster dropped mid transaction results in an error.
#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn concurrent_cluster_drop() {
    let server = test_util::TestHarness::default().start_blocking();
    let mut txn_client = server.connect(postgres::NoTls).unwrap();
    let mut drop_client = server.connect(postgres::NoTls).unwrap();

    txn_client
        .execute(
            "CREATE CLUSTER c REPLICAS (r1 (SIZE 'scale=1,workers=1'));",
            &[],
        )
        .expect("failed to create cluster");
    txn_client
        .execute("CREATE TABLE t (a INT);", &[])
        .expect("failed to create cluster");

    txn_client
        .execute("SET CLUSTER TO c", &[])
        .expect("success");
    txn_client.execute("BEGIN", &[]).expect("success");
    let _ = txn_client.query("SELECT * FROM t", &[]).expect("success");

    drop_client.execute("DROP CLUSTER c", &[]).expect("success");

    let err = txn_client
        .execute("SELECT * FROM t", &[])
        .expect_err("error");

    assert_eq!(
        err.as_db_error().unwrap().message(),
        "the transaction's active cluster has been dropped"
    );
}

// Test connection ID properties.
#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn test_connection_id() {
    let harness = test_util::TestHarness::default();
    let envid = harness.environment_id.organization_id().as_u128();
    let server = harness.start().await;
    let client = server.connect().await.unwrap();

    let pid: i32 = client
        .query_one("SELECT pg_backend_pid()", &[])
        .await
        .unwrap()
        .get(0);
    assert!(pid > 0);
    let envid_lower = u32::try_from(envid & 0xFFF).unwrap();
    let pid = u32::from_le_bytes(pid.to_le_bytes());
    let pid_envid = pid >> 19;
    assert_eq!(envid_lower, pid_envid);
}

// Test connection ID properties.
#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn test_github_25388() {
    let server = test_util::TestHarness::default()
        .unsafe_mode()
        .start()
        .await;
    server
        .enable_feature_flags(&["unsafe_enable_unsafe_functions"])
        .await;

    // TODO(peek-seq) The second part of this test no longer works with the new peek sequencing,
    // because we no longer check the catalog after optimization whether the original dependencies
    // still exist. This might be fine, because nothing bad happens: timestamp determination already
    // puts a a read hold on the index, so the index doesn't actually gets dropped in the
    // Controller, and therefore the peek actually succeeds. In other words, the old peek
    // sequencing's dependency check was overly cautious. I'm planning to revisit this later, and
    // probably delete the second part of the test.
    server
        .disable_feature_flags(&["enable_frontend_peek_sequencing"])
        .await;

    let client1 = server.connect().await.unwrap();

    client1
        .batch_execute("CREATE TABLE t (a int)")
        .await
        .unwrap();

    // May be flakey/racey due to various sleeps.
    Retry::default()
        .retry_async(|_| async {
            client1
                .batch_execute("DROP INDEX IF EXISTS idx")
                .await
                .unwrap();
            client1
                .batch_execute("CREATE INDEX idx ON t(a)")
                .await
                .unwrap();

            let client2 = server.connect().await.unwrap();
            mz_ore::task::spawn(|| "test", async move {
                tokio::time::sleep(Duration::from_secs(1)).await;
                client2.batch_execute("DROP INDEX idx").await.unwrap();
            });

            match client1
                .query("SUBSCRIBE (SELECT *, mz_unsafe.mz_sleep(2) FROM t)", &[])
                .await
            {
                Ok(_) => Err("unexpected query success".to_string()),
                Err(err) if err.to_string_with_causes().contains("was dropped") => Ok(()),
                Err(err) => Err(err.to_string_with_causes()),
            }
        })
        .await
        .unwrap();

    Retry::default()
        .retry_async(|_| async {
            client1
                .batch_execute("DROP INDEX IF EXISTS idx")
                .await
                .unwrap();
            client1
                .batch_execute("CREATE INDEX idx ON t(a)")
                .await
                .unwrap();

            let client2 = server.connect().await.unwrap();
            mz_ore::task::spawn(|| "test", async move {
                tokio::time::sleep(Duration::from_secs(1)).await;
                client2.batch_execute("DROP INDEX idx").await.unwrap();
            });

            match client1
                .query("SELECT *, mz_unsafe.mz_sleep(2) FROM t", &[])
                .await
            {
                Ok(_) => Err("unexpected query success".to_string()),
                Err(err) if err.to_string_with_causes().contains("was dropped") => Ok(()),
                Err(err) => Err(err.to_string_with_causes()),
            }
        })
        .await
        .unwrap();
}

#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn test_webhook_source_batch_interval() {
    let server = test_util::TestHarness::default().start().await;
    let client = server.connect().await.unwrap();

    #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize, Clone)]
    struct WebhookEvent {
        name: &'static str,
    }

    // Create a webhook source.
    client
        .execute(
            "CREATE SOURCE webhook_batch_interval_test FROM WEBHOOK BODY FORMAT JSON",
            &[],
        )
        .await
        .expect("failed to create source");

    let http_client = reqwest::Client::new();
    let webhook_url = format!(
        "http://{}/api/webhook/materialize/public/webhook_batch_interval_test",
        server.http_local_addr()
    );

    // Send one event to our webhook source.
    let start = Instant::now();
    let event = WebhookEvent { name: "first" };
    let resp = http_client
        .post(&webhook_url)
        .json(&event)
        .send()
        .await
        .expect("failed to POST event");
    let first_duration = start.elapsed();
    assert!(resp.status().is_success());
    // All is normal, the request should be fast.
    assert!(first_duration < Duration::from_secs(4));

    // Change our batch interval to be very high.
    let sys_client = server
        .connect()
        .internal()
        .user(&SYSTEM_USER.name)
        .await
        .unwrap();
    sys_client
        .execute(
            "ALTER SYSTEM SET user_storage_managed_collections_batch_duration TO '10s'",
            &[],
        )
        .await
        .expect("failed to set batch duration");

    // Send a second event...
    let start = Instant::now();
    let event = WebhookEvent { name: "second" };
    let resp = http_client
        .post(&webhook_url)
        .json(&event)
        .send()
        .await
        .expect("failed to POST event");
    let second_duration = start.elapsed();
    assert!(resp.status().is_success());

    // ... and then quickly follow up with a third.
    let start = Instant::now();
    let event = WebhookEvent { name: "third" };
    let resp = http_client
        .post(&webhook_url)
        .json(&event)
        .send()
        .await
        .expect("failed to POST event");
    let third_duration = start.elapsed();
    assert!(resp.status().is_success());

    // The second event should finish quickly like the first!
    assert!(second_duration < Duration::from_secs(4));
    // But the third will now be stuck behind the batch interval so it will take longer!
    assert!(third_duration > Duration::from_secs(7));
}

#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[cfg_attr(miri, ignore)] // too slow
async fn test_startup_cluster_notice_with_http_options() {
    let server = test_util::TestHarness::default().start().await;

    let http_client = reqwest::Client::new();

    let http_url = Url::parse(&format!("http://{}/api/sql", server.http_local_addr(),)).unwrap();

    let query = serde_json::json!({
        "query": "SHOW cluster"
    });

    // Specify a non-existant cluster should not show a notice.
    let options = serde_json::json!({
        "cluster": "i_do_not_exist"
    });
    let options = options.to_string();
    let mut non_existant_cluster_url = http_url.clone();
    non_existant_cluster_url
        .query_pairs_mut()
        .append_pair("options", &options);
    let result: serde_json::Value = http_client
        .post(non_existant_cluster_url)
        .json(&query)
        .send()
        .await
        .expect("failed to send SQL")
        .json()
        .await
        .expect("invalid JSON");
    let result = result
        .get("results")
        .expect("success")
        .as_array()
        .expect("results should be an array")
        .get(0)
        .expect("should have 1 element");

    let query_result = result.get("rows").expect("rows should exist");
    insta::assert_json_snapshot!(query_result, @r###"
    [
      [
        "i_do_not_exist"
      ]
    ]
    "###);

    // We have a single notice because we ran a query, not because we connected.
    let notices = result.get("notices").expect("notices should exist");
    insta::assert_json_snapshot!(notices, @r###"
    [
      {
        "message": "session default cluster \"i_do_not_exist\" does not exist",
        "code": "MZ005",
        "severity": "notice",
        "hint": "Pick an extant cluster with SET CLUSTER = name. Run SHOW CLUSTERS to see available clusters."
      },
      {
        "message": "cluster \"i_do_not_exist\" does not exist",
        "code": "MZ007",
        "severity": "notice",
        "hint": "Create the cluster with CREATE CLUSTER or pick an extant cluster with SET CLUSTER = name. List available clusters with SHOW CLUSTERS."
      }
    ]
    "###);
}

#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn test_startup_cluster_notice() {
    let server = test_util::TestHarness::default().start().await;

    let notices = Arc::new(Mutex::new(Vec::new()));
    let notices_ = Arc::clone(&notices);
    let client = server
        .connect()
        .notice_callback(move |notice| notices_.lock().expect("not poisoned").push(notice))
        .await
        .expect("success");
    let notices = { notices.lock().expect("not poisoned").clone() };
    assert!(notices.is_empty());
    drop(client);

    let sys_client = server
        .connect()
        .internal()
        .user(&SYSTEM_USER.name)
        .await
        .unwrap();
    sys_client
        .execute("DROP CLUSTER quickstart", &[])
        .await
        .expect("able to drop cluster");

    // Reconnecting without a system cluster should get us a notice.
    let notices = Arc::new(Mutex::new(Vec::new()));
    let notices_ = Arc::clone(&notices);
    let client = server
        .connect()
        .notice_callback(move |notice| notices_.lock().expect("not poisoned").push(notice))
        .await
        .expect("success");
    // Give the server a moment to flush all notices.
    tokio::time::sleep(Duration::from_secs(2)).await;

    let notices = { notices.lock().expect("not poisoned").clone() };
    insta::assert_debug_snapshot!(notices, @r###"
    [
        DbError {
            severity: "NOTICE",
            parsed_severity: None,
            code: SqlState(
                Other(
                    "MZ005",
                ),
            ),
            message: "system default cluster \"quickstart\" does not exist",
            detail: None,
            hint: Some(
                "Set a default cluster for the current role with `ALTER ROLE <role> SET cluster TO <cluster>;`.",
            ),
            position: None,
            where_: None,
            schema: None,
            table: None,
            column: None,
            datatype: None,
            constraint: None,
            file: None,
            line: None,
            routine: None,
        },
    ]
    "###);

    client
        .execute("ALTER ROLE materialize SET CLUSTER to non_existant", &[])
        .await
        .expect("able to set cluster");
    drop(client);

    // If we have a role default cluster, but it doesn't exist, we should still get a notice.
    let notices = Arc::new(Mutex::new(Vec::new()));
    let notices_ = Arc::clone(&notices);
    let _client = server
        .connect()
        .notice_callback(move |notice| notices_.lock().expect("not poisoned").push(notice))
        .await
        .expect("success");
    // Give the server a moment to flush all notices.
    tokio::time::sleep(Duration::from_secs(2)).await;

    let notices = { notices.lock().expect("not poisoned").clone() };
    insta::assert_debug_snapshot!(notices, @r###"
    [
        DbError {
            severity: "NOTICE",
            parsed_severity: None,
            code: SqlState(
                Other(
                    "MZ005",
                ),
            ),
            message: "role default cluster \"non_existant\" does not exist",
            detail: None,
            hint: Some(
                "Change the default cluster for the current role with `ALTER ROLE <role> SET cluster TO <cluster>;`.",
            ),
            position: None,
            where_: None,
            schema: None,
            table: None,
            column: None,
            datatype: None,
            constraint: None,
            file: None,
            line: None,
            routine: None,
        },
    ]
    "###);
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn test_durable_oids() {
    let data_dir = tempfile::tempdir().unwrap();
    let harness = test_util::TestHarness::default().data_directory(data_dir.path());

    let table_oid: u32 = {
        let server = harness.clone().start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();
        client
            .execute("CREATE TABLE t (a INT);", &[])
            .expect("failed to create table");
        client
            .query_one("SELECT oid FROM mz_tables WHERE name = 't'", &[])
            .expect("failed to select")
            .get(0)
    };

    {
        let server = harness.clone().start_blocking();
        let mut client = server.connect(postgres::NoTls).unwrap();

        let restarted_table_oid: u32 = client
            .query_one("SELECT oid FROM mz_tables WHERE name = 't'", &[])
            .expect("failed to select")
            .get(0);
        assert_eq!(table_oid, restarted_table_oid);

        client
            .execute("CREATE VIEW v AS SELECT 1;", &[])
            .expect("failed to create table");
        let view_oid: u32 = client
            .query_one("SELECT oid FROM mz_views WHERE name = 'v'", &[])
            .expect("failed to select")
            .get(0);
        assert_ne!(table_oid, view_oid);
    }
}

#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn test_double_encoded_json() {
    let server = test_util::TestHarness::default().start().await;
    let client = server.connect().await.expect("success");

    client
        .execute("CREATE TABLE t1 (a jsonb)", &[])
        .await
        .unwrap();
    client
        .execute("GRANT ALL PRIVILEGES ON TABLE t1 TO PUBLIC", &[])
        .await
        .unwrap();

    client
        .execute(r#"INSERT INTO t1 VALUES ('{ "type": "foo" }')"#, &[])
        .await
        .unwrap();
    client
        .execute(
            r#"INSERT INTO t1 VALUES ('" { \"type\": \"foo\" } "')"#,
            &[],
        )
        .await
        .unwrap();

    let http_client = reqwest::Client::new();
    let http_url = Url::parse(&format!("http://{}/api/sql", server.http_local_addr(),)).unwrap();
    let query = serde_json::json!({
        "query": "SELECT a FROM t1 ORDER BY a"
    });

    let result: serde_json::Value = http_client
        .post(http_url)
        .json(&query)
        .send()
        .await
        .expect("failed to send SQL")
        .json()
        .await
        .expect("invalid JSON");
    let rows = result
        .as_object()
        .unwrap()
        .get("results")
        .unwrap()
        .as_array()
        .unwrap()
        .get(0)
        .unwrap()
        .as_object()
        .unwrap()
        .get("rows")
        .unwrap();

    insta::assert_debug_snapshot!(rows, @r###"
    Array [
        Array [
            String(" { \"type\": \"foo\" } "),
        ],
        Array [
            Object {
                "type": String("foo"),
            },
        ],
    ]
    "###);

    let ws_url = server.ws_addr();
    let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap();
    let _ws_init = test_util::auth_with_ws(&mut ws, BTreeMap::default()).unwrap();

    let json = "{\"query\":\"SELECT a FROM t1 ORDER BY a;\"}";
    let json: serde_json::Value = serde_json::from_str(json).unwrap();
    ws.send(Message::Text(json.to_string().into())).unwrap();

    let mut read_msg = || -> WebSocketResponse {
        loop {
            let msg = ws.read().unwrap();
            if let Message::Text(text) = msg {
                return serde_json::from_str(&text).unwrap();
            }
        }
    };
    let _starting = read_msg();
    let _columns = read_msg();

    let row_1 = read_msg();
    insta::assert_debug_snapshot!(row_1, @r###"
    Row(
        [
            String(" { \"type\": \"foo\" } "),
        ],
    )
    "###);
    let row_2 = read_msg();
    insta::assert_debug_snapshot!(row_2, @r###"
    Row(
        [
            Object {
                "type": String("foo"),
            },
        ],
    )
    "###);
}

// Tests cert reloading of environmentd.
#[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
async fn test_cert_reloading() {
    let ca = Ca::new_root("test ca").unwrap();
    let (server_cert, server_key) = ca
        .request_cert("server", vec![IpAddr::V4(Ipv4Addr::LOCALHOST)])
        .unwrap();
    let metrics_registry = MetricsRegistry::new();

    let tenant_id = Uuid::new_v4();
    let email = "user@_.com".to_string();
    let password = Uuid::new_v4().to_string();
    let client_id = Uuid::new_v4();
    let secret = Uuid::new_v4();
    let initial_api_tokens = vec![ApiToken {
        client_id: client_id.clone(),
        secret: secret.clone(),
        description: None,
        created_at: Utc::now(),
    }];
    let roles = Vec::new();
    let users = BTreeMap::from([(
        email.clone(),
        UserConfig {
            id: Uuid::new_v4(),
            email,
            password,
            tenant_id,
            initial_api_tokens,
            roles,
            auth_provider: None,
            verified: None,
            metadata: None,
        },
    )]);

    let issuer = "frontegg-mock".to_owned();
    let encoding_key =
        EncodingKey::from_rsa_pem(&ca.pkey.private_key_to_pem_pkcs8().unwrap()).unwrap();
    let decoding_key = DecodingKey::from_rsa_pem(&ca.pkey.public_key_to_pem().unwrap()).unwrap();

    const EXPIRES_IN_SECS: i64 = 50;
    let frontegg_server = FronteggMockServer::start(
        None,
        issuer,
        encoding_key,
        decoding_key,
        users,
        BTreeMap::new(),
        None,
        SYSTEM_TIME.clone(),
        EXPIRES_IN_SECS,
        // Add a bit of delay so we can test connection de-duplication.
        Some(Duration::from_millis(100)),
        None,
    )
    .await
    .unwrap();

    let frontegg_auth = FronteggAuthentication::new(
        FronteggConfig {
            admin_api_token_url: frontegg_server.auth_api_token_url(),
            decoding_key: DecodingKey::from_rsa_pem(&ca.pkey.public_key_to_pem().unwrap()).unwrap(),
            tenant_id: Some(tenant_id),
            now: SYSTEM_TIME.clone(),
            admin_role: "mzadmin".to_string(),
            refresh_drop_lru_size: DEFAULT_REFRESH_DROP_LRU_CACHE_SIZE,
            refresh_drop_factor: DEFAULT_REFRESH_DROP_FACTOR,
        },
        mz_frontegg_auth::Client::default(),
        &metrics_registry,
    );
    let frontegg_user = "user@_.com";
    let frontegg_password = format!("mzp_{client_id}{secret}");

    let (mut reload_tx, reload_rx) = futures::channel::mpsc::channel(1);
    let reload_certs = Box::pin(reload_rx);

    let config = test_util::TestHarness::default()
        // Enable SSL on the main port. There should be a balancerd port with no SSL.
        .with_tls(server_cert.clone(), server_key.clone())
        .with_frontegg_auth(&frontegg_auth)
        .with_metrics_registry(metrics_registry);
    let envd_server = config.start_with_trigger(reload_certs).await;

    let body = r#"{"query": "select 12234"}"#;
    let ca_cert = reqwest::Certificate::from_pem(&ca.cert.to_pem().unwrap()).unwrap();
    let client = reqwest::Client::builder()
        .add_root_certificate(ca_cert)
        // No pool so that connections are never re-used which can use old ssl certs.
        .pool_max_idle_per_host(0)
        .tls_info(true)
        .build()
        .unwrap();

    let conn_str = Arc::new(format!(
        "user={frontegg_user} password={frontegg_password} host={} port={} sslmode=require",
        envd_server.sql_local_addr().ip(),
        envd_server.sql_local_addr().port()
    ));

    /// Asserts that the postgres connection provides the expected server-side certificate.
    async fn check_pgwire(conn_str: &str, ca_cert_path: &PathBuf, expected_cert: X509) {
        let tls = make_pg_tls(Box::new(move |b: &mut SslConnectorBuilder| {
            b.set_ca_file(ca_cert_path).unwrap();
            b.set_verify_callback(SslVerifyMode::all(), move |verify_success, x509store| {
                assert!(verify_success);
                for cert in x509store.chain().unwrap() {
                    // Expect exactly one cert to be the expected one.
                    if *cert == expected_cert {
                        return true;
                    }
                }
                false
            });
            Ok(())
        }));

        let (pg_client, conn) = tokio_postgres::connect(conn_str, tls.clone())
            .await
            .unwrap();
        task::spawn(|| "pg_client", async move {
            let _ = conn.await;
        });

        assert_eq!(
            pg_client
                .query_one("SELECT 2", &[])
                .await
                .unwrap()
                .get::<_, i32>(0),
            2
        );
    }

    // Various tests about reloading of certs.

    // Assert the current certificate is as expected.
    let https_url = format!(
        "https://{addr}/api/sql",
        addr = envd_server.http_local_addr(),
    );
    let resp = client
        .post(&https_url)
        .header("Content-Type", "application/json")
        .basic_auth(frontegg_user, Some(&frontegg_password))
        .body(body)
        .send()
        .await
        .unwrap();
    let tlsinfo = resp.extensions().get::<reqwest::tls::TlsInfo>().unwrap();
    let resp_x509 = X509::from_der(tlsinfo.peer_certificate().unwrap()).unwrap();
    let server_x509 = X509::from_pem(&std::fs::read(&server_cert).unwrap()).unwrap();
    assert_eq!(resp_x509, server_x509);
    assert_contains!(resp.text().await.unwrap(), "12234");
    check_pgwire(&conn_str, &ca.ca_cert_path(), server_x509.clone()).await;

    // Generate new certs. Install only the key, reload, and make sure the old cert is still in
    // use.
    let (next_cert, next_key) = ca
        .request_cert("next", vec![IpAddr::V4(Ipv4Addr::LOCALHOST)])
        .unwrap();
    let next_x509 = X509::from_pem(&std::fs::read(&next_cert).unwrap()).unwrap();
    assert_ne!(next_x509, server_x509);
    std::fs::copy(next_key, &server_key).unwrap();
    let (tx, rx) = oneshot::channel();
    reload_tx.try_send(Some(tx)).unwrap();
    let res = rx.await.unwrap();
    assert_err!(res);

    // We should still be on the old cert because now the cert and key mismatch.
    let resp = client
        .post(&https_url)
        .header("Content-Type", "application/json")
        .basic_auth(frontegg_user, Some(&frontegg_password))
        .body(body)
        .send()
        .await
        .unwrap();
    let tlsinfo = resp.extensions().get::<reqwest::tls::TlsInfo>().unwrap();
    let resp_x509 = X509::from_der(tlsinfo.peer_certificate().unwrap()).unwrap();
    assert_eq!(resp_x509, server_x509);
    check_pgwire(&conn_str, &ca.ca_cert_path(), server_x509.clone()).await;

    // Now move the cert too. Reloading should succeed and the response should have the new
    // cert.
    std::fs::copy(next_cert, &server_cert).unwrap();
    let (tx, rx) = oneshot::channel();
    reload_tx.try_send(Some(tx)).unwrap();
    let res = rx.await.unwrap();
    assert_ok!(res);
    let resp = client
        .post(&https_url)
        .header("Content-Type", "application/json")
        .basic_auth(frontegg_user, Some(&frontegg_password))
        .body(body)
        .send()
        .await
        .unwrap();
    let tlsinfo = resp.extensions().get::<reqwest::tls::TlsInfo>().unwrap();
    let resp_x509 = X509::from_der(tlsinfo.peer_certificate().unwrap()).unwrap();
    assert_eq!(resp_x509, next_x509);
    check_pgwire(&conn_str, &ca.ca_cert_path(), next_x509.clone()).await;
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn test_builtin_connection_alterations_are_preserved_across_restarts() {
    let data_dir = tempfile::tempdir().unwrap();
    let harness = test_util::TestHarness::default()
        .data_directory(data_dir.path())
        .unsafe_mode();

    {
        let server = harness.clone().start_blocking();
        let mut client = server
            .pg_config_internal()
            .user(&ANALYTICS_USER.name)
            .connect(postgres::NoTls)
            .unwrap();

        let row = client
            .query_one(
                "SELECT create_sql
            FROM mz_catalog.mz_connections
            WHERE name = 'mz_analytics'",
                &[],
            )
            .expect("success");
        let sql: String = row.get("create_sql");
        assert_eq!(
            sql,
            "CREATE CONNECTION \"mz_internal\".\"mz_analytics\" TO AWS (ASSUME ROLE ARN = '')"
        );
    }

    {
        let server = harness.clone().start_blocking();
        let mut client = server
            .pg_config_internal()
            .user(&ANALYTICS_USER.name)
            .connect(postgres::NoTls)
            .unwrap();

        let row = client
            .query_one(
                "SELECT create_sql
            FROM mz_catalog.mz_connections
            WHERE name = 'mz_analytics'",
                &[],
            )
            .expect("success");
        let sql: String = row.get("create_sql");
        assert_eq!(
            sql,
            "CREATE CONNECTION \"mz_internal\".\"mz_analytics\" TO AWS (ASSUME ROLE ARN = '')"
        );

        client
            .execute(
                "ALTER CONNECTION mz_internal.mz_analytics SET (ASSUME ROLE ARN = 'foo')",
                &[],
            )
            .expect("success");

        let row = client
            .query_one(
                "SELECT create_sql
            FROM mz_catalog.mz_connections
            WHERE name = 'mz_analytics'",
                &[],
            )
            .expect("success");
        let sql: String = row.get("create_sql");
        assert_eq!(
            sql,
            "CREATE CONNECTION \"mz_internal\".\"mz_analytics\" TO AWS (ASSUME ROLE ARN = 'foo')"
        );
    }

    {
        let server = harness.clone().start_blocking();
        let mut client = server
            .pg_config_internal()
            .user(&ANALYTICS_USER.name)
            .connect(postgres::NoTls)
            .unwrap();

        let row = client
            .query_one(
                "SELECT create_sql
            FROM mz_catalog.mz_connections
            WHERE name = 'mz_analytics'",
                &[],
            )
            .expect("success");
        let sql: String = row.get("create_sql");
        assert_eq!(
            sql,
            "CREATE CONNECTION \"mz_internal\".\"mz_analytics\" TO AWS (ASSUME ROLE ARN = 'foo')"
        );
    }
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
#[allow(clippy::disallowed_methods)]
fn test_webhook_request_compression() {
    let server = test_util::TestHarness::default().start_blocking();
    let mut client = server.connect(postgres::NoTls).unwrap();

    // Create a webhook source.
    client
        .execute(
            "CREATE CLUSTER webhook_cluster REPLICAS (r1 (SIZE 'scale=1,workers=1'));",
            &[],
        )
        .expect("failed to create cluster");
    client
        .execute(
            "CREATE SOURCE webhook_text IN CLUSTER webhook_cluster FROM WEBHOOK BODY FORMAT TEXT",
            &[],
        )
        .expect("failed to create source");

    let http_client = Client::new();
    let webhook_url = format!(
        "http://{}/api/webhook/materialize/public/webhook_text",
        server.http_local_addr(),
    );

    let og_body = "hello world!";

    let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
    encoder.write_all(og_body.as_bytes()).unwrap();
    let compressed_body = encoder.finish().unwrap();

    assert!(og_body.as_bytes().len() < compressed_body.len());

    let resp = http_client
        .post(&webhook_url)
        .header(CONTENT_ENCODING, "gzip")
        .body(compressed_body)
        .send()
        .expect("failed to POST event");
    assert!(resp.status().is_success(), "{resp:?}");

    // Wait for up to 10s for the webhook to get appended.
    let mut row = None;
    for _ in 0..10 {
        std::thread::sleep(std::time::Duration::from_secs(1));

        row = client
            .query_opt("SELECT body FROM webhook_text", &[])
            .unwrap();
        if row.is_some() {
            break;
        }
    }

    let rnd_body = row.as_ref().map(|r| r.get("body"));
    assert_eq!(rnd_body, Some(og_body));
}

// =============================================================================
// MCP (Model Context Protocol) integration tests
// =============================================================================

/// Helper to set up an MCP test server and run datadriven tests.
#[allow(clippy::disallowed_methods)]
fn run_mcp_datadriven(testdata_path: &str, harness: test_util::TestHarness) {
    run_mcp_datadriven_inner(testdata_path, harness, false)
}

/// Same as [`run_mcp_datadriven`], but additionally sets
/// `restrict_to_user_objects = true` as a default for the HTTP user, so
/// every MCP HTTP request session starts with the restriction active.
#[allow(clippy::disallowed_methods)]
fn run_mcp_datadriven_restricted(testdata_path: &str, harness: test_util::TestHarness) {
    run_mcp_datadriven_inner(testdata_path, harness, true)
}

#[allow(clippy::disallowed_methods)]
fn run_mcp_datadriven_inner(
    testdata_path: &str,
    harness: test_util::TestHarness,
    restrict_to_user_objects: bool,
) {
    let version_re = Regex::new(r#"\d+\.\d+\.\d+(\.\d+)?(-(dev|rc)(\.\d+)?)?"#).unwrap();

    datadriven::walk(testdata_path, |f| {
        let server = harness.clone().start_blocking();

        // Grant all privileges to default http user (same setup as test_http_sql).
        {
            let mut super_user = server
                .pg_config_internal()
                .user(&SYSTEM_USER.name)
                .connect(postgres::NoTls)
                .unwrap();
            super_user
                .batch_execute(&format!("CREATE ROLE {}", &HTTP_DEFAULT_USER.name))
                .unwrap();
            super_user
                .batch_execute(&format!(
                    "GRANT ALL PRIVILEGES ON SYSTEM TO {}",
                    &HTTP_DEFAULT_USER.name
                ))
                .unwrap();
            super_user
                .batch_execute(&format!(
                    "GRANT ALL PRIVILEGES ON CLUSTER quickstart TO {}",
                    &HTTP_DEFAULT_USER.name
                ))
                .unwrap();
            super_user
                .batch_execute(&format!(
                    "GRANT ALL PRIVILEGES ON DATABASE materialize TO {}",
                    &HTTP_DEFAULT_USER.name
                ))
                .unwrap();
            super_user
                .batch_execute(&format!(
                    "GRANT ALL PRIVILEGES ON SCHEMA materialize.public TO {}",
                    &HTTP_DEFAULT_USER.name
                ))
                .unwrap();

            if restrict_to_user_objects {
                // restrict_to_user_objects only takes effect on session start,
                // so it must be set as a role default; the per-request MCP
                // sessions pick it up automatically.
                super_user
                    .batch_execute("ALTER SYSTEM SET enable_rbac_checks TO true")
                    .unwrap();
                super_user
                    .batch_execute(&format!(
                        "ALTER ROLE {} SET restrict_to_user_objects = true",
                        &HTTP_DEFAULT_USER.name
                    ))
                    .unwrap();
            }
        }

        let agents_url = format!("http://{}/api/mcp/agent", server.http_local_addr());
        let developer_url = format!("http://{}/api/mcp/developer", server.http_local_addr());

        f.run(|tc| {
            let url = match tc.directive.as_str() {
                "mcp-agent" => &agents_url,
                "mcp-developer" => &developer_url,
                other => panic!("unknown directive: {}", other),
            };

            let json: serde_json::Value = serde_json::from_str(&tc.input).unwrap();
            let res = Client::new().post(url).json(&json).send().unwrap();

            let status = res.status();
            let body = res.text().unwrap();

            // Replace version numbers for stable output.
            let body = version_re.replace_all(&body, "<VERSION>").to_string();

            if body.is_empty() {
                format!("{}\n", status)
            } else {
                format!("{}\n{}\n", status, body)
            }
        });
    });
}

/// Tests the MCP agent endpoint with the query tool explicitly disabled.
#[mz_ore::test]
fn test_mcp_agent() {
    let harness = test_util::TestHarness::default()
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .with_system_parameter_default(
            "enable_mcp_agent_query_tool".to_string(),
            "false".to_string(),
        );
    run_mcp_datadriven("tests/testdata/mcp/agent", harness);
}

/// Tests the MCP agent endpoint with the query tool enabled (default behavior).
#[mz_ore::test]
fn test_mcp_agent_query_tool() {
    let harness = test_util::TestHarness::default()
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string());
    run_mcp_datadriven("tests/testdata/mcp/agent_query_tool", harness);
}

/// Tests the MCP agent endpoint under `restrict_to_user_objects = true`,
/// the setting clients use to scope sessions down to user objects only.
/// Covers a small, focused set of cases; broader agent coverage lives in
/// `test_mcp_agent` / `test_mcp_agent_query_tool`.
#[mz_ore::test]
fn test_mcp_agent_restricted() {
    let harness = test_util::TestHarness::default()
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string());
    run_mcp_datadriven_restricted("tests/testdata/mcp/agent_restricted", harness);
}

/// Tests that the MCP agent endpoint returns 503 when the feature flag is disabled.
#[mz_ore::test]
fn test_mcp_agent_disabled() {
    let harness = test_util::TestHarness::default()
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "false".to_string());
    run_mcp_datadriven("tests/testdata/mcp/agent_disabled", harness);
}

/// Tests the MCP developer endpoint.
#[mz_ore::test]
fn test_mcp_developer() {
    let harness = test_util::TestHarness::default()
        .with_mcp_routes(false, true)
        .with_system_parameter_default("enable_mcp_developer".to_string(), "true".to_string());
    run_mcp_datadriven("tests/testdata/mcp/developer", harness);
}

/// Tests the MCP developer endpoint when disabled (503).
#[mz_ore::test]
fn test_mcp_developer_disabled() {
    let harness = test_util::TestHarness::default()
        .with_mcp_routes(false, true)
        .with_system_parameter_default("enable_mcp_developer".to_string(), "false".to_string());
    run_mcp_datadriven("tests/testdata/mcp/developer_disabled", harness);
}

/// Regression test for database-issues#11320.
///
/// The developer endpoint validator allows unqualified `mz_*` table names as
/// a UX convenience. Before the fix, an attacker with CREATE privileges could
/// create `public.mz_leak` pointing at sensitive data and the session's
/// `search_path` would resolve the unqualified name to that view, bypassing
/// the system-catalog-only restriction. The fix sets a tight `search_path`
/// containing only system schemas before executing the query, so `mz_leak`
/// cannot resolve to a user-created object.
#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_mcp_developer_search_path_defense() {
    let server = test_util::TestHarness::default()
        .with_mcp_routes(false, true)
        .with_system_parameter_default("enable_mcp_developer".to_string(), "true".to_string())
        .start_blocking();

    let developer_url = format!("http://{}/api/mcp/developer", server.http_local_addr());

    // Set up a user view named `mz_leak` in the `public` schema and point the
    // HTTP user's search_path at `public` so an unqualified reference would
    // normally resolve to it.
    {
        let mut super_user = server
            .pg_config_internal()
            .user(&SYSTEM_USER.name)
            .connect(postgres::NoTls)
            .unwrap();

        super_user
            .batch_execute(&format!("CREATE ROLE {}", &HTTP_DEFAULT_USER.name))
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT ALL PRIVILEGES ON SYSTEM TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT ALL PRIVILEGES ON DATABASE materialize TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT ALL PRIVILEGES ON SCHEMA materialize.public TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();

        // Attacker-created view with the `mz_` prefix.
        super_user
            .batch_execute("CREATE VIEW public.mz_leak AS SELECT 'leaked_secret'::text AS payload")
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT SELECT ON public.mz_leak TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
        super_user
            .batch_execute(&format!(
                "ALTER ROLE {} SET search_path TO public",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
    }

    // Attempt to leak: unqualified `mz_leak` must NOT resolve to `public.mz_leak`.
    let (status, body) = mcp_post(
        &developer_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/call",
            "params": {
                "name": "query_system_catalog",
                "arguments": {"sql_query": "SELECT * FROM mz_leak"}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"].is_object(),
        "unqualified mz_leak should not resolve to the user view, got: {body}"
    );
    let result_text = body
        .get("result")
        .and_then(|r| r["content"].get(0))
        .and_then(|c| c["text"].as_str())
        .unwrap_or("");
    assert!(
        !result_text.contains("leaked_secret"),
        "user view contents must not leak through MCP, got: {body}"
    );

    // Explicit qualification with a non-system schema is still rejected by the
    // validator, regardless of search_path.
    let (status, body) = mcp_post(
        &developer_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 2,
            "method": "tools/call",
            "params": {
                "name": "query_system_catalog",
                "arguments": {"sql_query": "SELECT * FROM public.mz_leak"}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"]["message"]
            .as_str()
            .unwrap_or("")
            .contains("non-system"),
        "public.mz_leak should be rejected by the validator, got: {body}"
    );

    // Legitimate unqualified system queries must still work: the tight
    // search_path resolves `mz_tables` to `mz_catalog.mz_tables`.
    let (status, body) = mcp_post(
        &developer_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {
                "name": "query_system_catalog",
                "arguments": {"sql_query": "SELECT name FROM mz_databases WHERE name = 'materialize'"}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"].is_null(),
        "legitimate unqualified system query should succeed, got: {body}"
    );
    let result_text = body["result"]["content"][0]["text"].as_str().unwrap();
    assert!(
        result_text.contains("materialize"),
        "query should return the default database, got: {result_text}"
    );
}

/// Helper to POST a JSON-RPC request to an MCP endpoint and return the parsed response.
fn mcp_post(url: &str, json: serde_json::Value) -> (reqwest::StatusCode, serde_json::Value) {
    let res = Client::new().post(url).json(&json).send().unwrap();
    let status = res.status();
    let body: serde_json::Value = res.json().unwrap_or(serde_json::Value::Null);
    (status, body)
}

/// Tests get_data_products, get_data_product_details, and read_data_product against
/// a real data product (view + index + comment).
#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_mcp_agent_with_data_product() {
    let server = test_util::TestHarness::default()
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .start_blocking();

    let agents_url = format!("http://{}/api/mcp/agent", server.http_local_addr());

    // Set up a data product: create a view, index it, and add a comment on the index.
    {
        let mut super_user = server
            .pg_config_internal()
            .user(&SYSTEM_USER.name)
            .connect(postgres::NoTls)
            .unwrap();

        // Create the HTTP user and grant privileges.
        super_user
            .batch_execute(&format!("CREATE ROLE {}", &HTTP_DEFAULT_USER.name))
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT ALL PRIVILEGES ON SYSTEM TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT ALL PRIVILEGES ON CLUSTER quickstart TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT ALL PRIVILEGES ON DATABASE materialize TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT ALL PRIVILEGES ON SCHEMA materialize.public TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();

        // Create a materialized view and an index on it.
        super_user
            .batch_execute(
                "CREATE MATERIALIZED VIEW test_products IN CLUSTER quickstart AS SELECT 1::int AS id, 'widget'::text AS name",
            )
            .unwrap();
        super_user
            .batch_execute(
                "CREATE INDEX test_products_idx IN CLUSTER quickstart ON test_products (id)",
            )
            .unwrap();
        // Comment on the index enriches the data product description.
        super_user
            .batch_execute("COMMENT ON INDEX test_products_idx IS 'A test data product for integration testing'")
            .unwrap();
        // Grant SELECT on the view to the HTTP user so it appears in mz_show_my_object_privileges.
        super_user
            .batch_execute(&format!(
                "GRANT SELECT ON test_products TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
        // Grant USAGE on the cluster so it appears in mz_show_my_cluster_privileges.
        super_user
            .batch_execute(&format!(
                "GRANT USAGE ON CLUSTER quickstart TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();

        // Indexed regular view: cheap to query via the in-memory arrangement,
        // should appear as a data product. Comment on the view (not the index)
        // exercises the object-comment fallback path.
        super_user
            .batch_execute(
                "CREATE VIEW test_indexed_view AS SELECT 2::int AS id, 'gadget'::text AS name",
            )
            .unwrap();
        super_user
            .batch_execute(
                "CREATE INDEX test_indexed_view_idx IN CLUSTER quickstart ON test_indexed_view (id)",
            )
            .unwrap();
        super_user
            .batch_execute("COMMENT ON VIEW test_indexed_view IS 'View-level description'")
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT SELECT ON test_indexed_view TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();

        // Non-indexed regular view: would trigger full recompute on query,
        // must NOT appear as a data product.
        super_user
            .batch_execute(
                "CREATE VIEW test_unindexed_view AS SELECT 3::int AS id, 'sprocket'::text AS name",
            )
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT SELECT ON test_unindexed_view TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();

        // Materialized view without an explicit index: bounded query cost
        // (Persist-backed), should appear as a data product. Comment on the MV
        // itself exercises the object-comment fallback path.
        super_user
            .batch_execute(
                "CREATE MATERIALIZED VIEW test_unindexed_mv IN CLUSTER quickstart AS SELECT 4::int AS id, 'bolt'::text AS name",
            )
            .unwrap();
        super_user
            .batch_execute(
                "COMMENT ON MATERIALIZED VIEW test_unindexed_mv IS 'MV-level description'",
            )
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT SELECT ON test_unindexed_mv TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
    }

    // get_data_products should now return our data product.
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/call",
            "params": {
                "name": "get_data_products",
                "arguments": {}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    let result_text = body["result"]["content"][0]["text"].as_str().unwrap();
    let products: serde_json::Value = serde_json::from_str(result_text).unwrap();
    let has_object = |needle: &str| {
        products.as_array().unwrap().iter().any(|p| {
            p.as_array()
                .map(|arr| arr[0].as_str().unwrap_or("").contains(needle))
                .unwrap_or(false)
        })
    };
    // Indexed MV: appears.
    assert!(
        has_object("test_products"),
        "indexed MV should appear as a data product"
    );
    // Indexed regular view: appears (in-memory arrangement is cheap to query).
    assert!(
        has_object("test_indexed_view"),
        "indexed view should appear as a data product"
    );
    // Non-indexed MV: appears (Persist-backed, bounded query cost).
    assert!(
        has_object("test_unindexed_mv"),
        "non-indexed materialized view should appear as a data product"
    );
    // Non-indexed regular view: excluded (would trigger full recompute).
    assert!(
        !has_object("test_unindexed_view"),
        "non-indexed view must NOT appear as a data product"
    );

    // Description fallback: index comment is preferred, object comment is used
    // when no index comment exists.
    let find_product = |needle: &str| -> &serde_json::Value {
        products
            .as_array()
            .unwrap()
            .iter()
            .find(|p| {
                p.as_array()
                    .map(|arr| arr[0].as_str().unwrap_or("").contains(needle))
                    .unwrap_or(false)
            })
            .unwrap_or_else(|| panic!("{needle} should exist as a data product"))
    };
    let description_of = |needle: &str| -> String {
        find_product(needle).as_array().unwrap()[2]
            .as_str()
            .unwrap_or("")
            .to_string()
    };
    assert_eq!(
        description_of("test_products"),
        "A test data product for integration testing",
        "indexed MV with an index comment should use the index comment"
    );
    assert_eq!(
        description_of("test_indexed_view"),
        "View-level description",
        "indexed view without an index comment should fall back to the view comment"
    );
    assert_eq!(
        description_of("test_unindexed_mv"),
        "MV-level description",
        "non-indexed MV should use its own comment as description"
    );

    // Find our product
    let product = products
        .as_array()
        .unwrap()
        .iter()
        .find(|p| {
            p.as_array()
                .map(|arr| arr[0].as_str().unwrap_or("").contains("test_products"))
                .unwrap_or(false)
        })
        .expect("test_products data product should exist");
    let object_name = product.as_array().unwrap()[0].as_str().unwrap();
    assert!(
        object_name.contains("test_products"),
        "object_name should contain test_products, got: {}",
        object_name
    );

    // get_data_product_details with the exact name.
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 2,
            "method": "tools/call",
            "params": {
                "name": "get_data_product_details",
                "arguments": {"name": object_name}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(body["result"]["content"][0]["text"].as_str().is_some());
    assert!(body["error"].is_null());

    // Each row carries a `hydration` field (5th cell). For an MV that has
    // had time to hydrate on a single-replica `quickstart` cluster, expect
    // `hydrated: true` with 1/1 replicas.
    let rows_text = body["result"]["content"][0]["text"].as_str().unwrap();
    let rows: serde_json::Value = serde_json::from_str(rows_text).unwrap();
    let rows = rows.as_array().expect("details should return rows");
    assert!(!rows.is_empty(), "details should return at least one row");
    for row in rows {
        let row = row.as_array().expect("each row should be an array");
        assert_eq!(
            row.len(),
            5,
            "each details row should have 5 cells (object_name, cluster, description, schema, hydration), got: {:?}",
            row,
        );
        let hydration = &row[4];
        assert!(
            hydration.is_object(),
            "hydration cell should be a JSON object, got: {hydration}",
        );
        assert!(
            hydration.get("hydrated").is_some_and(|v| v.is_boolean()),
            "hydration.hydrated should be a bool, got: {hydration}",
        );
        let replica_count = hydration
            .get("replica_count")
            .and_then(|v| v.as_i64())
            .unwrap_or_else(|| {
                panic!("hydration.replica_count should be an int, got: {hydration}")
            });
        let hydrated_replica_count = hydration
            .get("hydrated_replica_count")
            .and_then(|v| v.as_i64())
            .unwrap_or_else(|| {
                panic!("hydration.hydrated_replica_count should be an int, got: {hydration}")
            });
        assert!(
            replica_count >= 0 && hydrated_replica_count >= 0,
            "replica counts must be non-negative, got: {hydration}",
        );
        assert!(
            hydrated_replica_count <= replica_count,
            "hydrated_replica_count ({hydrated_replica_count}) cannot exceed replica_count ({replica_count}): {hydration}",
        );
    }

    // get_data_product_details should also resolve the indexed view, proving
    // the filter change is applied consistently to mz_mcp_data_product_details.
    let indexed_view_name = find_product("test_indexed_view").as_array().unwrap()[0]
        .as_str()
        .unwrap()
        .to_string();
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 20,
            "method": "tools/call",
            "params": {
                "name": "get_data_product_details",
                "arguments": {"name": indexed_view_name}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"].is_null(),
        "indexed view should be resolvable via get_data_product_details, got: {body}"
    );
    assert!(body["result"]["content"][0]["text"].as_str().is_some());
    // Indexed view should also report a hydration object.
    let rows_text = body["result"]["content"][0]["text"].as_str().unwrap();
    let rows: serde_json::Value = serde_json::from_str(rows_text).unwrap();
    let rows = rows.as_array().expect("details should return rows");
    assert!(!rows.is_empty());
    for row in rows {
        let row = row.as_array().expect("each row should be an array");
        assert_eq!(row.len(), 5, "row should include hydration cell: {row:?}");
        assert!(row[4].is_object(), "hydration cell should be an object");
    }

    // read_data_product should return the row from the view.
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {
                "name": "read_data_product",
                "arguments": {"name": object_name, "limit": 10}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    let rows_text = body["result"]["content"][0]["text"].as_str().unwrap();
    let rows: serde_json::Value = serde_json::from_str(rows_text).unwrap();
    assert_eq!(rows.as_array().unwrap().len(), 1, "expected 1 row");
    // The row should have id=1, name=widget
    let row = &rows[0];
    assert_eq!(row[0].as_str().unwrap(), "1");
    assert_eq!(row[1].as_str().unwrap(), "widget");

    // read_data_product with cluster override.
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 4,
            "method": "tools/call",
            "params": {
                "name": "read_data_product",
                "arguments": {"name": object_name, "limit": 10, "cluster": "quickstart"}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(body["error"].is_null(), "cluster override should work");
    let rows_text = body["result"]["content"][0]["text"].as_str().unwrap();
    let rows: serde_json::Value = serde_json::from_str(rows_text).unwrap();
    assert_eq!(rows.as_array().unwrap().len(), 1);

    // DEX-27: an MV whose home cluster is NOT the HTTP session's default
    // (`quickstart`) must still be readable without the caller passing a
    // `cluster` argument. Before the auto-routing change the read would
    // execute against `quickstart`, missing the index. Reading without an
    // override must still succeed and return the row.
    {
        let mut super_user = server
            .pg_config_internal()
            .user(&SYSTEM_USER.name)
            .connect(postgres::NoTls)
            .unwrap();
        super_user
            .batch_execute(
                "CREATE CLUSTER dex27_other_cluster REPLICAS (r1 (SIZE 'scale=1,workers=1'))",
            )
            .unwrap();
        super_user
            .batch_execute(
                "CREATE MATERIALIZED VIEW test_off_default IN CLUSTER dex27_other_cluster \
                 AS SELECT 42::int AS id, 'off_default'::text AS name",
            )
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT SELECT ON test_off_default TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
        super_user
            .batch_execute(&format!(
                "GRANT USAGE ON CLUSTER dex27_other_cluster TO {}",
                &HTTP_DEFAULT_USER.name
            ))
            .unwrap();
    }
    // The MV was created after the cached `get_data_products` snapshot, so
    // build the qualified name directly rather than re-querying.
    let off_default_name = r#""materialize"."public"."test_off_default""#.to_string();
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 41,
            "method": "tools/call",
            "params": {
                "name": "read_data_product",
                "arguments": {"name": off_default_name, "limit": 10}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"].is_null(),
        "no-override read on off-default cluster should auto-route, got: {body}",
    );
    let rows_text = body["result"]["content"][0]["text"].as_str().unwrap();
    let rows: serde_json::Value = serde_json::from_str(rows_text).unwrap();
    assert_eq!(rows.as_array().unwrap().len(), 1);
    assert_eq!(rows[0][0].as_str().unwrap(), "42");
    assert_eq!(rows[0][1].as_str().unwrap(), "off_default");

    // read_data_product with limit 0 should return no rows.
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 5,
            "method": "tools/call",
            "params": {
                "name": "read_data_product",
                "arguments": {"name": object_name, "limit": 0}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(body["error"].is_null());
    let rows_text = body["result"]["content"][0]["text"].as_str().unwrap();
    let rows: serde_json::Value = serde_json::from_str(rows_text).unwrap();
    assert_eq!(
        rows.as_array().unwrap().len(),
        0,
        "limit 0 should return no rows"
    );

    // read_data_product with a large `limit` is not clamped or rejected: the
    // response is bounded only by the size cap, matching the SQL HTTP layer.
    // The test data product has 1 row, so LIMIT 9999 just returns that 1 row.
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 6,
            "method": "tools/call",
            "params": {
                "name": "read_data_product",
                "arguments": {"name": object_name, "limit": 9999}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"].is_null(),
        "large limit should succeed; size cap is the only guard"
    );
    let rows_text = body["result"]["content"][0]["text"].as_str().unwrap();
    let rows: serde_json::Value = serde_json::from_str(rows_text).unwrap();
    assert_eq!(rows.as_array().unwrap().len(), 1);

    // SQL injection attempt in data product name: should return DataProductNotFound, not execute.
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 7,
            "method": "tools/call",
            "params": {
                "name": "get_data_product_details",
                "arguments": {"name": "'; DROP TABLE test_products; --"}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"]["message"]
            .as_str()
            .unwrap()
            .contains("Data product not found"),
        "SQL injection in name should be safely handled"
    );

    // SQL injection attempt in read_data_product name.
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 8,
            "method": "tools/call",
            "params": {
                "name": "read_data_product",
                "arguments": {"name": "x' UNION SELECT * FROM mz_tables --"}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"]["message"]
            .as_str()
            .unwrap()
            .contains("Invalid data product name"),
        "SQL injection in read_data_product name should be rejected by validation"
    );

    // SQL injection attempt in cluster override.
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 9,
            "method": "tools/call",
            "params": {
                "name": "read_data_product",
                "arguments": {"name": object_name, "cluster": "'; DROP TABLE test_products; --"}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);
    // The cluster name is escaped, so this should fail as an invalid cluster, not execute injection.
    // It may be a query execution error (bad cluster name) which is fine - the key is no injection.
    assert!(
        body["error"].is_object(),
        "injection in cluster should produce an error, not succeed"
    );
}

/// Verifies that the `WWW-Authenticate` challenge on a 401 for an MCP
/// route includes both the Bearer (with `resource_metadata`) and Basic
/// challenges, so OAuth-aware clients (Claude Desktop, ChatGPT remote
/// MCP) can discover the protected resource metadata while existing
/// callers that send Basic still see a usable challenge. This is the
/// half of RFC 9728 that doesn't require an authorization server to be
/// configured — even with no `oidc_issuer` set, the Bearer challenge
/// itself fires (clients then probe the well-known URI and get a 404,
/// which is the documented fallback path).
#[mz_ore::test]
fn test_mcp_unauth_emits_bearer_and_basic_challenges() {
    // An OIDC listener is the authenticator kind that advertises OAuth
    // (see `McpOAuthDiscovery`); an unauthenticated MCP request on it 401s
    // with both the Bearer and Basic challenges. (A `None` listener would
    // default to anonymous_http_user and never challenge.) `with_oidc_auth`
    // must precede `with_mcp_routes` because it replaces the listener config.
    let server = test_util::TestHarness::default()
        .with_oidc_auth(None, None, None, None)
        .with_mcp_routes(true, true)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .with_system_parameter_default("enable_mcp_developer".to_string(), "true".to_string())
        .start_blocking();

    for endpoint in ["agent", "developer"] {
        let url = format!("http://{}/api/mcp/{endpoint}", server.http_local_addr());
        let res = Client::new()
            .post(&url)
            .header("Content-Type", "application/json")
            .body(r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#)
            .send()
            .unwrap();
        assert_eq!(
            res.status(),
            StatusCode::UNAUTHORIZED,
            "MCP {endpoint} without auth should 401",
        );
        let challenges: Vec<String> = res
            .headers()
            .get_all(http::header::WWW_AUTHENTICATE)
            .iter()
            .map(|v| v.to_str().unwrap().to_string())
            .collect();
        assert!(
            challenges.iter().any(|c| c.starts_with("Bearer ")
                && c.contains("resource_metadata=")
                && c.contains("/.well-known/oauth-protected-resource")
                && c.contains("scope=\"mcp.read\"")),
            "expected a Bearer challenge with resource_metadata AND scope on /api/mcp/{endpoint}, got: {challenges:?}",
        );
        assert!(
            challenges.iter().any(|c| c.starts_with("Basic ")),
            "expected a Basic challenge alongside Bearer so curl/legacy callers \
             still see a usable challenge, got: {challenges:?}",
        );
    }
}

/// MCP routes on a `Password`-authenticator listener must NOT advertise
/// OAuth — discovery is scoped to `Oidc` listeners by
/// [`oauth_metadata::McpOAuthDiscovery::for_authenticator`]. Regression
/// guard for the seam: even with MCP routes enabled, a Password listener
/// emits only `Basic` on a 401 and 404s the well-known discovery URI.
#[mz_ore::test]
fn test_mcp_unauth_on_password_listener_does_not_advertise_oauth() {
    let server = test_util::TestHarness::default()
        .with_password_auth(Password("mz_system_password".to_string()))
        .with_mcp_routes(true, true)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .with_system_parameter_default("enable_mcp_developer".to_string(), "true".to_string())
        .with_system_parameter_default(
            "oidc_issuer".to_string(),
            "https://issuer.test.example.com".to_string(),
        )
        .start_blocking();

    for endpoint in ["agent", "developer"] {
        let res = Client::new()
            .post(&format!(
                "http://{}/api/mcp/{endpoint}",
                server.http_local_addr()
            ))
            .header("Content-Type", "application/json")
            .body(r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#)
            .send()
            .unwrap();
        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
        let challenges: Vec<String> = res
            .headers()
            .get_all(http::header::WWW_AUTHENTICATE)
            .iter()
            .map(|v| v.to_str().unwrap().to_string())
            .collect();
        assert!(
            challenges.iter().any(|c| c.starts_with("Basic ")),
            "Password listener should still emit Basic on /api/mcp/{endpoint}: {challenges:?}",
        );
        assert!(
            !challenges.iter().any(|c| c.starts_with("Bearer ")),
            "Password listener MUST NOT advertise OAuth on /api/mcp/{endpoint}: {challenges:?}",
        );
    }

    // The discovery endpoint must 404 on a Password listener even though
    // `oidc_issuer` is set — there is no OAuth flow this listener honours.
    let res = Client::new()
        .get(&format!(
            "http://{}/.well-known/oauth-protected-resource",
            server.http_local_addr()
        ))
        .send()
        .unwrap();
    assert_eq!(
        res.status(),
        StatusCode::NOT_FOUND,
        "discovery doc must not be published on a Password listener",
    );
}

/// The SQL HTTP layer must NOT gain the Bearer challenge — only MCP
/// routes opt into OAuth discovery. Regression guard so we don't
/// accidentally widen the OAuth surface to the rest of the HTTP server.
#[mz_ore::test]
fn test_sql_http_unauth_keeps_only_basic_challenge() {
    let server = test_util::TestHarness::default()
        .with_password_auth(Password("mz_system_password".to_string()))
        .start_blocking();

    // `/` triggers the WWW-Authenticate path (see auth_middleware) so we
    // get a 401 with the challenge header.
    let res = Client::new()
        .get(&format!("http://{}/", server.http_local_addr()))
        .send()
        .unwrap();
    assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
    let challenges: Vec<String> = res
        .headers()
        .get_all(http::header::WWW_AUTHENTICATE)
        .iter()
        .map(|v| v.to_str().unwrap().to_string())
        .collect();
    assert!(
        challenges.iter().any(|c| c.starts_with("Basic ")),
        "SQL HTTP layer should still emit Basic challenge: {challenges:?}",
    );
    assert!(
        !challenges.iter().any(|c| c.starts_with("Bearer ")),
        "SQL HTTP layer must NOT advertise OAuth: {challenges:?}",
    );
}

/// The well-known protected resource metadata document is public (no
/// auth) and returns the configured authorization server when
/// `oidc_issuer` is set. Pinning the JSON shape protects MCP clients
/// from accidental contract drift.
///
/// Uses `with_oidc_auth` because the discovery handler only publishes on
/// an OIDC listener — the one authenticator kind that validates bearer
/// tokens against `oidc_issuer` (see `McpOAuthDiscovery`). The listener
/// never validates a token in this test, so no live IdP is needed.
#[mz_ore::test]
fn test_oauth_protected_resource_metadata_advertises_issuer() {
    let server = test_util::TestHarness::default()
        .with_oidc_auth(None, None, None, None)
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .with_system_parameter_default(
            "oidc_issuer".to_string(),
            "https://issuer.test.example.com".to_string(),
        )
        .start_blocking();

    let res = Client::new()
        .get(&format!(
            "http://{}/.well-known/oauth-protected-resource",
            server.http_local_addr()
        ))
        .send()
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);
    let body: serde_json::Value = res.json().unwrap();

    // `resource` is built from the request host and must point at the
    // MCP path prefix — clients use this as the `aud` value in their
    // RFC 8707 resource indicator.
    let resource = body["resource"].as_str().expect("resource is a string");
    assert!(
        resource.ends_with("/api/mcp"),
        "resource should point at /api/mcp prefix, got: {resource}",
    );

    let authorization_servers: Vec<String> = body["authorization_servers"]
        .as_array()
        .expect("authorization_servers is an array")
        .iter()
        .map(|v| v.as_str().unwrap().to_string())
        .collect();
    assert_eq!(
        authorization_servers,
        vec!["https://issuer.test.example.com".to_string()],
        "should advertise the configured oidc_issuer verbatim",
    );

    let bearer_methods: Vec<String> = body["bearer_methods_supported"]
        .as_array()
        .expect("bearer_methods_supported is an array")
        .iter()
        .map(|v| v.as_str().unwrap().to_string())
        .collect();
    assert_eq!(
        bearer_methods,
        vec!["header".to_string()],
        "we only accept the Authorization header, not a query param",
    );

    let scopes_supported: Vec<String> = body["scopes_supported"]
        .as_array()
        .expect("scopes_supported is an array")
        .iter()
        .map(|v| v.as_str().unwrap().to_string())
        .collect();
    assert_eq!(
        scopes_supported,
        vec!["mcp.read".to_string()],
        "should advertise the mcp.read scope so clients know what to request",
    );
}

/// Path-suffixed well-known URIs (RFC 9728 §3.1) MUST serve the same
/// document as the bare well-known URI. Strict clients look these up
/// first based on the resource path; if they 404 here, those clients
/// never make it to the working bare URI.
#[mz_ore::test]
fn test_oauth_protected_resource_metadata_path_suffixed_aliases() {
    let server = test_util::TestHarness::default()
        .with_oidc_auth(None, None, None, None)
        .with_mcp_routes(true, true)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .with_system_parameter_default("enable_mcp_developer".to_string(), "true".to_string())
        .with_system_parameter_default(
            "oidc_issuer".to_string(),
            "https://issuer.test.example.com".to_string(),
        )
        .start_blocking();

    for suffix in [
        "/.well-known/oauth-protected-resource",
        "/.well-known/oauth-protected-resource/api/mcp/agent",
        "/.well-known/oauth-protected-resource/api/mcp/developer",
    ] {
        let res = Client::new()
            .get(&format!("http://{}{suffix}", server.http_local_addr()))
            .send()
            .unwrap();
        assert_eq!(res.status(), StatusCode::OK, "expected 200 at {suffix}");
        let body: serde_json::Value = res.json().unwrap();
        let authorization_servers: Vec<String> = body["authorization_servers"]
            .as_array()
            .expect("authorization_servers is an array")
            .iter()
            .map(|v| v.as_str().unwrap().to_string())
            .collect();
        assert_eq!(
            authorization_servers,
            vec!["https://issuer.test.example.com".to_string()],
            "all three URIs must serve the same metadata document",
        );
    }
}

/// An `oidc_issuer` value that does not parse as a URL must NOT be
/// published verbatim — RFC 9728 §3 requires entries in
/// `authorization_servers` to be valid URLs, and publishing garbage
/// would steer every client into an unrecoverable error on their next
/// discovery hop. The honest response is a 503 plus a server-side
/// warning so operators can see and fix the misconfiguration.
#[mz_ore::test]
fn test_oauth_protected_resource_metadata_rejects_invalid_issuer() {
    let server = test_util::TestHarness::default()
        .with_oidc_auth(None, None, None, None)
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .with_system_parameter_default(
            "oidc_issuer".to_string(),
            // Whitespace + missing scheme: not a valid URI per `http::Uri`.
            "not a url".to_string(),
        )
        .start_blocking();

    let res = Client::new()
        .get(&format!(
            "http://{}/.well-known/oauth-protected-resource",
            server.http_local_addr()
        ))
        .send()
        .unwrap();
    assert_eq!(
        res.status(),
        StatusCode::SERVICE_UNAVAILABLE,
        "invalid oidc_issuer must surface as 503, not a published malformed doc",
    );
}

/// When no `oidc_issuer` is configured on a listener that *does*
/// validate tokens, the endpoint MUST 404. RFC 9728 requires
/// `authorization_servers` to be non-empty, and returning an honest
/// 404 lets clients fall back to whatever else they support instead of
/// being misled by a half-baked document.
#[mz_ore::test]
fn test_oauth_protected_resource_metadata_404_when_no_issuer() {
    let server = test_util::TestHarness::default()
        .with_oidc_auth(None, None, None, None)
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .start_blocking();

    let res = Client::new()
        .get(&format!(
            "http://{}/.well-known/oauth-protected-resource",
            server.http_local_addr()
        ))
        .send()
        .unwrap();
    assert_eq!(
        res.status(),
        StatusCode::NOT_FOUND,
        "no oidc_issuer => 404, not an empty metadata doc",
    );
}

/// End-to-end discovery flow: an unauthenticated MCP request emits a
/// 401 with a `resource_metadata` URL on the same host as the request,
/// and the discovery document at that path resolves to a 200 with the
/// configured issuer. Pinning the sequence guards against any one
/// piece being broken in a way the per-piece tests would miss.
///
/// The published `resource_metadata` URL uses the `https` scheme even
/// though the test listener is plain HTTP (we always advertise `https`;
/// see `oauth_metadata::PUBLISHED_SCHEME`), so we extract the path
/// suffix and re-issue the GET against the actual local HTTP listener.
#[mz_ore::test]
fn test_oauth_protected_resource_metadata_end_to_end_flow() {
    let server = test_util::TestHarness::default()
        .with_oidc_auth(None, None, None, None)
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .with_system_parameter_default(
            "oidc_issuer".to_string(),
            "https://issuer.test.example.com".to_string(),
        )
        .start_blocking();

    let local_addr = server.http_local_addr();
    let mcp_res = Client::new()
        .post(&format!("http://{local_addr}/api/mcp/agent"))
        .header("Content-Type", "application/json")
        .body(r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#)
        .send()
        .unwrap();
    assert_eq!(mcp_res.status(), StatusCode::UNAUTHORIZED);

    let bearer_challenge = mcp_res
        .headers()
        .get_all(http::header::WWW_AUTHENTICATE)
        .iter()
        .map(|v| v.to_str().unwrap().to_string())
        .find(|c| c.starts_with("Bearer "))
        .expect("a Bearer WWW-Authenticate challenge");
    let advertised_url = bearer_challenge
        .split(',')
        .find_map(|part| {
            part.trim()
                .strip_prefix("resource_metadata=\"")
                .and_then(|s| s.strip_suffix('"'))
        })
        .expect("resource_metadata parameter in Bearer challenge");
    assert!(
        advertised_url.starts_with("https://"),
        "discovery URL must always be https: {advertised_url}",
    );

    let path_suffix = advertised_url
        .splitn(4, '/')
        .nth(3)
        .expect("path on the advertised URL");
    let discovery_res = Client::new()
        .get(&format!("http://{local_addr}/{path_suffix}"))
        .send()
        .unwrap();
    assert_eq!(discovery_res.status(), StatusCode::OK);
    let body: serde_json::Value = discovery_res.json().unwrap();
    assert_eq!(
        body["authorization_servers"]
            .as_array()
            .and_then(|a| a.first())
            .and_then(|v| v.as_str()),
        Some("https://issuer.test.example.com"),
    );
}

/// A `None`-authenticator listener (anonymous_http_user) MUST 404 even
/// if an `oidc_issuer` is configured: publishing on a listener that
/// never validates tokens would mislead clients into a flow they cannot
/// complete.
#[mz_ore::test]
fn test_oauth_protected_resource_metadata_404_on_no_auth_listener() {
    let server = test_util::TestHarness::default()
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .with_system_parameter_default(
            "oidc_issuer".to_string(),
            "https://issuer.test.example.com".to_string(),
        )
        .start_blocking();

    let res = Client::new()
        .get(&format!(
            "http://{}/.well-known/oauth-protected-resource",
            server.http_local_addr()
        ))
        .send()
        .unwrap();
    assert_eq!(
        res.status(),
        StatusCode::NOT_FOUND,
        "None-authenticator listener must not publish OAuth metadata even with oidc_issuer set",
    );
}

/// Verifies that MCP requests update the Prometheus counters and that
/// tool calls record a histogram observation.
#[mz_ore::test]
fn test_mcp_metrics() {
    let server = test_util::TestHarness::default()
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .start_blocking();

    let agents_url = format!("http://{}/api/mcp/agent", server.http_local_addr());

    // Exercise three distinct request shapes:
    //   1. `initialize`: succeeds.
    //   2. `tools/list`: succeeds.
    //   3. `tools/call` for `read_data_product` with a nonexistent name:
    //      the request itself completes with an MCP error
    //      (`DataProductNotFound`), which both `requests_total` and
    //      `tool_calls_total` should reflect via the status label.

    let (status, _) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2025-11-25",
                "capabilities": {},
                "clientInfo": {"name": "metrics-test", "version": "0"}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);

    let (status, _) = mcp_post(
        &agents_url,
        serde_json::json!({"jsonrpc": "2.0", "id": 2, "method": "tools/list"}),
    );
    assert_eq!(status, StatusCode::OK);

    let (status, _) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {
                "name": "read_data_product",
                "arguments": {"name": "nonexistent_product"}
            }
        }),
    );
    assert_eq!(status, StatusCode::OK);

    // Helper: look up a counter value by exact (endpoint, method/tool, status)
    // label combination. Panics with a helpful message if not found, so test
    // failures point at the missing label set rather than a generic `None`.
    fn find_counter(
        gathered: &[prometheus::proto::MetricFamily],
        metric_name: &str,
        labels: &[(&str, &str)],
    ) -> f64 {
        let family = gathered
            .iter()
            .find(|m| m.name() == metric_name)
            .unwrap_or_else(|| panic!("metric family {} should be present", metric_name));
        for metric in family.get_metric() {
            let matches_all = labels.iter().all(|(name, want)| {
                metric
                    .get_label()
                    .iter()
                    .any(|l| l.name() == *name && l.value() == *want)
            });
            if matches_all {
                return metric.get_counter().value();
            }
        }
        panic!(
            "no {} entry matching labels {:?} in {:?}",
            metric_name,
            labels,
            family
                .get_metric()
                .iter()
                .map(|m| m
                    .get_label()
                    .iter()
                    .map(|l| (l.name().to_string(), l.value().to_string()))
                    .collect::<Vec<_>>())
                .collect::<Vec<_>>()
        );
    }

    let gathered = server.metrics_registry().gather();

    // requests_total: one entry per (endpoint, method, status). The
    // tools/call request failed at the tool level, so its status label is
    // the McpRequestError error_type, not "ok".
    assert_eq!(
        find_counter(
            &gathered,
            "mz_mcp_requests_total",
            &[
                ("endpoint_type", "agent"),
                ("method", "initialize"),
                ("status", "ok"),
            ],
        ),
        1.0,
    );
    assert_eq!(
        find_counter(
            &gathered,
            "mz_mcp_requests_total",
            &[
                ("endpoint_type", "agent"),
                ("method", "tools/list"),
                ("status", "ok"),
            ],
        ),
        1.0,
    );
    assert_eq!(
        find_counter(
            &gathered,
            "mz_mcp_requests_total",
            &[
                ("endpoint_type", "agent"),
                ("method", "tools/call"),
                ("status", "DataProductNotFound"),
            ],
        ),
        1.0,
    );

    // tool_calls_total: one entry per (endpoint, tool_name, status).
    assert_eq!(
        find_counter(
            &gathered,
            "mz_mcp_tool_calls_total",
            &[
                ("endpoint_type", "agent"),
                ("tool_name", "read_data_product"),
                ("status", "DataProductNotFound"),
            ],
        ),
        1.0,
    );

    // tool_call_duration_seconds: the single tools/call above should
    // produce exactly one histogram observation.
    let duration = gathered
        .iter()
        .find(|m| m.name() == "mz_mcp_tool_call_duration_seconds")
        .expect("mz_mcp_tool_call_duration_seconds should be present");
    let total_samples: u64 = duration
        .get_metric()
        .iter()
        .map(|m| m.get_histogram().get_sample_count())
        .sum();
    assert_eq!(
        total_samples, 1,
        "tool_call_duration_seconds should record exactly 1 sample",
    );
}

/// Tests runtime toggling of MCP feature flags.
#[mz_ore::test]
fn test_mcp_agent_runtime_flag_toggle() {
    let server = test_util::TestHarness::default()
        .with_mcp_routes(true, true)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .with_system_parameter_default("enable_mcp_developer".to_string(), "true".to_string())
        .start_blocking();

    let agents_url = format!("http://{}/api/mcp/agent", server.http_local_addr());
    let developer_url = format!("http://{}/api/mcp/developer", server.http_local_addr());

    let tools_list = serde_json::json!({
        "jsonrpc": "2.0",
        "id": 1,
        "method": "tools/list"
    });

    // Both endpoints should be enabled (feature flags set to true via system parameter defaults).
    let (status, _) = mcp_post(&agents_url, tools_list.clone());
    assert_eq!(status, StatusCode::OK);
    let (status, _) = mcp_post(&developer_url, tools_list.clone());
    assert_eq!(status, StatusCode::OK);

    // Disable MCP agent at runtime.
    server.disable_feature_flags(&["enable_mcp_agent"]);

    let res = Client::new()
        .post(&agents_url)
        .json(&tools_list)
        .send()
        .unwrap();
    assert_eq!(
        res.status(),
        reqwest::StatusCode::SERVICE_UNAVAILABLE,
        "agent should return 503 after disabling"
    );

    // Developer should still work.
    let (status, _) = mcp_post(&developer_url, tools_list.clone());
    assert_eq!(status, StatusCode::OK, "developer should still be enabled");

    // Re-enable MCP agent.
    server.enable_feature_flags(&["enable_mcp_agent"]);
    let (status, _) = mcp_post(&agents_url, tools_list.clone());
    assert_eq!(
        status,
        StatusCode::OK,
        "agent should work again after re-enabling"
    );

    // Test query tool toggling: enabled by default.
    let query_call = serde_json::json!({
        "jsonrpc": "2.0",
        "id": 2,
        "method": "tools/call",
        "params": {
            "name": "query",
            "arguments": {"cluster": "quickstart", "sql_query": "SELECT 1"}
        }
    });
    let (status, body) = mcp_post(&agents_url, query_call.clone());
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"].is_null(),
        "query tool should be enabled by default"
    );
    let result_text = body["result"]["content"][0]["text"].as_str().unwrap();
    assert!(
        result_text.contains("1"),
        "query should return result with 1"
    );

    // Disable query tool at runtime.
    server.disable_feature_flags(&["enable_mcp_agent_query_tool"]);

    let (status, body) = mcp_post(&agents_url, query_call.clone());
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"]["message"]
            .as_str()
            .unwrap_or("")
            .contains("query tool is not available"),
        "query tool should be disabled after disabling"
    );

    // Re-enable query tool.
    server.enable_feature_flags(&["enable_mcp_agent_query_tool"]);
    let (status, body) = mcp_post(&agents_url, query_call.clone());
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"].is_null(),
        "query tool should work after re-enabling"
    );
    let result_text = body["result"]["content"][0]["text"].as_str().unwrap();
    assert!(
        result_text.contains("1"),
        "query should return result with 1 after re-enabling"
    );
}

/// Tests that the MCP agent endpoint respects RBAC: data products are visible
/// to any user with SELECT on the view. Cluster USAGE is not required for
/// discovery — the view appears with a NULL cluster if no accessible index exists.
#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_mcp_agent_rbac() {
    let server = test_util::TestHarness::default()
        .with_mcp_routes(true, false)
        .with_system_parameter_default("enable_mcp_agent".to_string(), "true".to_string())
        .start_blocking();

    let agents_url = format!("http://{}/api/mcp/agent", server.http_local_addr());

    let mut super_user = server
        .pg_config_internal()
        .user(&SYSTEM_USER.name)
        .connect(postgres::NoTls)
        .unwrap();

    // Create the HTTP default user with basic system/database/schema privileges
    // but NO object-level grants yet.
    super_user
        .batch_execute(&format!("CREATE ROLE {}", &HTTP_DEFAULT_USER.name))
        .unwrap();
    super_user
        .batch_execute(&format!(
            "GRANT ALL PRIVILEGES ON SYSTEM TO {}",
            &HTTP_DEFAULT_USER.name
        ))
        .unwrap();
    super_user
        .batch_execute(&format!(
            "GRANT ALL PRIVILEGES ON DATABASE materialize TO {}",
            &HTTP_DEFAULT_USER.name
        ))
        .unwrap();
    super_user
        .batch_execute(&format!(
            "GRANT ALL PRIVILEGES ON SCHEMA materialize.public TO {}",
            &HTTP_DEFAULT_USER.name
        ))
        .unwrap();

    // Create a data product: a materialized view (no index or comment required).
    super_user
        .batch_execute("CREATE MATERIALIZED VIEW rbac_product IN CLUSTER quickstart AS SELECT 1::int AS id, 'secret'::text AS payload")
        .unwrap();

    let get_products = serde_json::json!({
        "jsonrpc": "2.0",
        "id": 1,
        "method": "tools/call",
        "params": {"name": "get_data_products", "arguments": {}}
    });

    let products_visible = |body: &serde_json::Value| -> bool {
        let text = body["result"]["content"][0]["text"]
            .as_str()
            .unwrap_or("[]");
        let products: serde_json::Value =
            serde_json::from_str(text).unwrap_or(serde_json::json!([]));
        products
            .as_array()
            .map(|arr| {
                arr.iter().any(|p| {
                    p.as_array()
                        .map(|row| row[0].as_str().unwrap_or("").contains("rbac_product"))
                        .unwrap_or(false)
                })
            })
            .unwrap_or(false)
    };

    // 1. No SELECT → product NOT visible.
    let (status, body) = mcp_post(&agents_url, get_products.clone());
    assert_eq!(status, StatusCode::OK);
    assert!(
        !products_visible(&body),
        "product should not be visible with no privileges"
    );

    // 2. Grant SELECT on view → product NOW visible (no index or cluster USAGE needed).
    super_user
        .batch_execute(&format!(
            "GRANT SELECT ON rbac_product TO {}",
            &HTTP_DEFAULT_USER.name
        ))
        .unwrap();
    let (status, body) = mcp_post(&agents_url, get_products.clone());
    assert_eq!(status, StatusCode::OK);
    assert!(
        products_visible(&body),
        "product should be visible with just SELECT (no index required)"
    );

    // Capture the fully-qualified product name for subsequent tests.
    let text = body["result"]["content"][0]["text"].as_str().unwrap();
    let products: serde_json::Value = serde_json::from_str(text).unwrap();
    let product_name = products
        .as_array()
        .unwrap()
        .iter()
        .find(|p| {
            p.as_array()
                .map(|row| row[0].as_str().unwrap_or("").contains("rbac_product"))
                .unwrap_or(false)
        })
        .unwrap()
        .as_array()
        .unwrap()[0]
        .as_str()
        .unwrap()
        .to_string();

    // 3. Revoke SELECT → product disappears.
    super_user
        .batch_execute(&format!(
            "REVOKE SELECT ON rbac_product FROM {}",
            &HTTP_DEFAULT_USER.name
        ))
        .unwrap();
    let (status, body) = mcp_post(&agents_url, get_products.clone());
    assert_eq!(status, StatusCode::OK);
    assert!(
        !products_visible(&body),
        "product should disappear after revoking SELECT"
    );

    // 4. read_data_product by name while lacking SELECT → DataProductNotFound.
    let (status, body) = mcp_post(
        &agents_url,
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 2,
            "method": "tools/call",
            "params": {"name": "read_data_product", "arguments": {"name": product_name}}
        }),
    );
    assert_eq!(status, StatusCode::OK);
    assert!(
        body["error"]["message"]
            .as_str()
            .unwrap_or("")
            .contains("Data product not found"),
        "read_data_product should fail with not-found when SELECT is revoked"
    );

    // 5. Re-grant SELECT → visible again.
    super_user
        .batch_execute(&format!(
            "GRANT SELECT ON rbac_product TO {}",
            &HTTP_DEFAULT_USER.name
        ))
        .unwrap();
    let (status, body) = mcp_post(&agents_url, get_products.clone());
    assert_eq!(status, StatusCode::OK);
    assert!(
        products_visible(&body),
        "product should reappear after re-granting SELECT"
    );
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
fn test_inject_audit_events() {
    let server = test_util::TestHarness::default().start_blocking();
    let mut pg_client = server.connect(postgres::NoTls).unwrap();

    // Inject two audit events via the HTTP API.
    let http_client = Client::new();
    let url = Url::parse(&format!(
        "http://{}/api/catalog/inject-audit-events",
        server.internal_http_local_addr()
    ))
    .unwrap();
    let res = http_client
        .post(url)
        .json(&serde_json::json!([
            {
                "event_type": "create",
                "object_type": "table",
                "details": {"IdNameV1": {"id": "u1", "name": "injected_table"}},
                "user": "mz_system"
            },
            {
                "event_type": "drop",
                "object_type": "table",
                "details": {"IdNameV1": {"id": "u1", "name": "injected_table"}},
                "user": null
            }
        ]))
        .send()
        .unwrap();
    assert_eq!(res.status(), StatusCode::OK);

    // Verify the injected events are the last two audit events.
    let rows = pg_client
        .query(
            "SELECT event_type, object_type, details->>'name' as name, occurred_at
             FROM mz_audit_events
             ORDER BY id DESC LIMIT 2",
            &[],
        )
        .unwrap();
    assert_eq!(rows.len(), 2);
    assert_eq!(rows[0].get::<_, String>("event_type"), "drop");
    assert_eq!(rows[0].get::<_, String>("object_type"), "table");
    assert_eq!(rows[0].get::<_, String>("name"), "injected_table");
    assert_eq!(rows[1].get::<_, String>("event_type"), "create");
    assert_eq!(rows[1].get::<_, String>("object_type"), "table");
    assert_eq!(rows[1].get::<_, String>("name"), "injected_table");
    let drop_ts: DateTime<Utc> = rows[0].get("occurred_at");
    let create_ts: DateTime<Utc> = rows[1].get("occurred_at");
    assert_eq!(drop_ts, create_ts);
}

#[mz_ore::test]
#[cfg_attr(miri, ignore)] // too slow
fn test_inject_audit_events_malformed() {
    let server = test_util::TestHarness::default().start_blocking();

    let http_client = Client::new();
    let url = Url::parse(&format!(
        "http://{}/api/catalog/inject-audit-events",
        server.internal_http_local_addr()
    ))
    .unwrap();

    // Missing required fields.
    let res = http_client
        .post(url.clone())
        .json(&serde_json::json!([
            {
                "event_type": "create"
            }
        ]))
        .send()
        .unwrap();
    assert_eq!(res.status(), StatusCode::UNPROCESSABLE_ENTITY);

    // Invalid event_type value.
    let res = http_client
        .post(url.clone())
        .json(&serde_json::json!([
            {
                "event_type": "bogus",
                "object_type": "table",
                "details": {"IdNameV1": {"id": "u1", "name": "t"}},
                "user": null
            }
        ]))
        .send()
        .unwrap();
    assert_eq!(res.status(), StatusCode::UNPROCESSABLE_ENTITY);

    // Not valid JSON at all.
    let res = http_client
        .post(url)
        .header("content-type", "application/json")
        .body("not json")
        .send()
        .unwrap();
    assert_eq!(res.status(), StatusCode::BAD_REQUEST);
}