Repository: risingwavelabs/risingwave
File: src/tests/simulation/tests/integration_tests/utils.rs (Rust)

Code before:
    WHERE database_id IS NOT NULL
)
SELECT d.name,
       COALESCE(r.event_type, 'DATABASE_RECOVERY_UNKNOWN') AS event_type
FROM rw_catalog.rw_databases d
LEFT JOIN ranked r
       ON d.id = r.database_id AND r.rn = 1
ORDER BY d.name;
"#;

pub(crate) async fn query_all_database_recovery_state(
    cluster: &mut Cluster,
) -> Result<HashMap<String, DatabaseRecoveryState>> {
    let mut session = cluster.start_session();
    let output = session.run(ALL_DATABASE_RECOVERY_STATUS_SQL).await?;
    drop(session);

    let mut result = HashMap::new();
    for line in output.lines().filter(|line| !line.trim().is_empty()) {
        let (name, event) = line
            .rsplit_once(' ')
            .unwrap_or((line, "DATABASE_RECOVERY_UNKNOWN"));
        result.insert(
            name.trim().to_owned(),
            DatabaseRecoveryState::from_event(event),
        );
    }
    Ok(result)
}

pub(crate) async fn cluster_fully_running(cluster: &mut Cluster) -> Result<bool> {
    let global_running = {
        let mut session = cluster.start_session();
        match session.run("select rw_recovery_status()").await {
            Ok(status) => status.trim() == "RUNNING",
            Err(_) => return Ok(false),
        }
    };

    let db_statuses = match query_all_database_recovery_state(cluster).await {
        Ok(statuses) => statuses,
        Err(_) => return Ok(false),
    };
    Ok(global_running && db_statuses.values().all(|s| s.is_running()))
}

const WAIT_ALL_DB_TIMEOUT: Duration = Duration::from_secs(100);
const WAIT_ALL_DB_INTERVAL: Duration = Duration::from_secs(10);

pub(crate) async fn wait_all_database_recovered(cluster: &mut Cluster) {
    let start = Instant::now();

Review comment (bug):
The parsing logic assumes the SQL query output format is \"<database_name> <event_type>\" separated by space. This is fragile as query output formatting can vary. Consider using a more robust parsing approach or documenting the expected format with a comment explaining the query output structure.
```suggestion
    // The SQL query is expected to output lines in the format: "<database_name> <event_type>", separated by whitespace.
    // This parsing logic assumes that neither database_name nor event_type contains whitespace.
    // If the format changes, this code may break.
    for line in output.lines().filter(|line| !line.trim().is_empty()) {
        let fields: Vec<&str> = line.split_whitespace().collect();
        if fields.len() == 2 {
            let name = fields[0];
            let event = fields[1];
            result.insert(
                name.trim().to_owned(),
                DatabaseRecoveryState::from_event(event),
            );
        } else {
            // If the line does not have exactly two fields, skip or handle as unknown.
            // Optionally, log or return an error here for debugging.
            continue;
        }
```

Code after:
            "SUCCESS" => Self::Success,
            "START" => Self::Start,
            "FAILURE" => Self::Failure,
            "UNKNOWN" => Self::Unknown,
            other => panic!("Unexpected last database event: {other}"),
        }
    }

    fn is_success(self) -> bool {
        matches!(self, Self::Success)
    }
}

#[derive(Debug)]
pub(crate) struct DatabaseRecoveryInfo {
    id: u32,
    name: String,
    last_database_event: DatabaseRecoveryEvent,
    last_global_event: GlobalRecoveryEvent,
    in_global_running: bool,
    in_global_recovering: bool,
    state: DatabaseRecoveryState,
}

#[derive(Deserialize)]
struct DatabaseRecoveryRow {
    #[serde(rename = "database_id")]
    database_id: u32,
    #[serde(rename = "database_name")]
    database_name: String,
    #[serde(rename = "last_database_event")]
    last_database_event: String,
    #[serde(rename = "last_global_event")]
    last_global_event: String,
    #[serde(rename = "recovery_state")]
    recovery_state: String,
    #[serde(rename = "in_global_running")]
    in_global_running: bool,
    #[serde(rename = "in_global_recovering")]
    in_global_recovering: bool,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum GlobalRecoveryEvent {
    Running,
    Recovering,
    Unknown,
}

impl GlobalRecoveryEvent {
    fn from_str(event: &str) -> Self {