rmill.java

与postgresql数据库结合的数据库水印软件包

/*
 * 
 * Copyright 2003,2004 The Watermill Team
 *
 * This file is part of Watermill.
 *
 *    Watermill is free software; you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation; either version 2 of the License, or
 *    (at your option) any later version.
 *
 *    Watermill is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with Watermill; if not, write to the Free Software
 *    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */
package Watermill.rmill;

import java.io.*;
import java.util.*;
import Watermill.relational.*;
import Watermill.kernel.*;
import org.gnu.readline.*;

import java.net.*;

public class Rmill {
    

    // channels used to communicate with psql
    private static BufferedReader psqlReader;
    private static BufferedReader psqlErrorReader;
    private static BufferedWriter psqlWriter;
    
    // channels used to communicate with users
    private static BufferedReader instream;
    private static BufferedWriter outstream;
    //private static BufferedWriter verbose;
    
    private static StringWriter buffer;
    private static CommandParser parser;
    private static Command command;

    private static Manager manager;
    private static Process psql;

    
    private static String fileToRead=null;
    
    final static int CMD = 0;
    final static int FILE = 1;
    final static int SVR = 2;

    private static int mode = CMD;
    private static boolean newCmd = true;

    private static File history = null;

    private static void frontEndMessage(){
	System.out.println("Error connection to the psql frontend.");
	System.exit(-1);
    }

    
    
    private static int parseArg(int ipos,String args[]){
	try {
	    int pos=ipos;
	    
	    if(args[pos].equals("--debug")){
		Msg.setDebug();
		return pos+1;
	    }
	    
	    if(args[pos].equals("--verbose")){
		Msg.setVerbose();
		return pos+1;
	    }
	    
	    if(args[pos].equals("--file")){
		fileToRead=args[pos+1];
		mode = FILE;
		return pos+2;
	    }
	    if(args[pos].equals("--server")){
		mode = SVR;
		return pos+1;
	    }
	    
	    System.out.println("Incorrect command line parameter.");
	    System.exit(-1);
	}
	catch (ArrayIndexOutOfBoundsException e){
	    System.out.println("Incorrect command line parameter.");
	    System.exit(-1);
	}
	return 0; // never reaching this point.
    }
    
    private static void parseCommandLine(String args[]){
	for(int pos=0;pos<args.length;){
	    pos=parseArg(pos,args);
	}
    }


    private static void run() {
	try {
	    if (mode == CMD)
		runCommand();
	    if (mode == FILE)
		runFile();
	    if (mode == SVR)
		runServer();
	}
	catch (EOFException e) {
	    // user entered logout command, capture and stop
	}
	catch (FileNotFoundException e) {
	    System.err.println("Unable to locate file :"+fileToRead);
	}
	catch (Exception e) {
	    Msg.debug("Program interrupted "+e);
	}
    }

    private static void runServer() throws Exception {
	System.out.println("Running rmill server on localhost:"+Constant.port);
	ServerSocket serversocket = new ServerSocket(Constant.port);
	while (true) {
	    try {
		Socket socket = serversocket.accept(); 
		instream = new BufferedReader(new InputStreamReader(socket.getInputStream()));
		outstream = new BufferedWriter(new OutputStreamWriter(socket.getOutputStream())); // should be dev null
		loop();
		socket.close();
	    }
	    catch (IOException e) {
		Msg.debug("Connexion abnormaly aborted");
	    }
	}
    }

    private static void runCommand() throws Exception{
	instream = new BufferedReader(new InputStreamReader(System.in));
	outstream = new BufferedWriter(new OutputStreamWriter(System.out));
	
	System.out.println("Welcome to Watermill-enabled psql, the PostgreSQL interactive terminal.");
	System.out.println("Type:  \\copyright for distribution terms");
	System.out.println("       \\h for help with SQL commands");
	System.out.println("       \\? for help on internal slash commands");
	System.out.println("       \\g or terminate with semicolon to execute query");
	System.out.println("       \\q to quit");
	
	while (true) 
	    loop();
    }

    private static void runFile() throws Exception {
	instream = new BufferedReader(new FileReader(fileToRead));
	outstream = new BufferedWriter(new OutputStreamWriter(System.out));
	while (true) 
	    loop();
    }
    
    private static void readPsqlAnswer(BufferedWriter outstream) {
	try {

	    int i =0;
	    while (!psqlReader.ready() & !psqlErrorReader.ready()) {}

	    while (psqlErrorReader.ready()) {
		outstream.write(psqlErrorReader.readLine());
		outstream.newLine();
		outstream.flush();
	    }

	    while (psqlReader.ready()) {
		outstream.write(psqlReader.readLine());
		outstream.newLine();
		outstream.flush();
	    }
	}
	catch(IOException e) {
		Msg.debug(e);
	}
    }

    private static Command getCommand(String message) throws Exception{
	parser.ReInit(new StringReader(message));
	try{
	    return parser.parse();
	}
	catch (ParseException e) {
	    Msg.debug("Not a Watermill command, going to psql");
	    return null;
	}
	catch (TokenMgrError e) {
	    Msg.debug("Not a Watermill command, going to psql");
	    return null;
	}
    }

    private static String getLine(BufferedReader instream) throws IOException {

	String line = null;

	if (mode == CMD) {
	    String prompt;
	    if (newCmd)
		prompt = "=>";
	    else
		prompt = "->";
	    line = Readline.readline(prompt);
	    if (line == null)
		line = "";
	}
	else 
	    line=instream.readLine();//} // reading it
    
	if (line.trim().startsWith("#")) // deleting comments.
	    line="";
	return line;
    }

    private static void loop() throws Exception {
	
	String line=getLine(instream);
	newCmd = true;

	try{
	    buffer.write(line);
	} 
	catch(Exception e){frontEndMessage();}
	    
	Msg.verbose("Current buffer:"+buffer.toString());
	
	if (line.trim().equals("\\q")) // \q is now a rmill command, do not send it to postgres
	    throw new Exception();
	
	if (line.trim().endsWith(";")|line.trim().startsWith("\\")){ // check if it is the last line
	    buffer.flush();
	    try {
		command = getCommand(buffer.toString());
		
		if (command == null) { // since this is not a Watermill command, we send it to postgres
		    try{ 
			String s = buffer.toString();
			Msg.debug("Sending:"+s+" to postgres");
				psqlWriter.write(s);
				psqlWriter.newLine();
				psqlWriter.flush();
				readPsqlAnswer(outstream);
		    }catch(IOException e2){frontEndMessage();} 
		}
		else if (command instanceof ExecuteCommand){ // this is a file command
		    ExecuteCommand executeCommand=(ExecuteCommand)command;
		    try {
			instream=new BufferedReader(new FileReader(executeCommand.filename));
		    }
		    catch (FileNotFoundException e) {
			outstream.write("Unable to find command file :"+executeCommand.filename);
		    }
		}
		else { // this is a regular command
		    command.execute(manager);
		    Msg.debug("Command successfully executed");
		}
	    }
	    catch (Exception e3) {
		frontEndMessage();
	    }
	    buffer = new StringWriter();
	} 
	else 
	    newCmd = false;
    }
    
    private static void init() throws Exception {
	
	try {
	    Readline.load(ReadlineLibrary.GnuReadline);
	}
	catch (UnsatisfiedLinkError e) {
	    System.err.println("couldn't load readline lib. Using simple stdin");
	}
	
	Readline.initReadline("rmill");
	 
	Runtime.getRuntime()                       
	    .addShutdownHook(new Thread() {        
		    public void run() {
			Readline.cleanup();
		    }
		});
	history = new File(".rmill_history");
	try {
	    if (history.exists())
		Readline.readHistoryFile(history.getName());
	}
	catch (Exception e) {
	    System.err.println("Error reading history file");
	}

	manager = new RManager();
	psql = null;
	parser = new CommandParser(new StringReader(""));
	command=new EmptyCommand();
	psql=Runtime.getRuntime().exec("psql -U "+Constant.masterDBuser+" "+Constant.masterDBname);
	
	
	psqlReader=new BufferedReader(new InputStreamReader(psql.getInputStream())); //stdout for psql
	psqlErrorReader=new BufferedReader(new InputStreamReader(psql.getErrorStream()));
	psqlWriter=new BufferedWriter(new OutputStreamWriter(psql.getOutputStream()));;

	buffer = new StringWriter();
	
    }
    
    private static void clean() throws Exception {
	try {
	    Msg.debug("Closing history file");
	    Readline.writeHistoryFile(history.getName());
	}
	catch (Exception e) {
	    System.err.println("Error writing history file!");
	}
	try {
	    Msg.debug("Ending psql process");
	    psqlWriter.write("\\q");
	    psqlWriter.newLine();
	    psqlWriter.flush();
	}
	catch (IOException e) {}

	try{
	    Msg.debug("Closing client streams");
	    instream.close();
	    outstream.close();
	}
	catch (IOException e) {}
	Readline.cleanup();
    }


    public static void main(String args[]){


	try {

	    init();
	    
	    parseCommandLine(args);
	    
	    run();
	
	    clean();

	}
	catch (Exception e) {
	    System.err.println("Unable to launch watermill");
	    System.err.println("An error has occured "+e);
	    e.printStackTrace();
	    
	}
	
	
    }
}