SQL: Add HasColumn() to check if a result contains a named column.

4 files changed, 62 insertions, 3 deletions

diff --git a/include/modules/sql.h b/include/modules/sql.h
index 14cd60a56..59955107d 100644
--- a/include/modules/sql.h
+++ b/include/modules/sql.h
@@ -122,6 +122,15 @@ class SQL::Result : public classbase
 	 * @param result A reference to the vector to store column names in.
 	 */
 	virtual void GetCols(std::vector<std::string>& result) = 0;
+
+	/**
+	 * Check if there's a column with the specified name in the result
+	 *
+	 * @param the column name
+	 * @param on success, this is the column index
+	 * @returns true, or false if the column is not found
+	 */
+	virtual bool HasColumn(const std::string& column, size_t& index) = 0;
 };
 
 /** SQL::Error holds the error state of a request.
diff --git a/src/modules/extra/m_mysql.cpp b/src/modules/extra/m_mysql.cpp
index 4ee6de527..a177951ce 100644
--- a/src/modules/extra/m_mysql.cpp
+++ b/src/modules/extra/m_mysql.cpp
@@ -203,6 +203,19 @@ class MySQLresult : public SQL::Result
 		result.assign(colnames.begin(), colnames.end());
 	}
 
+	bool HasColumn(const std::string& column, size_t& index)
+	{
+		for (size_t i = 0; i < colnames.size(); ++i)
+		{
+			if (colnames[i] == column)
+			{
+				index = i;
+				return true;
+			}
+		}
+		return false;
+	}
+
 	SQL::Field GetValue(int row, int column)
 	{
 		if ((row >= 0) && (row < rows) && (column >= 0) && (column < (int)fieldlists[row].size()))
diff --git a/src/modules/extra/m_pgsql.cpp b/src/modules/extra/m_pgsql.cpp
index 25ce6c7f1..8beea1265 100644
--- a/src/modules/extra/m_pgsql.cpp
+++ b/src/modules/extra/m_pgsql.cpp
@@ -87,6 +87,16 @@ class PgSQLresult : public SQL::Result
 	PGresult* res;
 	int currentrow;
 	int rows;
+	std::vector<std::string> colnames;
+
+	void getColNames()
+	{
+		colnames.resize(PQnfields(res));
+		for(unsigned int i=0; i < colnames.size(); i++)
+		{
+			colnames[i] = PQfname(res, i);
+		}
+	}
  public:
 	PgSQLresult(PGresult* result) : res(result), currentrow(0)
 	{
@@ -107,11 +117,25 @@ class PgSQLresult : public SQL::Result
 
 	void GetCols(std::vector<std::string>& result) CXX11_OVERRIDE
 	{
-		result.resize(PQnfields(res));
-		for(unsigned int i=0; i < result.size(); i++)
+		if (colnames.empty())
+			getColNames();
+		result = colnames;
+	}
+
+	bool HasColumn(const std::string& column, size_t& index)
+	{
+		if (colnames.empty())
+			getColNames();
+
+		for (size_t i = 0; i < colnames.size(); ++i)
 		{
-			result[i] = PQfname(res, i);
+			if (colnames[i] == column)
+			{
+				index = i;
+				return true;
+			}
 		}
+		return false;
 	}
 
 	SQL::Field GetValue(int row, int column)
diff --git a/src/modules/extra/m_sqlite3.cpp b/src/modules/extra/m_sqlite3.cpp
index 0f596a0f7..31821045d 100644
--- a/src/modules/extra/m_sqlite3.cpp
+++ b/src/modules/extra/m_sqlite3.cpp
@@ -82,6 +82,19 @@ class SQLite3Result : public SQL::Result
 	{
 		result.assign(columns.begin(), columns.end());
 	}
+
+	bool HasColumn(const std::string& column, size_t& index)
+	{
+		for (size_t i = 0; i < columns.size(); ++i)
+		{
+			if (columns[i] == column)
+			{
+				index = i;
+				return true;
+			}
+		}
+		return false;
+	}
 };
 
 class SQLConn : public SQL::Provider