PageRenderTime 40ms CodeModel.GetById 16ms app.highlight 21ms RepoModel.GetById 1ms app.codeStats 0ms

/mordor/pq/connection.cpp

http://github.com/mozy/mordor
C++ | 277 lines | 256 code | 20 blank | 1 comment | 36 complexity | 390ac4c1b1f7f4cc61c7796a361ff1ea MD5 | raw file
  1// Copyright (c) 2009 - Mozy, Inc.
  2
  3#include "connection.h"
  4
  5#include "mordor/assert.h"
  6#include "mordor/iomanager.h"
  7#include "mordor/log.h"
  8
  9#include "exception.h"
 10
 11#ifdef MSVC
 12#pragma comment(lib, "libpq")
 13#endif
 14
 15namespace Mordor {
 16namespace PQ {
 17
 18static Logger::ptr g_log = Log::lookup("mordor:pq");
 19
 20Connection::Connection(const std::string &conninfo, IOManager *ioManager,
 21    Scheduler *scheduler, bool connectImmediately)
 22: m_conninfo(conninfo)
 23, m_exceptioned(false)
 24{
 25#ifdef WINDOWS
 26    m_scheduler = scheduler;
 27#else
 28    m_scheduler = ioManager;
 29#endif
 30    if (connectImmediately)
 31        connect();
 32}
 33
 34ConnStatusType
 35Connection::status()
 36{
 37    if (!m_conn || m_exceptioned)
 38        return CONNECTION_BAD;
 39    return PQstatus(m_conn.get());
 40}
 41
 42void
 43Connection::connect()
 44{
 45    m_exceptioned = false;
 46#ifdef WINDOWS
 47    SchedulerSwitcher switcher(m_scheduler);
 48#else
 49    if (m_scheduler) {
 50        m_conn.reset(PQconnectStart(m_conninfo.c_str()), &PQfinish);
 51        if (!m_conn)
 52            MORDOR_THROW_EXCEPTION(std::bad_alloc());
 53        if (status() == CONNECTION_BAD)
 54            throwException(m_conn.get());
 55        if (PQsetnonblocking(m_conn.get(), 1))
 56            throwException(m_conn.get());
 57        int fd = PQsocket(m_conn.get());
 58        PostgresPollingStatusType whatToPoll = PGRES_POLLING_WRITING;
 59        while (true) {
 60            MORDOR_LOG_DEBUG(g_log) << m_conn.get() << " PQconnectPoll(): "
 61                << whatToPoll;
 62            switch (whatToPoll) {
 63                case PGRES_POLLING_READING:
 64                    m_scheduler->registerEvent(fd, SchedulerType::READ);
 65                    Scheduler::yieldTo();
 66                    break;
 67                case PGRES_POLLING_WRITING:
 68                    m_scheduler->registerEvent(fd, SchedulerType::WRITE);
 69                    Scheduler::yieldTo();
 70                    break;
 71                case PGRES_POLLING_FAILED:
 72                    throwException(m_conn.get());
 73                case PGRES_POLLING_OK:
 74                    MORDOR_LOG_INFO(g_log) << m_conn.get() << " PQconnectStart(\""
 75                        << m_conninfo << "\")";
 76                    return;
 77                default:
 78                    MORDOR_NOTREACHED();
 79            }
 80            whatToPoll = PQconnectPoll(m_conn.get());
 81        }
 82    } else
 83#endif
 84    {
 85        m_conn.reset(PQconnectdb(m_conninfo.c_str()), &PQfinish);
 86        if (!m_conn)
 87            MORDOR_THROW_EXCEPTION(std::bad_alloc());
 88        if (status() == CONNECTION_BAD)
 89            throwException(m_conn.get());
 90    }
 91    MORDOR_LOG_INFO(g_log) << m_conn.get() << " PQconnectdb(\"" << m_conninfo << "\")";
 92}
 93
 94void
 95Connection::reset()
 96{
 97    m_exceptioned = false;
 98#ifdef WINDOWS
 99    SchedulerSwitcher switcher(m_scheduler);
100#else
101    if (m_scheduler) {
102        if (!PQresetStart(m_conn.get()))
103            throwException(m_conn.get());
104        int fd = PQsocket(m_conn.get());
105        PostgresPollingStatusType whatToPoll = PGRES_POLLING_WRITING;
106        while (true) {
107            MORDOR_LOG_DEBUG(g_log) << m_conn.get() << " PQresetPoll(): "
108                << whatToPoll;
109            switch (whatToPoll) {
110                case PGRES_POLLING_READING:
111                    m_scheduler->registerEvent(fd, SchedulerType::READ);
112                    Scheduler::yieldTo();
113                    break;
114                case PGRES_POLLING_WRITING:
115                    m_scheduler->registerEvent(fd, SchedulerType::WRITE);
116                    Scheduler::yieldTo();
117                    break;
118                case PGRES_POLLING_FAILED:
119                    throwException(m_conn.get());
120                case PGRES_POLLING_OK:
121                    MORDOR_LOG_INFO(g_log) << m_conn.get() << " PQresetStart()";
122                    return;
123                default:
124                    MORDOR_NOTREACHED();
125            }
126            whatToPoll = PQresetPoll(m_conn.get());
127        }
128    } else
129#endif
130    {
131        PQreset(m_conn.get());
132        if (status() == CONNECTION_BAD)
133            throwException(m_conn.get());
134    }
135    MORDOR_LOG_INFO(g_log) << m_conn.get() << " PQreset()";
136}
137
138std::string escape(PGconn *conn, const std::string &string)
139{
140    std::string result;
141    result.resize(string.size() * 2);
142    int error = 0;
143    size_t resultSize = PQescapeStringConn(conn, &result[0],
144        string.c_str(), string.size(), &error);
145    if (error)
146        throwException(conn);
147    result.resize(resultSize);
148    return result;
149}
150
151std::string
152Connection::escape(const std::string &string)
153{
154    return PQ::escape(m_conn.get(), string);
155}
156
157static std::string escapeBinary(PGconn *conn, const std::string &blob)
158{
159    size_t length;
160    std::string resultString;
161    char *result = (char *)PQescapeByteaConn(conn,
162        (unsigned char *)blob.c_str(), blob.size(), &length);
163    if (!result)
164        throwException(conn);
165    try {
166        resultString.append(result, length);
167    } catch (...) {
168        PQfreemem(result);
169        throw;
170    }
171    PQfreemem(result);
172    return resultString;
173}
174
175std::string
176Connection::escapeBinary(const std::string &blob)
177{
178    return PQ::escapeBinary(m_conn.get(), blob);
179}
180
181#ifndef WINDOWS
182void flush(PGconn *conn, SchedulerType *scheduler)
183{
184    while (true) {
185        int result = PQflush(conn);
186        MORDOR_LOG_DEBUG(g_log) << conn << " PQflush(): " << result;
187        switch (result) {
188            case 0:
189                return;
190            case -1:
191                throwException(conn);
192            case 1:
193                scheduler->registerEvent(PQsocket(conn), SchedulerType::WRITE);
194                Scheduler::yieldTo();
195                continue;
196            default:
197                MORDOR_NOTREACHED();
198        }
199    }
200}
201
202PGresult *nextResult(PGconn *conn, SchedulerType *scheduler)
203{
204    while (true) {
205        if (!PQconsumeInput(conn))
206            throwException(conn);
207        if (PQisBusy(conn)) {
208            MORDOR_LOG_DEBUG(g_log) << conn << " PQisBusy()";
209            scheduler->registerEvent(PQsocket(conn), SchedulerType::READ);
210            Scheduler::yieldTo();
211            continue;
212        }
213        MORDOR_LOG_DEBUG(g_log) << conn << " PQconsumeInput()";
214        return PQgetResult(conn);
215    }
216}
217#endif
218
219PreparedStatement
220Connection::prepare(const std::string &command, const std::string &name, PreparedStatement::ResultFormat resultFormat)
221{
222    if (!name.empty()) {
223#ifdef WINDOWS
224        SchedulerSwitcher switcher(m_scheduler);
225#else
226        if (m_scheduler) {
227            if (!PQsendPrepare(m_conn.get(), name.c_str(), command.c_str(), 0, NULL))
228                throwException(m_conn.get());
229            flush(m_conn.get(), m_scheduler);
230            boost::shared_ptr<PGresult> result(nextResult(m_conn.get(), m_scheduler),
231                &PQclear);
232            while (result) {
233                ExecStatusType status = PQresultStatus(result.get());
234                MORDOR_LOG_DEBUG(g_log) << m_conn.get() << " PQresultStatus("
235                    << result.get() << "): " << PQresStatus(status);
236                if (status != PGRES_COMMAND_OK)
237                    throwException(result.get());
238                result.reset(nextResult(m_conn.get(), m_scheduler),
239                    &PQclear);
240            }
241            MORDOR_LOG_VERBOSE(g_log) << m_conn.get() << " PQsendPrepare(\""
242                << name << "\", \"" << command << "\")";
243            return PreparedStatement(m_conn, std::string(), name, m_scheduler, resultFormat);
244        } else
245#endif
246        {
247            boost::shared_ptr<PGresult> result(PQprepare(m_conn.get(),
248                name.c_str(), command.c_str(), 0, NULL), &PQclear);
249            if (!result)
250                throwException(m_conn.get());
251            ExecStatusType status = PQresultStatus(result.get());
252            MORDOR_LOG_DEBUG(g_log) << m_conn.get() << " PQresultStatus("
253                << result.get() << "): " << PQresStatus(status);
254            if (status != PGRES_COMMAND_OK)
255                throwException(result.get());
256            MORDOR_LOG_VERBOSE(g_log) << m_conn.get() << " PQprepare(\"" << name
257                << "\", \"" << command << "\")";
258            return PreparedStatement(m_conn, std::string(), name, m_scheduler, resultFormat);
259        }
260    } else {
261        return PreparedStatement(m_conn, command, name, m_scheduler, resultFormat);
262    }
263}
264
265Connection::CopyInParams
266Connection::copyIn(const std::string &table)
267{
268    return CopyInParams(table, m_conn, m_scheduler);
269}
270
271Connection::CopyOutParams
272Connection::copyOut(const std::string &table)
273{
274    return CopyOutParams(table, m_conn, m_scheduler);
275}
276
277}}