7 files changed, 151 insertions, 70 deletions

diff --git a/src/server/shared/Database/DatabaseWorkerPool.h b/src/server/shared/Database/DatabaseWorkerPool.h
index 5a5ced4a504..e0735a83fba 100644
--- a/src/server/shared/Database/DatabaseWorkerPool.h
+++ b/src/server/shared/Database/DatabaseWorkerPool.h
@@ -199,7 +199,12 @@ class DatabaseWorkerPool
 
         QueryResult Query(const char* sql)
         {
-            return GetConnection()->Query(sql);
+            ResultSet* result = GetConnection()->Query(sql);
+            if (!result || !result->GetRowCount())
+                return QueryResult(NULL);
+
+            result->NextRow();
+            return QueryResult(result);
         }
 
         QueryResult PQuery(const char* sql, ...)
@@ -291,7 +296,7 @@ class DatabaseWorkerPool
         PreparedQueryResult Query(PreparedStatement* stmt)
         {
             PreparedResultSet* ret = GetConnection()->Query(stmt);
-            if (!ret || !ret->num_rows)
+            if (!ret || !ret->GetRowCount())
                 return PreparedQueryResult(NULL);
 
             ret->NextRow();
diff --git a/src/server/shared/Database/MySQLConnection.cpp b/src/server/shared/Database/MySQLConnection.cpp
index e8f5bb08699..ef78d0afdb0 100644
--- a/src/server/shared/Database/MySQLConnection.cpp
+++ b/src/server/shared/Database/MySQLConnection.cpp
@@ -233,10 +233,10 @@ bool MySQLConnection::Execute(PreparedStatement* stmt)
     }
 }
 
-QueryResult MySQLConnection::Query(const char* sql)
+ResultSet* MySQLConnection::Query(const char* sql)
 {
     if (!sql)
-        return QueryResult(NULL);
+        return NULL;
 
     MYSQL_RES *result = NULL;
     MYSQL_FIELD *fields = NULL;
@@ -244,13 +244,9 @@ QueryResult MySQLConnection::Query(const char* sql)
     uint32 fieldCount = 0;
 
     if (!_Query(sql, &result, &fields, &rowCount, &fieldCount))
-        return QueryResult(NULL);
+        return NULL;
 
-    ResultSet *queryResult = new ResultSet(result, fields, rowCount, fieldCount);
-
-    queryResult->NextRow();
-
-    return QueryResult(queryResult);
+    return new ResultSet(result, fields, rowCount, fieldCount);
 }
 
 bool MySQLConnection::_Query(const char *sql, MYSQL_RES **pResult, MYSQL_FIELD **pFields, uint64* pRowCount, uint32* pFieldCount)
diff --git a/src/server/shared/Database/MySQLConnection.h b/src/server/shared/Database/MySQLConnection.h
index f8f8f79eb44..b707f8a4675 100644
--- a/src/server/shared/Database/MySQLConnection.h
+++ b/src/server/shared/Database/MySQLConnection.h
@@ -40,7 +40,7 @@ class MySQLConnection
     public:
         bool Execute(const char* sql);
         bool Execute(PreparedStatement* stmt);
-        QueryResult Query(const char* sql);
+        ResultSet* Query(const char* sql);
         PreparedResultSet* Query(PreparedStatement* stmt);
         bool _Query(const char *sql, MYSQL_RES **pResult, MYSQL_FIELD **pFields, uint64* pRowCount, uint32* pFieldCount);
 
diff --git a/src/server/shared/Database/SQLOperation.cpp b/src/server/shared/Database/SQLOperation.cpp
index 4abad405599..7ae690d2053 100644
--- a/src/server/shared/Database/SQLOperation.cpp
+++ b/src/server/shared/Database/SQLOperation.cpp
@@ -44,10 +44,14 @@ bool BasicStatementTask::Execute()
 {
     if (m_has_result)
     {
-        m_result.set(
-            m_conn->Query(m_sql)
-            );
-
+        ResultSet* result = m_conn->Query(m_sql);
+        if (!result || !result->GetRowCount())
+        {
+            m_result.set(QueryResult(NULL));
+            return false;
+        }
+        result->NextRow();
+        m_result.set(QueryResult(result));
         return true;
     }
 
@@ -62,15 +66,15 @@ bool SQLQueryHolder::SetQuery(size_t index, const char *sql)
         return false;
     }
 
-    if (m_queries[index].first != NULL)
-    {
-        sLog.outError("Attempt assign query to holder index (%zu) where other query stored (Old: [%s] New: [%s])",
-            index, m_queries[index].first, sql);
-        return false;
-    }
-
     /// not executed yet, just stored (it's not called a holder for nothing)
-    m_queries[index] = SQLResultPair(strdup(sql), QueryResult(NULL));
+    SQLElementData element;
+    element.type = SQL_ELEMENT_RAW;
+    element.element.query = strdup(sql);
+    
+    SQLResultSetUnion result;
+    result.qresult = NULL;
+
+    m_queries[index] = SQLResultPair(element, result);
     return true;
 }
 
@@ -99,26 +103,62 @@ bool SQLQueryHolder::SetPQuery(size_t index, const char *format, ...)
 
 QueryResult SQLQueryHolder::GetResult(size_t index)
 {
+    // Don't call to this function if the index is of a prepared statement
     if (index < m_queries.size())
     {
         /// the query strings are freed on the first GetResult or in the destructor
-        if (m_queries[index].first != NULL)
+        if (SQLElementData* data = &m_queries[index].first)
         {
-            free((void*)(const_cast<char*>(m_queries[index].first)));
-            m_queries[index].first = NULL;
+            free((void*)(const_cast<char*>(data->element.query)));
+            data->element.query = NULL;
         }
-        /// when you get a result aways remember to delete it!
-        return m_queries[index].second;
+
+        ResultSet* result = m_queries[index].second.qresult;
+        if (!result || !result->GetRowCount())
+            return QueryResult(NULL);
+
+        result->NextRow();
+        return QueryResult(result);
     }
     else
         return QueryResult(NULL);
 }
 
-void SQLQueryHolder::SetResult(size_t index, QueryResult result)
+PreparedQueryResult SQLQueryHolder::GetPreparedResult(size_t index)
+{
+    // Don't call to this function if the index is of a prepared statement
+    if (index < m_queries.size())
+    {
+        /// the query strings are freed on the first GetResult or in the destructor
+        if (SQLElementData* data = &m_queries[index].first)
+        {
+            delete data->element.stmt;
+            data->element.stmt = NULL;
+        }
+
+        PreparedResultSet* result = m_queries[index].second.presult;
+        if (!result || !result->GetRowCount())
+            return PreparedQueryResult(NULL);
+
+        result->NextRow();
+        return PreparedQueryResult(result);
+    }
+    else
+        return PreparedQueryResult(NULL);
+}
+
+void SQLQueryHolder::SetResult(size_t index, ResultSet* result)
 {
     /// store the result in the holder
     if (index < m_queries.size())
-        m_queries[index].second = result;
+        m_queries[index].second.qresult = result;
+}
+
+void SQLQueryHolder::SetPreparedResult(size_t index, PreparedResultSet* result)
+{
+    /// store the result in the holder
+    if (index < m_queries.size())
+        m_queries[index].second.presult = result;
 }
 
 SQLQueryHolder::~SQLQueryHolder()
@@ -127,8 +167,18 @@ SQLQueryHolder::~SQLQueryHolder()
     {
         /// if the result was never used, free the resources
         /// results used already (getresult called) are expected to be deleted
-        if (m_queries[i].first != NULL)
-            free((void*)(const_cast<char*>(m_queries[i].first)));
+        if (SQLElementData* data = &m_queries[i].first)
+        {
+            switch (data->type)
+            {
+                case SQL_ELEMENT_RAW:
+                    free((void*)(const_cast<char*>(data->element.query)));
+                    break;
+                case SQL_ELEMENT_PREPARED:
+                    delete data->element.stmt;
+                    break;
+            }
+        }
     }
 }
 
@@ -149,9 +199,26 @@ bool SQLQueryHolderTask::Execute()
     for (size_t i = 0; i < queries.size(); i++)
     {
         /// execute all queries in the holder and pass the results
-        char const *sql = queries[i].first;
-        if (sql)
-            m_holder->SetResult(i, m_conn->Query(sql));
+        if (SQLElementData* data = &queries[i].first)
+        {
+            switch (data->type)
+            {
+                case SQL_ELEMENT_RAW:
+                {
+                    char const *sql = data->element.query;
+                    if (sql)
+                        m_holder->SetResult(i, m_conn->Query(sql));
+                    break;
+                }
+                case SQL_ELEMENT_PREPARED:
+                {
+                    PreparedStatement* stmt = data->element.stmt;
+                    if (stmt)
+                        m_holder->SetPreparedResult(i, m_conn->Query(stmt));
+                    break;
+                }
+            }
+        }
     }
 
     m_result.set(m_holder);
diff --git a/src/server/shared/Database/SQLOperation.h b/src/server/shared/Database/SQLOperation.h
index 976e458a4f2..e5d7f7c7180 100644
--- a/src/server/shared/Database/SQLOperation.h
+++ b/src/server/shared/Database/SQLOperation.h
@@ -25,6 +25,37 @@
 #include "Common.h"
 #include "Callback.h"
 
+//- Forward declare (don't include header to prevent circular includes)
+class PreparedStatement;
+
+//- Union that holds element data
+union SQLElementUnion
+{
+    PreparedStatement* stmt;
+    const char* query;
+};
+
+//- Type specifier of our element data
+enum SQLElementDataType
+{
+    SQL_ELEMENT_RAW,
+    SQL_ELEMENT_PREPARED,
+};
+
+//- The element
+struct SQLElementData
+{
+    SQLElementUnion element;
+    SQLElementDataType type;
+};
+
+//- For ambigious resultsets
+union SQLResultSetUnion
+{
+    PreparedResultSet* presult;
+    ResultSet* qresult;
+};
+
 class MySQLConnection;
 
 class SQLOperation : public ACE_Method_Request
@@ -65,16 +96,19 @@ class SQLQueryHolder
 {
     friend class SQLQueryHolderTask;
     private:
-        typedef std::pair<const char*, QueryResult> SQLResultPair;
+        typedef std::pair<SQLElementData, SQLResultSetUnion> SQLResultPair;
         std::vector<SQLResultPair> m_queries;
     public:
         SQLQueryHolder() {}
         ~SQLQueryHolder();
         bool SetQuery(size_t index, const char *sql);
         bool SetPQuery(size_t index, const char *format, ...) ATTR_PRINTF(3,4);
+        bool SetPreparedQuery(size_t index, PreparedStatement* stmt);
         void SetSize(size_t size);
         QueryResult GetResult(size_t index);
-        void SetResult(size_t index, QueryResult result);
+        PreparedQueryResult GetPreparedResult(size_t index);
+        void SetResult(size_t index, ResultSet* result);
+        void SetPreparedResult(size_t index, PreparedResultSet* result);
 };
 
 typedef ACE_Future<SQLQueryHolder*> QueryResultHolderFuture;
diff --git a/src/server/shared/Database/Transaction.cpp b/src/server/shared/Database/Transaction.cpp
index 693bc1f96b2..0a7fe4dab60 100644
--- a/src/server/shared/Database/Transaction.cpp
+++ b/src/server/shared/Database/Transaction.cpp
@@ -22,8 +22,8 @@
 //- Append a raw ad-hoc query to the transaction
 void Transaction::Append(const char* sql)
 {
-    TransactionElementData data;
-    data.type = TRANSACTION_ELEMENT_RAW;
+    SQLElementData data;
+    data.type = SQL_ELEMENT_RAW;
     data.element.query = strdup(sql);
     m_queries.push(data);
 }
@@ -42,8 +42,8 @@ void Transaction::PAppend(const char* sql, ...)
 //- Append a prepared statement to the transaction
 void Transaction::Append(PreparedStatement* stmt)
 {
-    TransactionElementData data;
-    data.type = TRANSACTION_ELEMENT_PREPARED;
+    SQLElementData data;
+    data.type = SQL_ELEMENT_PREPARED;
     data.element.stmt = stmt;
     m_queries.push(data);
 }
@@ -52,13 +52,13 @@ void Transaction::Cleanup()
 {
     while (!m_queries.empty())
     {
-        TransactionElementData data = m_queries.front();
+        SQLElementData data = m_queries.front();
         switch (data.type)
         {
-            case TRANSACTION_ELEMENT_PREPARED:
+            case SQL_ELEMENT_PREPARED:
                 delete data.element.stmt;
             break;
-            case TRANSACTION_ELEMENT_RAW:
+            case SQL_ELEMENT_RAW:
                 free((void*)(data.element.query));
             break;
         }
@@ -68,17 +68,17 @@ void Transaction::Cleanup()
 
 bool TransactionTask::Execute()
 {
-    std::queue<TransactionElementData> &queries = m_trans->m_queries;
+    std::queue<SQLElementData> &queries = m_trans->m_queries;
     if (queries.empty())
         return false;
 
     m_conn->BeginTransaction();
     while (!queries.empty())
     {
-        TransactionElementData data = queries.front();
+        SQLElementData data = queries.front();
         switch (data.type)
         {
-            case TRANSACTION_ELEMENT_PREPARED:
+            case SQL_ELEMENT_PREPARED:
             {
                 PreparedStatement* stmt = data.element.stmt;
                 ASSERT(stmt);
@@ -91,7 +91,7 @@ bool TransactionTask::Execute()
                 delete data.element.stmt;
             }
             break;
-            case TRANSACTION_ELEMENT_RAW:
+            case SQL_ELEMENT_RAW:
             {
                 const char* sql = data.element.query;
                 ASSERT(sql);
diff --git a/src/server/shared/Database/Transaction.h b/src/server/shared/Database/Transaction.h
index 8d824c297d4..fa6cfeba747 100644
--- a/src/server/shared/Database/Transaction.h
+++ b/src/server/shared/Database/Transaction.h
@@ -24,27 +24,6 @@
 //- Forward declare (don't include header to prevent circular includes)
 class PreparedStatement;
 
-//- Union that holds element data
-union TransactionElementUnion
-{
-    PreparedStatement* stmt;
-    const char* query;
-};
-
-//- Type specifier of our element data
-enum TransactionElementDataType
-{
-    TRANSACTION_ELEMENT_RAW,
-    TRANSACTION_ELEMENT_PREPARED,
-};
-
-//- The transaction element
-struct TransactionElementData
-{
-    TransactionElementUnion element;
-    TransactionElementDataType type;
-};
-
 /*! Transactions, high level class. */
 class Transaction
 {
@@ -60,7 +39,7 @@ class Transaction
 
     protected:
         void Cleanup();
-        std::queue<TransactionElementData> m_queries;
+        std::queue<SQLElementData> m_queries;
 
     private:
         bool m_actioned;