package org.apache.torque.task;

 /*

  * Licensed to the Apache Software Foundation (ASF) under one

  * or more contributor license agreements.  See the NOTICE file

  * distributed with this work for additional information

  * regarding copyright ownership.  The ASF licenses this file

  * to you 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.

  */

 import java.io.BufferedOutputStream;

 import java.io.BufferedReader;

 import java.io.File;

 import java.io.FileInputStream;

 import java.io.FileOutputStream;

 import java.io.FileReader;

 import java.io.IOException;

 import java.io.InputStreamReader;

 import java.io.PrintStream;

 import java.io.Reader;

 import java.io.StringReader;

 import java.sql.Connection;

 import java.sql.DatabaseMetaData;

 import java.sql.Driver;

 import java.sql.ResultSet;

 import java.sql.ResultSetMetaData;

 import java.sql.SQLException;

 import java.sql.SQLWarning;

 import java.sql.Statement;

 import java.util.ArrayList;

 import java.util.HashMap;

 import java.util.Iterator;

 import java.util.List;

 import java.util.Map;

 import java.util.Properties;

 import org.apache.commons.lang.StringUtils;

 import org.apache.tools.ant.AntClassLoader;

 import org.apache.tools.ant.BuildException;

 import org.apache.tools.ant.Project;

 import org.apache.tools.ant.PropertyHelper;

 import org.apache.tools.ant.Task;

 import org.apache.tools.ant.types.EnumeratedAttribute;

 import org.apache.tools.ant.types.Path;

 import org.apache.tools.ant.types.Reference;

 /**

  * This task uses an SQL -> Database map in the form of a properties file to insert each SQL file listed into its

  * designated database.

  * 

  * @author <a href="mailto:jeff@custommonkey.org">Jeff Martin</a>

  * @author <a href="mailto:gholam@xtra.co.nz">Michael McCallum</A>

  * @author <a href="mailto:tim.stephenson@sybase.com">Tim Stephenson</A>

  * @author <a href="mailto:jvanzyl@apache.org">Jason van Zyl</A>

  * @author <a href="mailto:mpoeschl@marmot.at">Martin Poeschl</a>

  * @version $Id: TorqueSQLExec.java,v 1.1 2007-10-21 07:57:26 abyrne Exp $

  */

 public class TorqueSQLExec extends Task {

         private int goodSql = 0;

         private int totalSql = 0;

         private Path classpath;

         private AntClassLoader loader;

         /**

      *

      */

         public static class DelimiterType extends EnumeratedAttribute {

                 public static final String NORMAL = "normal";

                 public static final String ROW = "row";

                 public String[] getValues() {

                         return new String[] { NORMAL, ROW };

                 }

         }

         /** Database connection */

         private Connection conn = null;

         /** Autocommit flag. Default value is false */

         private boolean autocommit = false;

         /** SQL statement */

         private Statement statement = null;

         /** DB driver. */

         private String driver = null;

         /** DB url. */

         private String url = null;

         /** User name. */

         private String userId = null;

         /** Password */

         private String password = null;

         /** SQL Statement delimiter */

         private String delimiter = ";";

         /**

          * The delimiter type indicating whether the delimiter will only be recognized on a line by itself

          */

         private String delimiterType = DelimiterType.NORMAL;

         /** Print SQL results. */

         private boolean print = false;

         /** Print header columns. */

         private boolean showheaders = true;

         /** Results Output file. */

         private File output = null;

         /** RDBMS Product needed for this SQL. */

         private String rdbms = null;

         /** RDBMS Version needed for this SQL. */

         private String version = null;

         /** Action to perform if an error is found */

         private String onError = "abort";

         /** Encoding to use when reading SQL statements from a file */

         private String encoding = null;

         /** Src directory for the files listed in the sqldbmap. */

         private String srcDir;

         /** Properties file that maps an individual SQL file to a database. */

         private File sqldbmap;

         /**

          * Set the sqldbmap properties file.

          * 

          * @param sqldbmap

          *            filename for the sqldbmap

          */

         public void setSqlDbMap(String sqldbmap) {

                 this.sqldbmap = getProject().resolveFile(sqldbmap);

         }

         /**

          * Get the sqldbmap properties file.

          * 

          * @return filename for the sqldbmap

          */

         public File getSqlDbMap() {

                 return sqldbmap;

         }

         /**

          * Set the src directory for the sql files listed in the sqldbmap file.

          * 

          * @param srcDir

          *            sql source directory

          */

         public void setSrcDir(String srcDir) {

                 this.srcDir = getProject().resolveFile(srcDir).toString();

         }

         /**

          * Get the src directory for the sql files listed in the sqldbmap file.

          * 

          * @return sql source directory

          */

         public String getSrcDir() {

                 return srcDir;

         }

         /**

          * Set the classpath for loading the driver.

          * 

          * @param classpath

          *            the classpath

          */

         public void setClasspath(Path classpath) {

                 if (this.classpath == null) {

                         this.classpath = classpath;

                 } else {

                         this.classpath.append(classpath);

                 }

         }

         /**

          * Create the classpath for loading the driver.

          * 

          * @return the classpath

          */

         public Path createClasspath() {

                 if (this.classpath == null) {

                         this.classpath = new Path(getProject());

                 }

                 return this.classpath.createPath();

         }

         /**

          * Set the classpath for loading the driver using the classpath reference.

          * 

          * @param r

          *            reference to the classpath

          */

         public void setClasspathRef(Reference r) {

                 createClasspath().setRefid(r);

         }

         /**

          * Set the sql command to execute

          * 

          * @param sql

          *            sql command to execute

          * @deprecated This method has no effect and will be removed in a future version.

          */

         public void addText(String sql) {

         }

         /**

          * Set the JDBC driver to be used.

          * 

          * @param driver

          *            driver class name

          */

         public void setDriver(String driver) {

                 this.driver = driver;

         }

         /**

          * Set the DB connection url.

          * 

          * @param url

          *            connection url

          */

         public void setUrl(String url) {

                 this.url = url;

         }

         /**

          * Set the user name for the DB connection.

          * 

          * @param userId

          *            database user

          */

         public void setUserid(String userId) {

                 this.userId = userId;

         }

         /**

          * Set the file encoding to use on the sql files read in

          * 

          * @param encoding

          *            the encoding to use on the files

          */

         public void setEncoding(String encoding) {

                 this.encoding = encoding;

         }

         /**

          * Set the password for the DB connection.

          * 

          * @param password

          *            database password

          */

         public void setPassword(String password) {

                 this.password = password;

         }

         /**

          * Set the autocommit flag for the DB connection.

          * 

          * @param autocommit

          *            the autocommit flag

          */

         public void setAutocommit(boolean autocommit) {

                 this.autocommit = autocommit;

         }

         /**

          * Set the statement delimiter.

          * 

          * <p>

          * For example, set this to "go" and delimitertype to "ROW" for Sybase ASE or MS SQL Server.

          * </p>

          * 

          * @param delimiter

          */

         public void setDelimiter(String delimiter) {

                 this.delimiter = delimiter;

         }

         /**

          * Set the Delimiter type for this sql task. The delimiter type takes two values - normal and row. Normal means that

          * any occurence of the delimiter terminate the SQL command whereas with row, only a line containing just the

          * delimiter is recognized as the end of the command.

          * 

          * @param delimiterType

          */

         public void setDelimiterType(DelimiterType delimiterType) {

                 this.delimiterType = delimiterType.getValue();

         }

         /**

          * Set the print flag.

          * 

          * @param print

          */

         public void setPrint(boolean print) {

                 this.print = print;

         }

         /**

          * Set the showheaders flag.

          * 

          * @param showheaders

          */

         public void setShowheaders(boolean showheaders) {

                 this.showheaders = showheaders;

         }

         /**

          * Set the output file.

          * 

          * @param output

          */

         public void setOutput(File output) {

                 this.output = output;

         }

         /**

          * Set the rdbms required

          * 

          * @param vendor

          */

         public void setRdbms(String vendor) {

                 this.rdbms = vendor.toLowerCase();

         }

         /**

          * Set the version required

          * 

          * @param version

          */

         public void setVersion(String version) {

                 this.version = version.toLowerCase();

         }

         /**

          * Set the action to perform onerror

          * 

          * @param action

          */

         public void setOnerror(OnError action) {

                 this.onError = action.getValue();

         }

         /**

          * Load the sql file and then execute it

          * 

          * @throws BuildException

          */

         @SuppressWarnings("unchecked")

         public void execute() throws BuildException {

                 if (sqldbmap == null || getSqlDbMap().exists() == false) {

                         throw new BuildException("You haven't provided an sqldbmap, or " + "the one you specified doesn't exist: " + sqldbmap);

                 }

                 if (driver == null) {

                         throw new BuildException("Driver attribute must be set!", getLocation());

                 }

                 if (userId == null) {

                         throw new BuildException("User Id attribute must be set!", getLocation());

                 }

                 if (password == null) {

                         throw new BuildException("Password attribute must be set!", getLocation());

                 }

                 if (url == null) {

                         throw new BuildException("Url attribute must be set!", getLocation());

                 }

                 Properties map = new Properties();

                 try {

                         FileInputStream fis = new FileInputStream(getSqlDbMap());

                         map.load(fis);

                         fis.close();

                 } catch (IOException ioe) {

                         throw new BuildException("Cannot open and process the sqldbmap!");

                 }

                 Map<Object, Object> databases = new HashMap<Object, Object>();

                 Iterator<?> eachFileName = map.keySet().iterator();

                 while (eachFileName.hasNext()) {

                         String sqlfile = (String) eachFileName.next();

                         String database = map.getProperty(sqlfile);

                         List<Object> files = (List<Object>) databases.get(database);

                         if (files == null) {

                                 files = new ArrayList<Object>();

                                 databases.put(database, files);

                         }

                         // We want to make sure that the base schemas

                         // are inserted first.

                         if (sqlfile.indexOf("schema.sql") != -1) {

                                 files.add(0, sqlfile);

                         } else {

                                 files.add(sqlfile);

                         }

                 }

                 Iterator<?> eachDatabase = databases.keySet().iterator();

                 while (eachDatabase.hasNext()) {

                         String db = (String) eachDatabase.next();

                         List<Object> transactions = new ArrayList<Object>();

                         eachFileName = ((List<?>) databases.get(db)).iterator();

                         while (eachFileName.hasNext()) {

                                 String fileName = (String) eachFileName.next();

                                 File file = new File(srcDir, fileName);

                                 if (file.exists()) {

                                         Transaction transaction = new Transaction();

                                         transaction.setSrc(file);

                                         transactions.add(transaction);

                                 } else {

                                         System.out.println("File '" + file.getAbsolutePath() + "' in sqldbmap does not exist, so skipping it.");

                                 }

                         }

                         insertDatabaseSqlFiles(url, db, transactions);

                 }

         }

         /**

          * Take the base url, the target database and insert a set of SQL files into the target database.

          * 

          * @param url

          * @param database

          * @param transactions

          */

         private void insertDatabaseSqlFiles(String url, String database, List<?> transactions) {

                 url = StringUtils.replace(url, "@DB@", database);

                 System.out.println("Our new url -> " + url);

                 Driver driverInstance = null;

                 try {

                         Class<?> dc;

                         if (classpath != null) {

                                 log("Loading " + driver + " using AntClassLoader with classpath " + classpath, Project.MSG_VERBOSE);

                                 loader = new AntClassLoader(getProject(), classpath);

                                 dc = loader.loadClass(driver);

                         } else {

                                 log("Loading " + driver + " using system loader.", Project.MSG_VERBOSE);

                                 dc = Class.forName(driver);

                         }

                         driverInstance = (Driver) dc.newInstance();

                 } catch (ClassNotFoundException e) {

                         throw new BuildException("Class Not Found: JDBC driver " + driver + " could not be loaded", getLocation());

                 } catch (IllegalAccessException e) {

                         throw new BuildException("Illegal Access: JDBC driver " + driver + " could not be loaded", getLocation());

                 } catch (InstantiationException e) {

                         throw new BuildException("Instantiation Exception: JDBC driver " + driver + " could not be loaded", getLocation());

                 }

                 try {

                         log("connecting to " + url, Project.MSG_VERBOSE);

                         Properties info = new Properties();

                         info.put("user", userId);

                         info.put("password", password);

                         conn = driverInstance.connect(url, info);

                         if (conn == null) {

                                 // Driver doesn't understand the URL

                                 throw new SQLException("No suitable Driver for " + url);

                         }

                         if (!isValidRdbms(conn)) {

                                 return;

                         }

                         conn.setAutoCommit(autocommit);

                         statement = conn.createStatement();

                         PrintStream out = System.out;

                         try {

                                 if (output != null) {

                                         log("Opening PrintStream to output file " + output, Project.MSG_VERBOSE);

                                         out = new PrintStream(new BufferedOutputStream(new FileOutputStream(output)));

                                 }

                                 // Process all transactions

                                 for (Iterator<?> it = transactions.iterator(); it.hasNext();) {

                                         Transaction transaction = (Transaction) it.next();

                                         transaction.runTransaction(out);

                                         if (!autocommit) {

                                                 log("Commiting transaction", Project.MSG_VERBOSE);

                                                 conn.commit();

                                         }

                                 }

                         } finally {

                                 if (out != null && out != System.out) {

                                         out.close();

                                 }

                         }

                 } catch (IOException e) {

                         if (!autocommit && conn != null && onError.equals("abort")) {

                                 try {

                                         conn.rollback();

                                 } catch (SQLException ex) {

                                         // do nothing.

                                 }

                         }

                         throw new BuildException(e, getLocation());

                 } catch (SQLException e) {

                         if (!autocommit && conn != null && onError.equals("abort")) {

                                 try {

                                         conn.rollback();

                                 } catch (SQLException ex) {

                                         // do nothing.

                                 }

                         }

                         throw new BuildException(e, getLocation());

                 } finally {

                         try {

                                 if (statement != null) {

                                         statement.close();

                                 }

                                 if (conn != null) {

                                         conn.close();

                                 }

                         } catch (SQLException e) {

                         }

                 }

                 System.out.println(goodSql + " of " + totalSql + " SQL statements executed successfully");

         }

         /**

          * Read the statements from the .sql file and execute them. Lines starting with '//', '--' or 'REM ' are ignored.

          * 

          * @param reader

          * @param out

          * @throws SQLException

          * @throws IOException

          */

         protected void runStatements(Reader reader, PrintStream out) throws SQLException, IOException {

                 String sql = "";

                 String line = "";

                 BufferedReader in = new BufferedReader(reader);

                 PropertyHelper ph = PropertyHelper.getPropertyHelper(getProject());

                 try {

                         while ((line = in.readLine()) != null) {

                                 line = line.trim();

                                 line = ph.replaceProperties("", line, getProject().getProperties());

                                 if (line.startsWith("//") || line.startsWith("--")) {

                                         continue;

                                 }

                                 if (line.length() > 4 && line.substring(0, 4).equalsIgnoreCase("REM ")) {

                                         continue;

                                 }

                                 sql += " " + line;

                                 sql = sql.trim();

                                 // SQL defines "--" as a comment to EOL

                                 // and in Oracle it may contain a hint

                                 // so we cannot just remove it, instead we must end it

                                 if (line.indexOf("--") >= 0) {

                                         sql += "\n";

                                 }

                                 if (delimiterType.equals(DelimiterType.NORMAL) && sql.endsWith(delimiter) || delimiterType.equals(DelimiterType.ROW) && line.equals(delimiter)) {

                                         log("SQL: " + sql, Project.MSG_VERBOSE);

                                         execSQL(sql.substring(0, sql.length() - delimiter.length()), out);

                                         sql = "";

                                 }

                         }

                         // Catch any statements not followed by ;

                         if (!sql.equals("")) {

                                 execSQL(sql, out);

                         }

                 } catch (SQLException e) {

                         throw e;

                 }

         }

         /**

          * Verify if connected to the correct RDBMS

          * 

          * @param conn

          */

         protected boolean isValidRdbms(Connection conn) {

                 if (rdbms == null && version == null) {

                         return true;

                 }

                 try {

                         DatabaseMetaData dmd = conn.getMetaData();

                         if (rdbms != null) {

                                 String theVendor = dmd.getDatabaseProductName().toLowerCase();

                                 log("RDBMS = " + theVendor, Project.MSG_VERBOSE);

                                 if (theVendor == null || theVendor.indexOf(rdbms) < 0) {

                                         log("Not the required RDBMS: " + rdbms, Project.MSG_VERBOSE);

                                         return false;

                                 }

                         }

                         if (version != null) {

                                 String theVersion = dmd.getDatabaseProductVersion().toLowerCase();

                                 log("Version = " + theVersion, Project.MSG_VERBOSE);

                                 if (theVersion == null || !(theVersion.startsWith(version) || theVersion.indexOf(" " + version) >= 0)) {

                                         log("Not the required version: \"" + version + "\"", Project.MSG_VERBOSE);

                                         return false;

                                 }

                         }

                 } catch (SQLException e) {

                         // Could not get the required information

                         log("Failed to obtain required RDBMS information", Project.MSG_ERR);

                         return false;

                 }

                 return true;

         }

         /**

          * Exec the sql statement.

          * 

          * @param sql

          * @param out

          * @throws SQLException

          */

         protected void execSQL(String sql, PrintStream out) throws SQLException {

                 // Check and ignore empty statements

                 if ("".equals(sql.trim())) {

                         return;

                 }

                 try {

                         totalSql++;

                         if (!statement.execute(sql)) {

                                 log(statement.getUpdateCount() + " rows affected", Project.MSG_VERBOSE);

                         } else {

                                 if (print) {

                                         printResults(out);

                                 }

                         }

                         SQLWarning warning = conn.getWarnings();

                         while (warning != null) {

                                 log(warning + " sql warning", Project.MSG_VERBOSE);

                                 warning = warning.getNextWarning();

                         }

                         conn.clearWarnings();

                         goodSql++;

                 } catch (SQLException e) {

                         System.out.println("Failed to execute: " + sql);

                         if (!onError.equals("continue")) {

                                 throw e;

                         }

                         log(e.toString(), Project.MSG_ERR);

                 }

         }

         /**

          * print any results in the statement.

          * 

          * @param out

          * @throws SQLException

          */

         protected void printResults(PrintStream out) throws java.sql.SQLException {

                 ResultSet rs = null;

                 do {

                         rs = statement.getResultSet();

                         if (rs != null) {

                                 log("Processing new result set.", Project.MSG_VERBOSE);

                                 ResultSetMetaData md = rs.getMetaData();

                                 int columnCount = md.getColumnCount();

                                 StringBuffer line = new StringBuffer();

                                 if (showheaders) {

                                         for (int col = 1; col < columnCount; col++) {

                                                 line.append(md.getColumnName(col));

                                                 line.append(",");

                                         }

                                         line.append(md.getColumnName(columnCount));

                                         out.println(line);

                                         line.setLength(0);

                                 }

                                 while (rs.next()) {

                                         boolean first = true;

                                         for (int col = 1; col <= columnCount; col++) {

                                                 String columnValue = rs.getString(col);

                                                 if (columnValue != null) {

                                                         columnValue = columnValue.trim();

                                                 }

                                                 if (first) {

                                                         first = false;

                                                 } else {

                                                         line.append(",");

                                                 }

                                                 line.append(columnValue);

                                         }

                                         out.println(line);

                                         line.setLength(0);

                                 }

                         }

                 } while (statement.getMoreResults());

                 out.println();

         }

         /**

          * Enumerated attribute with the values "continue", "stop" and "abort" for the onerror attribute.

          */

         public static class OnError extends EnumeratedAttribute {

                 public static final String CONTINUE = "continue";

                 public static final String STOP = "stop";

                 public static final String ABORT = "abort";

                 public String[] getValues() {

                         return new String[] { CONTINUE, STOP, ABORT };

                 }

         }

         /**

          * Contains the definition of a new transaction element. Transactions allow several files or blocks of statements to

          * be executed using the same JDBC connection and commit operation in between.

          */

         public class Transaction {

                 private File tSrcFile = null;

                 private String tSqlCommand = "";

                 public void setSrc(File src) {

                         this.tSrcFile = src;

                 }

                 public void addText(String sql) {

                         this.tSqlCommand += sql;

                 }

                 private void runTransaction(PrintStream out) throws IOException, SQLException {

                         if (tSqlCommand.length() != 0) {

                                 log("Executing commands", Project.MSG_INFO);

                                 runStatements(new StringReader(tSqlCommand), out);

                         }

                         if (tSrcFile != null) {

                                 System.out.println("Executing file: " + tSrcFile.getAbsolutePath());

                                 Reader reader = (encoding == null) ? new FileReader(tSrcFile) : new InputStreamReader(new FileInputStream(tSrcFile), encoding);

                                 runStatements(reader, out);

                                 reader.close();

                         }

                 }

         }

 }