YamlWriter.java

  1. /*
  2.  *
  3.  * The DbUnit Database Testing Framework
  4.  * Copyright (C)2002-2004, DbUnit.org
  5.  *
  6.  * This library is free software; you can redistribute it and/or
  7.  * modify it under the terms of the GNU Lesser General Public
  8.  * License as published by the Free Software Foundation; either
  9.  * version 2.1 of the License, or (at your option) any later version.
  10.  *
  11.  * This library is distributed in the hope that it will be useful,
  12.  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13.  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14.  * Lesser General Public License for more details.
  15.  *
  16.  * You should have received a copy of the GNU Lesser General Public
  17.  * License along with this library; if not, write to the Free Software
  18.  * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  19.  *
  20.  */

  21. package org.dbunit.dataset.yaml;

  22. import org.dbunit.dataset.Column;
  23. import org.dbunit.dataset.DataSetException;
  24. import org.dbunit.dataset.IDataSet;
  25. import org.dbunit.dataset.ITable;
  26. import org.dbunit.dataset.ITableIterator;
  27. import org.dbunit.dataset.ITableMetaData;
  28. import org.yaml.snakeyaml.DumperOptions;
  29. import org.yaml.snakeyaml.Yaml;

  30. import java.io.Writer;
  31. import java.util.ArrayList;
  32. import java.util.LinkedHashMap;
  33. import java.util.List;

  34. /**
  35.  * @author Björn Beskow
  36.  * @version $Revision$ $Date$
  37.  */
  38. class YamlWriter
  39. {

  40.     private Yaml _yaml;
  41.     private Writer _out;
  42.     private boolean _useFlowStyle;

  43.     public YamlWriter(Writer out)
  44.     {
  45.         this(out, false);
  46.     }

  47.     public YamlWriter(Writer out, boolean useFlowStyle)
  48.     {
  49.         this._out = out;
  50.         this._useFlowStyle = useFlowStyle;
  51.     }

  52.     public void setUseFlowStyle(boolean useFlowStyle)
  53.     {
  54.         this._useFlowStyle = useFlowStyle;
  55.     }

  56.     public void write(IDataSet dataSet) throws DataSetException
  57.     {
  58.         LinkedHashMap<String, List<LinkedHashMap<String, Object>>> dataSetAsMap = dataSetAsMap(dataSet);
  59.         DumperOptions options = new DumperOptions();
  60.         if (_useFlowStyle)
  61.         {
  62.             options.setDefaultFlowStyle(DumperOptions.FlowStyle.FLOW);
  63.             options.setPrettyFlow(true);
  64.         } else
  65.         {
  66.             options.setDefaultFlowStyle(DumperOptions.FlowStyle.BLOCK);
  67.             options.setIndent(4);
  68.             options.setIndicatorIndent(2);
  69.         }
  70.         _yaml = new Yaml(options);
  71.         _yaml.dump(dataSetAsMap, _out);
  72.     }

  73.     private LinkedHashMap<String, List<LinkedHashMap<String, Object>>> dataSetAsMap(IDataSet dataSet) throws DataSetException
  74.     {
  75.         LinkedHashMap<String, List<LinkedHashMap<String, Object>>> result = new LinkedHashMap<>();
  76.         ITableIterator iterator = dataSet.iterator();
  77.         while (iterator.next())
  78.         {
  79.             ITableMetaData tableMetaData = iterator.getTableMetaData();
  80.             List<LinkedHashMap<String, Object>> tableRows = new ArrayList<>();
  81.             ITable table = iterator.getTable();
  82.             for (int row = 0; row < table.getRowCount(); row++)
  83.             {
  84.                 LinkedHashMap<String, Object> rowMap = new LinkedHashMap<>();
  85.                 for (Column column : tableMetaData.getColumns())
  86.                 {
  87.                     String columnName = column.getColumnName();
  88.                     Object value = table.getValue(row, columnName);
  89.                     if (value != null)
  90.                     {
  91.                         rowMap.put(columnName, value);
  92.                     }
  93.                 }
  94.                 tableRows.add(rowMap);
  95.             }
  96.             result.put(tableMetaData.getTableName(), tableRows);
  97.         }
  98.         return result;
  99.     }
  100. }