fix: handling timestamp escape clause

google/cloud/odbc/bq_driver/internal/odbc_sql_fetch.cc

@@ -237,6 +237,85 @@ StatusRecord WriteRowset(ResultSet const& result_set, int const rowset_size,
   return StatusRecord::Ok();
 }
 
+StatusRecord WriteRowset(StatementHandle& stmt_handle, int const rowset_size,
+                         DescriptorHandle& ard, DescriptorHandle& ird) {
+  if (rowset_size <= 0) {
+    LOG(ERROR) << "WriteRowset:: rowset_size should not be <= 0";
+    return StatusRecord{SQLStates::k_HY000(), "rowset_size should not be <= 0"};
+  }
+
+  int row_counter = 0;
+  SQLUSMALLINT* row_status_ptr = ird.GetHeaderRecord().array_status_ptr;
+
+  while (row_counter < rowset_size) {
+    ResultSet& result_set = stmt_handle.GetResultSet();
+    if (result_set.cursor >= static_cast<int>(result_set.rows.size())) {
+      StatusRecord next_page_status = FetchNextResultSet(stmt_handle);
+      if (!next_page_status.ok()) {
+        if (next_page_status.sql_state == SQLStates::k_SQL_NO_DATA()) {
+          break;
+        }
+        LOG(ERROR) << "WriteRowset::FetchNextResultSet:: "
+                   << next_page_status.message;
+        if (row_counter == 0) {
+          return next_page_status;
+        }
+        break;
+      }
+      ResultSet& updated_rs = stmt_handle.GetResultSet();
+      if (updated_rs.rows.empty()) {
+        break;
+      }
+      updated_rs.cursor++;
+    }
+
+    ResultSet& current_rs = stmt_handle.GetResultSet();
+    if (current_rs.cursor < 0 ||
+        current_rs.cursor >= static_cast<int>(current_rs.rows.size())) {
+      break;
+    }
+
+    StatusRecord status_record =
+        WriteDSRow(current_rs.rows[current_rs.cursor], current_rs.row_schema,
+                   ard, row_counter);
+    if (!status_record.ok()) {
+      LOG(ERROR) << "WriteRowset::WriteDSRow:: " << status_record.message;
+      return status_record;
+    }
+
+    if (row_status_ptr) {
+      row_status_ptr[row_counter] = SQL_ROW_SUCCESS;
+    }
+
+    row_counter++;
+    current_rs.cursor++;
+  }
+
+  ResultSet& final_rs = stmt_handle.GetResultSet();
+  if (row_counter > 0) {
+    final_rs.cursor--;
+  }
+
+  // Mark unused rows
+  if (row_status_ptr) {
+    for (int i = row_counter; i < rowset_size; i++) {
+      row_status_ptr[i] = SQL_ROW_NOROW;
+    }
+  }
+
+  SQLULEN* rows_processed_ptr = ird.GetHeaderRecord().rows_processed_ptr;
+  if (rows_processed_ptr) {
+    *rows_processed_ptr = row_counter;
+  }
+
+  if (row_counter == 0) {
+    return StatusRecord(
+        {SQLStates::k_SQL_NO_DATA(), "No more data to return."});
+  }
+
+  return StatusRecord::Ok();
+}
+
 StatusRecord FetchNextResultSet(StatementHandle& stmt_handle) {
   // We need to return only the top `SQL_ATTR_MAX_ROWS` number of rows
   auto max_rows_status = stmt_handle.GetAttribute(SQL_ATTR_MAX_ROWS);
@@ -258,21 +337,27 @@ StatusRecord FetchNextResultSet(StatementHandle& stmt_handle) {
   if (stmt_handle.WasHtapiEnabled()) {
     StatusRecord read_status = ReadNextResultsFromStream(stmt_handle);
     if (!read_status.ok()) {
-      LOG(ERROR) << "ReadNextResultsFromStream:: " << read_status.message;
+      if (read_status.sql_state != SQLStates::k_SQL_NO_DATA()) {
+        LOG(ERROR) << "ReadNextResultsFromStream:: " << read_status.message;
+      }
       return read_status;
     }
   } else {
     StatusRecord read_status = FetchNextPageResultSet(stmt_handle);
     if (!read_status.ok()) {
-      LOG(ERROR) << "FetchNextPageResultSet:: " << read_status.message;
+      if (read_status.sql_state != SQLStates::k_SQL_NO_DATA()) {
+        LOG(ERROR) << "FetchNextPageResultSet:: " << read_status.message;
+      }
       return read_status;
     }
   }
 #else
 
   StatusRecord read_status = FetchNextPageResultSet(stmt_handle);
   if (!read_status.ok()) {
-    LOG(ERROR) << "FetchNextPageResultSet:: " << read_status.message;
+    if (read_status.sql_state != SQLStates::k_SQL_NO_DATA()) {
+      LOG(ERROR) << "FetchNextPageResultSet:: " << read_status.message;
+    }
     return read_status;
   }
 #endif  // (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW)

google/cloud/odbc/bq_driver/internal/odbc_sql_fetch.h

@@ -21,7 +21,12 @@
 
 namespace google::cloud::odbc_bq_driver_internal {
 
-// Writes rowset_size number of rows to the columns bound by the application
+class StatementHandle;
+google::cloud::odbc_internal::StatusRecord WriteRowset(
+    StatementHandle& stmt_handle, int rowset_size, DescriptorHandle& ard,
+    DescriptorHandle& ird);
+
+// Overload for unit testing with ResultSet only
 google::cloud::odbc_internal::StatusRecord WriteRowset(
     ResultSet const& result_set, int rowset_size, DescriptorHandle& ard,
     DescriptorHandle& ird);

google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.cc

@@ -21,6 +21,7 @@
 #include "google/cloud/odbc/bq_driver/internal/odbc_sql_type_info.h"
 #include "google/cloud/odbc/bq_driver/internal/odbc_transactions.h"
 #include "google/cloud/odbc/bq_driver/internal/trace_utils.h"
+#include "google/cloud/odbc/bq_driver/internal/utils.h"
 #include "google/cloud/odbc/internal/status_record_or.h"
 
 namespace google::cloud::odbc_bq_driver_internal {
@@ -213,8 +214,14 @@ StatusRecord StatementHandle::PrepareQuery(std::string const& query) {
   }
   ConnectionHandle& conn_handle = *GetConnectionHandle();
 
+  std::string processed_query = query;
+  auto noscan_attr = GetAttribute(SQL_ATTR_NOSCAN);
+  if (!noscan_attr || *noscan_attr != SQL_NOSCAN_ON) {
+    processed_query = TranslateOdbcEscapeSequences(query);
+  }
+
   Job req;
-  req.configuration.query.query = query;
+  req.configuration.query.query = processed_query;
   req.configuration.query.use_query_cache = conn_handle.GetDsn().is_query_cache;
   req.configuration.dry_run = true;
   req.configuration.query.use_legacy_sql =
@@ -233,7 +240,7 @@ StatusRecord StatementHandle::PrepareQuery(std::string const& query) {
   // to be used during table creation. Subsequent operations on the table will
   // automatically use the KMS key without the application sending it.
   std::string kms_key_name = conn_handle.GetDsn().kms_key_name;
-  if (IsInsertQuery(query) || IsSelectQuery(query)) {
+  if (IsInsertQuery(processed_query) || IsSelectQuery(processed_query)) {
     if (!kms_key_name.empty()) {
       req.configuration.query.destination_encryption_configuration
           .kms_key_name = kms_key_name;
@@ -243,8 +250,8 @@ StatusRecord StatementHandle::PrepareQuery(std::string const& query) {
   if (!conn_handle.GetDsn().is_bq_legacy_sql) {
     // Detect POSITIONAL (`?`) and NAMED (`[:@]\w+`) parameter markers using
     // RE2 instead of a manual character scan.
-    bool has_positional = re2::RE2::PartialMatch(query, R"(\?)");
-    bool has_named = re2::RE2::PartialMatch(query, R"([:@]\w+)");
+    bool has_positional = re2::RE2::PartialMatch(processed_query, R"(\?)");
+    bool has_named = re2::RE2::PartialMatch(processed_query, R"([:@]\w+)");
     if (has_positional) {
       req.configuration.query.parameter_mode = "POSITIONAL";
     }
@@ -330,7 +337,7 @@ StatusRecord StatementHandle::PrepareQuery(std::string const& query) {
     conn_handle.SetSessionId(response->statistics.session_info.session_id);
   }
 
-  query_str_ = query;
+  query_str_ = processed_query;
   prepared_job_ = *response;
   return StatusRecord::Ok();
 }
@@ -551,6 +558,10 @@ StatusRecord StatementHandle::PopulateIpd(DescriptorHandle& handle,
 void StatementHandle::CloseCursor() {
   ResultSet result_set;
   result_set_ = result_set;
+#if (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW)
+  ClearReadRowsStream();
+  ClearReadRowsIterator();
+#endif
   if (StatementPrepared()) {
     SetStmtState(StmtStates::kStatementPrepared);
   } else {

google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h

@@ -141,6 +141,7 @@ class StatementHandle : public Handle {
       StreamRange<::google::cloud::bigquery::storage::v1::ReadRowsResponse>
           stream_range) {
     read_rows_stream_ = std::move(stream_range);
+    read_rows_iterator_.reset();
   }
 
   std::optional<