1.1 --- a/core/smn_database.cpp	Sun May 22 02:04:34 2011 -0700
     1.2 +++ b/core/smn_database.cpp	Sat May 14 20:21:37 2011 -0700
     1.3 @@ -36,9 +36,15 @@
     1.4  #include "PluginSys.h"
     1.5  #include "sm_stringutil.h"
     1.6  
     1.7 -HandleType_t hQueryType;
     1.8  HandleType_t hStmtType;
     1.9  
    1.10 +HandleType_t hCombinedQueryType;
    1.11 +typedef struct
    1.12 +{
    1.13 +	IQuery *query;
    1.14 +	IDatabase *db;
    1.15 +} CombinedQuery;
    1.16 +
    1.17  class DatabaseHelpers : 
    1.18  	public SMGlobalClass,
    1.19  	public IHandleTypeDispatch
    1.20 @@ -57,22 +63,23 @@
    1.21  		g_HandleSys.InitAccessDefaults(&tacc, NULL);
    1.22  		tacc.ident = g_pCoreIdent;
    1.23  
    1.24 -		hQueryType = g_HandleSys.CreateType("IQuery", this, 0, &tacc, &acc, g_pCoreIdent, NULL);
    1.25 -		hStmtType = g_HandleSys.CreateType("IPreparedQuery", this, hQueryType, &tacc, &acc, g_pCoreIdent, NULL);
    1.26 +		hCombinedQueryType = g_HandleSys.CreateType("IQuery", this, 0, &tacc, &acc, g_pCoreIdent, NULL);
    1.27 +		hStmtType = g_HandleSys.CreateType("IPreparedQuery", this, 0, &tacc, &acc, g_pCoreIdent, NULL);
    1.28  	}
    1.29  
    1.30  	virtual void OnSourceModShutdown()
    1.31  	{
    1.32  		g_HandleSys.RemoveType(hStmtType, g_pCoreIdent);
    1.33 -		g_HandleSys.RemoveType(hQueryType, g_pCoreIdent);
    1.34 +		g_HandleSys.RemoveType(hCombinedQueryType, g_pCoreIdent);
    1.35  	}
    1.36  
    1.37  	virtual void OnHandleDestroy(HandleType_t type, void *object)
    1.38  	{
    1.39 -		if (type == hQueryType)
    1.40 +		if (type == hCombinedQueryType)
    1.41  		{
    1.42 -			IQuery *query = (IQuery *)object;
    1.43 -			query->Destroy();
    1.44 +			CombinedQuery *combined = (CombinedQuery *)object;
    1.45 +			combined->query->Destroy();
    1.46 +			delete combined;
    1.47  		} else if (type == hStmtType) {
    1.48  			IPreparedQuery *query = (IPreparedQuery *)object;
    1.49  			query->Destroy();
    1.50 @@ -84,10 +91,37 @@
    1.51  inline HandleError ReadQueryHndl(Handle_t hndl, IPluginContext *pContext, IQuery **query)
    1.52  {
    1.53  	HandleSecurity sec;
    1.54 +	CombinedQuery *c;
    1.55  	sec.pOwner = pContext->GetIdentity();
    1.56  	sec.pIdentity = g_pCoreIdent;
    1.57  
    1.58 -	return g_HandleSys.ReadHandle(hndl, hQueryType, &sec, (void **)query);
    1.59 +	HandleError ret;
    1.60 +	
    1.61 +	if ((ret = g_HandleSys.ReadHandle(hndl, hStmtType, &sec, (void **)query)) != HandleError_None)
    1.62 +	{
    1.63 +		ret = g_HandleSys.ReadHandle(hndl, hCombinedQueryType, &sec, (void **)&c);
    1.64 +		if (ret == HandleError_None)
    1.65 +		{
    1.66 +			*query = c->query;
    1.67 +		}
    1.68 +	}
    1.69 +	return ret;
    1.70 +}
    1.71 +
    1.72 +inline HandleError ReadQueryAndDbHndl(Handle_t hndl, IPluginContext *pContext, IQuery **query, IDatabase **db)
    1.73 +{
    1.74 +	HandleSecurity sec;
    1.75 +	CombinedQuery *c;
    1.76 +	sec.pOwner = pContext->GetIdentity();
    1.77 +	sec.pIdentity = g_pCoreIdent;
    1.78 +
    1.79 +	HandleError ret = g_HandleSys.ReadHandle(hndl, hCombinedQueryType, &sec, (void **)&c);
    1.80 +	if (ret == HandleError_None)
    1.81 +	{
    1.82 +		*query = c->query;
    1.83 +		*db = c->db;
    1.84 +	}
    1.85 +	return ret;
    1.86  }
    1.87  
    1.88  inline HandleError ReadStmtHndl(Handle_t hndl, IPluginContext *pContext, IPreparedQuery **query)
    1.89 @@ -192,12 +226,17 @@
    1.90  		
    1.91  		if (m_pQuery)
    1.92  		{
    1.93 -			qh = g_HandleSys.CreateHandle(hQueryType, m_pQuery, me->GetIdentity(), g_pCoreIdent, NULL);
    1.94 +			CombinedQuery *c = new CombinedQuery;
    1.95 +			c->query = m_pQuery;
    1.96 +			c->db = m_pDatabase;
    1.97 +			
    1.98 +			qh = g_HandleSys.CreateHandle(hCombinedQueryType, c, me->GetIdentity(), g_pCoreIdent, NULL);
    1.99  			if (qh != BAD_HANDLE)
   1.100  			{
   1.101  				m_pQuery = NULL;
   1.102  			} else {
   1.103  				UTIL_Format(error, sizeof(error), "Could not alloc handle");
   1.104 +				delete c;
   1.105  			}
   1.106  		}
   1.107  
   1.108 @@ -537,42 +576,59 @@
   1.109  {
   1.110  	IDatabase *db = NULL;
   1.111  	IPreparedQuery *stmt = NULL;
   1.112 +	IQuery *query = NULL;
   1.113  	HandleError err;
   1.114  
   1.115 -	if ((err = ReadDbOrStmtHndl(params[1], pContext, &db, &stmt)) != HandleError_None)
   1.116 +	if (((err = ReadDbOrStmtHndl(params[1], pContext, &db, &stmt)) != HandleError_None)
   1.117 +		&& ((err = ReadQueryAndDbHndl(params[1], pContext, &query, &db)) != HandleError_None))
   1.118  	{
   1.119 -		return pContext->ThrowNativeError("Invalid statement or db Handle %x (error: %d)", params[1], err);
   1.120 +		return pContext->ThrowNativeError("Invalid statement, db, or query Handle %x (error: %d)", params[1], err);
   1.121  	}
   1.122  
   1.123 -	if (db)
   1.124 +
   1.125 +	if (stmt)
   1.126 +	{
   1.127 +		return stmt->GetAffectedRows();
   1.128 +	}
   1.129 +	else if (query)
   1.130 +	{
   1.131 +		return db->GetAffectedRowsForQuery(query);
   1.132 +	}
   1.133 +	else if (db)
   1.134  	{
   1.135  		return db->GetAffectedRows();
   1.136 -	} else if (stmt) {
   1.137 -		return stmt->GetAffectedRows();
   1.138  	}
   1.139  
   1.140 -	return pContext->ThrowNativeError("Unknown error reading db/stmt handles");
   1.141 +	return pContext->ThrowNativeError("Unknown error reading db/stmt/query handles");
   1.142  }
   1.143  
   1.144  static cell_t SQL_GetInsertId(IPluginContext *pContext, const cell_t *params)
   1.145  {
   1.146  	IDatabase *db = NULL;
   1.147 +	IQuery *query = NULL;
   1.148  	IPreparedQuery *stmt = NULL;
   1.149  	HandleError err;
   1.150  
   1.151 -	if ((err = ReadDbOrStmtHndl(params[1], pContext, &db, &stmt)) != HandleError_None)
   1.152 +	if (((err = ReadDbOrStmtHndl(params[1], pContext, &db, &stmt)) != HandleError_None)
   1.153 +		&& ((err = ReadQueryAndDbHndl(params[1], pContext, &query, &db)) != HandleError_None))
   1.154  	{
   1.155 -		return pContext->ThrowNativeError("Invalid statement or db Handle %x (error: %d)", params[1], err);
   1.156 +		return pContext->ThrowNativeError("Invalid statement, db, or query Handle %x (error: %d)", params[1], err);
   1.157  	}
   1.158  
   1.159 -	if (db)
   1.160 +	if (query)
   1.161 +	{
   1.162 +		return db->GetInsertIDForQuery(query);
   1.163 +	}
   1.164 +	else if (db)
   1.165  	{
   1.166  		return db->GetInsertID();
   1.167 -	} else if (stmt) {
   1.168 +	}
   1.169 +	else if (stmt)
   1.170 +	{
   1.171  		return stmt->GetInsertID();
   1.172  	}
   1.173  
   1.174 -	return pContext->ThrowNativeError("Unknown error reading db/stmt handles");
   1.175 +	return pContext->ThrowNativeError("Unknown error reading db/stmt/query handles");
   1.176  }
   1.177  
   1.178  static cell_t SQL_GetError(IPluginContext *pContext, const cell_t *params)
   1.179 @@ -682,10 +738,14 @@
   1.180  		return BAD_HANDLE;
   1.181  	}
   1.182  
   1.183 -	Handle_t hndl = g_HandleSys.CreateHandle(hQueryType, qr, pContext->GetIdentity(), g_pCoreIdent, NULL);
   1.184 +	CombinedQuery *c = new CombinedQuery;
   1.185 +	c->query = qr;
   1.186 +	c->db = db;
   1.187 +	Handle_t hndl = g_HandleSys.CreateHandle(hCombinedQueryType, c, pContext->GetIdentity(), g_pCoreIdent, NULL);
   1.188  	if (hndl == BAD_HANDLE)
   1.189  	{
   1.190  		qr->Destroy();
   1.191 +		delete c;
   1.192  		return BAD_HANDLE;
   1.193  	}
   1.194  

     2.1 --- a/extensions/mysql/mysql/MyBasicResults.cpp	Sun May 22 02:04:34 2011 -0700
     2.2 +++ b/extensions/mysql/mysql/MyBasicResults.cpp	Sat May 14 20:21:37 2011 -0700
     2.3 @@ -307,6 +307,8 @@
     2.4  : m_pParent(db), m_rs(res)
     2.5  {
     2.6  	m_pParent->IncReferenceCount();
     2.7 +	m_InsertID = m_pParent->GetInsertID();
     2.8 +	m_AffectedRows = m_pParent->GetAffectedRows();
     2.9  }
    2.10  
    2.11  IResultSet *MyQuery::GetResultSet()
    2.12 @@ -319,6 +321,16 @@
    2.13  	return &m_rs;
    2.14  }
    2.15  
    2.16 +unsigned int MyQuery::GetInsertID()
    2.17 +{
    2.18 +	return m_InsertID;
    2.19 +}
    2.20 +
    2.21 +unsigned int MyQuery::GetAffectedRows()
    2.22 +{
    2.23 +	return m_AffectedRows;
    2.24 +}
    2.25 +
    2.26  bool MyQuery::FetchMoreResults()
    2.27  {
    2.28  	if (m_rs.m_pRes == NULL)

     3.1 --- a/extensions/mysql/mysql/MyBasicResults.h	Sun May 22 02:04:34 2011 -0700
     3.2 +++ b/extensions/mysql/mysql/MyBasicResults.h	Sat May 14 20:21:37 2011 -0700
     3.3 @@ -87,9 +87,14 @@
     3.4  	IResultSet *GetResultSet();
     3.5  	bool FetchMoreResults();
     3.6  	void Destroy();
     3.7 +public: // Used by the driver to implement GetInsertIDForQuery()/GetAffectedRowsForQuery()
     3.8 +	unsigned int GetInsertID();
     3.9 +	unsigned int GetAffectedRows();
    3.10  private:
    3.11  	MyDatabase *m_pParent;
    3.12  	MyBasicResults m_rs;
    3.13 +	unsigned int m_InsertID;
    3.14 +	unsigned int m_AffectedRows;
    3.15  };
    3.16  
    3.17  #endif //_INCLUDE_SM_MYSQL_BASIC_RESULTS_H_

     4.1 --- a/extensions/mysql/mysql/MyDatabase.cpp	Sun May 22 02:04:34 2011 -0700
     4.2 +++ b/extensions/mysql/mysql/MyDatabase.cpp	Sat May 14 20:21:37 2011 -0700
     4.3 @@ -257,6 +257,16 @@
     4.4  	return new MyQuery(this, res);
     4.5  }
     4.6  
     4.7 +unsigned int MyDatabase::GetAffectedRowsForQuery(IQuery *query)
     4.8 +{
     4.9 +	return static_cast<MyQuery*>(query)->GetAffectedRows();
    4.10 +}
    4.11 +
    4.12 +unsigned int MyDatabase::GetInsertIDForQuery(IQuery *query)
    4.13 +{
    4.14 +	return static_cast<MyQuery*>(query)->GetInsertID();
    4.15 +}
    4.16 +
    4.17  IPreparedQuery *MyDatabase::PrepareQuery(const char *query, char *error, size_t maxlength, int *errCode)
    4.18  {
    4.19  	MYSQL_STMT *stmt = mysql_stmt_init(m_mysql);

     5.1 --- a/extensions/mysql/mysql/MyDatabase.h	Sun May 22 02:04:34 2011 -0700
     5.2 +++ b/extensions/mysql/mysql/MyDatabase.h	Sat May 14 20:21:37 2011 -0700
     5.3 @@ -60,6 +60,8 @@
     5.4  	IDBDriver *GetDriver();
     5.5  	bool DoSimpleQueryEx(const char *query, size_t len);
     5.6  	IQuery *DoQueryEx(const char *query, size_t len);
     5.7 +	unsigned int GetAffectedRowsForQuery(IQuery *query);
     5.8 +	unsigned int GetInsertIDForQuery(IQuery *query);
     5.9  public:
    5.10  	const DatabaseInfo &GetInfo();
    5.11  private:

     6.1 --- a/extensions/sqlite/driver/SqDatabase.cpp	Sun May 22 02:04:34 2011 -0700
     6.2 +++ b/extensions/sqlite/driver/SqDatabase.cpp	Sat May 14 20:21:37 2011 -0700
     6.3 @@ -186,6 +186,15 @@
     6.4  	return pQuery;
     6.5  }
     6.6  
     6.7 +unsigned int SqDatabase::GetAffectedRowsForQuery(IQuery *query)
     6.8 +{
     6.9 +	return static_cast<SqQuery*>(query)->GetAffectedRows();
    6.10 +}
    6.11 +unsigned int SqDatabase::GetInsertIDForQuery(IQuery *query)
    6.12 +{
    6.13 +	return static_cast<SqQuery*>(query)->GetInsertID();
    6.14 +}
    6.15 +
    6.16  IPreparedQuery *SqDatabase::PrepareQuery(const char *query, 
    6.17  										 char *error, 
    6.18  										 size_t maxlength, 

     7.1 --- a/extensions/sqlite/driver/SqDatabase.h	Sun May 22 02:04:34 2011 -0700
     7.2 +++ b/extensions/sqlite/driver/SqDatabase.h	Sat May 14 20:21:37 2011 -0700
     7.3 @@ -56,6 +56,8 @@
     7.4  	IDBDriver *GetDriver();
     7.5  	bool DoSimpleQueryEx(const char *query, size_t len);
     7.6  	IQuery *DoQueryEx(const char *query, size_t len);
     7.7 +	unsigned int GetAffectedRowsForQuery(IQuery *query);
     7.8 +	unsigned int GetInsertIDForQuery(IQuery *query);
     7.9  public:
    7.10  	sqlite3 *GetDb();
    7.11  private:

     8.1 --- a/plugins/include/dbi.inc	Sun May 22 02:04:34 2011 -0700
     8.2 +++ b/plugins/include/dbi.inc	Sat May 14 20:21:37 2011 -0700
     8.3 @@ -261,20 +261,20 @@
     8.4  /**
     8.5   * Returns the last query's insertion id.
     8.6   *
     8.7 - * @param hndl			A database OR statement Handle.
     8.8 + * @param hndl			A database, query, OR statement Handle.
     8.9   * @return				Last query's insertion id.
    8.10 - * @error				Invalid database or statement Handle.
    8.11 + * @error				Invalid database, query, or statement Handle.
    8.12   */
    8.13  native SQL_GetInsertId(Handle:hndl);
    8.14  
    8.15  /**
    8.16   * Returns the error reported by the last query.
    8.17   *
    8.18 - * @param hndl			A database OR statement Handle.
    8.19 + * @param hndl			A database, query, OR statement Handle.
    8.20   * @param error			Error buffer.
    8.21   * @param maxlength		Maximum length of the buffer.
    8.22   * @return				True if there was an error, false otherwise.
    8.23 - * @error				Invalid database or statement Handle.
    8.24 + * @error				Invalid database, query, or statement Handle.
    8.25   */
    8.26  native bool:SQL_GetError(Handle:hndl, String:error[], maxlength);
    8.27  

     9.1 --- a/public/IDBDriver.h	Sun May 22 02:04:34 2011 -0700
     9.2 +++ b/public/IDBDriver.h	Sat May 14 20:21:37 2011 -0700
     9.3 @@ -42,7 +42,7 @@
     9.4   */
     9.5  
     9.6  #define SMINTERFACE_DBI_NAME		"IDBI"
     9.7 -#define SMINTERFACE_DBI_VERSION		7
     9.8 +#define SMINTERFACE_DBI_VERSION		8
     9.9  
    9.10  namespace SourceMod
    9.11  {
    9.12 @@ -573,6 +573,33 @@
    9.13  		 * @return				IQuery pointer on success, NULL otherwise.
    9.14  		 */
    9.15  		virtual IQuery *DoQueryEx(const char *query, size_t len) =0;
    9.16 +		
    9.17 +		/**
    9.18 +		 * @brief Retrieves the number of affected rows from the last execute of
    9.19 +		 * the given query
    9.20 +		 *
    9.21 +		 * Note: This can only accept queries from this driver.
    9.22 +		 *
    9.23 +		 * This function is not thread safe and must be included in any locks.
    9.24 +		 *
    9.25 +		 * @param query			IQuery object from this driver
    9.26 +		 * @return				Rows affected from last execution of this query,
    9.27 +		 *						if applicable.
    9.28 +		 */
    9.29 +		virtual unsigned int GetAffectedRowsForQuery(IQuery *query) =0;
    9.30 +		
    9.31 +		/**
    9.32 +		 * @brief Retrieves the last insert id of the given query
    9.33 +		 *
    9.34 +		 * Note: This can only accept queries from this driver.
    9.35 +		 *
    9.36 +		 * This function is not thread safe and must be included in any locks.
    9.37 +		 *
    9.38 +		 * @param query			IQuery object from this driver
    9.39 +		 * @return				Insert Id from the last execution of this query,
    9.40 +		 *						if applicable.
    9.41 +		 */
    9.42 +		virtual unsigned int GetInsertIDForQuery(IQuery *query) =0;
    9.43  	};
    9.44  
    9.45  	/**