tinymux/mux/modules/sqlslave/sqlslave.cpp

/*! \file sqlslave.cpp
 * \brief SQLSlave Module
 *
 */

#include "autoconf.h"
#include "config.h"
#include "libmux.h"
#include "modules.h"
#include "autoconf.h"
#if defined(HAVE_MYSQL_H)
#include <mysql.h>
#endif // HAVE_MYSQL_H
#include "sql.h"

#include <atomic>
#include <new>
#include <string>

class CQueryServer : public mux_IQueryControl
{
public:
    // mux_IUnknown
    //
    virtual MUX_RESULT QueryInterface(MUX_IID iid, void **ppv);
    virtual uint32_t     AddRef(void);
    virtual uint32_t     Release(void);

    // mux_IQueryControl
    //
    virtual MUX_RESULT Connect(const UTF8 *pServer, const UTF8 *pDatabase, const UTF8 *pUser, const UTF8 *pPassword);
    virtual MUX_RESULT Advise(mux_IQuerySink *pIQuerySink);
    virtual MUX_RESULT Query(uint32_t iQueryHandle, const UTF8 *pDatabaseName, const UTF8 *pQuery);

    CQueryServer(void);
    MUX_RESULT FinalConstruct(void);
    virtual ~CQueryServer();

private:
    std::atomic<uint32_t> m_cRef;
    mux_IQuerySink *m_pIQuerySink;
    mux_ILog       *m_pILog;
#if defined(HAVE_MYSQL)
    MYSQL          *m_database;
#endif // HAVE_MYSQL
    // Module-owned copies of the connection parameters.  The caller retains
    // ownership of the buffers it passes to Connect(), so we must not alias or
    // free them here.
    //
    std::string     m_sServer;
    std::string     m_sDatabase;
    std::string     m_sUser;
    std::string     m_sPassword;

    // Returns true if a live MySQL session is established.
    //
    bool ConnectionHelper();

    // Best-effort logging of a MySQL error.  Silently does nothing if no log
    // interface is available (e.g. when running in the slave process).
    //
    void LogError(const UTF8 *pContext);
};

static std::atomic<int32_t> g_cComponents(0);
static std::atomic<int32_t> g_cServerLocks(0);

#define NUM_CLASSES 1
static MUX_CLASS_INFO sum_classes[NUM_CLASSES] =
{
    { CID_QueryServer }
};

// The following four functions are for access by dlopen.
//
extern "C" MUX_RESULT DCL_EXPORT DCL_API mux_CanUnloadNow(void)
{
    if (  0 == g_cComponents
       && 0 == g_cServerLocks)
    {
        return MUX_S_OK;
    }
    else
    {
        return MUX_S_FALSE;
    }
}

extern "C" MUX_RESULT DCL_EXPORT DCL_API mux_GetClassObject(MUX_CID cid, MUX_IID iid, void **ppv)
{
    MUX_RESULT mr = MUX_E_CLASSNOTAVAILABLE;

    if (CID_QueryServer == cid)
    {
        CQueryServerFactory *pQueryServerFactory = NULL;
        try
        {
            pQueryServerFactory = new CQueryServerFactory;
        }
        catch (const std::bad_alloc &)
        {
            ; // Handled by the NULL check below.
        }

        if (NULL == pQueryServerFactory)
        {
            return MUX_E_OUTOFMEMORY;
        }

        mr = pQueryServerFactory->QueryInterface(iid, ppv);
        pQueryServerFactory->Release();
    }
    return mr;
}

extern "C" MUX_RESULT DCL_EXPORT DCL_API mux_Register(void)
{
    // Advertise our components.
    //
    MUX_RESULT mr = mux_RegisterClassObjects(NUM_CLASSES, sum_classes, NULL);
#if defined(HAVE_MYSQL)
    if (MUX_SUCCEEDED(mr))
    {
        if (mysql_library_init(0, NULL, NULL))
        {
            mr = MUX_E_FAIL;
        }
    }
#endif
    return mr;
}

extern "C" MUX_RESULT DCL_EXPORT DCL_API mux_Unregister(void)
{
#if defined(HAVE_MYSQL)
    mysql_library_end();
#endif
    return mux_RevokeClassObjects(NUM_CLASSES, sum_classes);
}

// QueryServer component which is not directly accessible.
//
CQueryServer::CQueryServer(void) : m_cRef(1), m_pIQuerySink(NULL), m_pILog(NULL)
{
#if defined(HAVE_MYSQL)
    m_database = NULL;
#endif // HAVE_MYSQL

    g_cComponents++;
}

MUX_RESULT CQueryServer::FinalConstruct(void)
{
    // Acquire a logging interface on a best-effort basis.  CID_Log is provided
    // by the engine in the main process; when this component runs in the slave
    // process the interface may be unavailable, in which case we simply do not
    // log.  Failure here must not prevent the component from being created.
    //
    mux_CreateInstance(CID_Log, NULL, UseSameProcess, IID_ILog,
                       reinterpret_cast<void **>(&m_pILog));

    return MUX_S_OK;
}

CQueryServer::~CQueryServer()
{
    if (NULL != m_pIQuerySink)
    {
        m_pIQuerySink->Release();
        m_pIQuerySink = NULL;
    }

    if (NULL != m_pILog)
    {
        m_pILog->Release();
        m_pILog = NULL;
    }

#if defined(HAVE_MYSQL)
    if (NULL != m_database)
    {
        mysql_close(m_database);
        m_database = NULL;
    }
#endif // HAVE_MYSQL

    g_cComponents--;
}

MUX_RESULT CQueryServer::QueryInterface(MUX_IID iid, void **ppv)
{
    if (mux_IID_IUnknown == iid)
    {
        *ppv = static_cast<mux_IQueryControl *>(this);
    }
    else if (IID_IQueryControl == iid)
    {
        *ppv = static_cast<mux_IQueryControl *>(this);
    }
    else
    {
        *ppv = NULL;
        return MUX_E_NOINTERFACE;
    }
    reinterpret_cast<mux_IUnknown *>(*ppv)->AddRef();
    return MUX_S_OK;
}

uint32_t CQueryServer::AddRef(void)
{
    return m_cRef.fetch_add(1, std::memory_order_relaxed) + 1;
}

uint32_t CQueryServer::Release(void)
{
    uint32_t prev = m_cRef.fetch_sub(1, std::memory_order_acq_rel);
    if (1 == prev)
    {
        delete this;
        return 0;
    }
    return prev - 1;
}

MUX_RESULT CQueryServer::Connect(const UTF8 *pServer, const UTF8 *pDatabase, const UTF8 *pUser, const UTF8 *pPassword)
{
    // mysql_real_connect() dereferences each of these, and ConnectionHelper()
    // inspects pServer[0].  Reject null arguments rather than crash later.
    //
    if (  NULL == pServer
       || NULL == pDatabase
       || NULL == pUser
       || NULL == pPassword)
    {
        return MUX_E_INVALIDARG;
    }

    // Save copies of the new Server/Database/User/Password values.  These are
    // used later if reconnection is necessary.  We must copy rather than alias
    // the caller's buffers: in-process the caller retains ownership, and across
    // the proxy/stub boundary these arrive in transient (often stack) storage.
    //
    m_sServer.assign(reinterpret_cast<const char *>(pServer));
    m_sDatabase.assign(reinterpret_cast<const char *>(pDatabase));
    m_sUser.assign(reinterpret_cast<const char *>(pUser));
    m_sPassword.assign(reinterpret_cast<const char *>(pPassword));

#if defined(HAVE_MYSQL)
    // Close any existing session.
    //
    if (NULL != m_database)
    {
        mysql_close(m_database);
        m_database = NULL;
    }

    m_database = mysql_init(NULL);

    if (NULL != m_database)
    {
        ConnectionHelper();
    }
#endif // HAVE_MYSQL
    return MUX_S_OK;
}

void CQueryServer::LogError(const UTF8 *pContext)
{
#if defined(HAVE_MYSQL)
    if (  NULL != m_pILog
       && NULL != m_database)
    {
        bool fStarted = false;
        if (  MUX_SUCCEEDED(m_pILog->start_log(&fStarted, LOG_ALWAYS, T("SQL"), T("ERR")))
           && fStarted)
        {
            m_pILog->log_text(pContext);
            m_pILog->log_text(T(": "));
            m_pILog->log_text(reinterpret_cast<const UTF8 *>(mysql_error(m_database)));
            m_pILog->end_log();
        }
    }
#else
    UNUSED_PARAMETER(pContext);
#endif
}

bool CQueryServer::ConnectionHelper()
{
#if defined(HAVE_MYSQL)
    if (m_sServer.empty())
    {
        return false;
    }

#ifdef MYSQL_OPT_RECONNECT
    // As of MySQL 5.0.3, the default is no longer to reconnect.
    //
    my_bool reconnect = 1;
    if (0 != mysql_options(m_database, MYSQL_OPT_RECONNECT, reinterpret_cast<const char *>(&reconnect)))
    {
        LogError(T("mysql_options(MYSQL_OPT_RECONNECT)"));
    }
#endif
    if (0 != mysql_options(m_database, MYSQL_SET_CHARSET_NAME, "utf8"))
    {
        LogError(T("mysql_options(MYSQL_SET_CHARSET_NAME)"));
    }

    if (mysql_real_connect(m_database, m_sServer.c_str(), m_sUser.c_str(),
         m_sPassword.c_str(), m_sDatabase.c_str(), 0, NULL, 0) != 0)
    {
#ifdef MYSQL_OPT_RECONNECT
        // Before MySQL 5.0.19, mysql_real_connect sets the option
        // back to default, so we set it again.
        //
        if (0 != mysql_options(m_database, MYSQL_OPT_RECONNECT, reinterpret_cast<const char *>(&reconnect)))
        {
            LogError(T("mysql_options(MYSQL_OPT_RECONNECT)"));
        }
#endif
        return true;
    }

    // The connection attempt failed.  Surface the reason rather than leaving
    // m_database silently unconnected.
    //
    LogError(T("mysql_real_connect"));
    return false;
#else
    return false;
#endif
}

MUX_RESULT CQueryServer::Advise(mux_IQuerySink *pIQuerySink)
{
    if (NULL != m_pIQuerySink)
    {
        m_pIQuerySink->Release();
        m_pIQuerySink = NULL;
    }

    if (NULL == pIQuerySink)
    {
        return MUX_E_INVALIDARG;
    }

    m_pIQuerySink = pIQuerySink;
    return MUX_S_OK;
}

MUX_RESULT CQueryServer::Query(uint32_t iQueryHandle, const UTF8 *pDatabaseName, const UTF8 *pQuery)
{
    UNUSED_PARAMETER(pDatabaseName);

    if (NULL == m_pIQuerySink)
    {
        return MUX_E_NOTREADY;
    }

    uint32_t iError = QS_SUCCESS;

    QUEUE_INFO qiResultsSet;
    Pipe_InitializeQueueInfo(&qiResultsSet);

#if defined(HAVE_MYSQL)
    if (NULL == m_database)
    {
        iError = QS_NO_SESSION;
    }
    else if (mysql_ping(m_database) != 0)
    {
        // The session is gone.  Attempt our own reconnection and distinguish a
        // failed connection (bad credentials, missing database) from a server
        // that is simply unreachable, so softcode can tell them apart.
        //
        if (!ConnectionHelper())
        {
            iError = QS_CONNECT_FAILED;
        }
        else if (mysql_ping(m_database) != 0)
        {
            iError = QS_SQL_UNAVAILABLE;
        }
    }

    if (  QS_SUCCESS == iError
       && mysql_real_query(m_database, reinterpret_cast<const char *>(pQuery), strlen(reinterpret_cast<const char *>(pQuery))) != 0)
    {
        LogError(T("mysql_real_query"));
        iError = QS_QUERY_ERROR;
    }

    MYSQL_RES *result = NULL;
    MYSQL_ROW  row;

    int nFields = 0;
    if (iError == QS_SUCCESS)
    {
        size_t nRows = 0;
        result = mysql_store_result(m_database);
        if (NULL == result)
        {
            Pipe_AppendBytes(&qiResultsSet, sizeof(nFields), &nFields);
            Pipe_AppendBytes(&qiResultsSet, sizeof(nRows), &nRows);
        }
        else
        {
            nFields = mysql_num_fields(result);
            Pipe_AppendBytes(&qiResultsSet, sizeof(nFields), &nFields);

            row = mysql_fetch_row(result);
            while (row)
            {
                nRows++;

                int loop;
                for (loop = 0; loop < nFields; loop++)
                {
                    const char *p;
                    if (NULL != row[loop])
                    {
                        p = row[loop];
                    }
                    else
                    {
                        p = "";
                    }
                    size_t n = strlen(p)+1;
                    Pipe_AppendBytes(&qiResultsSet, sizeof(n), &n);
                    Pipe_AppendBytes(&qiResultsSet, n, p);
                }
                row = mysql_fetch_row(result);
            }
            mysql_free_result(result);
            Pipe_AppendBytes(&qiResultsSet, sizeof(nRows), &nRows);
        }

        // Drain any remaining result sets from stored procedures.
        // mysql_next_result() returns 0 for another result set, -1 when there
        // are no more, and a positive value on error.
        //
        int iNext;
        while (0 == (iNext = mysql_next_result(m_database)))
        {
            MYSQL_RES *extra = mysql_store_result(m_database);
            if (extra)
            {
                mysql_free_result(extra);
            }
        }
        if (0 < iNext)
        {
            // A later statement (e.g. inside a stored procedure) failed.  Do
            // not let it masquerade as overall success.
            //
            LogError(T("mysql_next_result"));
            iError = QS_QUERY_ERROR;
        }
    }
#else // HAVE_MYSQL
    iError = QS_NO_SESSION;
#endif // HAVE_MYSQL

    MUX_RESULT mr = m_pIQuerySink->Result(iQueryHandle, iError, &qiResultsSet);
    Pipe_EmptyQueue(&qiResultsSet);
    return mr;
}

// Factory for CQueryServer component which is not directly accessible.
//
CQueryServerFactory::CQueryServerFactory(void) : m_cRef(1)
{
}

CQueryServerFactory::~CQueryServerFactory()
{
}

MUX_RESULT CQueryServerFactory::QueryInterface(MUX_IID iid, void **ppv)
{
    if (mux_IID_IUnknown == iid)
    {
        *ppv = static_cast<mux_IClassFactory *>(this);
    }
    else if (mux_IID_IClassFactory == iid)
    {
        *ppv = static_cast<mux_IClassFactory *>(this);
    }
    else
    {
        *ppv = NULL;
        return MUX_E_NOINTERFACE;
    }
    reinterpret_cast<mux_IUnknown *>(*ppv)->AddRef();
    return MUX_S_OK;
}

uint32_t CQueryServerFactory::AddRef(void)
{
    return m_cRef.fetch_add(1, std::memory_order_relaxed) + 1;
}

uint32_t CQueryServerFactory::Release(void)
{
    uint32_t prev = m_cRef.fetch_sub(1, std::memory_order_acq_rel);
    if (1 == prev)
    {
        delete this;
        return 0;
    }
    return prev - 1;
}

MUX_RESULT CQueryServerFactory::CreateInstance(mux_IUnknown *pUnknownOuter, MUX_IID iid, void **ppv)
{
    // Disallow attempts to aggregate this component.
    //
    if (NULL != pUnknownOuter)
    {
        return MUX_E_NOAGGREGATION;
    }

    CQueryServer *pQueryServer = NULL;
    try
    {
        pQueryServer = new CQueryServer;
    }
    catch (const std::bad_alloc &)
    {
        ; // Handled by the NULL check below.
    }

    MUX_RESULT mr;
    if (NULL == pQueryServer)
    {
        return MUX_E_OUTOFMEMORY;
    }
    else
    {
        mr = pQueryServer->FinalConstruct();
        if (MUX_FAILED(mr))
        {
            pQueryServer->Release();
            return mr;
        }
    }

    mr = pQueryServer->QueryInterface(iid, ppv);
    pQueryServer->Release();
    return mr;
}

MUX_RESULT CQueryServerFactory::LockServer(bool bLock)
{
    if (bLock)
    {
        g_cServerLocks++;
    }
    else
    {
        g_cServerLocks--;
    }
    return MUX_S_OK;
}