001    /**
002     * Copyright (c) 2000-present Liferay, Inc. All rights reserved.
003     *
004     * This library is free software; you can redistribute it and/or modify it under
005     * the terms of the GNU Lesser General Public License as published by the Free
006     * Software Foundation; either version 2.1 of the License, or (at your option)
007     * any later version.
008     *
009     * This library is distributed in the hope that it will be useful, but WITHOUT
010     * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
011     * FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
012     * details.
013     */
014    
015    package com.liferay.portal.tools.samplesqlbuilder;
016    
017    import com.liferay.portal.dao.db.MySQLDB;
018    import com.liferay.portal.freemarker.FreeMarkerUtil;
019    import com.liferay.portal.kernel.dao.db.DB;
020    import com.liferay.portal.kernel.dao.db.DBFactoryUtil;
021    import com.liferay.portal.kernel.io.CharPipe;
022    import com.liferay.portal.kernel.io.OutputStreamWriter;
023    import com.liferay.portal.kernel.io.unsync.UnsyncBufferedReader;
024    import com.liferay.portal.kernel.io.unsync.UnsyncBufferedWriter;
025    import com.liferay.portal.kernel.io.unsync.UnsyncTeeWriter;
026    import com.liferay.portal.kernel.util.FileUtil;
027    import com.liferay.portal.kernel.util.GetterUtil;
028    import com.liferay.portal.kernel.util.SortedProperties;
029    import com.liferay.portal.kernel.util.StringBundler;
030    import com.liferay.portal.kernel.util.StringPool;
031    import com.liferay.portal.kernel.util.StringUtil;
032    import com.liferay.portal.tools.ToolDependencies;
033    
034    import java.io.File;
035    import java.io.FileInputStream;
036    import java.io.FileOutputStream;
037    import java.io.FileReader;
038    import java.io.FileWriter;
039    import java.io.IOException;
040    import java.io.Reader;
041    import java.io.Writer;
042    
043    import java.nio.channels.FileChannel;
044    
045    import java.util.ArrayList;
046    import java.util.HashMap;
047    import java.util.List;
048    import java.util.Map;
049    import java.util.Properties;
050    
051    /**
052     * @author Brian Wing Shun Chan
053     * @author Shuyang Zhou
054     */
055    public class SampleSQLBuilder {
056    
057            public static void main(String[] args) {
058                    ToolDependencies.wireBasic();
059    
060                    Reader reader = null;
061    
062                    try {
063                            Properties properties = new SortedProperties();
064    
065                            reader = new FileReader(args[0]);
066    
067                            properties.load(reader);
068    
069                            DataFactory dataFactory = new DataFactory(properties);
070    
071                            new SampleSQLBuilder(properties, dataFactory);
072                    }
073                    catch (Exception e) {
074                            e.printStackTrace();
075                    }
076                    finally {
077                            if (reader != null) {
078                                    try {
079                                            reader.close();
080                                    }
081                                    catch (IOException ioe) {
082                                            ioe.printStackTrace();
083                                    }
084                            }
085                    }
086            }
087    
088            public SampleSQLBuilder(Properties properties, DataFactory dataFactory)
089                    throws Exception {
090    
091                    _dbType = properties.getProperty("sample.sql.db.type");
092    
093                    _csvFileNames = StringUtil.split(
094                            properties.getProperty("sample.sql.output.csv.file.names"));
095                    _optimizeBufferSize = GetterUtil.getInteger(
096                            properties.getProperty("sample.sql.optimize.buffer.size"));
097                    _outputDir = properties.getProperty("sample.sql.output.dir");
098                    _script = properties.getProperty("sample.sql.script");
099    
100                    _dataFactory = dataFactory;
101    
102                    // Generic
103    
104                    Reader reader = generateSQL();
105    
106                    File tempDir = new File(_outputDir, "temp");
107    
108                    tempDir.mkdirs();
109    
110                    try {
111    
112                            // Specific
113    
114                            compressSQL(reader, tempDir);
115    
116                            // Merge
117    
118                            boolean outputMerge = GetterUtil.getBoolean(
119                                    properties.getProperty("sample.sql.output.merge"));
120    
121                            if (outputMerge) {
122                                    File sqlFile = new File(
123                                            _outputDir, "sample-" + _dbType + ".sql");
124    
125                                    FileUtil.delete(sqlFile);
126    
127                                    mergeSQL(tempDir, sqlFile);
128                            }
129                            else {
130                                    File outputDir = new File(_outputDir, "output");
131    
132                                    FileUtil.deltree(outputDir);
133    
134                                    if (!tempDir.renameTo(outputDir)) {
135    
136                                            // This will only happen when temp and output directories
137                                            // are on different file systems
138    
139                                            FileUtil.copyDirectory(tempDir, outputDir);
140                                    }
141                            }
142                    }
143                    finally {
144                            FileUtil.deltree(tempDir);
145                    }
146    
147                    StringBundler sb = new StringBundler();
148    
149                    for (String key : properties.stringPropertyNames()) {
150                            if (!key.startsWith("sample.sql")) {
151                                    continue;
152                            }
153    
154                            String value = properties.getProperty(key);
155    
156                            sb.append(key);
157                            sb.append(StringPool.EQUAL);
158                            sb.append(value);
159                            sb.append(StringPool.NEW_LINE);
160                    }
161    
162                    FileUtil.write(
163                            new File(_outputDir, "benchmarks-actual.properties"),
164                            sb.toString());
165            }
166    
167            protected void compressSQL(
168                            DB db, File directory, Map<String, Writer> insertSQLWriters,
169                            Map<String, StringBundler> sqls, String insertSQL)
170                    throws IOException {
171    
172                    String tableName = insertSQL.substring(0, insertSQL.indexOf(' '));
173    
174                    int index = insertSQL.indexOf(" values ") + 8;
175    
176                    StringBundler sb = sqls.get(tableName);
177    
178                    if ((sb == null) || (sb.index() == 0)) {
179                            sb = new StringBundler();
180    
181                            sqls.put(tableName, sb);
182    
183                            sb.append("insert into ");
184                            sb.append(insertSQL.substring(0, index));
185                            sb.append("\n");
186                    }
187                    else {
188                            sb.append(",\n");
189                    }
190    
191                    String values = insertSQL.substring(index, insertSQL.length() - 1);
192    
193                    sb.append(values);
194    
195                    if (sb.index() >= _optimizeBufferSize) {
196                            sb.append(";\n");
197    
198                            insertSQL = db.buildSQL(sb.toString());
199    
200                            sb.setIndex(0);
201    
202                            writeToInsertSQLFile(
203                                    directory, tableName, insertSQLWriters, insertSQL);
204                    }
205            }
206    
207            protected void compressSQL(Reader reader, File dir) throws Exception {
208                    DB db = DBFactoryUtil.getDB(_dbType);
209    
210                    if (db instanceof MySQLDB) {
211                            db = new SampleMySQLDB();
212                    }
213    
214                    Map<String, Writer> insertSQLWriters = new HashMap<String, Writer>();
215                    Map<String, StringBundler> insertSQLs =
216                            new HashMap<String, StringBundler>();
217                    List<String> miscSQLs = new ArrayList<String>();
218    
219                    try (UnsyncBufferedReader unsyncBufferedReader =
220                                    new UnsyncBufferedReader(reader)) {
221    
222                            String s = null;
223    
224                            while ((_freemarkerException == null) &&
225                                       ((s = unsyncBufferedReader.readLine()) != null)) {
226    
227                                    s = s.trim();
228    
229                                    if (s.length() > 0) {
230                                            if (s.startsWith("insert into ")) {
231                                                    compressSQL(
232                                                            db, dir, insertSQLWriters, insertSQLs,
233                                                            s.substring(12));
234                                            }
235                                            else {
236                                                    miscSQLs.add(s);
237                                            }
238                                    }
239                            }
240                    }
241    
242                    if (_freemarkerException != null) {
243                            throw new Exception(
244                                    "Unable to process freemarker template ", _freemarkerException);
245                    }
246    
247                    for (Map.Entry<String, StringBundler> entry : insertSQLs.entrySet()) {
248                            String tableName = entry.getKey();
249                            StringBundler sb = entry.getValue();
250    
251                            if (sb.index() == 0) {
252                                    continue;
253                            }
254    
255                            String insertSQL = db.buildSQL(sb.toString());
256    
257                            writeToInsertSQLFile(dir, tableName, insertSQLWriters, insertSQL);
258    
259                            try (Writer insertSQLWriter = insertSQLWriters.remove(tableName)) {
260                                    insertSQLWriter.write(";\n");
261                            }
262                    }
263    
264                    try (Writer miscSQLWriter = new FileWriter(new File(dir, "misc.sql"))) {
265                            for (String miscSQL : miscSQLs) {
266                                    miscSQL = db.buildSQL(miscSQL);
267    
268                                    miscSQLWriter.write(miscSQL);
269                                    miscSQLWriter.write(StringPool.NEW_LINE);
270                            }
271                    }
272            }
273    
274            protected Writer createFileWriter(File file) throws IOException {
275                    FileOutputStream fileOutputStream = new FileOutputStream(file);
276    
277                    Writer writer = new OutputStreamWriter(fileOutputStream);
278    
279                    return createUnsyncBufferedWriter(writer);
280            }
281    
282            protected Writer createUnsyncBufferedWriter(Writer writer) {
283                    return new UnsyncBufferedWriter(writer, _WRITER_BUFFER_SIZE) {
284    
285                            @Override
286                            public void flush() {
287    
288                                    // Disable FreeMarker from flushing
289    
290                            }
291    
292                    };
293            }
294    
295            protected Reader generateSQL() {
296                    final CharPipe charPipe = new CharPipe(_PIPE_BUFFER_SIZE);
297    
298                    Thread thread = new Thread() {
299    
300                            @Override
301                            public void run() {
302                                    Writer sampleSQLWriter = null;
303                                    Map<String, Object> context = null;
304    
305                                    try {
306                                            sampleSQLWriter = new UnsyncTeeWriter(
307                                                    createUnsyncBufferedWriter(charPipe.getWriter()),
308                                                    createFileWriter(new File(_outputDir, "sample.sql")));
309    
310                                            context = getContext();
311    
312                                            FreeMarkerUtil.process(_script, context, sampleSQLWriter);
313                                    }
314                                    catch (Exception e) {
315                                            _freemarkerException = e;
316                                    }
317                                    finally {
318                                            for (String csvFileName : _csvFileNames) {
319                                                    Writer csvWriter = (Writer)context.get(
320                                                            csvFileName + "CSVWriter");
321    
322                                                    try {
323                                                            csvWriter.close();
324                                                    }
325                                                    catch (IOException ioe) {
326                                                            ioe.printStackTrace();
327                                                    }
328                                            }
329    
330                                            if (sampleSQLWriter != null) {
331                                                    try {
332                                                            sampleSQLWriter.close();
333                                                    }
334                                                    catch (IOException ioe) {
335                                                            ioe.printStackTrace();
336                                                    }
337                                            }
338    
339                                            charPipe.close();
340                                    }
341                            }
342    
343                    };
344    
345                    thread.start();
346    
347                    return charPipe.getReader();
348            }
349    
350            protected Map<String, Object> getContext() throws Exception {
351                    Map<String, Object> context = new HashMap<String, Object>();
352    
353                    context.put("dataFactory", _dataFactory);
354    
355                    for (String csvFileName : _csvFileNames) {
356                            Writer csvWriter = createFileWriter(
357                                    new File(_outputDir, csvFileName + ".csv"));
358    
359                            context.put(csvFileName + "CSVWriter", csvWriter);
360                    }
361    
362                    return context;
363            }
364    
365            protected Properties getProperties(String[] args) throws Exception {
366                    Reader reader = null;
367    
368                    try {
369                            Properties properties = new SortedProperties();
370    
371                            reader = new FileReader(args[0]);
372    
373                            properties.load(reader);
374    
375                            return properties;
376                    }
377                    finally {
378                            if (reader != null) {
379                                    try {
380                                            reader.close();
381                                    }
382                                    catch (IOException ioe) {
383                                            ioe.printStackTrace();
384                                    }
385                            }
386                    }
387            }
388    
389            protected void mergeSQL(File inputDir, File outputSQLFile)
390                    throws IOException {
391    
392                    FileOutputStream outputSQLFileOutputStream = new FileOutputStream(
393                            outputSQLFile);
394    
395                    try (FileChannel outputFileChannel =
396                                    outputSQLFileOutputStream.getChannel()) {
397    
398                            File miscSQLFile = null;
399    
400                            for (File inputFile : inputDir.listFiles()) {
401                                    String inputFileName = inputFile.getName();
402    
403                                    if (inputFileName.equals("misc.sql")) {
404                                            miscSQLFile = inputFile;
405    
406                                            continue;
407                                    }
408    
409                                    mergeSQL(inputFile, outputFileChannel);
410                            }
411    
412                            if (miscSQLFile != null) {
413                                    mergeSQL(miscSQLFile, outputFileChannel);
414                            }
415                    }
416            }
417    
418            protected void mergeSQL(File inputFile, FileChannel outputFileChannel)
419                    throws IOException {
420    
421                    FileInputStream inputFileInputStream = new FileInputStream(inputFile);
422    
423                    try (FileChannel inputFileChannel = inputFileInputStream.getChannel()) {
424                            inputFileChannel.transferTo(
425                                    0, inputFileChannel.size(), outputFileChannel);
426                    }
427    
428                    inputFile.delete();
429            }
430    
431            protected void writeToInsertSQLFile(
432                            File dir, String tableName, Map<String, Writer> insertSQLWriters,
433                            String insertSQL)
434                    throws IOException {
435    
436                    Writer insertSQLWriter = insertSQLWriters.get(tableName);
437    
438                    if (insertSQLWriter == null) {
439                            File file = new File(dir, tableName + ".sql");
440    
441                            insertSQLWriter = createFileWriter(file);
442    
443                            insertSQLWriters.put(tableName, insertSQLWriter);
444                    }
445    
446                    insertSQLWriter.write(insertSQL);
447            }
448    
449            private static final int _PIPE_BUFFER_SIZE = 16 * 1024 * 1024;
450    
451            private static final int _WRITER_BUFFER_SIZE = 16 * 1024;
452    
453            private final String[] _csvFileNames;
454            private final DataFactory _dataFactory;
455            private final String _dbType;
456            private volatile Exception _freemarkerException;
457            private final int _optimizeBufferSize;
458            private final String _outputDir;
459            private final String _script;
460    
461    }