java - 使用 Spring 的 ResourceDatabasePopulator 导入 MySQL 表定义转储

Question

我们有一个 Spring 测试执行监听器，它通过 Spring 的 ResourceDatabasePopulator 导入转储的 MySQL 表定义:

public class DBSetupExecutionListener extends AbstractTestExecutionListener {
    ...

    @Override
    public synchronized void beforeTestClass(TestContext testContext) {
        ...
        // import the table definitions from a previously dumped file
        ResourceDatabasePopulator populator = new ResourceDatabasePopulator();
        populator.addScript(new ClassPathResource("some-table-definitions.sql"));
        populator.execute(dataSource);
        ...
    }

    ...
}

导入的 MySQL 转储文件确实如下所示(简化):

/*!40101 SET @OLD_CHARACTER_SET_CLIENT=@@CHARACTER_SET_CLIENT */;
/*!40101 SET @OLD_CHARACTER_SET_RESULTS=@@CHARACTER_SET_RESULTS */;
/*!40101 SET @OLD_COLLATION_CONNECTION=@@COLLATION_CONNECTION */;
/*!40101 SET NAMES utf8 */;
/*!40014 SET @OLD_FOREIGN_KEY_CHECKS=@@FOREIGN_KEY_CHECKS, FOREIGN_KEY_CHECKS=0 */;
/*!40101 SET @OLD_SQL_MODE=@@SQL_MODE, SQL_MODE='NO_AUTO_VALUE_ON_ZERO' */;
/*!40111 SET @OLD_SQL_NOTES=@@SQL_NOTES, SQL_NOTES=0 */;

DROP TABLE IF EXISTS `someTable`;

CREATE TABLE `someTable`(
  `id` bigint(20) NOT NULL AUTO_INCREMENT,
  `other_table_id` bigint(20) NOT NULL,
   ...
   PRIMARY KEY (`id`)
   KEY `FK_OTHERTABLE_ID`(`other_table_id`),
   CONSTRAINT `FK_OTHERTABLE_ID` FOREIGN KEY (`other_table_id`) REFERENCES `otherTable`(`id`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8;

DROP TABLE IF EXISTS `otherTable`;

CREATE TABLE `otherTable`(
  `id` bigint(20) NOT NULL AUTO_INCREMENT,
  ...
  PRIMARY KEY (`id`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8;

...

/*!40111 SET SQL_NOTES=@OLD_SQL_NOTES */;
/*!40101 SET SQL_MODE=@OLD_SQL_MODE */;
/*!40014 SET FOREIGN_KEY_CHECKS=@OLD_FOREIGN_KEY_CHECKS */;
/*!40101 SET CHARACTER_SET_CLIENT=@OLD_CHARACTER_SET_CLIENT */;
/*!40101 SET CHARACTER_SET_RESULTS=@OLD_CHARACTER_SET_RESULTS */;
/*!40101 SET COLLATION_CONNECTION=@OLD_COLLATION_CONNECTION */;

通过 command line 导入转储时工作没有任何问题，ResourceDatabasePopulator 忽略预处理器指令，因此在创建实际引用表之前无法首先创建引用表。更改表定义的顺序确实可以解决问题，但对于多个表来说有点乏味 - 特别是当您将当前表定义转储到文件时。

由于导入转储的 SQL 表定义是一项非常常见的任务，我想我在这里做错了什么。有没有办法告诉ResourceDatabasePopulator遵守预处理指令，从而防止在执行脚本时进行外键检查？

Answer 1

我听从了 @M.Deinum 的建议，实现了我自己版本的 ResourceDatabasePopulator。由于成员范围有限并且没有任何 getter 来检索这些字段，不幸的是，原始 ResourceDatabasePopulator 实现中存在一些复制和粘贴。这个类绝对不是为支持子类化而设计的......

下面是在我的用例中适用的当前解决方案。它并不完美，可能是针对 MySQL 语句量身定制的，但至少它完成了工作。

import java.io.BufferedReader;
import java.io.IOException;
import java.lang.invoke.MethodHandles;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.EncodedResource;
import org.springframework.jdbc.datasource.init.DatabasePopulatorUtils;
import org.springframework.jdbc.datasource.init.ResourceDatabasePopulator;
import org.springframework.jdbc.datasource.init.ScriptUtils;
import org.springframework.jdbc.datasource.init.UncategorizedScriptException;
import org.springframework.util.StringUtils;

/**
 * Populates, initializes, or cleans up a database using SQL scripts defined in
 * external resources.
 * <p>
 * This implementation will clean up a given script by removing comments and processing instructions
 * contained in the script and perform a table reordering based on defined foreign keys in a
 * best-effort attempt.
 *
 * <ul>
 * <li>Call {@link #addScript} to add a single SQL script location.
 * <li>Call {@link #addScripts} to add multiple SQL script locations.
 * <li>Consult the setter methods in this class for further configuration options.
 * <li>Call {@link #populate} or {@link #execute} to initialize or clean up the
 * database using the configured scripts.
 * </ul>
 *
 * @author Keith Donald
 * @author Dave Syer
 * @author Juergen Hoeller
 * @author Chris Beams
 * @author Oliver Gierke
 * @author Sam Brannen
 * @author Chris Baldwin
 * @author Roman Vottner
 *
 * @since 3.0
 * @see DatabasePopulatorUtils
 * @see ScriptUtils
 */
public class OrderedResourceDatabasePopulator extends ResourceDatabasePopulator {

  private static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

  protected String sqlScriptEncoding;

  /**
   * Construct a new {@code OrderedResourceDatabasePopulator} with default settings.
   */
  public OrderedResourceDatabasePopulator() {
    /* no-op */
  }

  /**
   * Construct a new {@code OrderedResourceDatabasePopulator} with default settings
   * for the supplied scripts.
   * 
   * @param scripts the scripts to execute to initialize or clean up the database
   *                (never {@code null})
   */
  public OrderedResourceDatabasePopulator(Resource... scripts) {
    this();
    setScripts(cleanResource(scripts));
  }

  /**
   * Construct a new {@code OrderedResourceDatabasePopulator} with the supplied values.
   *
   * @param continueOnError flag to indicate that all failures in SQL should be logged but not cause
   *                        a failure
   * @param ignoreFailedDrops flag to indicate that a failed SQL {@code DROP} statement can be
   *                          ignored
   * @param sqlScriptEncoding the encoding for the supplied SQL scripts; may be {@code null} or
   *                           <em>empty</em> to indicate platform encoding
   * @param scripts the scripts to execute to initialize or clean up the database (never {@code
   *                null})
   */
  public OrderedResourceDatabasePopulator(boolean continueOnError, boolean ignoreFailedDrops,
                                          String sqlScriptEncoding, Resource... scripts) {

    super(continueOnError, ignoreFailedDrops, sqlScriptEncoding, scripts);
  }

  /**
   * Add a script to execute to initialize or clean up the database.
   *
   * @param script the path to an SQL script (never {@code null})
   */
  @Override
  public void addScript(Resource script) {
    super.addScript(cleanResource(script)[0]);
  }

  /**
   * Add multiple scripts to execute to initialize or clean up the database.
   * 
   * @param scripts the scripts to execute (never {@code null})
   */
  @Override
  public void addScripts(Resource... scripts) {
    super.addScripts(cleanResource(scripts));
  }

  /**
   * Set the scripts to execute to initialize or clean up the database, replacing any previously 
   * added scripts.
   * 
   * @param scripts the scripts to execute (never {@code null})
   */
  @Override
  public void setScripts(Resource... scripts) {
    super.setScripts(cleanResource(scripts));
  }

  /**
   * Specify the encoding for the configured SQL scripts, if different from the platform encoding.
   * 
   * @param sqlScriptEncoding the encoding used in scripts; may be {@code null} or empty to indicate
   *                          platform encoding
   * 
   * @see #addScript(Resource)
   */
  @Override
  public void setSqlScriptEncoding(String sqlScriptEncoding) {
    super.setSqlScriptEncoding(sqlScriptEncoding);
    this.sqlScriptEncoding = StringUtils.hasText(sqlScriptEncoding) ? sqlScriptEncoding : null;
  }

  private Resource[] cleanResource(Resource... scripts) {
    List<Resource> cleaned = new ArrayList<>();
    for (Resource script : scripts) {
      EncodedResource encodedScript = new EncodedResource(script, this.sqlScriptEncoding);

      StringBuilder sb = new StringBuilder();
      try (BufferedReader in = new BufferedReader(encodedScript.getReader())) {
        String line;

        // parse script
        Map<String, Table> tables = new LinkedHashMap<>();
        Table curTable = null;
        while ((line = in.readLine()) != null) {
          if (!line.startsWith("#") && !line.startsWith("/*")) {
            line = line.replaceAll("`", "");
            String lowerCaseLine = line.toLowerCase();

            // parse table definitions
            if (lowerCaseLine.startsWith("drop table") || "".equals(line.trim())) {
              if (curTable != null) {
                curTable = null;
              }
              continue;
            }

            curTable = parseTable(line, tables, curTable);
          }
        }

        // check the order of table definitions and reorder them to ensure referenced tables are
        // defined before table specifying the foreign key
        Map<String, Table> orderedTables = new LinkedHashMap<>();
        List<Table> backlog = new ArrayList<>();
        for (String tableName : tables.keySet()) {
          // check for available foreign key definitions
          List<String> referencedTables = tables.get(tableName).getReferencedTables();
          // if no foreign keys are defined in the table we can add the table to the current set
          if (referencedTables.isEmpty()) {
            orderedTables.put(tableName, tables.get(tableName));
            processBacklog(backlog, orderedTables);
            continue;
          }

          // foreign keys are defined. If any of the referenced tables is not yet in the set move it
          // to the backlog table and try it later
          boolean dependenciesAvailable = true;
          for (String dependency : referencedTables) {
            if (!orderedTables.containsKey(dependency)) {
              backlog.add(tables.get(tableName));
              dependenciesAvailable = false;
              break;
            }
          }

          // if all referred tables are available in the set add the current table also to the
          // ordered set
          if (dependenciesAvailable) {
            orderedTables.put(tableName, tables.get(tableName));
            processBacklog(backlog, orderedTables);
            continue;
          }

          // probe tables kept in the backlog again
          processBacklog(backlog, orderedTables);
        }

        while (!backlog.isEmpty()) {
          processBacklog(backlog, orderedTables);
        }

        ArrayList<Table> reverse = new ArrayList<>(orderedTables.values());
        ListIterator<Table> reverseIter = reverse.listIterator(reverse.size());
        sb.append("DROP TABLE IF EXISTS ");
        int originalSize = sb.length();
        while (reverseIter.hasPrevious()) {
          if (sb.length() > originalSize) {
            sb.append(", ");
          }
          sb.append(reverseIter.previous().getName());
        }
        sb.append(";\n\n");

        for (Table table : orderedTables.values()) {
          sb.append(table.getCreateStatement()).append("\n");
        }

      } catch (IOException ex) {
        throw new UncategorizedScriptException(ex.getMessage(), ex);
      }

      if (LOG.isTraceEnabled()) {
        LOG.trace("SQL script after cleaning: \n{}", sb.toString());
      }
      Resource cleanedResource = new ByteArrayResource(sb.toString().getBytes());
      cleaned.add(cleanedResource);
    }

    return cleaned.toArray(new Resource[cleaned.size()]);
  }

  private Table parseTable(String line, Map<String, Table> tables, Table curTable) {
    String lowerCaseLine = line.toLowerCase();
    if (lowerCaseLine.startsWith("create table")) {
      String name = line.substring("create table ".length(), line.indexOf("(")).trim();
      curTable = new Table(name);
      tables.put(name, curTable);
    } else if (curTable != null) {
      if (lowerCaseLine.contains("primary key")) {
        curTable.setPrimaryKey(line);
      } else if (lowerCaseLine.contains("unique key")) {
        curTable.addUniqueKey(line);
      } else if (lowerCaseLine.contains("foreign key")) {
        curTable.addForeignKey(line);
      } else if (lowerCaseLine.contains(" key ")) {
        curTable.addIndex(line);
      } else if (lowerCaseLine.contains(" charset=") || lowerCaseLine.contains("engine=")) {
        curTable.setMetaData(line);
      } else {
        curTable.addColumn(line);
      }
    }
    return curTable;
  }

  private void processBacklog(List<Table> backlog, Map<String, Table> orderedTables) {
    Iterator<Table> iter = backlog.iterator();
    while (iter.hasNext()) {
      Table table = iter.next();
      boolean allDependenciesAvailable = true;
      for (String dependency : table.getReferencedTables()) {
        if (!orderedTables.containsKey(dependency)) {
          allDependenciesAvailable = false;
        }

        if (allDependenciesAvailable) {
          orderedTables.put(table.getName(), table);
          iter.remove();
        }
      }
    }
  }

  private class Table {

    private final String name;
    private List<String> columns = new ArrayList<>();
    private String primaryKey;
    private List<String> uniqueKeys = new ArrayList<>();
    private List<String> indices = new ArrayList<>();
    private List<String> foreignKeys = new ArrayList<>();
    private String metaData;

    public Table(String name) {
      this.name = name;
    }

    String getName() {
      return this.name;
    }

    void addColumn(String column) {
      this.columns.add(column);
    }

    void setPrimaryKey(String primaryKey) {
      this.primaryKey = primaryKey;
    }

    void addUniqueKey(String uniqueKey) {
      this.uniqueKeys.add(uniqueKey);
    }

    void addIndex(String index) {
      this.indices.add(index);
    }

    void addForeignKey(String foreignKey) {
      this.foreignKeys.add(foreignKey);
    }

    List<String> getReferencedTables() {
      List<String> referencedTables = new ArrayList<>();
      for (String foreignKey : foreignKeys) {
        int start = foreignKey.toLowerCase().indexOf("references ") + "references ".length();
        String table = foreignKey.substring(start, foreignKey.indexOf(" ", start));
        referencedTables.add(table);
      }
      return referencedTables;
    }

    void setMetaData(String metaData) {
      this.metaData = metaData;
    }

    String getCreateStatement() {
      StringBuilder sb = new StringBuilder();
      sb.append("CREATE TABLE ").append(this.name).append(" (");
      for (String column : this.columns) {
        sb.append("\n");
        sb.append("  ").append(column);
      }
      if (null != primaryKey) {
        sb.append("\n");
        sb.append("  ").append(this.primaryKey);
      }
      if (!uniqueKeys.isEmpty()) {
        for (String uniqueKey : uniqueKeys) {
          sb.append("\n");
          sb.append("  ").append(uniqueKey);
        }
      }
      if (!indices.isEmpty()) {
        for (String index : indices) {
          sb.append("\n");
          sb.append("  ").append(index);
        }
      }
      if (!foreignKeys.isEmpty()) {
        for (String foreignKey : foreignKeys) {
          sb.append("\n");
          sb.append("  ").append(foreignKey);
        }
      }
      sb.append("\n");
      if (metaData == null || !metaData.contains(")")) {
        sb.append(") ");
      }
      if (metaData != null) {
        sb.append(metaData);
      }
      sb.append("\n");

      return sb.toString();
    }
  }
}

此填充器将首先清除给定脚本中任何不需要的字符、注释或处理指令，然后构建虚拟表结构，以便在声明外部引用的表之前对作为外部引用目标的表进行排序。表定义虚拟化后，填充器将重新创建脚本，方法是首先按虚拟模型的相反顺序为每个表添加 DROP TABLE 语句，然后添加 CREATE TABLE 语句按计算顺序排列的语句。

请注意，此版本尚不支持脚本中包含的其他 SQL 命令(例如 INSERT，...)。如果需要，请随意根据您的喜好自定义此代码。

为了使用重新排序的 SQL 脚本，只需将 ResourceDatabasePopulator 替换为 OrderedResourceDatabasePopulator

@Override
public synchronized void beforeTestClass(TestContext testContext) {
    ...
    // import the table definitions from a previously dumped file
    ResourceDatabasePopulator populator = new OrderedResourceDatabasePopulator();
    populator.addScript(new ClassPathResource("some-table-definitions.sql"));
    populator.execute(dataSource);
    ...
}