Add query id to error messages when possible

Author: slabko

driver/api/impl/impl.cpp

@@ -505,6 +505,18 @@ SQLRETURN GetStmtAttr(
             CASE_NUM(SQL_ATTR_RETRIEVE_DATA, SQLULEN, SQL_RD_ON);
             CASE_NUM(SQL_ATTR_USE_BOOKMARKS, SQLULEN, SQL_UB_OFF);
 
+            case SQL_CH_STMT_ATTR_LAST_QUERY_ID: {
+                auto & query_id = statement.getQueryId();
+                if (query_id.empty())
+                    throw SqlException("Invalid cursor state", "24000");
+
+                // Note that the driver always returns the query ID as a UTF-8 string, and the driver manager does not
+                // alter or convert this value. Therefore, the user should convert the value to the expected format.
+                // Since this is for internal use, it's simpler to keep it this way. If we decide to make it public,
+                // we should convert the value to a SQLGUID.
+                return fillOutputString<PTChar>(query_id, out_value, out_value_max_length, out_value_length, true);
+            }
+
             case SQL_ATTR_FETCH_BOOKMARK_PTR:
             case SQL_ATTR_KEYSET_SIZE:
             case SQL_ATTR_SIMULATE_CURSOR:

driver/driver.cpp

@@ -140,3 +140,16 @@ void Driver::writeLogSessionEnd(std::ostream & stream) {
     }
     stream << " ====================" << std::endl;
 }
+
+std::string Driver::addContextInfoToExceptionMessage(const std::string & message, SQLHANDLE handle, SQLSMALLINT handle_type) const
+{
+    // Add query_id to exceptions related to statements
+    if (handle_type == SQL_HANDLE_STMT) {
+        auto * statement  = static_cast<Statement *>(handle);
+        const auto query_id = statement->getQueryId();
+        if (!query_id.empty()) {
+            return "Error while processing query " + query_id + ": " + message;
+        }
+    }
+    return message;
+}

driver/driver.h

@@ -93,6 +93,8 @@ class Driver
     virtual void onAttrChange(int attr) final override;
 
 private:
+    std::string addContextInfoToExceptionMessage(const std::string & message, SQLHANDLE handle, SQLSMALLINT handle_type) const;
+
     template <typename Callable>
     static inline SQLRETURN doCall(Callable & callable,
         typename std::enable_if<
@@ -235,27 +237,31 @@ inline SQLRETURN Driver::call(Callable && callable, SQLHANDLE handle, SQLSMALLIN
                     return doCall(callable, descendant, skip_diag);
                 }
                 catch (const SqlException & ex) {
-                    LOG(ex.getSQLState() << " (" << ex.what() << ")" << "[rc: " << ex.getReturnCode() << "]");
+                    auto error_message = addContextInfoToExceptionMessage(ex.what(), handle, handle_type);
+                    LOG(ex.getSQLState() << " (" << error_message << ")" << "[rc: " << ex.getReturnCode() << "]");
                     if (!skip_diag)
-                        descendant.fillDiag(ex.getReturnCode(), ex.getSQLState(), ex.what(), 1);
+                        descendant.fillDiag(ex.getReturnCode(), ex.getSQLState(), error_message, 1);
                     return ex.getReturnCode();
                 }
                 catch (const Poco::Exception & ex) {
-                    LOG("HY000 (" << ex.displayText() << ")");
+                    auto error_message = addContextInfoToExceptionMessage(ex.displayText(), handle, handle_type);
+                    LOG("HY000 (" << error_message << ")");
                     if (!skip_diag)
-                        descendant.fillDiag(SQL_ERROR, "HY000", ex.displayText(), 1);
+                        descendant.fillDiag(SQL_ERROR, "HY000", error_message, 1);
                     return SQL_ERROR;
                 }
                 catch (const std::exception & ex) {
-                    LOG("HY000 (" << ex.what() << ")");
+                    auto error_message = addContextInfoToExceptionMessage(ex.what(), handle, handle_type);
+                    LOG("HY000 (" << error_message << ")");
                     if (!skip_diag)
-                        descendant.fillDiag(SQL_ERROR, "HY000", ex.what(), 1);
+                        descendant.fillDiag(SQL_ERROR, "HY000", error_message, 1);
                     return SQL_ERROR;
                 }
                 catch (...) {
+                    auto error_message = addContextInfoToExceptionMessage("Unknown exception", handle, handle_type);
                     LOG("HY000 (Unknown exception)");
                     if (!skip_diag)
-                        descendant.fillDiag(SQL_ERROR, "HY000", "Unknown exception", 2);
+                        descendant.fillDiag(SQL_ERROR, "HY000", error_message, 2);
                     return SQL_ERROR;
                 }
             };

driver/platform/platform.h

@@ -117,3 +117,7 @@
 #define CH_SQL_ATTR_DRIVERLOG            (SQL_ATTR_TRACE + CH_SQL_OFFSET)
 #define CH_SQL_ATTR_DRIVERLOGFILE        (SQL_ATTR_TRACEFILE + CH_SQL_OFFSET)
 
+// User range is between SQL_DRIVER_STMT_ATTR_BASE (0x00004000) and 0x00007FFF, see
+// https://learn.microsoft.com/en-us/sql/odbc/reference/develop-app/driver-specific-data-types-descriptor-information-diagnostic
+#define SQL_CH_STMT_ATTR_BASE (SQL_DRIVER_STMT_ATTR_BASE + 1)
+#define SQL_CH_STMT_ATTR_LAST_QUERY_ID (SQL_CH_STMT_ATTR_BASE + 1)

driver/statement.cpp

@@ -99,6 +99,7 @@ Statement::HttpRequestData Statement::prepareHttpRequest()
 
 void Statement::requestNextPackOfResultSets(std::unique_ptr<ResultMutator> && mutator) {
     result_reader.reset();
+    query_id.clear();
 
     const auto param_set_array_size = getEffectiveDescriptor(SQL_ATTR_APP_PARAM_DESC).getAttrAs<SQLULEN>(SQL_DESC_ARRAY_SIZE, 1);
     if (next_param_set_idx >= param_set_array_size)
@@ -167,6 +168,13 @@ void Statement::requestNextPackOfResultSets(std::unique_ptr<ResultMutator> && mu
         }
     }
 
+    // response->get accepts default value as `const std::string &`
+    // and so we cannot pass an rvalue to it (e.g.`std::string()`),
+    // it will be deleted when the function returns.
+    // We use a static variable that will never be deleted.
+    static const std::string empty_query_id{};
+    query_id = response->get("X-ClickHouse-Query-Id", empty_query_id);
+
     Poco::Net::HTTPResponse::HTTPStatus status = response->getStatus();
     if (status != Poco::Net::HTTPResponse::HTTP_OK) {
         std::stringstream error_message;
@@ -355,6 +363,10 @@ ResultSet & Statement::getResultSet() {
     return result_reader->getResultSet();
 }
 
+const std::string & Statement::getQueryId() const {
+    return query_id;
+}
+
 bool Statement::advanceToNextResultSet() {
     if (!is_executed)
         return false;
@@ -382,6 +394,7 @@ void Statement::closeCursor() {
     }
 
     result_reader.reset();
+    query_id.clear();
     in = nullptr;
     response.reset();
 

driver/statement.h

@@ -44,9 +44,12 @@ class Statement
     /// Indicates whether there is an result set available for reading.
     bool hasResultSet() const;
 
-    /// Get the current resukt set.
+    /// Get the current result set.
     ResultSet & getResultSet();
 
+    /// Get the current query id, or an empty string if query id is not applicable or not available.
+    const std::string & getQueryId() const;
+
     /// Make the next result set current, if any.
     bool advanceToNextResultSet();
 
@@ -113,5 +116,6 @@ class Statement
     std::unique_ptr<Poco::Net::HTTPResponse> response;
     std::istream* in = nullptr;
     std::unique_ptr<ResultReader> result_reader;
+    std::string query_id;
     std::size_t next_param_set_idx = 0;
 };

driver/test/CMakeLists.txt

@@ -85,6 +85,7 @@ function (declare_odbc_test_targets libname UNICODE)
         performance_it.cpp
         type_info_it.cpp
         authentication_it.cpp
+        error_handling_it.cpp
     )
 
     if (CH_ODBC_ENABLE_CODE_COVERAGE)

driver/test/error_handling_it.cpp

@@ -0,0 +1,86 @@
+#include "driver/test/client_utils.h"
+#include "driver/test/client_test_base.h"
+#include "driver/test/result_set_reader.hpp"
+
+#include <gtest/gtest.h>
+#include <gmock/gmock.h>
+
+class ErrorHandlingTest
+    : public ClientTestBase
+{
+public:
+    std::string getQueryId()
+    {
+        char query_id_data[74] = {};
+        SQLINTEGER query_id_len{};
+        SQLGetStmtAttr(hstmt, SQL_CH_STMT_ATTR_LAST_QUERY_ID, query_id_data, std::size(query_id_data), &query_id_len);
+
+#ifdef _win_
+        return toUTF8(reinterpret_cast<PTChar*>(query_id_data));
+#else
+        // unixODBC usually converts all strings from the driver encoding to the application encoding,
+        // similar to how Windows behaves. However, `SQLGetStmtAttr` appears to be an exception in unixODBC,
+        // as it passes the result in the driver encoding unchanged.
+        // See: https://github.com/lurcher/unixODBC/issues/22
+        //
+        // To solve this issue and ensure the result reaches the application in the correct encoding,
+        // we would need to implement both `SQLGetStmtAttr` and `SQLGetStmtAttrW` in the driver.
+        // This would introduce a number of workarounds and non-uniform solutions.
+        // Currently, there is only one string attribute, `SQL_CH_STMT_ATTR_LAST_QUERY_ID`, which is internal,
+        // so this added complexity in the driver is not currently justified.
+        //
+        // Since we know that the Query ID is 36 bytes long, we can infer the encoding of the result
+        // and re-encode it accordingly here in the test.
+        if (query_id_len == 36)
+            return toUTF8(query_id_data);
+        else
+            return toUTF8(reinterpret_cast<char16_t*>(query_id_data));
+#endif
+    }
+};
+
+/**
+ * Given a query with a syntax error, we receive the error description with the query id.
+ */
+TEST_F(ErrorHandlingTest, ServerExceptionInBeginning)
+{
+    static const size_t stream_size = 100000;
+    auto query = fromUTF8<PTChar>("SELECT * FROM this.table.cannot.exist");
+
+    ODBC_CALL_ON_STMT_THROW(hstmt, SQLPrepare(hstmt, ptcharCast(query.data()), SQL_NTS));
+    ASSERT_EQ(SQLExecute(hstmt), SQL_ERROR);
+    auto error_message = extract_diagnostics(hstmt, SQL_HANDLE_STMT);
+    auto query_id = getQueryId();
+
+    auto expect_error_message = "1:[HY000][1]Error while processing query " + query_id + ": HTTP status code: 400";
+    ASSERT_EQ(error_message.substr(0, expect_error_message.size()), expect_error_message);
+}
+
+/**
+ * Given a query with an exception in the middle of the data stream, we receive the error description with the query id.
+ * (We do not yet parse error messages, we only report "Incomplete input stream, expected at least" at the moment)
+ */
+TEST_F(ErrorHandlingTest, ServerExceptionInMiddleOfStream)
+{
+    static const size_t stream_size = 100000;
+    auto query = fromUTF8<PTChar>(
+        "SELECT throwIf(number >= "
+        + std::to_string(stream_size) + ") FROM numbers("
+        + std::to_string(stream_size) + " + 1)");
+
+    ODBC_CALL_ON_STMT_THROW(hstmt, SQLPrepare(hstmt, ptcharCast(query.data()), SQL_NTS));
+    ODBC_CALL_ON_STMT_THROW(hstmt, SQLExecute(hstmt));
+
+    std::string error_message = "";
+    SQLRETURN res = SQL_SUCCESS;
+    while (res == SQL_SUCCESS)
+    {
+        res = SQLFetch(hstmt);
+        if (res != SQL_SUCCESS)
+            error_message = extract_diagnostics(hstmt, SQL_HANDLE_STMT);
+    }
+
+    auto query_id = getQueryId();
+    auto expect_error_message = "1:[HY000][1]Error while processing query " + query_id + ": ";
+    ASSERT_EQ(error_message.substr(0, expect_error_message.size()), expect_error_message);
+}

driver/utils/amortized_istream_reader.h

@@ -23,9 +23,10 @@ class AmortizedIStreamReader
 
     ~AmortizedIStreamReader() {
         // Put back any pre-read characters, just in case...
+        // (it should be done in reverse order)
         if (available() > 0) {
-            for (std::size_t i = buffer_.size() - 1; i >= offset_; --i) {
-                raw_stream_.putback(buffer_[i]);
+            for (auto it = buffer_.rbegin(); it < buffer_.rend() - offset_; ++it) {
+                raw_stream_.putback(*it);
             }
         }
     }