org.unitils.dbmaintainer.structure.impl.DtdDataSetStructureGenerator Maven / Gradle / Ivy

Go to download
/*
 * Copyright 2008,  Unitils.org
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.unitils.dbmaintainer.structure.impl;

import static org.unitils.thirdparty.org.apache.commons.io.IOUtils.closeQuietly;
import static org.unitils.thirdparty.org.apache.commons.io.IOUtils.write;

import org.apache.commons.lang.StringUtils;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.FilteredDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.filter.IncludeTableFilter;
import org.dbunit.dataset.xml.FlatDtdWriter;
import org.unitils.core.UnitilsException;
import org.unitils.dbmaintainer.structure.DataSetStructureGenerator;
import org.unitils.dbmaintainer.util.BaseDatabaseAccessor;
import org.unitils.thirdparty.org.apache.commons.dbutils.DbUtils;
import org.unitils.util.PropertyUtils;

import java.io.File;
import java.io.FileWriter;
import java.io.StringWriter;
import java.io.Writer;
import java.sql.Connection;
import java.util.Properties;
import java.util.Set;

/**
 * Implementation of {@link DataSetStructureGenerator} for the DbUnit {@link FlatXmlDataSet} XML test data files format
 * 
 * todo test and fix for hsqldb (see sample project)
 *
 * @author Filip Neven
 * @author Tim Ducheyne
 */
public class DtdDataSetStructureGenerator extends BaseDatabaseAccessor implements DataSetStructureGenerator {

    /* Property key of the filename of the generated DTD  */
    public static final String PROPKEY_DTD_FILENAME = "dtdGenerator.dtd.filename";

    /* The DTD file name */
    private String dtdFileName;


    /**
     * Initializes the generator by retrieving the name for the DTD file.
     *
     * @param configuration The config, not null
     */
    @Override
    protected void doInit(Properties configuration) {
        dtdFileName = PropertyUtils.getString(PROPKEY_DTD_FILENAME, configuration);
    }


    /**
     * Generates the DTD, and writes it to the file specified by the property {@link #PROPKEY_DTD_FILENAME}.
     * The DTD will contain the structure of the database. All tables will be written as optional elements and
     * all columns will be optional attributes.
     */
    public void generateDataSetStructure() {
        Writer writer = null;
        try {
            // creates the DTD file
            File dtdFile = new File(dtdFileName);
            File parentDirectory = dtdFile.getParentFile();
            if (parentDirectory == null) {
                throw new UnitilsException("Error generating DTD file. Could not find parent directory for DTD file: " + dtdFileName);
            }
            parentDirectory.mkdirs();

            // generate content as a string
            String dtdContent = generateDtdContent();
            // make all elements optional
            dtdContent = StringUtils.replace(dtdContent, "#REQUIRED\n", "#IMPLIED\n");

            // write the content to the file
            writer = new FileWriter(dtdFile);
            write(dtdContent, writer);

        } catch (UnitilsException e) {
            throw e;
        } catch (Exception e) {
            throw new UnitilsException("Error generating DTD file: " + dtdFileName, e);
        } finally {
            closeQuietly(writer);
        }
    }


    /**
     * Generates the actual content of the DTD file as an in-memory string.
     *
     * @return the DTD content, not null
     */
    protected String generateDtdContent() {
        Connection conn = null;
        try {
            conn = sqlHandler.getDataSource().getConnection();
            IDatabaseConnection dbUnitDatabaseConnection = new DatabaseConnection(conn, defaultDbSupport.getSchemaName());

            StringWriter stringWriter = new StringWriter();

            FlatDtdWriter datasetWriter = new FlatDtdWriter(stringWriter);
            datasetWriter.setContentModel(FlatDtdWriter.CHOICE);

            // create a dataset for the database content
            // filter out all system table names
            Set tableNames = defaultDbSupport.getTableNames();
            IDataSet actualDataSet = dbUnitDatabaseConnection.createDataSet();
            IDataSet filteredActualDataSet = new FilteredDataSet(new IncludeTableFilter(tableNames.toArray(new String[0])), actualDataSet);

            datasetWriter.write(filteredActualDataSet);
            return stringWriter.toString();

        } catch (Exception e) {
            throw new UnitilsException("Error generating content for DTD file.", e);
        } finally {
            DbUtils.closeQuietly(conn);
        }
    }
}