dbml_language_server/

formatter.rs

// src/formatter.rs
use crate::ast::*;

/// Formatting options for DBML documents
#[derive(Debug, Clone)]
pub struct FormattingOptions {
    /// Size of a tab in spaces
    pub tab_size: u32,
    /// Prefer spaces over tabs
    pub insert_spaces: bool,
    /// Trim trailing whitespace on lines
    pub trim_trailing_whitespace: bool,
    /// Insert a final newline at the end of the file
    pub insert_final_newline: bool,
    /// Trim all newlines after the final newline at the end of the file
    pub trim_final_newlines: bool,
}

impl Default for FormattingOptions {
    fn default() -> Self {
        Self {
            tab_size: 2,
            insert_spaces: true,
            trim_trailing_whitespace: true,
            insert_final_newline: true,
            trim_final_newlines: true,
        }
    }
}

impl FormattingOptions {
    /// Get the indentation string for the given level
    pub fn indent(&self, level: usize) -> String {
        let single_indent = if self.insert_spaces {
            " ".repeat(self.tab_size as usize)
        } else {
            "\t".to_string()
        };
        single_indent.repeat(level)
    }
}

/// Format a DBML document with proper indentation and spacing
pub fn format_document(ast: &Document, options: &FormattingOptions) -> String {
    let mut output = String::new();
    let mut first = true;

    for item in &ast.items {
        if !first {
            output.push_str("\n\n");
        }
        first = false;

        match item {
            DocumentItem::Table(table) => {
                output.push_str(&format_table(table, options));
            }
            DocumentItem::Enum(enum_def) => {
                output.push_str(&format_enum(enum_def, options));
            }
            DocumentItem::Ref(ref_def) => {
                output.push_str(&format_ref(ref_def));
            }
            DocumentItem::Project(project) => {
                output.push_str(&format_project(project, options));
            }
        }
    }

    // Apply final newline settings
    if options.trim_final_newlines {
        output = output.trim_end().to_string();
    }
    
    if options.insert_final_newline && !output.ends_with('\n') {
        output.push('\n');
    }

    // Trim trailing whitespace on each line if requested
    if options.trim_trailing_whitespace {
        output = output
            .lines()
            .map(|line| line.trim_end())
            .collect::<Vec<_>>()
            .join("\n");
        
        // Re-add final newline if it was removed during line processing
        if options.insert_final_newline && !output.ends_with('\n') {
            output.push('\n');
        }
    }

    output
}

fn format_table(table: &Table, options: &FormattingOptions) -> String {
    let mut output = String::new();

    // Table declaration
    output.push_str("Table ");
    if let Some(schema) = &table.schema {
        output.push_str(&schema.name);
        output.push('.');
    }
    output.push_str(&table.name.name);

    if let Some(alias) = &table.alias {
        output.push_str(" as ");
        output.push_str(&alias.name);
    }

    output.push_str(" {\n");

    // Table items
    for item in &table.items {
        match item {
            TableItem::Column(column) => {
                output.push_str(&format_column(column, 1, options));
            }
            TableItem::Indexes(indexes_block) => {
                output.push_str("\n");
                output.push_str(&options.indent(1));
                output.push_str("indexes {\n");
                for index in &indexes_block.indexes {
                    output.push_str(&format_index(index, 2, options));
                }
                output.push_str(&options.indent(1));
                output.push_str("}\n");
            }
            TableItem::Note(note) => {
                output.push_str(&options.indent(1));
                output.push_str("Note: ");
                output.push_str(&format_string_literal(note));
                output.push('\n');
            }
        }
    }

    // Table settings
    if !table.settings.is_empty() {
        output.push('\n');
        for setting in &table.settings {
            output.push_str(&options.indent(1));
            output.push_str(&format!("{}: {}\n", setting.key, setting.value));
        }
    }

    output.push_str("}\n");
    output
}

fn format_column(column: &Column, indent_level: usize, options: &FormattingOptions) -> String {
    let mut output = String::new();

    output.push_str(&options.indent(indent_level));
    output.push_str(&column.name.name);
    output.push(' ');
    output.push_str(&column.col_type);

    if !column.settings.is_empty() {
        output.push_str(" [");
        let settings: Vec<String> = column
            .settings
            .iter()
            .map(|s| format_column_setting(s))
            .collect();
        output.push_str(&settings.join(", "));
        output.push(']');
    }

    output.push('\n');
    output
}

fn format_column_setting(setting: &ColumnSetting) -> String {
    match setting {
        ColumnSetting::PrimaryKey => "pk".to_string(),
        ColumnSetting::NotNull => "not null".to_string(),
        ColumnSetting::Null => "null".to_string(),
        ColumnSetting::Unique => "unique".to_string(),
        ColumnSetting::Increment => "increment".to_string(),
        ColumnSetting::Default(val) => format!("default: {}", format_default_value(val)),
        ColumnSetting::Note(note) => format!("note: {}", format_string_literal(note)),
        ColumnSetting::Ref(inline_ref) => {
            let rel = format_relationship(&inline_ref.relationship);
            format!(
                "ref: {} {}.{}",
                rel, inline_ref.target_table.name, inline_ref.target_column.name
            )
        }
    }
}

fn format_default_value(val: &DefaultValue) -> String {
    match val {
        DefaultValue::String(s) => format!("'{}'", s.replace('\'', "\\'")),
        DefaultValue::Number(n) => n.clone(),
        DefaultValue::Boolean(b) => b.to_string(),
        DefaultValue::Expression(e) => format!("`{}`", e),
    }
}

fn format_string_literal(lit: &StringLiteral) -> String {
    // Use single quotes for DBML strings
    format!("'{}'", lit.value.replace('\'', "\\'"))
}

fn format_relationship(rel: &RelationshipType) -> String {
    match rel {
        RelationshipType::OneToOne => "-".to_string(),
        RelationshipType::OneToMany => "<".to_string(),
        RelationshipType::ManyToOne => ">".to_string(),
        RelationshipType::ManyToMany => "<>".to_string(),
    }
}

fn format_index(index: &Index, indent_level: usize, options: &FormattingOptions) -> String {
    let mut output = String::new();

    output.push_str(&options.indent(indent_level));
    output.push('(');

    let columns: Vec<String> = index
        .columns
        .iter()
        .map(|col| match col {
            IndexColumn::Simple(ident) => ident.name.clone(),
            IndexColumn::Expression(expr) => format!("`{}`", expr),
        })
        .collect();

    output.push_str(&columns.join(", "));
    output.push(')');

    if !index.settings.is_empty() {
        output.push_str(" [");
        let settings: Vec<String> = index
            .settings
            .iter()
            .map(|s| format_index_setting(s))
            .collect();
        output.push_str(&settings.join(", "));
        output.push(']');
    }

    output.push('\n');
    output
}

fn format_index_setting(setting: &IndexSetting) -> String {
    match setting {
        IndexSetting::PrimaryKey => "pk".to_string(),
        IndexSetting::Unique => "unique".to_string(),
        IndexSetting::Name(name) => format!("name: '{}'", name.replace('\'', "\\'")),
        IndexSetting::Type(idx_type) => format!("type: {}", idx_type),
    }
}

fn format_enum(enum_def: &Enum, options: &FormattingOptions) -> String {
    let mut output = String::new();

    output.push_str("enum ");
    output.push_str(&enum_def.name.name);
    output.push_str(" {\n");

    for member in &enum_def.members {
        output.push_str(&options.indent(1));
        output.push_str(&member.name.name);

        if let Some(note) = &member.note {
            output.push_str(" [note: ");
            output.push_str(&format_string_literal(note));
            output.push(']');
        }

        output.push('\n');
    }

    output.push_str("}\n");
    output
}

fn format_ref(ref_def: &Ref) -> String {
    let mut output = String::new();

    output.push_str("Ref");

    if let Some(name) = &ref_def.name {
        output.push(' ');
        output.push_str(&name.name);
    }

    output.push_str(": ");
    output.push_str(&ref_def.from_table.name);
    output.push('.');

    if ref_def.from_columns.len() == 1 {
        output.push_str(&ref_def.from_columns[0].name);
    } else {
        output.push('(');
        let cols: Vec<String> = ref_def
            .from_columns
            .iter()
            .map(|c| c.name.clone())
            .collect();
        output.push_str(&cols.join(", "));
        output.push(')');
    }

    output.push(' ');
    output.push_str(&format_relationship(&ref_def.relationship));
    output.push(' ');

    output.push_str(&ref_def.to_table.name);
    output.push('.');

    if ref_def.to_columns.len() == 1 {
        output.push_str(&ref_def.to_columns[0].name);
    } else {
        output.push('(');
        let cols: Vec<String> = ref_def.to_columns.iter().map(|c| c.name.clone()).collect();
        output.push_str(&cols.join(", "));
        output.push(')');
    }

    // Add referential actions if present
    let mut settings = Vec::new();
    if let Some(on_delete) = &ref_def.on_delete {
        settings.push(format!("delete: {}", format_referential_action(on_delete)));
    }
    if let Some(on_update) = &ref_def.on_update {
        settings.push(format!("update: {}", format_referential_action(on_update)));
    }

    if !settings.is_empty() {
        output.push_str(" [");
        output.push_str(&settings.join(", "));
        output.push(']');
    }

    output.push('\n');
    output
}

fn format_referential_action(action: &ReferentialAction) -> String {
    match action {
        ReferentialAction::Cascade => "cascade".to_string(),
        ReferentialAction::Restrict => "restrict".to_string(),
        ReferentialAction::NoAction => "no action".to_string(),
        ReferentialAction::SetNull => "set null".to_string(),
        ReferentialAction::SetDefault => "set default".to_string(),
    }
}

fn format_project(project: &Project, options: &FormattingOptions) -> String {
    let mut output = String::new();

    output.push_str("Project");

    if let Some(name) = &project.name {
        output.push(' ');
        output.push_str(&name.name);
    }

    output.push_str(" {\n");

    for setting in &project.settings {
        match setting {
            ProjectSetting::DatabaseType(db_type) => {
                output.push_str(&options.indent(1));
                output.push_str("database_type: '");
                output.push_str(&db_type.replace('\'', "\\'"));
                output.push_str("'\n");
            }
            ProjectSetting::Note(note) => {
                output.push_str(&options.indent(1));
                output.push_str("Note: ");
                output.push_str(&format_string_literal(note));
                output.push('\n');
            }
        }
    }

    output.push_str("}\n");
    output
}