View Javadoc

1   /**
2    * Copyright 2008 Bryan Ray
3    * 
4    * Licensed under the Apache License, Version 2.0 (the "License"); 
5    * you may not use this file except in compliance with the License.
6    * You may obtain a copy of the License at 
7    * 
8    * http://www.apache.org/licenses/LICENSE-2.0 
9    * 
10   * Unless required by applicable law or agreed to in writing, software 
11   * distributed under the License is distributed on an "AS IS" BASIS,
12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
13   * See the License for the specific language governing permissions and 
14   * limitations under the License. 
15   */
16  package uk.org.bryanray.testtoys.recorder;
17  
18  import java.io.File;
19  import java.io.FileInputStream;
20  import java.io.FileNotFoundException;
21  import java.io.FileOutputStream;
22  import java.io.IOException;
23  import java.io.InputStream;
24  import java.io.OutputStream;
25  import java.sql.Connection;
26  import java.sql.DriverManager;
27  import java.sql.SQLException;
28  import java.util.ArrayList;
29  import java.util.HashMap;
30  import java.util.Iterator;
31  import java.util.List;
32  import java.util.Map;
33  import java.util.Properties;
34  import java.util.regex.Matcher;
35  import java.util.regex.Pattern;
36  
37  import org.apache.commons.cli.CommandLine;
38  import org.apache.commons.cli.CommandLineParser;
39  import org.apache.commons.cli.GnuParser;
40  import org.apache.commons.cli.HelpFormatter;
41  import org.apache.commons.cli.Option;
42  import org.apache.commons.cli.OptionBuilder;
43  import org.apache.commons.cli.Options;
44  import org.apache.commons.cli.ParseException;
45  import org.dbunit.database.DatabaseConfig;
46  import org.dbunit.database.DatabaseConnection;
47  import org.dbunit.database.IDatabaseConnection;
48  import org.dbunit.database.QueryDataSet;
49  import org.dbunit.dataset.IDataSet;
50  import org.dbunit.dataset.xml.FlatXmlDataSet;
51  
52  import uk.org.bryanray.testtoys.util.ExceptionUtil;
53  import uk.org.bryanray.testtoys.util.TemplateUtil;
54  
55  /**
56   * <DL>
57   * <DT><B>Project:</B></DT>
58   * <DD>TestDataLoader</DD>
59   * <DT><B>Filename:</B></DT>
60   * <DD>Recorder.java</DD>
61   * <DT><B>Creation Date:</B></DT>
62   * <DD>18 May 2008</DD>
63   * </DL>
64   * 
65   * @author Bryan Ray
66   */
67  public class Recorder {
68  
69  	private Properties databaseProperties = new Properties(System
70  			.getProperties());
71  
72  	private String databasePropertyPrefix = "db1.";
73  
74  	private List<String> queryList;
75  	private Map<String, String> dbUnitProperties;
76  	private Map<String, String> dbUnitFeatures;
77  
78  	/**
79  	 * A regular expression that is used to get the table name from a SQL
80  	 * 'select' statement. This pattern matches a string that starts with any
81  	 * characters, followed by the case-insensitive word 'from', followed by a
82  	 * table name of the form 'foo' or 'schema.foo', followed by any number of
83  	 * remaining characters.
84  	 */
85  	private static final Pattern TABLE_MATCH_PATTERN = Pattern.compile(
86  			".*\\s+from\\s+(\\w+(\\.\\w+)?).*", Pattern.CASE_INSENSITIVE);
87  
88  	public static void main(String[] args) {
89  		Recorder recorder = new Recorder();
90  
91  		// Read configuration file
92  		Options commandOptions = getCommandOptions();
93  		CommandLine commandLine = parseCommandArguments(args, commandOptions);
94  		if (commandLine.hasOption("configfile")) {
95  			String configFileName = commandLine.getOptionValue("configfile");
96  			try {
97  				recorder.readProperties(configFileName);
98  			} catch (FileNotFoundException fnfe) {
99  				System.err.println("Config file could not be read.");
100 				return;
101 			}
102 			Map<String, String> dbUnitFeatures = new HashMap<String, String>();
103 			dbUnitFeatures.put(
104 					"http://www.dbunit.org/features/qualifiedTableNames",
105 					"true");
106 			recorder.setDbUnitFeatures(dbUnitFeatures);
107 
108 			recorder.initialiseQueries();
109 
110 			File outputFile = new File(recorder.getOutputFile());
111 			try {
112 				OutputStream outputStream = new FileOutputStream(outputFile);
113 				recorder.generateDataSet(outputStream);
114 			} catch (Exception e) {
115 				throw ExceptionUtil.convertException(e);
116 			}
117 		} else {
118 			HelpFormatter formatter = new HelpFormatter();
119 			formatter.printHelp("generator", commandOptions);
120 		}
121 	}
122 
123 	private void initialiseQueries() {
124 		queryList = new ArrayList<String>();
125 		for (Iterator<Map.Entry<Object, Object>> propertiesIterator = databaseProperties
126 				.entrySet().iterator(); propertiesIterator.hasNext();) {
127 			Map.Entry<Object, Object> entry = propertiesIterator.next();
128 			String propertyName = (String) entry.getKey();
129 			if (propertyName.startsWith("query.")) {
130 				String propertyValue = (String) entry.getValue();
131 				queryList.add(propertyValue);
132 			}
133 		}
134 	}
135 
136 	private void readProperties(String propertiesArgument)
137 			throws FileNotFoundException {
138 		String[] propertiesFiles = propertiesArgument.split(",");
139 		for (int i = 0; i < propertiesFiles.length; i++) {
140 			File propertiesFile = new File(propertiesFiles[i]);
141 			InputStream inStream = new FileInputStream(propertiesFile);
142 			this.initialiseProperties(inStream);
143 		}
144 	}
145 
146 	private String getOutputFile() {
147 		return databaseProperties.getProperty("outputfile");
148 	}
149 
150 	/**
151 	 * @param commandArguments
152 	 *            The command arguments passed to {@link #main(String[]).
153 	 * @param options
154 	 *            The options builder from {@link #getCommandOptions().
155 	 */
156 	private static CommandLine parseCommandArguments(String[] commandArguments,
157 			Options options) {
158 		CommandLineParser parser = new GnuParser();
159 		try {
160 			return parser.parse(options, commandArguments);
161 		} catch (ParseException pe) {
162 			throw ExceptionUtil.convertException(pe);
163 		}
164 	}
165 
166 	/**
167 	 * @return A command line object that follows the builder pattern.
168 	 */
169 	@SuppressWarnings("static-access")
170 	private static Options getCommandOptions() {
171 		Options options = new Options();
172 		Option configFile = OptionBuilder.withArgName("file").hasArg()
173 				.withDescription("use given configuration file").create(
174 						"configfile");
175 		options.addOption(configFile);
176 		return options;
177 	}
178 
179 	/**
180 	 * Loads properties into {@link Recorder} instance, from the input stream.
181 	 * 
182 	 * @param inStream
183 	 *            The input stream that contains the properties.
184 	 */
185 	private void initialiseProperties(InputStream inStream) {
186 		if (null != inStream) {
187 			try {
188 				databaseProperties.load(inStream);
189 			} catch (IOException ioe) {
190 				throw new RuntimeException(ioe);
191 			}
192 		}
193 	}
194 
195 	/**
196 	 * Returns the value for the property from the initialised properties.
197 	 * 
198 	 * @param name
199 	 *            The property name.
200 	 * @return The property value.
201 	 */
202 	private String dbProperty(String name) {
203 		String property = databaseProperties.getProperty(databasePropertyPrefix
204 				+ name);
205 		return TemplateUtil.substituteTemplatePlaceholders(property,
206 				databaseProperties);
207 	}
208 
209 	/**
210 	 * Creates a database connection from the properties.
211 	 * 
212 	 * @return A database connection.
213 	 */
214 	private Connection createConnection() throws SQLException,
215 			ClassNotFoundException {
216 		String driverClassName = dbProperty("driver");
217 		Class.forName(driverClassName);
218 		String url = dbProperty("url");
219 		Connection connection = DriverManager.getConnection(url,
220 				dbProperty("username"), dbProperty("password"));
221 		return connection;
222 	}
223 
224 	private void generateDataSet(OutputStream outputStream) {
225 		Connection connection = null;
226 		try {
227 			connection = createConnection();
228 			IDatabaseConnection dbConnection = new DatabaseConnection(
229 					connection);
230 			configConnection(dbConnection);
231 			if (null != queryList && !queryList.isEmpty()) {
232 				// Partial database export
233 				QueryDataSet partialDataSet = new QueryDataSet(dbConnection);
234 				addQueries(partialDataSet);
235 				FlatXmlDataSet.write(partialDataSet, outputStream);
236 			} else {
237 				// Full database export
238 				IDataSet fullDataSet = dbConnection.createDataSet();
239 				FlatXmlDataSet.write(fullDataSet, outputStream);
240 			}
241 		} catch (Exception e) {
242 			throw ExceptionUtil.convertException(e);
243 		} finally {
244 			if (connection != null) {
245 				try {
246 					connection.close();
247 				} catch (SQLException se) {
248 					throw ExceptionUtil.convertException(se);
249 				}
250 			}
251 		}
252 	}
253 
254 	private void addQueries(QueryDataSet dataSet) {
255 		if (queryList == null)
256 			return;
257 		for (Iterator<String> queryIterator = queryList.iterator(); queryIterator.hasNext();) {
258 			String query = queryIterator.next();
259 			Matcher m = TABLE_MATCH_PATTERN.matcher(query);
260 			if (!m.matches()) {
261 				System.out.println("Unable to parse query. Ignoring '" + query
262 						+ "'.");
263 			} else {
264 				String table = m.group(1);
265 				// TODO Only add if the table has not been added
266 				dataSet.addTable(table, query);
267 			}
268 		}
269 	}
270 
271 	// Connection configuration
272 	public void setDbUnitFeatures(Map<String, String> dbUnitFeatures) {
273 		this.dbUnitFeatures = dbUnitFeatures;
274 	}
275 
276 	public void setDbUnitProperties(Map<String, String> dbUnitProperties) {
277 		this.dbUnitProperties = dbUnitProperties;
278 	}
279 
280 	private void configConnection(IDatabaseConnection conn) {
281 		DatabaseConfig config = conn.getConfig();
282 		if (dbUnitProperties != null) {
283 			for (Iterator<Map.Entry<String, String>> dbUnitPropertiesIterator = dbUnitProperties
284 					.entrySet().iterator(); dbUnitPropertiesIterator.hasNext();) {
285 				Map.Entry<String, String> entry = (Map.Entry<String, String>) dbUnitPropertiesIterator
286 						.next();
287 				String name = (String) entry.getKey();
288 				Object value = entry.getValue();
289 				config.setProperty(name, value);
290 			}
291 		}
292 		if (dbUnitFeatures != null) {
293 			for (Iterator<Map.Entry<String, String>> dbUnitFeaturesIterator = dbUnitFeatures
294 					.entrySet().iterator(); dbUnitFeaturesIterator.hasNext();) {
295 				Map.Entry<String, String> entry = (Map.Entry<String, String>) dbUnitFeaturesIterator
296 						.next();
297 				String name = (String) entry.getKey();
298 				boolean value = Boolean.valueOf((String) entry.getValue())
299 						.booleanValue();
300 				config.setFeature(name, value);
301 			}
302 		}
303 	}
304 }