import java.sql.Connection;
import java.sql.Statement;
import java.util.Arrays;
import java.util.Random;

import javax.sql.XAConnection;
import javax.transaction.xa.XAResource;
import javax.transaction.xa.Xid;

import org.postgresql.xa.PGXAConnection;
import org.postgresql.xa.PGXADataSource;

public class XATest3 {

	/**
	 * @param args
	 */
	public static void main(String[] args) {
		PGXADataSource dataSource = new PGXADataSource();
		dataSource.setDatabaseName("postgres");		
		dataSource.setUser("postgres");				
		dataSource.setServerName("localhost");		
		
		try {

			XAConnection connection = dataSource.getXAConnection();
			System.out.println("Connected!");
			XAResource resource = connection.getXAResource();
			CustomXid identifier = new CustomXid(10);

			// Apparently the application server gets the connection
			// handle once and reuses it because there was no
			// trace of getConnection() calls in the log.
			Connection jdbcConnection = connection.getConnection();

			// TM
			resource.start(identifier, XAResource.TMNOFLAGS);

			// application, EJB call
			Statement statement = jdbcConnection.createStatement();
			String sql = "INSERT INTO foo values ('ejb')";
			statement.executeUpdate(sql);

			// TM
		    resource.end(identifier, XAResource.TMSUCCESS);
		    resource.commit(identifier, true);

			// application, java block
			jdbcConnection.setAutoCommit(false);
			statement = jdbcConnection.createStatement();
			sql = "INSERT INTO foo values ('javablock, autocommit off')";
			statement.executeUpdate(sql);
			jdbcConnection.commit();

			// TM
			resource.start(identifier, XAResource.TMNOFLAGS);

			// application, EJB call
			statement = jdbcConnection.createStatement();
			sql = "INSERT INTO foo values ('ejb again')";
			statement.executeUpdate(sql);

			// TM
		    resource.end(identifier, XAResource.TMSUCCESS);
		    resource.commit(identifier, true);

			// application, 2nd java block
			jdbcConnection.setAutoCommit(false);
			statement = jdbcConnection.createStatement();
			sql = "INSERT INTO foo values ('javablock, autocommit still off')";
			statement.executeUpdate(sql);
			jdbcConnection.commit();
		    
		} catch(Throwable tw)
		{
			tw.printStackTrace();
		}
	}

	
    static class CustomXid implements Xid {
        private static Random rand = new Random(System.currentTimeMillis());
        byte[] gtrid = new byte[Xid.MAXGTRIDSIZE];
        byte[] bqual = new byte[Xid.MAXBQUALSIZE];

        CustomXid(int i)
        {
            rand.nextBytes(gtrid);
            gtrid[0] = (byte)i;
            gtrid[1] = (byte)i;
            gtrid[2] = (byte)i;
            gtrid[3] = (byte)i;
            gtrid[4] = (byte)i;
            bqual[0] = 4;
            bqual[1] = 5;
            bqual[2] = 6;
        }

        public int getFormatId() {
            return 0;
        }


        public byte[] getGlobalTransactionId() {
            return gtrid;
        }

        public byte[] getBranchQualifier() {
            return bqual;
        }
        public boolean equals(Object o) {
            Xid other = (Xid)o;
            if (other.getFormatId() != this.getFormatId())
                return false;
            if (!Arrays.equals(other.getBranchQualifier(), this.getBranchQualifier()))
                return false;
            if (!Arrays.equals(other.getGlobalTransactionId(), this.getGlobalTransactionId()))
                return false;

            return true;
        }
    }
	
}