/* $Id: PostgresqlConnector.java,v 1.28 2006/01/12 15:13:08 psionides Exp $ */
   package net.sourceforge.jdbdump.connect.connectors;
   
   import java.sql.DatabaseMetaData;
   import java.sql.PreparedStatement;
   import java.sql.ResultSet;
   import java.sql.SQLException;
   import java.sql.Statement;
   import java.util.TreeMap;
  import java.util.Vector;
  
  import net.sourceforge.jdbdump.connect.DatabaseConnectionData;
  import net.sourceforge.jdbdump.connect.DatabaseConnector;
  import net.sourceforge.jdbdump.dump.Column;
  import net.sourceforge.jdbdump.dump.Constraint;
  import net.sourceforge.jdbdump.dump.Dump;
  import net.sourceforge.jdbdump.dump.Table;
  
  import org.apache.log4j.Logger;
  
  /***
   * This is a class handling connection to the database, an implementation of the abstract
   * class DatabaseConnector. It is a specific connector, providing connection to
   * PostgreSQL databases.
   * 
   * @see DatabaseConnector
   * @author jsuder
   */
  
  public class PostgresqlConnector extends DatabaseConnector {
  
  	/*** A log4j logger for this class. */
  	private static Logger logger = Logger.getLogger(PostgresqlConnector.class);
  
  	/*** A ResultSet storing table data when it is being downloaded. */
  	protected ResultSet dataSet = null;
  
  	/*** The query used to download next parts of a table into dataSet. */
  	protected String query = null;
  	
  	/*** Number of current record in a table while it is being downloaded. */
  	protected int position;
  
  	/*** Number of columns in the current table. */
  	protected int columns;
  
  	/*** Maximum number of table records downloaded at a time. */
  	public static final int FETCH_SIZE = 1000;
  	
  	/***
  	 * Opens a connection to the database, using the provided connection data.
  	 * @param url a correct connection URL
  	 * @param user user's login in the database
  	 * @param pass user's password
  	 * @throws SQLException if the connection is not possible
  	 * @throws ClassNotFoundException if there is a problem with loading a driver
  	 */
  
  	public void connect(String url, String user, String pass) throws SQLException, ClassNotFoundException {
  		connect(url, user, pass, "org.postgresql.Driver");
  	}
  
  	/***
  	 * Creates a backup of the structure of the entire database. It doesn't backup
  	 * any data.
  	 * 
  	 * @todo dump other things like procedures, sequences, etc.
  	 * @return a Dump object containing information about database structure
  	 * @throws SQLException if the backup process is interrupted by a critical error
  	 */
  
  	public Dump dump() {
  		Dump dump = new Dump();
  		dumpTables(dump);
  
  		// TODO: dump other things like procedures, sequences, etc.
  
  		dump.setDatabaseReader(this);
  		return dump;
  	}
  
  	/***
  	 * Downloads the table's primary key from the database. This method implements
  	 * the empty method dumpTablePrimaryKeys() in the superclass, so is will be called
  	 * by superclass's dumpTable() after it retrieves all the columns.
  	 * 
  	 * @param table a table which should be checked for primary keys
  	 * @throws SQLException if the backup process is interrupted by a critical error
  	 */
  
  	protected void dumpTablePrimaryKeys(Table table) throws SQLException {
  		ResultSet rsKeys = meta.getPrimaryKeys(null, null, table.getName());
  
  		if (!rsKeys.next()) {
  			return;
  		}
  		
  		TreeMap<String, String> pkeys = new TreeMap<String, String>(); 
  		
 		String pkname = ""; 
 		do {
 			pkeys.put(rsKeys.getString("KEY_SEQ"), rsKeys.getString("COLUMN_NAME"));
 			pkname = rsKeys.getString("PK_NAME");
 		} while (rsKeys.next());
 		
 		String conText = "CONSTRAINT " + pkname + " PRIMARY KEY (";
 		boolean first = true;
 		for (String key : pkeys.keySet()) {
 			if (!first) {
 				conText = conText + ", ";
 			} else {
 				first = false;
 			}
 			conText = conText + pkeys.get(key);
 		}
 		conText = conText + ")";
 		
 		table.addConstraint(new Constraint("PRIMARY", Constraint.ConstraintType.PRIMARY_KEY, conText));
 	}
 
 	/***
 	 * Downloads the table's foreign (imported) keys from the database. This method implements
 	 * the empty method dumpTableForeignKeys() in the superclass, so is will be called
 	 * by superclass's dumpTable() after it retrieves all the columns.
 	 * 
 	 * @param table a table which should be checked for foreign keys
 	 * @throws SQLException if the backup process is interrupted by a critical error
 	 */
 
 	protected void dumpTableForeignKeys(Table table) throws SQLException {
 		ResultSet rsKeys = meta.getImportedKeys(null, null, table.getName());
 
 		/* the main map stores: foreign key name -> internal map
 		  the internal map stores: sequence number in a foreign key -> name of the 
 		     column in foreign key, name of the column in primary key, and
 		     primary key table name */
 		TreeMap<String, TreeMap<String, String[]>> fkeys
 			= new TreeMap<String, TreeMap<String, String[]>>(); 
 
 		// read all the key data into the map
 		while (rsKeys.next()) {
 			TreeMap<String, String[]> map = fkeys.get(rsKeys.getString("FK_NAME"));
 			if (map == null) {
 				map = new TreeMap<String, String[]>();
 				fkeys.put(rsKeys.getString("FK_NAME"), map);
 			}
 			String[] tab = new String[3];
 			tab[0] = rsKeys.getString("FKCOLUMN_NAME");
 			tab[1] = rsKeys.getString("PKTABLE_NAME");
 			tab[2] = rsKeys.getString("PKCOLUMN_NAME");
 			map.put(rsKeys.getString("KEY_SEQ"), tab);
 		}
 		
 		// create proper CONSTRAINT commands from the key data
 		for (String fkey : fkeys.keySet()) {
 			String fkname = fkey;
 			String pktable = "";
 			
 			TreeMap<String, String[]> map = fkeys.get(fkey);
 			boolean first = true;
 			String fkcolumns = "", pkcolumns = "";
 			for (String keyseq : map.keySet()) {
 				String[] tab = map.get(keyseq);
 				pktable = tab[1];
 				if (!first) {
 					fkcolumns = fkcolumns + ", ";
 					pkcolumns = pkcolumns + ", ";
 				} else {
 					first = false;
 				}
 				fkcolumns = fkcolumns + tab[0];
 				pkcolumns = pkcolumns + tab[2];
 			}
 				
 			String conText = "CONSTRAINT " + fkname + " FOREIGN KEY (" + fkcolumns
 				+ ") REFERENCES " + pktable + " (" + pkcolumns + ")";
 			table.addConstraint(new Constraint(fkname,
 				Constraint.ConstraintType.FOREIGN_KEY, conText));
 		}
 	}
 
 	/***
 	 * Downloads the structure of a single table column from the database. This method
 	 * overrides the method in superclass and will be called by superclass's dumpTable()
 	 * for each column in a table. The specific behaviour which is needed for PostgreSQL
 	 * is renaming the 'bpchar' data type to 'char'.
 	 * 
 	 * @param column a column which should have its parameters downloaded
 	 * @param rsColumns a ResultSet returned by DatabaseMetaData.getColumns(),
 	 *     set on this column's record
 	 * @throws SQLException if the backup process is interrupted by a critical error
 	 */
 
 	protected void dumpColumn(Column column, ResultSet rsColumns) throws SQLException {
 		super.dumpColumn(column, rsColumns);
 		if (column.getType().equals("bpchar")) {
 			column.setType("char");
 		}
 	}
 
 	/***
 	 * Restores the structure of the database from a Dump object loaded from a previously
 	 * created backup file. It doesn't restore any data.
 	 * 
 	 * @todo restore other things like procedures, sequences, etc.
 	 * @param dump a Dump object containing information about database structure
 	 * @throws SQLException if the restore process is interrupted by a critical error
 	 */
 
 	public void restore(Dump dump) throws SQLException {
 		clearDatabase(dump);
 		for (Table table : dump.getTables()) {
 			restoreTableStructure(table);
 			restoreTableData(table);
 		}
 		for (Table table : dump.getTables()) {
 			restoreTableConstraints(table);
 		}
 		
 		// TODO: restore other things like procedures, sequences, etc.
 	}
 
 	/***
 	 * Clears tables from the database if they already exist, before they are
 	 * inserted from a dump. This method is called by restore().
 	 * 
 	 * @param dump the dump which contains names of the tables which should be cleared
 	 * @throws SQLException if the clear process is interrupted by a critical error
 	 */
 
 	protected void clearDatabase(Dump dump) throws SQLException {
 		try {
 			Statement st = connection.createStatement();
 			for (Table table : dump.getTables()) {
 				try {
 					clearConstraints(table);
 				} catch (SQLException e) {
 					// ignore that - this means that table doesn't exist, which is not an error
 					logger.debug(e);
 				}
 			}
 			for (Table table : dump.getTables()) {
 				try {
 					logger.debug("DROP TABLE " + table.getName());
 					st.executeUpdate("DROP TABLE " + table.getName());
 				} catch (SQLException e) {
 					// ignore that - this means that table doesn't exist, which is not an error
 					logger.debug(e);
 				}
 			}
 		} catch (SQLException e) {
 			logger.error("Couldn't clear database: " + e);
 		}
 	}
 	
 	/***
 	 * Removes constraints from the database's tables (if they exist), before they are
 	 * dropped from the database and then inserted from a dump. This method is called by
 	 * clearDatabase() and is needed because it may not be possible to drop some of the
 	 * tables if they have constraints binding them to other tables.  
 	 *  
 	 * @param table the table which contains constraints which should be removed
 	 * @throws SQLException if the clear process is interrupted by a critical error
 	 */
 
 	protected void clearConstraints(Table table) throws SQLException {
 		Statement st = connection.createStatement();
 		
 		String query = "ALTER TABLE " + table.getName() + " DROP CONSTRAINT ";
 		Vector<Constraint> constraints = table.getConstraints();
 		for (int i=0; i<constraints.size(); i++) {
 			Constraint con = constraints.get(i);
 			if (con.getType() == Constraint.ConstraintType.PRIMARY_KEY) continue;
 
 			logger.debug(query + con.getName());
 			st.executeUpdate(query + con.getName());
 		}
 		
 		st.close();
 	}
 
 	/***
 	 * Creates a single table using information from a dump. It restores all the
 	 * fields, also adds a primary key if there is one, but doesn't upload any data yet
 	 * and doesn't add constraints other than the primary key.
 	 * 
 	 * @param table the table which should be added to the database
 	 * @throws SQLException if the create process is interrupted by a critical error
 	 */
 
 	protected void restoreTableStructure(Table table) throws SQLException {
 		try {
 			Statement st = connection.createStatement();
 			
 			// prepare the query
 			StringBuffer query = new StringBuffer("CREATE TABLE " + table.getName() + " (");
 	
 			Vector<Column> columns = table.getColumns();
 	
 			// get a list of constraints, but remove everything which is not a primary key
 			Vector<Constraint> constraints = new Vector<Constraint>();
 			for (int i=0; i<table.getConstraints().size(); i++) {
 				Constraint con = table.getConstraints().get(i);
 				if (con.getType() == Constraint.ConstraintType.PRIMARY_KEY) {
 					constraints.add(con);
 				}
 			}
 	
 			// add all the columns
 			for (int i=0; i<columns.size(); i++) {
 				Column col = columns.get(i); 
 				String type = col.getType();
 				
 				// type name may consist of many words, but it may be necessary 
 				// to add a size parameter after the first word
 				String[] typeWords = type.split(" ");
 				type = typeWords[0].toUpperCase();
 				
 				// insert the first word
 				query.append(col.getName() + " " + typeWords[0]);
 	
 				// for various number or character types, add size parameter 
 				if (type.equals("BIT") || type.endsWith("CHAR") || type.equals("INTERVAL")
 					|| type.startsWith("TIME")) {
 					query.append(" (" + col.getLength() + ")");
 				} else if ((typeWords.length > 1) && (typeWords[1].toUpperCase().equals("VARYING"))) {
 					// in 'character varying' or 'bit varying', the size has to be after the
 					// SECOND word from the name...
 					query.append(" VARYING ");
 					typeWords[1] = "";
 					query.append(" (" + col.getLength() + ")");
 				} else if (type.equals("DECIMAL") || type.equals("NUMERIC")) {
 					query.append(" (" + col.getLength() + ", " + col.getDecimalDigits() + ")");
 				}
 	
 				// add the rest of the words
 				for (int j=1; j<typeWords.length; j++) {
 					query.append(" " + typeWords[j]);
 				}
 				
 				// add NULL/NOT NULL parameter
 				if (col.getNullable() == DatabaseMetaData.columnNullable) {
 					query.append(" NULL");
 				} else if ((col.getNullable() == DatabaseMetaData.columnNoNulls)
 						&& (col.getNullable2().equals("NO"))) {
 					query.append(" NOT NULL");
 				}
 	
 				// add DEFAULT parameter
 				if (col.getDefault() != null) {
 					query.append(" DEFAULT " + col.getDefault());
 				}
 	
 				// add other parameters
 				if (col.getRemarks() != null && col.getRemarks().length() > 0) {
 					query.append(" " + col.getRemarks());
 				}
 	
 				// add a comma if it is not the end of the query
 				if ((i != columns.size()-1) || constraints.size() > 0) {
 					query.append(", ");
 				}
 			}
 	
 			// now add the primary key
 			for (int i=0; i<constraints.size(); i++) {
 				Constraint con = constraints.get(i);
 				String text = con.getText();
 				query.append(text);
 				if (i != constraints.size()-1) {
 					query.append(", ");
 				}
 			}
 	
 			// finish building the query
 			query.append(")");
 			
 			// finally, execute the query
 			logger.debug(query+"");
 			st.executeUpdate(query+"");
 			st.close();
 		} catch (SQLException e) {
 			logger.error("Couldn't restore the structure of table "
 				+ table.getName() + ": " + e);
 		}
 	}
 	
 	/***
 	 * Restores the table's constraints (e.g. foreign keys). This method is called
 	 * at the end of restoring process, after the data is uploaded, because otherwise
 	 * we would have to find a correct order in which to download the data to prevent
 	 * any conflicts.
 	 * 
 	 * @param table the table whose constraints should be added
 	 * @return SQLException if the constraints can't be added for some reason
 	 */
 
 	protected void restoreTableConstraints(Table table) throws SQLException {
 		try {
 			Statement st = connection.createStatement();
 			
 			String query = "ALTER TABLE " + table.getName() + " ADD ";
 			Vector<Constraint> constraints = table.getConstraints();
 			for (int i=0; i<constraints.size(); i++) {
 				Constraint con = constraints.get(i);
 				if (con.getType() != Constraint.ConstraintType.PRIMARY_KEY) {
 					String text = con.getText();
 					logger.debug(query + text);
 					st.executeUpdate(query + text);
 				}
 			}
 			
 			st.close();
 		} catch (SQLException e) {
 			logger.error("Couldn't restore constraints of table " + table.getName() + ": " + e);
 		}
 	}
 
 	/***
 	 * Restores the table's data into the database. The data are loaded from
 	 * the backup file record after record just before they are uploaded into
 	 * the database.
 	 * 
 	 * @param table the table whose data should be uploaded
 	 * @throws SQLException when the data can't be uploaded
 	 */
 
 	protected void restoreTableData(Table table) throws SQLException {
 		try {
 			table.initializeData();
 	
 			// prepare the query
 			StringBuffer query = new StringBuffer("INSERT INTO " + table.getName() + " VALUES (");
 			int columns = table.getColumns().size();
 			for (int i=0; i<columns; i++) {
 				Column col = table.getColumns().get(i);
 				if (i != 0) {
 					query.append(", ?");
 				} else {
 					query.append("?");
 				}
 			}
 			query.append(")");
 			logger.debug(query+"");
 			
 			// create a PreparedStatement to speed up the insert process
 			PreparedStatement stat = connection.prepareStatement(query+""); 
 	
 			// insert all records until the end of the table
 			while (true) {
 				try {
 					String[] record = table.getDataLine();
 					if (record == null) {
 						// end of the table
 						break;
 					}
 					
 					// execute the PreparedStatement
 					for (int i=0; i<columns; i++) {
 						stat.setString(i+1, record[i]);
 					}
 					stat.executeUpdate();
 				} catch (Exception e) {
 					logger.error("Can't read record of data from the dump: " + e);
 					return;
 				}
 			}
 		} catch (SQLException e) {
 			logger.error("Couldn't restore data of table " + table.getName() + ": " + e);
 		}
 	}
 
 	/***
 	 * Creates a PostgreSQL-specific JDBC connection url (it has the form:
 	 * jdbc:postgresql://host:port/database).
 	 * @param data an object providing all the data necessary to open the connection 
 	 * @return a connection url string generated from the data 
 	 */
 
 	public String createURL(DatabaseConnectionData data) {
 		String port = data.getPort();
 		String portStr;
 
 		try {
 			int portNr = Integer.parseInt(port);
 			portStr = ":" + port;
 		} catch (NumberFormatException e) {
 			portStr = "";
 		}
 		
 		return "jdbc:postgresql://" + data.getServer() + portStr + "/" + data.getDatabaseName();
 	}
 
 	/***
 	 * Prepares the given table for downloading its data. That means, it runs a query
 	 * "SELECT * FROM tablename" and stores the received result set.
 	 * @param tableName name of the table from which data should be downloaded
 	 * @throw SQLException if the resultset can't be returned
 	 */
 
 	public void initializeTableData(String tableName) throws SQLException {
 		Statement statement = connection.createStatement();
 
 		query = "SELECT * FROM " + tableName;
 		dataSet = statement.executeQuery(query + " LIMIT " + FETCH_SIZE);
 		columns = dataSet.getMetaData().getColumnCount();
 		position = FETCH_SIZE;
 	}
 
 	/***
 	 * Downloads one record of data from the previously initialized table in database.
 	 * @return next record from the table, or null if there are no more records
 	 * @throws SQLException if data record can't be downloaded because of a connection error
 	 *     or an error in the database
 	 */
 
 	public String[] getTableDataLine() throws SQLException { 
 		if (dataSet.next()) {
 			String[] dataLine = new String[columns];
 			for (int i=0; i<columns; i++) {
 				dataLine[i] = dataSet.getString(i+1);
 			}
 			return dataLine;
 		} else {
 			Statement statement = connection.createStatement();
 			dataSet = statement.executeQuery(query + " LIMIT " + FETCH_SIZE + " OFFSET " + position);
 			position += FETCH_SIZE;
 			
 			if (dataSet.next()) {
 				String[] dataLine = new String[columns];
 				for (int i=0; i<columns; i++) {
 					dataLine[i] = dataSet.getString(i+1);
 				}
 				return dataLine;
 			} else {
 				logger.info("getTableDataLine(): no more data lines in the table.");
 				return null; // end of data
 			}
 		}
 	}
 
 }