pg_tviews 0.1.0-beta.11

Transactional materialized views with incremental refresh for PostgreSQL
use crate::error::TViewResult;
use pgrx::prelude::*;
use std::collections::HashMap;

/// Infer column types from `PostgreSQL` catalog
///
/// # Errors
/// Returns error if column type lookup fails or column not found in table
pub fn infer_column_types(
    table_name: &str,
    columns: &[String],
) -> TViewResult<HashMap<String, String>> {
    let mut types = HashMap::new();

    for col in columns {
        // Query PostgreSQL catalog for column type
        let type_query = format!(
            "SELECT format_type(atttypid, atttypmod)
             FROM pg_attribute
             WHERE attrelid = '{table_name}'::regclass
               AND attname = '{col}'
               AND attnum > 0
               AND NOT attisdropped"
        );

        let col_type = crate::utils::spi_get_string(&type_query)
            .map_err(|e| crate::error::TViewError::SpiError {
                query: type_query.clone(),
                error: e.to_string(),
            })?
            .ok_or_else(|| crate::error::TViewError::CatalogError {
                operation: format!("find column '{col}' in table '{table_name}'"),
                pg_error: "Column not found".to_string(),
            })?;

        types.insert(col.clone(), col_type);
    }

    Ok(types)
}

/// Check if a table exists in the database
///
/// # Errors
/// Returns error if `pg_class` query fails
pub fn table_exists(table_name: &str) -> TViewResult<bool> {
    let query = format!(
        "SELECT COUNT(*) = 1 FROM pg_class
         WHERE relname = '{table_name}' AND relkind = 'r'"
    );

    Spi::get_one::<bool>(&query)
        .map_err(|e| crate::error::TViewError::SpiError {
            query,
            error: e.to_string(),
        })
        .map(|opt| opt.unwrap_or(false))
}

#[cfg(any(test, feature = "pg_test"))]
#[pg_schema]
mod tests {
    use super::*;

    #[pg_test]
    fn test_infer_column_types() {
        // Create test table
        Spi::run(
            "CREATE TABLE test_types (
            pk INTEGER PRIMARY KEY,
            id UUID NOT NULL,
            name TEXT,
            is_active BOOLEAN,
            created_at TIMESTAMPTZ DEFAULT NOW(),
            tags TEXT[],
            data JSONB
        )",
        )
        .unwrap();

        let columns = vec![
            "pk".to_string(),
            "id".to_string(),
            "name".to_string(),
            "is_active".to_string(),
            "created_at".to_string(),
            "tags".to_string(),
            "data".to_string(),
        ];

        let types = infer_column_types("test_types", &columns).unwrap();

        assert_eq!(types.get("pk"), Some(&"integer".to_string()));
        assert_eq!(types.get("id"), Some(&"uuid".to_string()));
        assert_eq!(types.get("name"), Some(&"text".to_string()));
        assert_eq!(types.get("is_active"), Some(&"boolean".to_string()));
        assert_eq!(
            types.get("created_at"),
            Some(&"timestamp with time zone".to_string())
        );
        assert_eq!(types.get("tags"), Some(&"text[]".to_string()));
        assert_eq!(types.get("data"), Some(&"jsonb".to_string()));
    }

    #[pg_test]
    fn test_infer_column_types_missing_table() {
        let columns = vec!["id".to_string()];
        let result = infer_column_types("nonexistent_table", &columns);
        assert!(result.is_err());
    }

    #[pg_test]
    fn test_infer_column_types_missing_column() {
        // Create test table
        Spi::run("CREATE TABLE test_missing_col (id UUID)").unwrap();

        let columns = vec!["id".to_string(), "missing_col".to_string()];
        let result = infer_column_types("test_missing_col", &columns);
        assert!(result.is_err());
    }

    #[pg_test]
    fn test_table_exists() {
        // Create test table
        Spi::run("CREATE TABLE test_exists (id UUID)").unwrap();

        assert_eq!(table_exists("test_exists"), Ok(true));
        assert_eq!(table_exists("nonexistent_table"), Ok(false));
    }
}