chore: response-driven SQLExecDirect execution

bazel/deps.bzl

@@ -70,19 +70,19 @@ 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",
         ],
     )
 
     maybe(
         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",
     )

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)

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<DSResults> 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<DSResults> 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<MaxRetriesOption>(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<DSResults> 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<MaxRetriesOption>(
       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<DSResults> 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<DSResults> 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<DSResults> 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<GetQueryResults> 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("");

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

@@ -61,6 +61,16 @@ odbc_internal::StatusRecordOr<DSResults> 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

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

@@ -765,20 +765,18 @@ StatusRecordOr<Section> ParseConnectionString(std::string& str) {
 std::string GetPathToOdbcIni() {
 #ifdef _WIN32
   // 64-bit
-  absl::optional<std::string> 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<std::string> path = google::cloud::internal::GetEnv("ODBCINI");
+  auto path = google::cloud::internal::GetEnv("ODBCINI");
   if (path) {
     return *path;
   }
-  absl::optional<std::string> 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<std::string> path =
-      google::cloud::internal::GetEnv("GOOGLEBIGQUERYODBCINI");
+  auto path = google::cloud::internal::GetEnv("GOOGLEBIGQUERYODBCINI");
   if (path) {
     return *path;
   }