/mordor/pq/copy.cpp

  1// Copyright (c) 2010 - Mozy, Inc.
  2
  3#include "mordor/assert.h"
  4#include "mordor/iomanager.h"
  5#include "mordor/log.h"
  6#include "mordor/streams/buffer.h"
  7#include "mordor/streams/stream.h"
  8
  9#include "connection.h"
 10#include "exception.h"
 11
 12namespace Mordor {
 13namespace PQ {
 14
 15static Logger::ptr g_log = Log::lookup("mordor:pq");
 16
 17Connection::CopyParams::CopyParams(const std::string &table,
 18    boost::shared_ptr<PGconn> conn, SchedulerType *scheduler)
 19    : m_table(table),
 20      m_scheduler(scheduler),
 21      m_conn(conn),
 22      m_binary(false),
 23      m_csv(false),
 24      m_header(false),
 25      m_delimiter('\0'),
 26      m_quote('\0'),
 27      m_escape('\0')
 28{}
 29
 30Connection::CopyParams &
 31Connection::CopyParams::columns(const std::vector<std::string> &columns)
 32{
 33    m_columns = columns;
 34    return *this;
 35}
 36
 37Connection::CopyParams &
 38Connection::CopyParams::binary()
 39{
 40    MORDOR_ASSERT(!m_csv);
 41    m_binary = true;
 42    return *this;
 43}
 44
 45Connection::CopyParams &
 46Connection::CopyParams::csv()
 47{
 48    MORDOR_ASSERT(!m_binary);
 49    m_csv = true;
 50    return *this;
 51}
 52
 53Connection::CopyParams &
 54Connection::CopyParams::delimiter(char delimiter)
 55{
 56    MORDOR_ASSERT(!m_binary);
 57    m_delimiter = delimiter;
 58    return *this;
 59}
 60
 61Connection::CopyParams &
 62Connection::CopyParams::nullString(const std::string &nullString)
 63{
 64    MORDOR_ASSERT(!m_binary);
 65    m_nullString = nullString;
 66    return *this;
 67}
 68
 69Connection::CopyParams &
 70Connection::CopyParams::header()
 71{
 72    MORDOR_ASSERT(m_csv);
 73    m_header = true;
 74    return *this;
 75}
 76
 77Connection::CopyParams &
 78Connection::CopyParams::quote(char quote)
 79{
 80    MORDOR_ASSERT(m_csv);
 81    m_quote = quote;
 82    return *this;
 83}
 84
 85Connection::CopyParams &
 86Connection::CopyParams::escape(char escape)
 87{
 88    MORDOR_ASSERT(m_csv);
 89    m_escape = escape;
 90    return *this;
 91}
 92
 93Connection::CopyParams &
 94Connection::CopyParams::notNullQuoteColumns(const std::vector<std::string> &columns)
 95{
 96    MORDOR_ASSERT(m_csv);
 97    m_notNullQuoteColumns = columns;
 98    return *this;
 99}
100
101class CopyInStream : public Stream
102{
103public:
104    CopyInStream(boost::shared_ptr<PGconn> conn, SchedulerType *scheduler)
105        : m_conn(conn),
106          m_scheduler(scheduler)
107    {}
108
109    ~CopyInStream()
110    {
111        boost::shared_ptr<PGconn> sharedConn = m_conn.lock();
112        if (sharedConn) {
113            PGconn *conn = sharedConn.get();
114            try {
115                putCopyEnd(conn, "COPY IN aborted");
116            } catch (...) {
117            }
118        }
119    }
120
121    bool supportsWrite() { return true; }
122
123    void close(CloseType type)
124    {
125        MORDOR_ASSERT(type & WRITE);
126        boost::shared_ptr<PGconn> sharedConn = m_conn.lock();
127        if (sharedConn) {
128            PGconn *conn = sharedConn.get();
129            putCopyEnd(conn, NULL);
130            m_conn.reset();
131        }
132    }
133
134    size_t write(const void *buffer, size_t length)
135    {
136        boost::shared_ptr<PGconn> sharedConn = m_conn.lock();
137        MORDOR_ASSERT(sharedConn);
138        PGconn *conn = sharedConn.get();
139        int status = 0;
140        length = std::min<size_t>(length, 0x7fffffff);
141#ifdef WINDOWS
142        SchedulerSwitcher switcher(m_scheduler);
143#endif
144        while (status == 0) {
145            status = PQputCopyData(conn, (const char *)buffer, (int)length);
146            switch (status) {
147                case 1:
148                    return length;
149                case -1:
150                    throwException(conn);
151#ifndef WINDOWS
152                case 0:
153                    MORDOR_ASSERT(m_scheduler);
154                    m_scheduler->registerEvent(PQsocket(conn),
155                        SchedulerType::WRITE);
156                    Scheduler::yieldTo();
157                    break;
158#endif
159                default:
160                    MORDOR_NOTREACHED();
161            }
162        }
163        MORDOR_NOTREACHED();
164    }
165
166private:
167    void putCopyEnd(PGconn *conn, const char *error) {
168#ifdef WINDOWS
169        SchedulerSwitcher switcher(m_scheduler);
170#endif
171        int status = 0;
172        while (status == 0) {
173            status = PQputCopyEnd(conn, error);
174            switch (status) {
175                case 1:
176                    break;
177                case -1:
178                    throwException(conn);
179#ifndef WINDOWS
180                case 0:
181                    MORDOR_ASSERT(m_scheduler);
182                    m_scheduler->registerEvent(PQsocket(conn),
183                        SchedulerType::WRITE);
184                    Scheduler::yieldTo();
185                    break;
186#endif
187                default:
188                    MORDOR_NOTREACHED();
189            }
190        }
191#ifndef WINDOWS
192        if (m_scheduler)
193            PQ::flush(conn, m_scheduler);
194#endif
195        boost::shared_ptr<PGresult> result;
196#ifndef WINDOWS
197        if (m_scheduler)
198            result.reset(nextResult(conn, m_scheduler), &PQclear);
199        else
200#endif
201            result.reset(PQgetResult(conn), &PQclear);
202        while (result) {
203            ExecStatusType status = PQresultStatus(result.get());
204            MORDOR_LOG_DEBUG(g_log) << conn << " PQresultStatus("
205                << result.get() << "): " << PQresStatus(status);
206            if (status != PGRES_COMMAND_OK)
207                throwException(result.get());
208#ifndef WINDOWS
209            if (m_scheduler)
210                result.reset(nextResult(conn, m_scheduler), &PQclear);
211            else
212#endif
213                result.reset(PQgetResult(conn), &PQclear);
214        }
215        MORDOR_LOG_VERBOSE(g_log) << conn << " PQputCopyEnd(\""
216            << (error ? error : "") << "\")";
217    }
218
219private:
220    boost::weak_ptr<PGconn> m_conn;
221    SchedulerType *m_scheduler;
222};
223
224class CopyOutStream : public Stream
225{
226public:
227    CopyOutStream(boost::shared_ptr<PGconn> conn, SchedulerType *scheduler)
228        : m_conn(conn),
229          m_scheduler(scheduler)
230    {}
231
232    bool supportsRead() { return true; }
233
234    size_t read(Buffer &buffer, size_t length)
235    {
236        if (m_readBuffer.readAvailable()) {
237            length = (std::min)(m_readBuffer.readAvailable(), length);
238            buffer.copyIn(m_readBuffer, length);
239            m_readBuffer.consume(length);
240            return length;
241        }
242        boost::shared_ptr<PGconn> sharedConn = m_conn.lock();
243        if (!sharedConn)
244            return 0;
245        PGconn *conn = sharedConn.get();
246#ifndef WINDOWS
247        SchedulerSwitcher switcher(m_scheduler);
248#endif
249        int status = 0;
250        do {
251            char *data = NULL;
252            status = PQgetCopyData(conn, &data,
253#ifdef WINDOWS
254                0
255#else
256                m_scheduler ? 1 : 0
257#endif
258                );
259            switch (status) {
260                case 0:
261#ifdef WINDOWS
262                    MORDOR_NOTREACHED();
263#else
264                    MORDOR_ASSERT(m_scheduler);
265                    m_scheduler->registerEvent(PQsocket(conn),
266                        SchedulerType::READ);
267                    Scheduler::yieldTo();
268                    continue;
269#endif
270                case -1:
271                    break;
272                case -2:
273                    throwException(conn);
274                default:
275                    MORDOR_ASSERT(status > 0);
276                    try {
277                        m_readBuffer.copyIn(data, status);
278                    } catch (...) {
279                        PQfreemem(data);
280                        throw;
281                    }
282                    PQfreemem(data);
283                    break;
284            }
285        } while (false);
286
287        if (status == -1) {
288            m_conn.reset();
289            boost::shared_ptr<PGresult> result;
290#ifndef WINDOWS
291            if (m_scheduler)
292                result.reset(nextResult(conn, m_scheduler), &PQclear);
293            else
294#endif
295                result.reset(PQgetResult(conn), &PQclear);
296            while (result) {
297                ExecStatusType status = PQresultStatus(result.get());
298                MORDOR_LOG_DEBUG(g_log) << conn << " PQresultStatus("
299                    << result.get() << "): " << PQresStatus(status);
300                if (status != PGRES_COMMAND_OK)
301                    throwException(result.get());
302#ifndef WINDOWS
303                if (m_scheduler)
304                    result.reset(nextResult(conn, m_scheduler), &PQclear);
305                else
306#endif
307                    result.reset(PQgetResult(conn), &PQclear);
308            }
309        }
310
311        length = (std::min)(m_readBuffer.readAvailable(), length);
312        buffer.copyIn(m_readBuffer, length);
313        m_readBuffer.consume(length);
314        return length;
315    }
316
317private:
318    boost::weak_ptr<PGconn> m_conn;
319    SchedulerType *m_scheduler;
320    Buffer m_readBuffer;
321};
322
323Stream::ptr
324Connection::CopyParams::execute(bool out)
325{
326    PGconn *conn = m_conn.get();
327    std::ostringstream os;
328    os << "COPY " << m_table << " ";
329    if (!m_columns.empty()) {
330        os << "(";
331        for (std::vector<std::string>::const_iterator it(m_columns.begin());
332            it != m_columns.end();
333            ++it) {
334            if (it != m_columns.begin())
335                os << ", ";
336            os << *it;
337        }
338        os << ") ";
339    }
340    os << (out ? "TO STDOUT" : "FROM STDIN");
341    if (m_binary) {
342        os << " BINARY";
343    } else {
344        if (m_delimiter != '\0')
345            os << " DELIMITER '"
346                << PQ::escape(conn, std::string(1, m_delimiter)) << '\'';
347        if (!m_nullString.empty())
348            os << " NULL '" << PQ::escape(conn, m_nullString) << '\'';
349        if (m_csv) {
350            os << " CSV";
351            if (m_header)
352                os << " HEADER";
353            if (m_quote != '\0')
354                os << " QUOTE '"
355                    << PQ::escape(conn, std::string(1, m_quote)) << '\'';
356            if (m_escape != '\0')
357                os << " ESCAPE '"
358                    << PQ::escape(conn, std::string(1, m_escape)) << '\'';
359            if (!m_notNullQuoteColumns.empty()) {
360                os << (out ? " FORCE QUOTE" : " FORCE NOT NULL ");
361                for (std::vector<std::string>::const_iterator it(m_notNullQuoteColumns.begin());
362                    it != m_notNullQuoteColumns.end();
363                    ++it) {
364                    if (it != m_notNullQuoteColumns.begin())
365                        os << ", ";
366                    os << *it;
367                }
368            }
369        }
370    }
371
372    boost::shared_ptr<PGresult> result, next;
373    const char *api = NULL;
374#ifdef WINDOWS
375    SchedulerSwitcher switcher(m_scheduler);
376#else
377    if (m_scheduler) {
378        api = "PQsendQuery";
379        if (!PQsendQuery(conn, os.str().c_str()))
380            throwException(conn);
381        flush(conn, m_scheduler);
382        next.reset(nextResult(conn, m_scheduler), &PQclear);
383        while (next) {
384            result = next;
385            if (PQresultStatus(result.get()) ==
386                (out ? PGRES_COPY_OUT : PGRES_COPY_IN))
387                break;
388            next.reset(nextResult(conn, m_scheduler), &PQclear);
389            if (next) {
390                ExecStatusType status = PQresultStatus(next.get());
391                MORDOR_LOG_VERBOSE(g_log) << conn << "PQresultStatus(" <<
392                    next.get() << "): " << PQresStatus(status);
393                switch (status) {
394                    case PGRES_COMMAND_OK:
395                    case PGRES_TUPLES_OK:
396                        break;
397                    default:
398                        throwException(next.get());
399                        MORDOR_NOTREACHED();
400                }
401            }
402        }
403    } else
404#endif
405    {
406        api = "PQexec";
407        result.reset(PQexec(conn, os.str().c_str()), &PQclear);
408    }
409    if (!result)
410        throwException(conn);
411    ExecStatusType status = PQresultStatus(result.get());
412    MORDOR_ASSERT(api);
413    MORDOR_LOG_VERBOSE(g_log) << conn << " " << api << "(\"" << os.str()
414        << "\"), PQresultStatus(" << result.get() << "): "
415        << PQresStatus(status);
416    switch (status) {
417        case PGRES_COMMAND_OK:
418        case PGRES_TUPLES_OK:
419            MORDOR_NOTREACHED();
420        case PGRES_COPY_IN:
421            MORDOR_ASSERT(!out);
422            return Stream::ptr(new CopyInStream(m_conn, m_scheduler));
423        case PGRES_COPY_OUT:
424            MORDOR_ASSERT(out);
425            return Stream::ptr(new CopyOutStream(m_conn, m_scheduler));
426        default:
427            throwException(result.get());
428            MORDOR_NOTREACHED();
429    }
430}
431
432Stream::ptr
433Connection::CopyInParams::operator()()
434{
435    return execute(false);
436}
437
438Stream::ptr
439Connection::CopyOutParams::operator()()
440{
441    return execute(true);
442}
443
444}}