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                    File tempDir = new File(_outputDir, "temp");
105    
106                    tempDir.mkdirs();
107    
108                    Reader reader = generateSQL();
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                                    String insertSQL = db.buildSQL(sb.toString());
253    
254                                    writeToInsertSQLFile(
255                                            dir, tableName, insertSQLWriters, insertSQL);
256                            }
257    
258                            try (Writer insertSQLWriter = insertSQLWriters.remove(tableName)) {
259                                    insertSQLWriter.write(";\n");
260                            }
261                    }
262    
263                    try (Writer miscSQLWriter = new FileWriter(new File(dir, "misc.sql"))) {
264                            for (String miscSQL : miscSQLs) {
265                                    miscSQL = db.buildSQL(miscSQL);
266    
267                                    miscSQLWriter.write(miscSQL);
268                                    miscSQLWriter.write(StringPool.NEW_LINE);
269                            }
270                    }
271            }
272    
273            protected Writer createFileWriter(File file) throws IOException {
274                    FileOutputStream fileOutputStream = new FileOutputStream(file);
275    
276                    Writer writer = new OutputStreamWriter(fileOutputStream);
277    
278                    return createUnsyncBufferedWriter(writer);
279            }
280    
281            protected Writer createUnsyncBufferedWriter(Writer writer) {
282                    return new UnsyncBufferedWriter(writer, _WRITER_BUFFER_SIZE) {
283    
284                            @Override
285                            public void flush() {
286    
287                                    // Disable FreeMarker from flushing
288    
289                            }
290    
291                    };
292            }
293    
294            protected Reader generateSQL() {
295                    final CharPipe charPipe = new CharPipe(_PIPE_BUFFER_SIZE);
296    
297                    Thread thread = new Thread() {
298    
299                            @Override
300                            public void run() {
301                                    Writer sampleSQLWriter = null;
302                                    Map<String, Object> context = null;
303    
304                                    try {
305                                            sampleSQLWriter = new UnsyncTeeWriter(
306                                                    createUnsyncBufferedWriter(charPipe.getWriter()),
307                                                    createFileWriter(new File(_outputDir, "sample.sql")));
308    
309                                            context = getContext();
310    
311                                            FreeMarkerUtil.process(_script, context, sampleSQLWriter);
312                                    }
313                                    catch (Exception e) {
314                                            _freemarkerException = e;
315                                    }
316                                    finally {
317                                            for (String csvFileName : _csvFileNames) {
318                                                    Writer csvWriter = (Writer)context.get(
319                                                            csvFileName + "CSVWriter");
320    
321                                                    try {
322                                                            csvWriter.close();
323                                                    }
324                                                    catch (IOException ioe) {
325                                                            ioe.printStackTrace();
326                                                    }
327                                            }
328    
329                                            if (sampleSQLWriter != null) {
330                                                    try {
331                                                            sampleSQLWriter.close();
332                                                    }
333                                                    catch (IOException ioe) {
334                                                            ioe.printStackTrace();
335                                                    }
336                                            }
337    
338                                            charPipe.close();
339                                    }
340                            }
341    
342                    };
343    
344                    thread.start();
345    
346                    return charPipe.getReader();
347            }
348    
349            protected Map<String, Object> getContext() throws Exception {
350                    Map<String, Object> context = new HashMap<String, Object>();
351    
352                    context.put("dataFactory", _dataFactory);
353    
354                    for (String csvFileName : _csvFileNames) {
355                            Writer csvWriter = createFileWriter(
356                                    new File(_outputDir, csvFileName + ".csv"));
357    
358                            context.put(csvFileName + "CSVWriter", csvWriter);
359                    }
360    
361                    return context;
362            }
363    
364            protected Properties getProperties(String[] args) throws Exception {
365                    Reader reader = null;
366    
367                    try {
368                            Properties properties = new SortedProperties();
369    
370                            reader = new FileReader(args[0]);
371    
372                            properties.load(reader);
373    
374                            return properties;
375                    }
376                    finally {
377                            if (reader != null) {
378                                    try {
379                                            reader.close();
380                                    }
381                                    catch (IOException ioe) {
382                                            ioe.printStackTrace();
383                                    }
384                            }
385                    }
386            }
387    
388            protected void mergeSQL(File inputDir, File outputSQLFile)
389                    throws IOException {
390    
391                    FileOutputStream outputSQLFileOutputStream = new FileOutputStream(
392                            outputSQLFile);
393    
394                    try (FileChannel outputFileChannel =
395                                    outputSQLFileOutputStream.getChannel()) {
396    
397                            File miscSQLFile = null;
398    
399                            for (File inputFile : inputDir.listFiles()) {
400                                    String inputFileName = inputFile.getName();
401    
402                                    if (inputFileName.equals("misc.sql")) {
403                                            miscSQLFile = inputFile;
404    
405                                            continue;
406                                    }
407    
408                                    mergeSQL(inputFile, outputFileChannel);
409                            }
410    
411                            if (miscSQLFile != null) {
412                                    mergeSQL(miscSQLFile, outputFileChannel);
413                            }
414                    }
415            }
416    
417            protected void mergeSQL(File inputFile, FileChannel outputFileChannel)
418                    throws IOException {
419    
420                    FileInputStream inputFileInputStream = new FileInputStream(inputFile);
421    
422                    try (FileChannel inputFileChannel = inputFileInputStream.getChannel()) {
423                            inputFileChannel.transferTo(
424                                    0, inputFileChannel.size(), outputFileChannel);
425                    }
426    
427                    inputFile.delete();
428            }
429    
430            protected void writeToInsertSQLFile(
431                            File dir, String tableName, Map<String, Writer> insertSQLWriters,
432                            String insertSQL)
433                    throws IOException {
434    
435                    Writer insertSQLWriter = insertSQLWriters.get(tableName);
436    
437                    if (insertSQLWriter == null) {
438                            File file = new File(dir, tableName + ".sql");
439    
440                            insertSQLWriter = createFileWriter(file);
441    
442                            insertSQLWriters.put(tableName, insertSQLWriter);
443                    }
444    
445                    insertSQLWriter.write(insertSQL);
446            }
447    
448            private static final int _PIPE_BUFFER_SIZE = 16 * 1024 * 1024;
449    
450            private static final int _WRITER_BUFFER_SIZE = 16 * 1024;
451    
452            private final String[] _csvFileNames;
453            private final DataFactory _dataFactory;
454            private final String _dbType;
455            private volatile Exception _freemarkerException;
456            private final int _optimizeBufferSize;
457            private final String _outputDir;
458            private final String _script;
459    
460    }