source: OAIHarvester/trunk/OAIHarvester/src/test/java/eu/clarin/cmdi/oai/harvester/HarvesterTest.java @ 1128

package eu.clarin.cmdi.oai.harvester;

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.Date;
import java.util.List;
import java.util.zip.ZipEntry;
import java.util.zip.ZipException;
import java.util.zip.ZipOutputStream;

import javax.xml.stream.XMLOutputFactory;
import javax.xml.stream.XMLStreamConstants;
import javax.xml.stream.XMLStreamReader;
import javax.xml.stream.XMLStreamWriter;

import eu.clarin.cmdi.oai.harvester.HarvestHandlerAdapter;
import eu.clarin.cmdi.oai.harvester.impl.SimpleHarvester;
import eu.clarin.cmdi.oai.harvester.util.XMLStreamCopier;


public class HarvesterTest {
    private static XMLOutputFactory factory =
        XMLOutputFactory.newInstance();

    private class MyHarvestHandler extends HarvestHandlerAdapter {
        private ZipOutputStream output;
        private int reqNum;
        private String prefix;

        public MyHarvestHandler(File file) {
            try {
                FileOutputStream fos = new FileOutputStream(file);
                output = new ZipOutputStream(fos);
                output.setMethod(ZipOutputStream.DEFLATED);
            } catch (Exception e) {
                throw new RuntimeException("xxx", e);
            }
        }

        public void close() {
            try {
                if (output != null) {
                    output.flush();
                    output.close();
                    output = null;
                }
            } catch (ZipException e) {
                /* IGNORE */
            } catch (Exception e) {
                throw new RuntimeException("xxx", e);
            }
        }

        @Override
        public void onIdentify(Repository repository) {
            System.err.println("onIdentify");
            System.err.println("     repositoryName: " +
                    repository.getRepositoryName());
            System.err.println("            baseURL: " +
                    repository.getBaseURL());
            System.err.println("    protocolVersion: " +
                    repository.getProtocolVersion());
            for (String adminEmail : repository.getAdminEmail()) {
                System.err.println("         adminEmail: " + adminEmail);
            }
            System.err.println("  earliestTimestamp: " +
                    repository.getEarliestTimestamp());
            System.err.println("      deletedPolicy: " +
                    repository.getDeletedNotion());
            System.err.println("        granularity: " +
                    repository.getGranularity());
            StringBuilder sb = new StringBuilder();
            final int mask = repository.getCompressionMask();
            if (mask != 0) {
                if ((mask & Repository.COMPRESSION_METHOD_DEFLATE) > 0) {
                    sb.append("DEFLATE");
                }
                if ((mask & Repository.COMPRESSION_METHOD_GZIP) > 0) {
                    if (sb.length() > 0) {
                        sb.append(", ");
                    }
                    sb.append("GZIP");
                }
            } else {
                sb.append("N/A");
            }
            System.err.println("        compression: " + sb.toString());
        }
        
        @Override
        public void onListMetadataFormats(List<MetadataFormat> metadataFormats) {
            System.err.println("onListMetadataFormats");
            for (MetadataFormat format : metadataFormats) {
                System.err.println("==> MetadataFormat");
                System.err.println("    Prefix: " + format.getPrefix());
                System.err.println("    Schema: " + format.getSchema());
                System.err.println("    Namespace: " +format.getNamespace());
            }
        }

        @Override
        public void onStartListingRecords(String prefix) {
            System.err.println("onStartListingRecords: " + prefix);
            this.prefix = prefix;
        }

        @Override
        public void onFinishListingRecords(String prefix, Date latestDatestamp) {
            System.err.println("onFinishListingRecords: " + prefix +
                ", latest = " + latestDatestamp);
            this.prefix = null;
        }


        @Override
        public void onRecordMetadata(Header header, XMLStreamReader reader) {
            System.err.println("==> Record");
            System.err.println("    Identifier: " + header.getIdentifier());
            System.err.println("    Datestamp: " + header.getDatestamp());
            if (header.getSets() != null) {
                System.err.println("    Sets: " + header.getSets());
            }
            String filename = header.getIdentifier();
            filename = replaceBadChars(filename) + ".xml";
            if (prefix != null) {
                filename = replaceBadChars(prefix) + "/" + filename; 
            }
            try {
                ZipEntry entry = new ZipEntry(filename);
                output.putNextEntry(entry);
                XMLStreamWriter writer =
                    factory.createXMLStreamWriter(output);
                XMLStreamCopier.copy(reader, writer);
                writer.flush();
                writer.close();
                output.closeEntry();
                output.flush();
                reader.close();
            } catch (Exception e) {
                e.printStackTrace();
                throw new RuntimeException("record failed");
            }
        }

        @Override
        public InputStream wrap(InputStream stream) throws IOException {
            final int size = 16*1024*1024;
            BufferedInputStream in =
                new BufferedInputStream(stream, size);
            in.mark(size);
            String x = Integer.toHexString(reqNum++);
            while (x.length() < 4) {
                x = "0" + x;
            }
            final String filename = "request/0x" + x + ".xml";
            System.err.println("FN: " + filename);
            ZipEntry entry = new ZipEntry(filename);
            output.putNextEntry(entry);
            byte[] buffer = new byte[8192];
            int r = -1;
            do {
                r = in.read(buffer, 0, buffer.length - 1);
                if (r > 0) {
                    output.write(buffer, 0, r);
                }
            } while (r > 0);
            output.closeEntry();
            output.flush();
            in.reset();
            return in;
        }

        private String replaceBadChars(String s) {
            return s.replace(':', '_').replace('.', '_').replace('\\', '_')
                    .replace('/', '_');
        }
    };
    
    public void run(String[] args) {
        if (args.length < 2) {
            System.err.println("Usage: [baseURI] [outfile]");
            System.exit(1);
        }
        final String repos = args[0];
        MyHarvestHandler handler = new MyHarvestHandler(new File(args[1]));
        try {
            Harvester harvester = SimpleHarvester.newInstance();
            HarvestJob job = harvester.createJob(repos, handler);
//            Calendar cal = Calendar.getInstance(TimeZone.getDefault());
//            cal.clear();
//            cal.set(2010, Calendar.DECEMBER, 15);
//            job.setFrom(cal.getTime());
            job.setMetadataPrefixes(Arrays.asList("oai_dc"));
            job.run();

            System.err.println("==> " + job.getState() + ": "+
                    job.getRecordCount() + " record(s) [resumed: " +
                    job.getResumptionCount() + " time(s)] / " +
                    job.getRequestCount() + " request(s) / " +
                    job.getBytesTransferred() + " bytes transferred");
            System.err.println("          " +
                    (job.getTotalTime() / 1000f) + " total / " +
                    (job.getNetworkTime() / 1000f) + " network / " +
                    (job.getWaitTime() / 1000f) + " wait / " +
                    (job.getProcessingTime() / 1000f) + " processing");
        } catch (HarvesterProtocolErrorException e) {
            System.err.println("Protocol error: ");
            if (e.getErrors() != null) {
                for (ProtocolError error : e.getErrors()) {
                    System.err.println("  " + error);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            handler.close();
        }
    }

    public static void main(String[] args) {
        HarvesterTest main = new HarvesterTest();
        main.run(args);
    }

    @SuppressWarnings("unused")
    private static void dump(XMLStreamReader reader) {
        StringBuilder sb = new StringBuilder("Main: ");
        switch (reader.getEventType()) {
        case XMLStreamConstants.START_DOCUMENT:
            sb.append("START_DOC");
            break;
        case XMLStreamConstants.END_DOCUMENT:
            sb.append("END_DOC");
            break;
        case XMLStreamConstants.COMMENT:
            sb.append("COMMENT[");
            sb.append(reader.getTextCharacters(), reader.getTextStart(),
                    reader.getTextLength());
            sb.append("]");
            break;
        case XMLStreamConstants.START_ELEMENT:
            sb.append("START[");
            sb.append(reader.getNamespaceURI());
            sb.append(",");
            sb.append(reader.getLocalName());
            sb.append("]");
            break;
        case XMLStreamConstants.END_ELEMENT:
            sb.append("END[");
            sb.append(reader.getNamespaceURI());
            sb.append(",");
            sb.append(reader.getLocalName());
            sb.append("]");
            break;
        case XMLStreamConstants.CHARACTERS:
            sb.append("CHARACTERS[");
            String text = reader.getText();
            text = text.replace("\n", "\\n").replace("\r", "\\r")
                    .replace("\t", "\\t");
            sb.append(text);
            sb.append("]");
            break;
        case XMLStreamConstants.CDATA:
            sb.append("CDATA[");
            sb.append(reader.getText());
            sb.append("]");
            break;
        default:
            sb.append(Integer.toString(reader.getEventType()));
        }
        System.err.println(sb.toString());
    }

} // class HarvesterTest