From c3518fd4e0e0a0bc15d28da2f5d8cf6d64087bf8 Mon Sep 17 00:00:00 2001 From: Sachin Purohit Date: Tue, 15 Sep 2026 23:26:07 +0000 Subject: [PATCH 1/2] deps: update google-cloud-cpp to v3.10.0 --- bazel/deps.bzl | 12 ++++++------ google/cloud/odbc/CMakeLists.txt | 2 +- google/cloud/odbc/bq_driver/internal/utils.cc | 13 +++++-------- google/cloud/odbc/testing/utils/env_vars.cc | 3 +-- 4 files changed, 13 insertions(+), 17 deletions(-) diff --git a/bazel/deps.bzl b/bazel/deps.bzl index 4169b38510..a9a5da608b 100644 --- a/bazel/deps.bzl +++ b/bazel/deps.bzl @@ -70,10 +70,10 @@ def cpp_bigquery_odbc_deps(name = None): maybe( http_archive, name = "bazel_skylib", - sha256 = "bc283cdfcd526a52c3201279cda4bc298652efa898b10b4db0837dc51652756f", + sha256 = "37cdfbc6faefea94f7b37760a305c98c08981116c2bc9e821e3b423221fad8c8", urls = [ - "https://mirror.bazel.build/github.com/bazelbuild/bazel-skylib/releases/download/1.7.1/bazel-skylib-1.7.1.tar.gz", - "https://github.com/bazelbuild/bazel-skylib/releases/download/1.7.1/bazel-skylib-1.7.1.tar.gz", + "https://mirror.bazel.build/github.com/bazelbuild/bazel-skylib/releases/download/1.9.2/bazel-skylib-1.9.2.tar.gz", + "https://github.com/bazelbuild/bazel-skylib/releases/download/1.9.2/bazel-skylib-1.9.2.tar.gz", ], ) @@ -81,8 +81,8 @@ def cpp_bigquery_odbc_deps(name = None): http_archive, name = "com_google_cloud_cpp", urls = [ - "https://github.com/googleapis/google-cloud-cpp/archive/85bfaff9038e0791c7f48a995253dc20c307ac78.tar.gz", + "https://github.com/googleapis/google-cloud-cpp/archive/refs/tags/v3.10.0.tar.gz", ], - sha256 = "65cf83a2c6fc730ca50420ccbe4d4dc5691bbd4365a0020dca918d9c5a443988", - strip_prefix = "google-cloud-cpp-85bfaff9038e0791c7f48a995253dc20c307ac78", + sha256 = "8ffe5675b72822789704bca9e045a4a7e1fc047199f1f41d5df1549318563020", + strip_prefix = "google-cloud-cpp-3.10.0", ) diff --git a/google/cloud/odbc/CMakeLists.txt b/google/cloud/odbc/CMakeLists.txt index ec04725342..4b3d077180 100644 --- a/google/cloud/odbc/CMakeLists.txt +++ b/google/cloud/odbc/CMakeLists.txt @@ -147,7 +147,7 @@ if (NOT google_cloud_cpp_bigquery_rest_FOUND OR NOT FetchContent_Declare( google-cloud-cpp - URL https://github.com/googleapis/google-cloud-cpp/archive/85bfaff9038e0791c7f48a995253dc20c307ac78.tar.gz + URL https://github.com/googleapis/google-cloud-cpp/archive/refs/tags/v3.10.0.tar.gz PATCH_COMMAND ${CMAKE_COMMAND} -P ${CMAKE_SOURCE_DIR}/cmake/PatchGoogleCloudCpp.cmake) FetchContent_MakeAvailable(google-cloud-cpp) diff --git a/google/cloud/odbc/bq_driver/internal/utils.cc b/google/cloud/odbc/bq_driver/internal/utils.cc index 201a0bf3da..26537b99ca 100644 --- a/google/cloud/odbc/bq_driver/internal/utils.cc +++ b/google/cloud/odbc/bq_driver/internal/utils.cc @@ -765,20 +765,18 @@ StatusRecordOr
ParseConnectionString(std::string& str) { std::string GetPathToOdbcIni() { #ifdef _WIN32 // 64-bit - absl::optional path = "SOFTWARE\\ODBC\\ODBC.INI"; + std::string path = "SOFTWARE\\ODBC\\ODBC.INI"; #ifndef _WIN64 // 32-bit path = "SOFTWARE\\WOW6432Node\\ODBC\\ODBC.INI"; #endif // _WIN64 - if (path) { - return *path; - } + return path; #else - absl::optional path = google::cloud::internal::GetEnv("ODBCINI"); + auto path = google::cloud::internal::GetEnv("ODBCINI"); if (path) { return *path; } - absl::optional home = google::cloud::internal::GetEnv("HOME"); + auto home = google::cloud::internal::GetEnv("HOME"); if (home) { return *home + "/.odbc.ini"; } @@ -788,8 +786,7 @@ std::string GetPathToOdbcIni() { std::string GetOdbcTraceConfigPath() { #ifndef _WIN32 - absl::optional path = - google::cloud::internal::GetEnv("GOOGLEBIGQUERYODBCINI"); + auto path = google::cloud::internal::GetEnv("GOOGLEBIGQUERYODBCINI"); if (path) { return *path; } diff --git a/google/cloud/odbc/testing/utils/env_vars.cc b/google/cloud/odbc/testing/utils/env_vars.cc index e3d1bb987f..b6c69a84af 100644 --- a/google/cloud/odbc/testing/utils/env_vars.cc +++ b/google/cloud/odbc/testing/utils/env_vars.cc @@ -18,8 +18,7 @@ namespace google::cloud::odbc_testing_utils { std::string GetRequiredEnvVar(std::string const& var) { - absl::optional optional_env_var = - ::google::cloud::internal::GetEnv(var.c_str()); + auto optional_env_var = ::google::cloud::internal::GetEnv(var.c_str()); if (!optional_env_var) { throw std::runtime_error(var + " environment variable is not set"); } From 6ddb426a36dd45bd437f6a2fb601ff9894b8a317 Mon Sep 17 00:00:00 2001 From: Sachin Purohit Date: Tue, 15 Sep 2026 23:26:15 +0000 Subject: [PATCH 2/2] chore: removing dry run for queries without positional params --- .../internal/odbc_sql_execute_utils.cc | 111 ++++++++--- .../internal/odbc_sql_execute_utils.h | 10 + .../cloud/odbc/bq_driver/odbc_sql_requests.cc | 177 +++++++++++------- .../odbc_driver_tests/connection_test.cc | 8 +- .../odbc_driver_tests/statement_test.cc | 34 ++++ 5 files changed, 245 insertions(+), 95 deletions(-) diff --git a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc index dbfb02792f..d3604a8e15 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc +++ b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.cc @@ -39,14 +39,15 @@ using ::google::cloud::bigquery::storage::v1::ReadRowsRequest; using ::google::cloud::bigquery::storage::v1::ReadRowsResponse; using ::google::cloud::bigquery::storage::v1::ReadSession; using ::google::cloud::bigquery::storage::v1::DataFormat::ARROW; -using ::google::cloud::bigquery_v2_minimal_internal::Job; using ::google::cloud::bigquery_v2_minimal_internal::JobReference; using ::google::cloud::bigquery_v2_minimal_internal::QueryRequest; #endif // (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW) using ::google::cloud::bigquery_v2_minimal_internal::GetQueryResults; using ::google::cloud::bigquery_v2_minimal_internal::GetQueryResultsRequest; +using ::google::cloud::bigquery_v2_minimal_internal::Job; using ::google::cloud::bigquery_v2_minimal_internal::PostQueryRequest; using ::google::cloud::bigquery_v2_minimal_internal::QueryParameter; +using ::google::cloud::bigquery_v2_minimal_internal::TableReference; using google::cloud::odbc_bigquery_client_interface::MaxRetriesOption; using google::cloud::odbc_bq_driver_internal::DescriptorRecord; using google::cloud::odbc_bq_driver_internal::DoubleStrToInt; @@ -222,6 +223,7 @@ StatusRecordOr ExecuteScript( } DSResults results; + results.job_ref = pq_status->job_reference; if (pq_status->job_complete && pq_status->page_token.empty()) { // we have gotten all the results results.num_dml_affected_rows = pq_status->num_dml_affected_rows; @@ -241,18 +243,43 @@ StatusRecordOr ExecuteScript( results.data_source_results = *gq_status; } - // Retrieve job information + if (!conn_handle->IsSessionStarted() && + !pq_status->session_info.session_id.empty()) { + conn_handle->SetSessionId(pq_status->session_info.session_id); + } + + auto populate_status = + PopulateScriptChildJobs(stmt_handle, results, post_query_request); + if (!populate_status.ok()) { + return populate_status; + } + + return results; +} + +StatusRecord PopulateScriptChildJobs( + StatementHandle& stmt_handle, DSResults& results, + PostQueryRequest const& post_query_request) { + ConnectionHandle* conn_handle = stmt_handle.GetConnectionHandle(); + if (!conn_handle || !conn_handle->IsConnected()) { + return StatusRecord::Ok(); + } + auto bq_client = conn_handle->GetClient(); + if (!bq_client || !results.job_ref.has_value()) { + return StatusRecord::Ok(); + } + Options list_job_options; list_job_options.set(conn_handle->GetDsn().max_retries); - auto all_jobs_status = - bq_client->ListAllJobs(pq_status->job_reference.project_id, - pq_status->job_reference.job_id, list_job_options); + auto all_jobs_status = bq_client->ListAllJobs( + results.job_ref->project_id, results.job_ref->job_id, list_job_options); if (!all_jobs_status) { - LOG(ERROR) << "ExecuteScript::ListAllJobs:: " + LOG(ERROR) << "PopulateScriptChildJobs::ListAllJobs:: " << all_jobs_status.GetStatusRecord().message; return all_jobs_status.GetStatusRecord(); } + int statement_jobs_count = 0; for (auto const& job_status : all_jobs_status.GetValue()) { if (job_status.statistics.job_query_stats.statement_type != "CREATE_PROCEDURE" && @@ -261,51 +288,54 @@ StatusRecordOr ExecuteScript( stmt_handle.SetJobData( job_status.job_reference.job_id, job_status.statistics.job_query_stats.statement_type); + statement_jobs_count++; } } - // Fetch query results if job data is available - if (!stmt_handle.HasJobData()) { - return results; + if (statement_jobs_count <= 1) { + return StatusRecord::Ok(); } + + // If there are multiple statement jobs, results currently contains the + // output of the LAST statement executed in the script. + // We need to fetch and return the results of the FIRST statement. auto job_status = stmt_handle.GetNextJobData(); if (!job_status.Ok()) { - LOG(ERROR) << "ExecuteScript::GetNextJobData:: " + LOG(ERROR) << "PopulateScriptChildJobs::GetNextJobData:: " << job_status.GetStatusRecord().message; return job_status.GetStatusRecord(); } auto job_data = job_status.GetValue(); - std::string job_id = job_data.first; - std::string statement_type = job_data.second; + std::string first_job_id = job_data.first; + std::string first_statement_type = job_data.second; Options query_results_options; query_results_options.set( conn_handle->GetDsn().max_retries); auto gq_status = bq_client->GetAllQueryResults( - pq_status->job_reference.project_id, job_id, - pq_status->job_reference.location, + results.job_ref->project_id, first_job_id, results.job_ref->location, post_query_request.query_request().timeout(), query_results_options); if (!gq_status) { - LOG(ERROR) << "ExecuteScript::GetAllQueryResults:: " + LOG(ERROR) << "PopulateScriptChildJobs::GetAllQueryResults:: " << gq_status.GetStatusRecord().message; return gq_status.GetStatusRecord(); } - // Assign DML row counts - if (statement_type == "INSERT" || statement_type == "UPDATE" || - statement_type == "DELETE") { + if (first_statement_type == "UPDATE" || first_statement_type == "INSERT" || + first_statement_type == "DELETE") { results.num_dml_affected_rows = gq_status->num_dml_affected_rows; } results.data_source_results = *gq_status; stmt_handle.SetDSResults(results); - if (!conn_handle->IsSessionStarted() && - !pq_status->session_info.session_id.empty()) { - conn_handle->SetSessionId(pq_status->session_info.session_id); - } + // Unbind IRD and populate with the first child statement's schema + DescriptorHandle& ird = stmt_handle.GetDescriptorHandle(DescriptorType::kIRD); + ird.UnbindAllDescriptorRecordsFrom(0); + TableReference table_fields; + StatementHandle::PopulateIrd(ird, gq_status->schema, table_fields); - return results; + return StatusRecord::Ok(); } #if (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW) @@ -1113,6 +1143,35 @@ StatusRecordOr FetchBQData( conn_handle.SetSessionId(pq_status->session_info.session_id); } + // If we skipped the dry-run job (e.g. SQLExecDirect without positional + // params), populate the statement handle's prepared job metadata, IRD, and + // location so that subsequent processing remains identical to the dry-run + // flow. + PostQueryRequest actual_post_query_request = post_query_request; + if (!stmt_handle.GetPreparedJob().has_value()) { + Job executed_job; + executed_job.job_reference = pq_status->job_reference; + executed_job.statistics.job_query_stats.statement_type = + pq_status->statement_type; + executed_job.statistics.job_query_stats.schema = pq_status->schema; + stmt_handle.SetPreparedJob(executed_job); + + if (!pq_status->schema.fields.empty()) { + DescriptorHandle& ird = + stmt_handle.GetDescriptorHandle(DescriptorType::kIRD); + ird.SetConnectionHandle(&conn_handle); + ird.ClearDescriptorRecordsMap(); + TableReference table_fields; + StatementHandle::PopulateIrd(ird, pq_status->schema, table_fields); + } + if (!pq_status->job_reference.location.empty()) { + auto query_req = actual_post_query_request.query_request(); + query_req.set_location(pq_status->job_reference.location); + actual_post_query_request.set_query_request(query_req); + stmt_handle.SetPostQueryRequest(actual_post_query_request); + } + } + DSResults results; results.num_dml_affected_rows = pq_status->num_dml_affected_rows; results.job_ref = pq_status->job_reference; @@ -1134,7 +1193,7 @@ StatusRecordOr FetchBQData( if (!read_status.ok()) { LOG(WARNING) << "FetchBQDataReadFromJob failed: " << read_status.message << ", falling back to FetchBQDataRead."; - read_status = FetchBQDataRead(stmt_handle, post_query_request); + read_status = FetchBQDataRead(stmt_handle, actual_post_query_request); if (!read_status.ok()) { return read_status; } @@ -1150,7 +1209,7 @@ StatusRecordOr FetchBQData( results.data_source_results = *pq_status; } else { auto gq_status = - FetchNextPageOfQueryResults(stmt_handle, post_query_request); + FetchNextPageOfQueryResults(stmt_handle, actual_post_query_request); if (!gq_status) { LOG(ERROR) << "FetchBQData::FetchNextPageOfQueryResults:: " << gq_status.GetStatusRecord().message; @@ -1213,7 +1272,7 @@ StatusRecordOr FetchNextPageOfQueryResults( auto* connection_handle = stmt_handle.GetConnectionHandle(); Options options; - auto job_client = stmt_handle.GetConnectionHandle()->GetClient(); + auto job_client = connection_handle->GetClient(); LOG(INFO) << "FetchNextPageOfQueryResults:: Request body: " << get_query_results_request.DebugString(""); diff --git a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h index ffd2665145..e0526e1747 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h +++ b/google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h @@ -61,6 +61,16 @@ odbc_internal::StatusRecordOr ExecuteScript( google::cloud::bigquery_v2_minimal_internal::PostQueryRequest const& post_query_request); +/* + * @brief If the executed query was a multi-statement script, populates child + * jobs into the statement handle for SQLMoreResults and retrieves results of + * the first statement. + */ +odbc_internal::StatusRecord PopulateScriptChildJobs( + StatementHandle& stmt_handle, DSResults& results, + google::cloud::bigquery_v2_minimal_internal::PostQueryRequest const& + post_query_request); + #if (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW) /* * @brief Reads the next set of rows from the stream cached in the statement diff --git a/google/cloud/odbc/bq_driver/odbc_sql_requests.cc b/google/cloud/odbc/bq_driver/odbc_sql_requests.cc index d2dc2bdbd9..d3235ad0d6 100644 --- a/google/cloud/odbc/bq_driver/odbc_sql_requests.cc +++ b/google/cloud/odbc/bq_driver/odbc_sql_requests.cc @@ -18,8 +18,10 @@ #include "google/cloud/odbc/bq_driver/internal/odbc_internal_commons.h" #include "google/cloud/odbc/bq_driver/internal/odbc_sql_execute_utils.h" #include "google/cloud/odbc/bq_driver/internal/odbc_stmt_handle.h" +#include "google/cloud/odbc/bq_driver/internal/odbc_transactions.h" #include "google/cloud/odbc/bq_driver/internal/odbc_type_utils.h" #include "google/cloud/odbc/bq_driver/internal/trace_utils.h" +#include "google/cloud/odbc/bq_driver/internal/utils.h" #include "google/cloud/odbc/bq_driver/odbc_descriptor.h" #include "google/cloud/odbc/bq_driver/odbc_utils.h" #include "google/cloud/odbc/internal/status_record_or.h" @@ -34,6 +36,7 @@ using ::google::cloud::bigquery_v2_minimal_internal::QueryParameter; using ::google::cloud::bigquery_v2_minimal_internal::QueryRequest; using google::cloud::bigquery_v2_minimal_internal::TableReference; using google::cloud::odbc_bq_driver::ToCharStr; +using google::cloud::odbc_bq_driver_internal::BeginTransactionIfNeeded; using google::cloud::odbc_bq_driver_internal::CancelBQJob; using google::cloud::odbc_bq_driver_internal::ConnectionHandle; using google::cloud::odbc_bq_driver_internal::ConstructBasicPostQueryRequest; @@ -46,6 +49,7 @@ using google::cloud::odbc_bq_driver_internal::ExecuteScript; using google::cloud::odbc_bq_driver_internal::FetchBQData; using google::cloud::odbc_bq_driver_internal::IntValueToOutputBufferResponse; using google::cloud::odbc_bq_driver_internal::LogAndReturnCode; +using google::cloud::odbc_bq_driver_internal::PopulateScriptChildJobs; using google::cloud::odbc_bq_driver_internal::ResultSet; using google::cloud::odbc_bq_driver_internal::StatementHandle; using google::cloud::odbc_bq_driver_internal::StmtStates; @@ -215,6 +219,26 @@ SQLRETURN HandleAsyncExecute(StatementHandle& handle_ref) { return LogAndReturnCode(handle_ref, status_record); } +// Helper function to process DSResults and apply SQL_ATTR_MAX_ROWS truncation. +StatusRecordOr ProcessQueryResultsWithMaxRows( + StatementHandle& stmt_handle, DSResults const& results) { + auto rs_status_record_or = ProcessQueryResults(results); + if (!rs_status_record_or) { + return rs_status_record_or.GetStatusRecord(); + } + auto max_rows_status = stmt_handle.GetAttribute(SQL_ATTR_MAX_ROWS); + if (!max_rows_status) { + return max_rows_status.GetStatusRecord(); + } + SQLULEN max_rows = *max_rows_status; + ResultSet result_set = std::move(*rs_status_record_or); + if (max_rows > 0 && max_rows < result_set.rows.size()) { + result_set.rows.erase(result_set.rows.begin() + max_rows, + result_set.rows.end()); + } + return result_set; +} + // @brief This function synchronously processes current execute requests // assuming PrepareQuery was called // @param stmt_handle The statement handle @@ -236,17 +260,20 @@ StatusRecord ActuallyProcessExecute(StatementHandle& stmt_handle, } int query_timeout = *query_timeout_status; - // Ensure a prepared job exists - if (!stmt_handle.GetPreparedJob().has_value()) { - LOG(ERROR) - << "ActuallyProcessExecute::Internal state error when executing query"; - return StatusRecord{SQLStates::k_HY000(), - "Internal state error when executing query"}; + std::string location; + std::string statement_type; + auto prepared_job = stmt_handle.GetPreparedJob(); + if (prepared_job.has_value()) { + location = prepared_job->job_reference.location; + statement_type = prepared_job->statistics.job_query_stats.statement_type; + } else { + location = conn_handle.GetDsn().session_location; } - Job prepared_job = stmt_handle.GetPreparedJob().value(); // We assume that the dry run job would detect the `location` properly. // The execution utils `FetchBQData` and others will use it through the + // `PostQueryRequest`. `SetPostQueryRequest` called subsequently caches it in + // the statement_handle, so it will can be used for next pages as well. std::vector basic_query_params = stmt_handle.GetQueryParameters(); DescriptorHandle& apd = stmt_handle.GetDescriptorHandle(DescriptorType::kAPD); @@ -259,8 +286,6 @@ StatusRecord ActuallyProcessExecute(StatementHandle& stmt_handle, SQLUSMALLINT* param_status_ptr = ipd.GetHeaderRecord().array_status_ptr; if (rows_processed_ptr) *rows_processed_ptr = 0; - std::string statement_type = - prepared_job.statistics.job_query_stats.statement_type; std::string sub_statement_type; StatusRecordOr ds_status_record_or; std::int64_t total_dml_affected_rows = 0; @@ -269,15 +294,15 @@ StatusRecord ActuallyProcessExecute(StatementHandle& stmt_handle, bool any_success = false; StatusRecord last_error; DSResults combined_ds_results; + PostQueryRequest post_request; for (SQLULEN i = 0; i < array_size; i++) { // We assume that the dry run job would detect the `location` properly. // The execution utils `FetchBQData` and others will use it through the // `PostQueryRequest`. `SetPostQueryRequest` called subsequently caches it // in the statement_handle, so it will can be used for next pages as well. - PostQueryRequest post_request = - ConstructBasicPostQueryRequest(conn_handle, query_str, query_timeout, - prepared_job.job_reference.location); + post_request = ConstructBasicPostQueryRequest(conn_handle, query_str, + query_timeout, location); std::vector query_params = basic_query_params; if (!query_params.empty()) { StatusRecord status = ConstructPositionalQueryParams( @@ -312,9 +337,13 @@ StatusRecord ActuallyProcessExecute(StatementHandle& stmt_handle, // Execute the script or fetch data based on statement type if (statement_type == "SCRIPT") { ds_status_record_or = ExecuteScript(stmt_handle, post_request); + } else if (!prepared_job.has_value() || statement_type == "SELECT") { + // It doesn't make sense to read from HTAPI if it is not a select + // statement. We get an error otherwise: + // "Cannot set destination table in jobs with DDL statements" + ds_status_record_or = FetchBQData(stmt_handle, post_request, true); } else { - bool const with_htapi = (statement_type == "SELECT"); - ds_status_record_or = FetchBQData(stmt_handle, post_request, with_htapi); + ds_status_record_or = FetchBQData(stmt_handle, post_request); } if (!ds_status_record_or) { @@ -342,6 +371,22 @@ StatusRecord ActuallyProcessExecute(StatementHandle& stmt_handle, combined_ds_results.num_dml_affected_rows = total_dml_affected_rows; stmt_handle.SetDSResults(combined_ds_results); + if (statement_type.empty() && stmt_handle.GetPreparedJob().has_value()) { + statement_type = + stmt_handle.GetPreparedJob()->statistics.job_query_stats.statement_type; + } + + if (statement_type == "SCRIPT" && !stmt_handle.HasJobData()) { + StatusRecord script_status = + PopulateScriptChildJobs(stmt_handle, combined_ds_results, post_request); + if (!script_status.ok()) { + stmt_handle.SetStmtState(failure_state); + LOG(ERROR) << "ActuallyProcessExecute::PopulateScriptChildJobs:: " + << script_status.message; + return script_status; + } + } + // If the statement was a script, retrieve sub-statement type if (statement_type == "SCRIPT" && stmt_handle.HasJobData()) { auto job_status = stmt_handle.GetNextJobData(); @@ -353,41 +398,34 @@ StatusRecord ActuallyProcessExecute(StatementHandle& stmt_handle, sub_statement_type = job_status.GetValue().second; } - // Process DSResults into a ResultSet - auto rs_status_record_or = ProcessQueryResults(combined_ds_results); + // Process DSResults into a ResultSet with MAX_ROWS applied + auto rs_status_record_or = + ProcessQueryResultsWithMaxRows(stmt_handle, combined_ds_results); if (!rs_status_record_or) { stmt_handle.SetStmtState(failure_state); - LOG(ERROR) << "ActuallyProcessExecute::ProcessQueryResults::" + LOG(ERROR) << "ActuallyProcessExecute::ProcessQueryResultsWithMaxRows::" << rs_status_record_or.GetStatusRecord().message; return rs_status_record_or.GetStatusRecord(); } - - // 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); - if (!max_rows_status) { - return max_rows_status.GetStatusRecord(); - } - SQLULEN max_rows = *max_rows_status; ResultSet& result_set = *rs_status_record_or; - auto& rs_rows = result_set.rows; - if (max_rows > 0 && max_rows < rs_rows.size()) { - rs_rows.erase(rs_rows.begin() + max_rows, rs_rows.end()); - } // Determine execution state based on statement type if (statement_type == "SELECT" || (statement_type == "SCRIPT" && sub_statement_type == "SELECT")) { stmt_handle.SetStmtState(StmtStates::kStatementExecutedWithRs); stmt_handle.SetResultSet(result_set); - } else if ((statement_type == "UPDATE" || statement_type == "INSERT" || - statement_type == "DELETE") && - combined_ds_results.num_dml_affected_rows == 0) { - stmt_handle.SetStmtState(StmtStates::kStatementExecutedWithoutRs); - // Note: The message is not supposed to be propagated to the application in - // case of SQL_NO_DATA - LOG(WARNING) << "ActuallyProcessExecute::No data found"; - return StatusRecord{SQLStates::k_SQL_NO_DATA(), "No data found"}; } else { + stmt_handle.GetDescriptorHandle(DescriptorType::kIRD) + .UnbindAllDescriptorRecordsFrom(0); + if ((statement_type == "UPDATE" || statement_type == "INSERT" || + statement_type == "DELETE" || statement_type == "MERGE") && + combined_ds_results.num_dml_affected_rows == 0) { + stmt_handle.SetStmtState(StmtStates::kStatementExecutedWithoutRs); + // Note: The message is not supposed to be propagated to the application + // in case of SQL_NO_DATA + LOG(WARNING) << "ActuallyProcessExecute::No data found"; + return StatusRecord{SQLStates::k_SQL_NO_DATA(), "No data found"}; + } stmt_handle.SetStmtState(StmtStates::kStatementExecutedWithoutRs); } @@ -402,19 +440,39 @@ StatusRecord ActuallyProcessExecDirect(StatementHandle& stmt_handle) { stmt_handle.SetStmtState(StmtStates::kStatementStillExecuting); std::string query_str = stmt_handle.GetQueryString(); - // We need to call `PrepareQuery` because: - // 1) We need to get `statement_type` during `ActuallyProcessExecute` - // through - // `Job::statistics.job_query_stats.statement_type`. This is not possible - // through `PostQueryResults` - // 2) For positional params, we need to get `QueryParameter`s before - // SQLExecDirect is called. - StatusRecord prepare_status = stmt_handle.PrepareQuery(query_str); - if (!prepare_status.ok()) { - LOG(ERROR) << "ActuallyProcessExecDirect::PrepareQuery:: " - << prepare_status.message; - return prepare_status; + ConnectionHandle& conn_handle = *(stmt_handle.GetConnectionHandle()); + + // For queries without bound positional parameters, bypass the PrepareQuery + // dry run to save an extra round-trip to BigQuery. We start an explicit + // transaction if needed, clear any prior prepared state and descriptors, and + // execute directly via FetchBQData (which populates the statement handle on + // completion). For parameterized queries, PrepareQuery is still required to + // resolve parameter types and descriptors from the dry run. + bool has_bound_params = stmt_handle.GetDescriptorHandle(DescriptorType::kAPD) + .HasDescriptorRecord(1); + + if (!has_bound_params) { + StatusRecord transaction_status = BeginTransactionIfNeeded(conn_handle); + if (!transaction_status.ok()) { + LOG(ERROR) << "ActuallyProcessExecDirect::BeginTransactionIfNeeded:: " + << transaction_status.message; + return transaction_status; + } + stmt_handle.SetNullPreparedJob(); + stmt_handle.SetQueryParameters({}); + stmt_handle.GetDescriptorHandle(DescriptorType::kIPD) + .ClearDescriptorRecordsMap(); + stmt_handle.GetDescriptorHandle(DescriptorType::kIRD) + .UnbindAllDescriptorRecordsFrom(0); + } else { + StatusRecord prepare_status = stmt_handle.PrepareQuery(query_str); + if (!prepare_status.ok()) { + LOG(ERROR) << "ActuallyProcessExecDirect::PrepareQuery:: " + << prepare_status.message; + return prepare_status; + } } + return ActuallyProcessExecute(stmt_handle, StmtStates::kStatementNotPrepared); } @@ -484,7 +542,7 @@ StatusRecord ActuallyGetMoreResults(StatementHandle& stmt_handle) { if (!job_status.Ok()) { return job_status.GetStatusRecord(); } - auto [job_id, statement_type] = stmt_handle.GetNextJobData().GetValue(); + auto [job_id, statement_type] = job_status.GetValue(); // Fetch query results from BigQuery. Options options; @@ -509,23 +567,13 @@ StatusRecord ActuallyGetMoreResults(StatementHandle& stmt_handle) { stmt_handle.SetDSResults(results); - // Process query results into a result set if it's a SELECT statement. - auto rs_status_record_or = ProcessQueryResults(results); - auto max_rows_status = stmt_handle.GetAttribute(SQL_ATTR_MAX_ROWS); - if (!max_rows_status) { - return max_rows_status.GetStatusRecord(); - } - SQLULEN max_rows = *max_rows_status; - ResultSet& result_set = *rs_status_record_or; - auto& rs_rows = result_set.rows; - if (max_rows > 0 && max_rows < rs_rows.size()) { - rs_rows.erase(rs_rows.begin() + max_rows, rs_rows.end()); - } - + // Process query results into a result set with MAX_ROWS applied. + auto rs_status_record_or = + ProcessQueryResultsWithMaxRows(stmt_handle, results); if (!rs_status_record_or || statement_type != "SELECT") { stmt_handle.SetStmtState(StmtStates::kStatementExecutedWithoutRs); } else { - stmt_handle.SetResultSet(result_set); + stmt_handle.SetResultSet(std::move(*rs_status_record_or)); stmt_handle.SetStmtState(StmtStates::kStatementExecutedWithRs); } @@ -535,8 +583,7 @@ StatusRecord ActuallyGetMoreResults(StatementHandle& stmt_handle) { // TODO(b/413273776): Handle PopulateIrd call in SQLMoreResults TableReference table_fields; - google::cloud::odbc_bq_driver_internal::StatementHandle::PopulateIrd( - ird, ds_status_record_or->schema, table_fields); + StatementHandle::PopulateIrd(ird, ds_status_record_or->schema, table_fields); return StatusRecord::Ok(); } diff --git a/google/cloud/odbc/integration_tests/odbc_driver_tests/connection_test.cc b/google/cloud/odbc/integration_tests/odbc_driver_tests/connection_test.cc index 7fb05d387c..77f8b7e077 100644 --- a/google/cloud/odbc/integration_tests/odbc_driver_tests/connection_test.cc +++ b/google/cloud/odbc/integration_tests/odbc_driver_tests/connection_test.cc @@ -597,6 +597,9 @@ TEST(MultipleConnectionTest, SQLDriverConnect) { } } +// SQL_ATTR_ANSI_APP is a unixODBC extension, so it's not expected to work on +// Windows +#ifndef _WIN32 TEST(ConnectionTest, VerifySQLANSIAttributes) { auto conn = std::make_shared(); SQLRETURN status; @@ -611,14 +614,11 @@ TEST(ConnectionTest, VerifySQLANSIAttributes) { sizeof(conn->outdsn), &buflen, SQL_DRIVER_COMPLETE); CheckError(status, "SQLDriverConnectA", conn, true); -#ifndef WIN32 - // SQL_ATTR_ANSI_APP is a unixODBC extension, so it's not expected to work on - // Windows status = SQLSetConnectAttr(conn->hdbc, SQL_ATTR_ANSI_APP, ToSqlPointer(SQL_AA_FALSE), 0); EXPECT_EQ(status, SQL_SUCCESS); -#endif } +#endif // _WIN32 TEST(ConnectionTest, SQLDriverConnectA) { auto conn = std::make_shared(); diff --git a/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc b/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc index 7e1db5cda9..916d906d99 100644 --- a/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc +++ b/google/cloud/odbc/integration_tests/odbc_driver_tests/statement_test.cc @@ -4623,4 +4623,38 @@ TEST(StatementTest, SQLExecute_RebindAfterExecute) { EXPECT_EQ(Disconnect(conn), SQL_SUCCESS); } +TEST(StatementTest, ScriptQueryWithSelect) { + auto conn = std::make_shared(); + ASSERT_EQ(Connect(kDefaultConnectionString, conn), SQL_SUCCESS); + + std::string query = + "SET @@query_label = 'k1:v1,k2:2,k3:true'; SELECT 1 AS col1, 'abc' AS " + "col2 LIMIT 1001"; + SQLRETURN ret = SQLExecDirect(conn->hstmt, (SQLCHAR*)query.c_str(), SQL_NTS); + CheckError(ret, "SQLExecDirect", conn); + ASSERT_EQ(ret, SQL_SUCCESS); + + SQLSMALLINT col_count = 0; + ASSERT_EQ(SQLNumResultCols(conn->hstmt, &col_count), SQL_SUCCESS); + EXPECT_EQ(col_count, 2); + + ASSERT_EQ(SQLFetch(conn->hstmt), SQL_SUCCESS); + + SQLBIGINT col1_val = 0; + SQLLEN ind1 = 0; + ASSERT_EQ(SQLGetData(conn->hstmt, 1, SQL_C_SBIGINT, &col1_val, + sizeof(col1_val), &ind1), + SQL_SUCCESS); + EXPECT_EQ(col1_val, 1); + + char col2_val[64] = {}; + SQLLEN ind2 = 0; + ASSERT_EQ( + SQLGetData(conn->hstmt, 2, SQL_C_CHAR, col2_val, sizeof(col2_val), &ind2), + SQL_SUCCESS); + EXPECT_STREQ(col2_val, "abc"); + + EXPECT_EQ(SQLFetch(conn->hstmt), SQL_NO_DATA); + EXPECT_EQ(Disconnect(conn), SQL_SUCCESS); +} } // namespace google::cloud::odbc_tests