bun-mariadb/src/lib/schema/sync-indexes.ts

import type {
    BUN_MARIADB_IndexSchemaType,
    BUN_MARIADB_TableSchemaType,
    BunMariaDBConfig,
} from "../../types";
import isVectorField from "./is-vector-field";
import MariaDBQuoteGen from "./mariadb-quote-gen";
import runSchemaQuery, { querySchemaRows } from "./run-schema-query";
import schemaCondition from "./schema-condition";

function isVectorIndexDef(
    index: BUN_MARIADB_IndexSchemaType,
    table: BUN_MARIADB_TableSchemaType,
): boolean {
    if (index.indexType === "VECTOR") return true;
    if (table.isVector) return true;

    const firstFieldName = index.indexTableFields?.[0];
    if (!firstFieldName) return false;

    const field = table.fields?.find((f) => f.fieldName === firstFieldName);
    return isVectorField(field);
}

function vectorDistanceMetric(index: BUN_MARIADB_IndexSchemaType): string {
    if (index.vectorDistanceMetric === "cosine") return "cosine";
    if (index.vectorDistanceMetric === "euclidean") return "euclidean";
    return "euclidean";
}

export default async function syncIndexes({
    table,
    config,
}: {
    table: BUN_MARIADB_TableSchemaType;
    config?: BunMariaDBConfig;
}): Promise<void> {
    const schemaCond = schemaCondition(config);
    const rows = await querySchemaRows<{
        INDEX_NAME: string;
        COLUMN_NAME: string;
        INDEX_TYPE: string;
    }>({
        query: `SELECT INDEX_NAME, COLUMN_NAME, INDEX_TYPE FROM information_schema.STATISTICS WHERE ${schemaCond.where} AND TABLE_NAME = ? AND INDEX_NAME <> 'PRIMARY' ORDER BY INDEX_NAME, SEQ_IN_INDEX`,
        values: [...schemaCond.values, table.tableName],
        config,
    });

    /**
     * Indexes required by foreign keys / unique constraints cannot be dropped
     * freely. Skip those when cleaning up schema indexes.
     */
    const protectedConstraintRows = await querySchemaRows<{
        CONSTRAINT_NAME: string;
        CONSTRAINT_TYPE: string;
    }>({
        query: `SELECT CONSTRAINT_NAME, CONSTRAINT_TYPE FROM information_schema.TABLE_CONSTRAINTS WHERE ${schemaCond.where} AND TABLE_NAME = ? AND CONSTRAINT_TYPE IN ('FOREIGN KEY', 'UNIQUE')`,
        values: [...schemaCond.values, table.tableName],
        config,
    });
    const protectedIndexNames = new Set(
        protectedConstraintRows.map((r) => r.CONSTRAINT_NAME),
    );

    // Column-level UNIQUE creates an index often named after the column
    for (const field of table.fields || []) {
        if (field.unique && field.fieldName) {
            protectedIndexNames.add(field.fieldName);
        }
    }
    for (const constraint of table.uniqueConstraints || []) {
        if (constraint.constraintName) {
            protectedIndexNames.add(constraint.constraintName);
        }
    }

    const existingIndexesMap = new Map<
        string,
        { columns: string[]; type: string }
    >();

    for (const row of rows) {
        if (!existingIndexesMap.has(row.INDEX_NAME)) {
            existingIndexesMap.set(row.INDEX_NAME, {
                columns: [],
                type: row.INDEX_TYPE,
            });
        }
        existingIndexesMap.get(row.INDEX_NAME)!.columns.push(row.COLUMN_NAME);
    }

    for (const [indexName, details] of existingIndexesMap.entries()) {
        if (protectedIndexNames.has(indexName)) {
            continue;
        }

        const schemaIndex = table.indexes?.find((i) => i.indexName === indexName);

        if (!schemaIndex) {
            console.log(`Dropping index: ${indexName}`);
            try {
                await runSchemaQuery({
                    query: `DROP INDEX ${MariaDBQuoteGen(indexName)} ON ${MariaDBQuoteGen(table.tableName)}`,
                    config,
                });
            } catch (err: any) {
                if (
                    String(err?.message || "").includes(
                        "needed in a foreign key constraint",
                    )
                ) {
                    console.warn(
                        `Skipping drop of index ${indexName}: required by a foreign key constraint`,
                    );
                    continue;
                }
                throw err;
            }
        } else {
            const schemaColumns = schemaIndex.indexTableFields || [];
            const columnsMatch =
                details.columns.length === schemaColumns.length &&
                details.columns.every((col, idx) => col === schemaColumns[idx]);

            if (!columnsMatch) {
                console.log(`Recreating changed index: ${indexName}`);
                await runSchemaQuery({
                    query: `DROP INDEX ${MariaDBQuoteGen(indexName)} ON ${MariaDBQuoteGen(table.tableName)}`,
                    config,
                });
                existingIndexesMap.delete(indexName);
            }
        }
    }

    for (const index of table.indexes || []) {
        if (
            !index.indexName ||
            !index.indexTableFields ||
            index.indexTableFields.length === 0
        ) {
            continue;
        }

        if (!existingIndexesMap.has(index.indexName)) {
            if (isVectorIndexDef(index, table)) {
                console.log(`Creating Vector index: ${index.indexName}`);
                const targetField = MariaDBQuoteGen(index.indexTableFields[0]!);
                const distanceMetric = vectorDistanceMetric(index);

                await runSchemaQuery({
                    query: `ALTER TABLE ${MariaDBQuoteGen(table.tableName)} ADD VECTOR INDEX ${MariaDBQuoteGen(index.indexName)} (${targetField}) M=8 DISTANCE=${distanceMetric}`,
                    config,
                });
            } else {
                console.log(`Creating standard index: ${index.indexName}`);
                const fields = index.indexTableFields
                    .map((field) => MariaDBQuoteGen(field))
                    .join(", ");
                const typeUpper = index.indexType?.toUpperCase();
                const isSpecialType =
                    typeUpper === "FULLTEXT" || typeUpper === "SPATIAL";
                const indexPrefix = isSpecialType ? `${typeUpper} ` : "";
                const indexSuffix =
                    !isSpecialType &&
                    (typeUpper === "BTREE" || typeUpper === "HASH")
                        ? ` USING ${typeUpper}`
                        : "";

                await runSchemaQuery({
                    query: `CREATE ${indexPrefix}INDEX ${MariaDBQuoteGen(index.indexName)} ON ${MariaDBQuoteGen(table.tableName)} (${fields})${indexSuffix}`,
                    config,
                });
            }
        }
    }
}