/home/mhx/git/github/libmoost/include/moost/psql/pgq.hpp

Go to the documentation of this file.

/* vim:set ts=3 sw=3 sts=3 et: */
#ifndef MOOST_PSQL_PGQ_CONSUMER_HPP__
#define MOOST_PSQL_PGQ_CONSUMER_HPP__

#include <string>
#include <sstream>
#include <cstring>

#include <pqxx/pqxx>

namespace moost {
namespace psql {

class pgq_consumer
{
public:
    pgq_consumer(std::string const &dbconn,
                 std::string const &queue_name,
                 std::string const &consumer_name)
        : m_dbconn(dbconn)
        , m_queue_name(queue_name)
        , m_consumer_name(consumer_name)
    {
    }

    bool register_consumer() const
    {
        pqxx::connection conn(m_dbconn);
        std::stringstream query;
        query << "select pgq.register_consumer("
              << conn.quote(m_queue_name) << ','
              << conn.quote(m_consumer_name) << ");";

        pqxx::work transaction(conn, "RegisterPgqConsumer");
        pqxx::result res = transaction.exec(query);
        // If no exception has been thrown, the consumer is registered now.

        bool const return_value = res.size() == 1 && res[0][0].as(0);

        transaction.commit();

        return return_value;
    }

    bool unregister_consumer() const
    {
        pqxx::connection conn(m_dbconn);

        bool return_value = false;

        try
        {
            std::stringstream query;
            query << "select pgq.unregister_consumer("
                  << conn.quote(m_queue_name) << ','
                  << conn.quote(m_consumer_name) << ");";

            pqxx::work transaction(conn, "UnregisterPgqConsumer");
            pqxx::result res = transaction.exec(query);

            return_value = res.size() == 1 && res[0][0].as(0);

            transaction.commit();
        }
        catch (pqxx::sql_error const & e)
        {
            if (!std::strstr(e.what(), "consumer not registered on queue"))
            {
                throw;
            }
        }

        return return_value;
    }

    bool is_registered() const
    {
        pqxx::connection conn(m_dbconn);
        std::stringstream query;
        query << "select 1 as ok from pgq.get_consumer_info() where queue_name="
              << conn.quote(m_queue_name) << " and consumer_name="
              << conn.quote(m_consumer_name) << ';';

        pqxx::work transaction(conn, "CheckPgqConsumer");
        pqxx::result res = transaction.exec(query);

        return res.size() == 1 && res[0][0].as(0);
    }

    template<class FunctorType>
    void poll(char const * columns, FunctorType const & functor) const
    {
        pqxx::connection conn(m_dbconn);

        for(;; )
        {
            std::stringstream query;
            query << "select next_batch from pgq.next_batch("
                  << conn.quote(m_queue_name) << ','
                  << conn.quote(m_consumer_name) << ");";

            pqxx::work transaction1(conn, "PollPgq1");
            pqxx::result res = transaction1.exec(query);

            if (res.empty() || res.front()[0].is_null())
            {
                transaction1.commit();
                return;
            }

            long next_batch;
            res.front()[0].to(next_batch);
            transaction1.commit();

            query.str(std::string());
            query << "select " << columns << " from pgq.get_batch_events(" << next_batch << ");";

            pqxx::work transaction2(conn, "PollPgq2");

            functor(transaction2.exec(query));

            query.str(std::string());
            query << "select pgq.finish_batch(" << next_batch << ");";
            transaction2.exec(query);
            transaction2.commit();
        }
    }

    std::string const & dbconn() const
    {
        return m_dbconn;
    }

    std::string const & queue_name() const
    {
        return m_queue_name;
    }

    std::string const & consumer_name() const
    {
        return m_consumer_name;
    }

private:
    std::string const m_dbconn;
    std::string const m_queue_name;
    std::string const m_consumer_name;
};

} // namespace
} // namespace

#endif // ifndef MOOST_PSQL_PGQ_CONSUMER_HPP__