path: root/src/io/DataWriter.java

package io;

import java.io.Closeable;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.HashSet;
import java.util.Scanner;
import java.util.Set;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.apache.commons.io.Charsets;
import org.apache.commons.io.IOUtils;
import org.json.JSONException;
import org.json.JSONObject;
import provider.ResultListener;

/**
 * This class writes the JSON objects for tweets and users to (separate) files.
 *
 * @author Maurice Laveaux
 */
public class DataWriter implements ResultListener, Closeable {

    /**
     * The writer for the tweet stream.
     */
    private final Store m_tweet;

    /**
     * the writer for the profile stream.
     */
    private final Store m_profile;

    /**
     * the buffer of tweet ids that already exist.
     */
    private final Set<Long> m_tweetIdSet;

    /**
     * the buffer of profile ids that already exist.
     */
    private final Set<Long> m_profileIdSet;

    public final static String CFG_PROFILE_FILENAME = "profiles-filename";
    public final static String CFG_TWEETS_FILENAME = "tweets-filename";

    /**
     * Creates an instance, specifying the profile and file names to which
     * tweets and users are written.
     *
     * @param profilesName The file to write the profiles to.
     * @param tweetsName The file to write the tweets to.
     */
    public DataWriter(final String profilesName, final String tweetsName) {
        m_profile = getStore(profilesName);
        m_profileIdSet = new HashSet<>();
        m_tweet = getStore(tweetsName);
        m_tweetIdSet = new HashSet<>();
    }

    /**
     * Opens a stream to every single file that data will be streamed to.
     *
     * @throws java.io.IOException if the files cannot be read or written.
     */
    public void open() throws IOException {
        try {
            m_profile.open();
            m_tweet.open();
            readIds(m_profileIdSet, m_profile);
            readIds(m_tweetIdSet, m_tweet);
        } catch (IOException ex) {
            m_profile.close();
            m_tweet.close();
            throw ex;
        }
    }

    /**
     * For a given filename, get an instance that provides I/O.
     *
     * @param filename File to read/write.
     * @return A Store instance providing read and write access.
     */
    protected Store getStore(String filename) {
        return new SimpleFileStore(filename);
    }

    public void close() {
        m_tweet.close();
        m_profile.close();
    }

    @Override
    public void profileGenerated(JSONObject obj) {
        writeObject(obj, m_profile.getOutputStream(), m_profileIdSet);
    }

    @Override
    public void tweetGenerated(JSONObject obj) {
        writeObject(obj, m_tweet.getOutputStream(), m_tweetIdSet);
    }

    /**
     * Read the current existing tweetName and profileName filenames and fill
     * the existing id set.
     *
     * @param is An input stream that provides JSON objects.
     * @return The set of ids, may be empty if the fill does not exist.
     */
    private void readIds(Set<Long> idSet, Store store) throws IOException {
        InputStream is = null;
        try {
            is = store.getInputStream();
            Scanner reader = new Scanner(is);
            // parse each line into a JSONObject, read the id and add it to
            // the set of ids.
            while (reader.hasNext()) {
                JSONObject obj = new JSONObject(reader.nextLine());
                long id = obj.getLong("id");
                idSet.add(id);
            }
        } catch (FileNotFoundException ex) {
            // ignore, file will be created if necessary.
        } catch (JSONException ex) {
            getLogger().log(Level.WARNING, store.getFileName()
                    + ": Corrupt file?", ex);
            throw new IOException(ex);
        } finally {
            IOUtils.closeQuietly(is);
        }
    }

    /**
     * Writes the JSONObject to a writer and update the idSet.
     *
     * @param obj The object to write.
     * @param output The stream to write objects to.
     * @param idSet The id set to add the obj id to.
     */
    private void writeObject(JSONObject obj, OutputStream output,
            Set<Long> idSet) {
        try {
            long id = obj.getLong("id");

            if (!idSet.contains(id)) {
                // Write a single profile into the profile file.
                try {
                    output.write((obj.toString() + "\n").getBytes(Charsets.UTF_8));
                    idSet.add(id);
                } catch (IOException ex) {
                    getLogger().log(Level.WARNING, "Cannot write to file", ex);
                }
            }
        } catch (JSONException ex) {
            getLogger().log(Level.WARNING, "ID not found?!", ex);
        }
    }

    private Logger getLogger() {
        return Logger.getLogger(getClass().getName());
    }

    abstract class Store {

        private final String filename;
        protected OutputStream os;

        Store(String filename) {
            this.filename = filename;
        }

        /**
         * @return The file name associated with this store.
         */
        public String getFileName() {
            return filename;
        }

        /**
         * Opens the writable stream (if not open yet), must be called before
         * writing a stream.
         *
         * @throws IOException
         */
        abstract public void open() throws IOException;

        /**
         * Properly closes a file, flushing any buffers.
         */
        public void close() {
            try {
                if (os != null) {
                    os.close();
                    os = null;
                }
            } catch (IOException ex) {
                Logger.getLogger(getClass().getName()).log(Level.WARNING,
                        "Could not close " + getFileName(), ex);
            }
        }

        /**
         * Returns a previously opened writable file stream.
         *
         * @return A stream to which JSON objects can be written (one per line).
         * If no writable stream is open, {@code null} is returned.
         */
        public OutputStream getOutputStream() {
            return os;
        }

        /**
         * Opens a new input stream for the file.
         *
         * @return A stream from which JSON objects can be read (one per line).
         * @throws IOException if no input file can be retrieved.
         */
        abstract public InputStream getInputStream() throws IOException;
    }

    class SimpleFileStore extends Store {

        SimpleFileStore(String filename) {
            super(filename);
        }

        @Override
        public void open() throws IOException {
            if (os == null) {
                os = new FileOutputStream(getFileName(), true);
            }
        }

        @Override
        public InputStream getInputStream() throws FileNotFoundException {
            return new FileInputStream(getFileName());
        }
    }
}