dbx/crates/dbx-core/src/db/starrocks.rs

use mysql_async::prelude::*;
use std::collections::HashSet;

use crate::models::connection::{ConnectionConfig, DatabaseType};
use crate::types::{ObjectInfo, TableInfo};

use super::mysql::{
    get_conn_with_timeout, get_str_by_name, list_routine_objects, list_tables_show_with_status, quote_value,
    table_infos_to_objects, MySqlPool, TableStatusMeta,
};

pub use super::mysql_compatible::{
    get_columns_show_from as get_catalog_columns, list_catalog_indexes, list_catalogs,
    list_databases_show_from as list_catalog_databases, list_indexes_with_ddl_fallback as list_indexes,
    list_tables_show_from as list_catalog_tables, show_create_table_ddl_from as get_catalog_table_ddl,
};

pub fn is_config(config: &ConnectionConfig) -> bool {
    is_profile(&config.db_type, config.driver_profile.as_deref())
}

pub fn is_profile(db_type: &DatabaseType, driver_profile: Option<&str>) -> bool {
    *db_type == DatabaseType::StarRocks
        || (*db_type == DatabaseType::Mysql
            && driver_profile.is_some_and(|profile| profile.eq_ignore_ascii_case("starrocks")))
}

fn materialized_views_sql(database: &str) -> String {
    format!(
        "SELECT TABLE_NAME FROM information_schema.materialized_views WHERE TABLE_SCHEMA = {}",
        quote_value(database)
    )
}

pub(crate) fn materialized_view_definition_sql(database: &str, name: &str) -> String {
    format!(
        "SELECT MATERIALIZED_VIEW_DEFINITION \
         FROM information_schema.materialized_views \
         WHERE TABLE_SCHEMA = {} AND TABLE_NAME = {} \
         LIMIT 1",
        quote_value(database),
        quote_value(name)
    )
}

async fn list_materialized_view_names(pool: &MySqlPool, database: &str) -> Result<HashSet<String>, String> {
    let sql = materialized_views_sql(database);
    let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?;
    let result = conn.query_iter(&sql).await.map_err(|error| error.to_string())?;
    let rows: Vec<mysql_async::Row> = result.collect_and_drop().await.map_err(|error| error.to_string())?;
    Ok(rows
        .iter()
        .filter_map(|row| {
            let name = get_str_by_name(row, "TABLE_NAME").trim().to_string();
            (!name.is_empty()).then_some(name)
        })
        .collect())
}

fn merge_materialized_views(
    tables: &mut Vec<TableInfo>,
    materialized_view_names: Result<HashSet<String>, String>,
    database: &str,
) {
    let materialized_view_names = match materialized_view_names {
        Ok(names) => names,
        Err(error) => {
            log::warn!("Skipping materialized view classification for StarRocks database `{database}`: {error}");
            return;
        }
    };

    let known_names: HashSet<String> = tables.iter().map(|table| table.name.clone()).collect();
    for table in tables.iter_mut() {
        if materialized_view_names.contains(&table.name) {
            table.table_type = "MATERIALIZED_VIEW".to_string();
        }
    }

    let mut sorted_names: Vec<&String> = materialized_view_names.iter().collect();
    sorted_names.sort();
    for name in sorted_names {
        if !known_names.contains(name.as_str()) {
            tables.push(TableInfo {
                name: name.clone(),
                table_type: "MATERIALIZED_VIEW".to_string(),
                comment: None,
                parent_schema: None,
                parent_name: None,
            });
        }
    }
}

async fn list_tables_with_status(
    pool: &MySqlPool,
    database: &str,
) -> Result<(Vec<TableInfo>, std::collections::HashMap<String, TableStatusMeta>), String> {
    let (tables, materialized_view_names) =
        tokio::join!(list_tables_show_with_status(pool, database), list_materialized_view_names(pool, database));
    let (mut tables, status) = tables?;
    merge_materialized_views(&mut tables, materialized_view_names, database);
    Ok((tables, status))
}

pub async fn list_tables(pool: &MySqlPool, database: &str) -> Result<Vec<TableInfo>, String> {
    list_tables_with_status(pool, database).await.map(|(tables, _)| tables)
}

pub async fn list_table_objects(pool: &MySqlPool, database: &str) -> Result<Vec<ObjectInfo>, String> {
    let (tables, routines) =
        tokio::join!(list_tables_with_status(pool, database), list_routine_objects(pool, database));
    let (tables, status) = tables?;
    let mut objects = table_infos_to_objects(tables, &status, database);

    match routines {
        Ok(routines) => objects.extend(routines),
        Err(error) => log::warn!("Skipping routines for database `{database}` in object browser: {error}"),
    }

    Ok(objects)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn profile_matches_starrocks_only() {
        assert!(is_profile(&DatabaseType::StarRocks, None));
        assert!(is_profile(&DatabaseType::Mysql, Some("STARROCKS")));
        assert!(!is_profile(&DatabaseType::Mysql, Some("doris")));
        assert!(!is_profile(&DatabaseType::Postgres, Some("starrocks")));
    }

    #[test]
    fn materialized_views_are_classified_without_duplicates() {
        let mut tables = vec![table("orders", "BASE TABLE"), table("orders_view", "VIEW"), table("orders_mv", "VIEW")];

        merge_materialized_views(&mut tables, Ok(HashSet::from(["orders_mv".to_string()])), "analytics");

        assert_eq!(
            tables.iter().map(|table| (table.name.as_str(), table.table_type.as_str())).collect::<Vec<_>>(),
            vec![("orders", "BASE TABLE"), ("orders_view", "VIEW"), ("orders_mv", "MATERIALIZED_VIEW")]
        );
    }

    #[test]
    fn missing_materialized_views_are_appended_deterministically() {
        let mut tables = vec![table("orders", "BASE TABLE")];
        let names = HashSet::from(["orders_mv".to_string(), "daily_orders_mv".to_string()]);

        merge_materialized_views(&mut tables, Ok(names), "analytics");

        assert_eq!(
            tables.iter().map(|table| (table.name.as_str(), table.table_type.as_str())).collect::<Vec<_>>(),
            vec![
                ("orders", "BASE TABLE"),
                ("daily_orders_mv", "MATERIALIZED_VIEW"),
                ("orders_mv", "MATERIALIZED_VIEW"),
            ]
        );
    }

    #[test]
    fn lookup_failure_preserves_base_table_types() {
        let mut tables = vec![table("orders_mv", "VIEW")];

        merge_materialized_views(&mut tables, Err("permission denied".to_string()), "analytics");

        assert_eq!(tables[0].table_type, "VIEW");
    }

    #[test]
    fn materialized_view_queries_escape_database_and_name() {
        assert_eq!(
            materialized_views_sql("tenant's analytics"),
            "SELECT TABLE_NAME FROM information_schema.materialized_views WHERE TABLE_SCHEMA = 'tenant\\'s analytics'"
        );
        assert_eq!(
            materialized_view_definition_sql("tenant's analytics", "weird'name"),
            "SELECT MATERIALIZED_VIEW_DEFINITION FROM information_schema.materialized_views WHERE TABLE_SCHEMA = 'tenant\\'s analytics' AND TABLE_NAME = 'weird\\'name' LIMIT 1"
        );
    }

    #[test]
    fn object_conversion_preserves_materialized_view_type() {
        let objects = table_infos_to_objects(
            vec![table("orders", "BASE TABLE"), table("orders_view", "VIEW"), table("orders_mv", "MATERIALIZED_VIEW")],
            &std::collections::HashMap::new(),
            "analytics",
        );

        assert_eq!(
            objects.iter().map(|object| (object.name.as_str(), object.object_type.as_str())).collect::<Vec<_>>(),
            vec![("orders", "TABLE"), ("orders_view", "VIEW"), ("orders_mv", "MATERIALIZED_VIEW")]
        );
    }

    fn table(name: &str, table_type: &str) -> TableInfo {
        TableInfo {
            name: name.to_string(),
            table_type: table_type.to_string(),
            comment: None,
            parent_schema: None,
            parent_name: None,
        }
    }
}