hamelin_lib 0.21.6

//! Match catalog table references against replacement catalog keys.

use std::collections::HashMap;

use crate::tree::ast::dataset_identifier::{DatasetIdentifier, UnqualifiedDatasetIdentifier};
use crate::tree::ast::identifier::SimpleIdentifier;

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum CatalogKeyMatchingError {
    #[error("ambiguous catalog replacement for {0}")]
    AmbiguousReplacement(String),
}

/// Candidate alias strings that may appear in table references for a catalog key.
pub fn catalog_key_aliases(
    catalog_key: &str,
    default_space: Option<&SimpleIdentifier>,
) -> Vec<String> {
    let mut aliases = vec![catalog_key.to_string()];
    if let Ok(id) = DatasetIdentifier::from_ref_str(catalog_key) {
        aliases.push(id.to_string());
        if let Some(qualified) = id.maybe_qualify(default_space) {
            aliases.push(qualified.to_string());
        }
        if let DatasetIdentifier::Qualified(q) = id {
            if default_space.is_some_and(|space| space.as_str() == q.space.as_str()) {
                let unqual = UnqualifiedDatasetIdentifier {
                    namespace: q.namespace,
                    table: q.table,
                };
                aliases.push(unqual.to_string());
            }
        }
    }
    aliases.sort();
    aliases.dedup();
    aliases
}

/// Candidate lookup keys for a table reference.
pub fn catalog_keys_for_dataset(
    dataset: &DatasetIdentifier,
    default_space: Option<&SimpleIdentifier>,
) -> Vec<String> {
    let mut keys = vec![dataset.to_string()];
    if let Some(qualified) = dataset.maybe_qualify(default_space) {
        keys.push(qualified.to_string());
    }
    keys
}

/// Find the CTE name for `dataset` when it matches one of the catalog replacements.
pub fn find_catalog_replacement(
    dataset: &DatasetIdentifier,
    replacements: &HashMap<String, String>,
    default_space: Option<&SimpleIdentifier>,
) -> Result<Option<String>, CatalogKeyMatchingError> {
    let ref_keys = catalog_keys_for_dataset(dataset, default_space);
    let mut matched_cte: Option<String> = None;

    for (catalog_key, cte_name) in replacements {
        let aliases = catalog_key_aliases(catalog_key, default_space);
        if !ref_keys
            .iter()
            .any(|key| aliases.iter().any(|alias| alias == key))
        {
            continue;
        }

        match &matched_cte {
            None => matched_cte = Some(cte_name.clone()),
            Some(existing) if existing == cte_name => {}
            Some(_) => {
                return Err(CatalogKeyMatchingError::AmbiguousReplacement(
                    dataset.to_string(),
                ));
            }
        }
    }

    Ok(matched_cte)
}

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

    #[test]
    fn qualified_key_does_not_alias_bare_table_without_default_space() {
        let aliases = catalog_key_aliases("events:foo", None);
        assert!(!aliases.contains(&"foo".to_string()));
    }

    #[test]
    fn qualified_key_aliases_bare_table_when_default_space_matches() {
        let default_space = SimpleIdentifier::new("events");
        let aliases = catalog_key_aliases("events:foo", Some(&default_space));
        assert!(aliases.contains(&"foo".to_string()));
    }

    #[test]
    fn qualified_key_does_not_alias_bare_table_for_other_default_space() {
        let default_space = SimpleIdentifier::new("signals");
        let aliases = catalog_key_aliases("events:foo", Some(&default_space));
        assert!(!aliases.contains(&"foo".to_string()));
    }

    #[test]
    fn unqualified_ref_with_default_space_does_not_match_other_space_key() {
        let dataset = DatasetIdentifier::from_ref_str("foo").unwrap();
        let mut replacements = HashMap::new();
        replacements.insert("events:foo".to_string(), "events_foo".to_string());
        replacements.insert("signals:foo".to_string(), "signals_foo".to_string());
        let default_space = SimpleIdentifier::new("events");

        let matched =
            find_catalog_replacement(&dataset, &replacements, Some(&default_space)).unwrap();
        assert_eq!(matched, Some("events_foo".to_string()));
    }
}