GZCSVFileWriter.cpp

/*
** libutil++
** $Id: GZCSVFileWriter.cpp 1653 2016-02-28 19:54:59Z sella $
** Copyright (c) 2011-2016 Digital Genesis, LLC. All Rights Reserved.
** Released under the LGPL Version 2.1 License.
** http://www.digitalgenesis.com
*/

static const char rcsid[] __attribute__((used)) = "$Id: GZCSVFileWriter.cpp 1653 2016-02-28 19:54:59Z sella $";

#include "GZCSVFileWriter.h"

#include <fcntl.h>

using namespace sella::file;

GZCSVFileWriter::GZCSVFileWriter(const std::string & file, char delimiter, char seperator) THROWS(IOException) :
        filename(file),

        delimiter(delimiter),
        seperator(seperator),

        descriptor(::open(file.c_str(), O_WRONLY | O_LARGEFILE | O_CREAT, S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH | S_IWOTH))
{
        /* Ensure the file was opened for writing, and zero all buffers and state machines. */
        if (this->descriptor <= 0) {
                THROW(IOException, "failed to open file '%s' for writing", this->filename.c_str());
        }
        memset(&(this->stream), 0, sizeof(this->stream));
        memset(&(this->buffer), 0, sizeof(this->buffer));

        /* Initialize for the writing of GZIP compressed files. */
        if (::deflateInit2(&(this->stream), 9 /* Best Compression */, Z_DEFLATED, 16 + MAX_WBITS /* GZIP */, 9 /* Maximum Memory and Speed */, Z_DEFAULT_STRATEGY) != Z_OK) {
                ::close(this->descriptor);

                THROW(IOException, "failed to initilize zlib instance for writing file '%s'", this->filename.c_str());
        }
}

GZCSVFileWriter::~GZCSVFileWriter() NO_EXCEPTIONS {
        if (this->descriptor != 0) {
                ::deflateEnd(&(this->stream));

                this->close();
        }
}

void GZCSVFileWriter::write(const std::vector<std::string> & data) THROWS(IOException) {
        size_t length = 0;

        /* Write all fields, one at a time, to the compression stream. */
        for (auto iter = data.begin(); iter != data.end(); ++iter) {
                ssize_t offset = 0;

                /* Append the field delimiter if this is not the first field of the record. */
                if (iter != data.begin()) {
                        /* Purge the buffer if required to make room for the field delimiter. */
                        if (sizeof(this->delimiter) > (sizeof(this->buffer[1]) - length)) {
                                this->write(this->buffer[1], length, Z_NO_FLUSH);

                                length = 0;
                        }
                        this->buffer[1][length++] = this->delimiter;
                }

                do { /* Append the field value to the buffer. */
                        /* Copy as much of the field value as possible to the buffer. */
                        if ((iter->length() - offset) > (sizeof(this->buffer[1]) - length)) {
                                ::memcpy(&(this->buffer[1][length]), &(iter->c_str()[offset]), sizeof(this->buffer[1]) - length);

                                this->write(this->buffer[1], sizeof(this->buffer[1]), Z_NO_FLUSH);

                                offset += (sizeof(this->buffer[1]) - length);
                                length = 0;
                        } else {
                                ::memcpy(&(this->buffer[1][length]), &(iter->c_str()[offset]), iter->length() - offset);

                                length += (iter->length() - offset);
                                offset += (iter->length() - offset);
                        }
                } while (offset != static_cast<ssize_t>(iter->length()));
        }

        /* Purge the buffer if required to make room for the record seperator. */
        if (sizeof(this->seperator) > (sizeof(this->buffer[1]) - length)) {
                this->write(this->buffer[1], length, Z_NO_FLUSH);

                length = 0;
        }
        this->buffer[1][length++] = this->seperator;

        this->write(this->buffer[1], length, Z_NO_FLUSH);
}

void GZCSVFileWriter::write(uint8_t * input, ssize_t length, int flush) THROWS(IOException) {
        ssize_t expect;

        this->stream.next_in = input;
        this->stream.avail_in = length;

        do {
                this->stream.next_out = this->buffer[0];
                this->stream.avail_out = sizeof(this->buffer[0]);

                /* Compress the supplied data into the buffer. */
                switch (::deflate(&(this->stream), flush)) {
                        case Z_STREAM_ERROR:
                                THROW(IOException, "zlib compression failure when writing buffer to compressed file '%s'", this->filename.c_str());
                }

                /* Write the compressed data to the file. */
                if ((expect = (sizeof(this->buffer[0]) - this->stream.avail_out)) > 0) {
                        while ((length = ::write(this->descriptor, this->buffer[0], sizeof(this->buffer[0]) - this->stream.avail_out)) != expect) {
                                if (length == -1) {
                                        THROW(IOException, "failure writing buffer to compressed file '%s'", this->filename.c_str());
                                }
                                expect -= length;
                        }
                }
        } while (this->stream.avail_out == 0);

        /* Ensure all data was compressed and written to the file. */
        if (this->stream.avail_in != 0) {
                THROW(IOException, "failed to write full buffer to compressed file '%s'", this->filename.c_str());
        }
}

void GZCSVFileWriter::close(const std::string & rename) THROWS(IOException) {
        uint8_t buffer[4];

        /* Close the compression stream out. */
        this->write(buffer, 0, Z_FINISH);
        
        ::close(this->descriptor);

        if (rename.length() > 0) {
                ::rename(this->filename.c_str(), rename.c_str());
        }
        this->filename.clear();
        this->descriptor = 0;
}

/*
** vim: noet ts=3 sw=3
*/