druid解析数据库（多表）主键、外键、唯一键

用的是h2，对sql文件语法有要求。
像下面这种格式就可以：

-- 创建 departments 表

CREATE TABLE departments (

    did INT PRIMARY KEY,

    name VARCHAR(255) -- 部门名称

);

-- 创建 example 表，并设置外键约束和唯一约束

CREATE TABLE example (

    id INT PRIMARY KEY,

    name VARCHAR(255) UNIQUE, -- 姓名，加上唯一约束

    age INT,                  -- 年龄

    department_did INT,

    FOREIGN KEY (department_did) REFERENCES departments(did)

);
实体类：

ForeignKeyInfo.java:

package com.example.drcbackend.entity.dto;

import lombok.Data;

@Data
public class ForeignKeyInfo {
    private String columnName;
    private String referencedTableName;
    private String referencedColumnName;

    public ForeignKeyInfo(String columnName, String referencedTableName, String referencedColumnName) {
        this.columnName = columnName;
        this.referencedTableName = referencedTableName;
        this.referencedColumnName = referencedColumnName;
    }
}

TableSchema.java:

 

package com.example.drcbackend.entity.dto;

import lombok.Data;
import java.util.List;

@Data
public class TableSchema {
    private String name;
    private List<String> primaryKeyColumns;
    private List<ForeignKeyInfo> foreignKeyColumns;
    private List<String> uniqueColumns;
}
核心解析部分：

package com.example.drcbackend.util;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLName;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
import com.alibaba.druid.sql.parser.SQLStatementParser;
import com.example.drcbackend.entity.dto.ForeignKeyInfo;
import com.example.drcbackend.entity.dto.TableSchema;

import java.util.ArrayList;
import java.util.List;

public class SqlParserUtil {//TODO 同一个sql文件只有reRun才能继续检测，造孽啊

    public static List<TableSchema> extractAllSchemas(String sql) {
        SQLStatementParser parser = new MySqlStatementParser(sql);
        List<SQLStatement> stmts = parser.parseStatementList();
        List<TableSchema> schemas = new ArrayList<>();

        for (SQLStatement stmt : stmts) {
            if (stmt instanceof SQLCreateTableStatement) {
                SQLCreateTableStatement create = (SQLCreateTableStatement) stmt;
                TableSchema schema = parseCreateStatement(create);
                schemas.add(schema);
            }
        }

        printSchemas(schemas);  // 调试输出
        return schemas;
    }

    private static TableSchema parseCreateStatement(SQLCreateTableStatement create) {
        TableSchema schema = new TableSchema();
        schema.setName(create.getName().getSimpleName());

        List<String> primaryKeys = new ArrayList<>();
        List<ForeignKeyInfo> foreignKeys = new ArrayList<>();
        List<String> uniqueColumns = new ArrayList<>();

        for (SQLTableElement element : create.getTableElementList()) {
            if (element instanceof SQLPrimaryKey) {
                // 处理表级主键约束（如：PRIMARY KEY (id, name)）
                SQLPrimaryKey pk = (SQLPrimaryKey) element;
                for (SQLSelectOrderByItem col : pk.getColumns()) {
                    primaryKeys.add(extractColumnName(col.getExpr()));
                }
            } else if (element instanceof SQLForeignKeyConstraint) {
                // 处理外键约束
                SQLForeignKeyConstraint fk = (SQLForeignKeyConstraint) element;
                String refTable = fk.getReferencedTableName().getSimpleName();
                List<SQLName> srcCols = fk.getReferencingColumns();
                List<SQLName> refCols = fk.getReferencedColumns();

                for (int i = 0; i < srcCols.size(); i++) {
                    foreignKeys.add(new ForeignKeyInfo(
                            srcCols.get(i).getSimpleName(),
                            refTable,
                            refCols.get(i).getSimpleName()
                    ));
                }
            } else if (element instanceof SQLUnique) {
                // 处理唯一约束（如：UNIQUE (name)）
                SQLUnique unique = (SQLUnique) element;
                for (SQLSelectOrderByItem col : unique.getColumns()) {
                    uniqueColumns.add(extractColumnName(col.getExpr()));
                }
            } else if (element instanceof SQLColumnDefinition) {
                // 处理列定义中的主键和唯一约束
                SQLColumnDefinition column = (SQLColumnDefinition) element;
                String columnName = column.getName().getSimpleName(); // 修正获取列名的方式

                // 检测列级主键
                if (isPrimaryKeyColumn(column)) {
                    primaryKeys.add(columnName);
                }

                // 检测列级唯一约束
                if (isUniqueColumn(column)) {
                    uniqueColumns.add(columnName);
                }
            }
        }

        schema.setPrimaryKeyColumns(primaryKeys);
        schema.setForeignKeyColumns(foreignKeys);
        schema.setUniqueColumns(uniqueColumns);
        return schema;
    }

    private static boolean isPrimaryKeyColumn(SQLColumnDefinition column) {
        return column.getConstraints().stream().anyMatch(
                c -> c instanceof SQLColumnPrimaryKey || "PRIMARY KEY".equalsIgnoreCase(c.toString())
        );
    }

    private static boolean isUniqueColumn(SQLColumnDefinition column) {
        return column.getConstraints().stream().anyMatch(
                c -> c instanceof SQLColumnUniqueKey || "UNIQUE".equalsIgnoreCase(c.toString())
        );
    }

    private static String extractColumnName(SQLExpr expr) {
        if (expr instanceof SQLIdentifierExpr) {
            return ((SQLIdentifierExpr) expr).getName();
        } else if (expr instanceof SQLPropertyExpr) {
            return ((SQLPropertyExpr) expr).getName();
        }
        return expr.toString();  // 其他情况返回字符串表示
    }

    private static void printSchemas(List<TableSchema> schemas) {
        for (TableSchema schema : schemas) {
            System.out.println("\n表名: " + schema.getName());
            System.out.println("主键列: " + schema.getPrimaryKeyColumns());

            System.out.println("外键列:");
            for (ForeignKeyInfo fk : schema.getForeignKeyColumns()) {
                System.out.printf("  %s → %s(%s)\n",
                        fk.getColumnName(),
                        fk.getReferencedTableName(),
                        fk.getReferencedColumnName()
                );
            }

            System.out.println("唯一列: " + schema.getUniqueColumns());
        }
        System.out.println("\n解析完成");
    }
}