impl(bq_driver): add max_retries impl

Author: NeerajDwivedii

google/cloud/odbc/bq_client_interface/jobs.cc

@@ -120,8 +120,10 @@ StatusRecordOr<Job> GetJob(JobClient& job_client, std::string const& project_id,
   get_job_request.set_job_id(job_id);
   get_job_request.set_location(location);
   LOG(INFO) << "GetJob:: Request body: " << get_job_request.DebugString("");
-  auto response = RetryLoop(
-      [&] { return job_client.GetJob(get_job_request, options); }, "GetJob");
+  auto max_retries = options.get<MaxRetriesOption>();
+  auto response =
+      RetryLoop([&] { return job_client.GetJob(get_job_request, options); },
+                "GetJob", max_retries);
 
   if (!response.ok()) {
     LOG(WARNING) << "GetJob:: Request failed: " << response.status();
@@ -248,8 +250,10 @@ StatusRecordOr<Job> InsertJob(JobClient& job_client,
   request.set_json_filter_keys(CreateKeysToFilterOut(job));
 
   LOG(INFO) << "InsertJob:: Request body: " << request.DebugString("");
-  auto response = RetryLoop(
-      [&] { return job_client.InsertJob(request, options); }, "InsertJob");
+  auto max_retries = options.get<MaxRetriesOption>();
+  auto response =
+      RetryLoop([&] { return job_client.InsertJob(request, options); },
+                "InsertJob", max_retries);
 
   if (!response.ok()) {
     LOG(WARNING) << "InsertJob:: Request failed: " << response.status();
@@ -277,8 +281,10 @@ StatusRecordOr<Job> CancelJob(JobClient& job_client,
     request.set_location(location);
   }
   LOG(INFO) << "CancelJob:: Request body: " << request.DebugString("");
-  auto response = RetryLoop(
-      [&] { return job_client.CancelJob(request, options); }, "CancelJob");
+  auto max_retries = options.get<MaxRetriesOption>();
+  auto response =
+      RetryLoop([&] { return job_client.CancelJob(request, options); },
+                "CancelJob", max_retries);
 
   if (!response.ok()) {
     LOG(WARNING) << "CancelJob:: Request failed: " << response.status();
@@ -301,8 +307,10 @@ StatusRecordOr<PostQueryResults> Query(JobClient& job_client,
   post_query_request.set_json_filter_keys(CreateKeysToFilterOut(query_request));
 
   LOG(INFO) << "Query:: Request body: " << post_query_request.DebugString("");
-  auto response = RetryLoop(
-      [&] { return job_client.Query(post_query_request, options); }, "Query");
+  auto max_retries = options.get<MaxRetriesOption>();
+  auto response =
+      RetryLoop([&] { return job_client.Query(post_query_request, options); },
+                "Query", max_retries);
 
   if (!response.ok()) {
     LOG(WARNING) << "Query:: Request failed: " << response.status();
@@ -406,11 +414,12 @@ StatusRecordOr<GetQueryResults> FilterQueryResults(
 
   LOG(INFO) << "FilterQueryResults:: Request body: "
             << get_query_results_request.DebugString("");
+  auto max_retries = options.get<MaxRetriesOption>();
   auto response = RetryLoop(
       [&] {
         return job_client.QueryResults(get_query_results_request, options);
       },
-      "QueryResults");
+      "QueryResults", max_retries);
 
   if (!response.ok()) {
     LOG(WARNING) << "FilterQueryResults:: Request failed: "

google/cloud/odbc/bq_client_interface/tables.cc

@@ -47,8 +47,10 @@ StatusRecordOr<Table> GetTable(TableClient& table_client,
 
   LOG(INFO) << "GetTable:: Request body: " << request.DebugString("");
 
-  auto response = RetryLoop(
-      [&] { return table_client.GetTable(request, options); }, "GetTable");
+  auto max_retries = options.get<MaxRetriesOption>();
+  auto response =
+      RetryLoop([&] { return table_client.GetTable(request, options); },
+                "GetTable", max_retries);
 
   if (!response.ok()) {
     LOG(WARNING) << "GetTable:: Request failed: " << response.status();

google/cloud/odbc/bq_client_interface/utils.h

@@ -22,6 +22,11 @@
 #include <string>
 #include <thread>
 
+namespace google::cloud::odbc_bigquery_client_interface {
+struct MaxRetriesOption {
+  using Type = int;
+};
+
 // URL-encodes a value for safe use as a single path segment.
 inline std::string UrlEncodeSegment(std::string const& value) {
   std::ostringstream os;
@@ -38,11 +43,10 @@ inline std::string UrlEncodeSegment(std::string const& value) {
 
 template <typename Functor>
 auto RetryLoop(Functor&& functor, std::string const& operation_name,
-               int max_retries = 6, int initial_delay_ms = 500,
+               int max_retries, int initial_delay_ms = 500,
                int max_delay_ms = 20000,
                double backoff_multiplier = 2.0) -> decltype(functor()) {
   int attempt = 0;
-
   using ReturnType = decltype(functor());
   ReturnType response;
 
@@ -85,5 +89,5 @@ auto RetryLoop(Functor&& functor, std::string const& operation_name,
 
   return response;
 }
-
+}  // namespace google::cloud::odbc_bigquery_client_interface
 #endif  // CPP_BIGQUERY_ODBC_GOOGLE_CLOUD_ODBC_BQ_CLIENT_INTERFACE_UTILS_H

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

@@ -62,6 +62,7 @@ std::string AdvanceOptions::allow_large_results_;
 std::string AdvanceOptions::use_default_large_results_;
 std::string AdvanceOptions::encryption_type_ = kDefaultEncryptionType;
 std::string AdvanceOptions::max_threads_ = std::to_string(kDefaultMaxThreads);
+std::string AdvanceOptions::max_retries_ = std::to_string(kDefaultMaxRetries);
 
 std::string const kLanguageDialect = "SQLDialect";
 std::string const kLargeResultsDatasetId = "LargeResultsDatasetId";
@@ -82,6 +83,7 @@ std::string const kUseDefaultLargeResultsDataset =
     "UseDefaultLargeResultsDataset";
 std::string const kEncryptionType = "EncryptionType";
 std::string const kMaxThreads = "MaxThreads";
+std::string const kMaxRetries = "MaxRetries";
 
 // Control dimensions and positions
 int const kHeight = 20;
@@ -341,6 +343,21 @@ void AdvanceOptions::CreateAdditionalControls(HFONT h_font) {
       h_max_threads_edit, GWL_STYLE,
       GetWindowLongPtr(h_max_threads_edit, GWL_STYLE) | ES_RIGHT | ES_NUMBER);
 
+  // max retries
+  HWND h_max_retries_label =
+      CreateLabel(adv_hwnd, "Max Retries:", kXAxis, kYAxis + 390, kWidth * 7,
+                  kHeight, WS_VISIBLE | SS_LEFT);
+  SendMessage(h_max_retries_label, WM_SETFONT, (WPARAM)h_font, TRUE);
+  HWND h_max_retries_edit =
+      CreateEditBox(adv_hwnd, kinputComboBoxXAxis, kYAxis + 390, kEditBoxWidth,
+                    kEditBoxHeight, kIdcMaxRetriesEdit);
+  SendMessage(h_max_retries_edit, WM_SETFONT, (WPARAM)h_font, TRUE);
+  SetWindowSubclass(GetDlgItem(adv_hwnd, kIdcMaxRetriesEdit), InputSubclassProc,
+                    0, 0);
+  SetWindowText(h_max_retries_edit, max_retries_.c_str());
+  SetWindowLongPtr(
+      h_max_retries_edit, GWL_STYLE,
+      GetWindowLongPtr(h_max_retries_edit, GWL_STYLE) | ES_RIGHT | ES_NUMBER);
   // TODO(b/497725655): Enable UI feature after public release
   // HWND h_variables_checkbox = CreateCheckBox(
   //     adv_hwnd, "Use SQL_WVARCHAR instead of SQL_VARCHAR", kXAxis, kYAxis +
@@ -351,11 +368,11 @@ void AdvanceOptions::CreateAdditionalControls(HFONT h_font) {
   // SetWindowSubclass(GetDlgItem(adv_hwnd, kIdcVariableCheckbox),
   //                   CheckboxSubclassProc, 0, 0);
   HWND h_additional_projects_label =
-      CreateLabel(adv_hwnd, "Additional projects:", kXAxis, kYAxis + 390,
+      CreateLabel(adv_hwnd, "Additional projects:", kXAxis, kYAxis + 420,
                   kWidth * 5, kHeight, WS_VISIBLE | SS_LEFT);
   SendMessage(h_additional_projects_label, WM_SETFONT, (WPARAM)h_font, TRUE);
   HWND h_additional_projects_edit =
-      CreateScrollableEditBox(adv_hwnd, kXAxis, kYAxis + 410, kWidth + 380,
+      CreateScrollableEditBox(adv_hwnd, kXAxis, kYAxis + 440, kWidth + 380,
                               kHeight + 32, kIdcAdditionalProjectsEdit);
   SendMessage(h_additional_projects_edit, WM_SETFONT, (WPARAM)h_font, TRUE);
 
@@ -364,11 +381,11 @@ void AdvanceOptions::CreateAdditionalControls(HFONT h_font) {
                     InputSubclassProc, 0, 0);
 
   HWND h_query_properties_label =
-      CreateLabel(adv_hwnd, "Query properties:", kXAxis, kYAxis + 470,
+      CreateLabel(adv_hwnd, "Query properties:", kXAxis, kYAxis + 500,
                   kWidth * 5, kHeight, WS_VISIBLE | SS_LEFT);
   SendMessage(h_query_properties_label, WM_SETFONT, (WPARAM)h_font, TRUE);
   HWND h_query_properties_edit =
-      CreateScrollableEditBox(adv_hwnd, kXAxis, kYAxis + 485, kWidth + 385,
+      CreateScrollableEditBox(adv_hwnd, kXAxis, kYAxis + 520, kWidth + 385,
                               kHeight + 13, kIdcQueryPropertiesEdit);
   SendMessage(h_query_properties_edit, WM_SETFONT, (WPARAM)h_font, TRUE);
 
@@ -392,12 +409,12 @@ void AdvanceOptions::CreateAdditionalControls(HFONT h_font) {
 }
 
 void AdvanceOptions::CreateButtons(HFONT h_font) {
-  HWND h_ok_button = CreateButton(adv_hwnd, "OK", kOkButtonX + 2, kButtonY + 15,
+  HWND h_ok_button = CreateButton(adv_hwnd, "OK", kOkButtonX + 2, kButtonY + 38,
                                   kButtonWidth, kButtonHeight, kIdcOKButton);
   SendMessage(h_ok_button, WM_SETFONT, (WPARAM)h_font, TRUE);
 
   HWND h_cancel_button =
-      CreateButton(adv_hwnd, "Cancel", kCancelButtonX, kButtonY + 15,
+      CreateButton(adv_hwnd, "Cancel", kCancelButtonX, kButtonY + 38,
                    kButtonWidth, kButtonHeight, kIdcCancelButton);
   SendMessage(h_cancel_button, WM_SETFONT, (WPARAM)h_font, TRUE);
 }
@@ -565,6 +582,20 @@ LRESULT CALLBACK AdvanceOptions::AdvanceOptProc(HWND hwnd, UINT u_msg,
             ShowErrorWindow(hwnd, err_msg);
             return true;
           }
+          HWND h_max_retries_edit = GetDlgItem(hwnd, kIdcMaxRetriesEdit);
+          char max_retries_buff[256] = {0};
+          GetWindowText(h_max_retries_edit, max_retries_buff,
+                        sizeof(max_retries_buff));
+          if (isValidUint32(max_retries_buff)) {
+            max_retries_ = max_retries_buff;
+          } else {
+            std::string err_msg =
+                "Invalid number of max retries: Valid values are in range [0," +
+                std::to_string(UINT32_MAX) + "]";
+            ShowErrorWindow(hwnd, err_msg);
+            return true;
+          }
+
           HWND h_additional_projects_edit =
               GetDlgItem(hwnd, kIdcAdditionalProjectsEdit);
           char additional_projects_buffer[1024] = {0};
@@ -742,6 +773,7 @@ void AdvanceOptions::SetValues(Section const& attribute_map) {
   temp_expiration_ =
       GetValueOrDefault(attribute_map, kLargeResultsTempTableExpirationTime);
   max_threads_ = GetValueOrDefault(attribute_map, kMaxThreads);
+  max_retries_ = GetValueOrDefault(attribute_map, kMaxRetries);
   session_location_ = GetValueOrDefault(attribute_map, kSessionLocation);
   additional_projects_ = GetValueOrDefault(attribute_map, kAdditionalProjects);
   query_properties_ = GetValueOrDefault(attribute_map, kQueryProperties);
@@ -766,6 +798,7 @@ void AdvanceOptions::ResetToDefaults() {
   rows_per_block_ = kDefaultRowsPerBlock;
   default_string_length_ = kDefaultStringLength;
   max_threads_ = std::to_string(kDefaultMaxThreads);
+  max_retries_ = std::to_string(kDefaultMaxRetries);
   session_location_.clear();
   additional_projects_.clear();
   query_properties_.clear();
@@ -798,7 +831,7 @@ void AdvanceOptions::Show(HWND hwnd) {
   RegisterClass(&wc_adv);
 
   int window_width = 462;
-  int window_height = 620;
+  int window_height = 650;
   int screen_width = GetSystemMetrics(SM_CXSCREEN);
   int screen_height = GetSystemMetrics(SM_CYSCREEN);
   int x_pos = (screen_width - window_width) / 2;

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

@@ -42,6 +42,7 @@ static int const KIdcLargeResultHeader = 145;
 static int const kIdcEncryptionKeyComboBox = 146;
 static int const kIdcHyperlink2 = 147;
 static int const kIdcMaxThreadsEdit = 148;
+static int const kIdcMaxRetriesEdit = 149;
 
 class AdvanceOptions {
  public:
@@ -96,6 +97,7 @@ class AdvanceOptions {
     return use_default_large_results_;
   }
   inline std::string const& GetMaxThreads() const { return max_threads_; }
+  inline std::string const& GetMaxRetries() const { return max_retries_; }
   void SetValues(Section const& attributes_map);
   void ResetToDefaults();
 
@@ -124,6 +126,7 @@ class AdvanceOptions {
   static std::string use_default_large_results_;
   static std::string encryption_type_;
   static std::string max_threads_;
+  static std::string max_retries_;
 
   static LRESULT CALLBACK AdvanceOptProc(HWND hwnd, UINT uMsg, WPARAM w_param,
                                          LPARAM l_param);

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

@@ -83,7 +83,8 @@ TEST_F(AdvanceOptionsTest, SetValuesValidinput) {
                            {"AdditionalProjects", "projectA,projectB"},
                            {"QueryProperties", "property1=value1"},
                            {"HTAPI_ActivationThreshold", "10000"},
-                           {"MaxThreads", "10"}};
+                           {"MaxThreads", "10"},
+                           {"MaxRetries", "9"}};
 
   AdvanceOptions options;
   options.SetValues(attribute_map);
@@ -98,7 +99,8 @@ TEST_F(AdvanceOptionsTest, SetValuesValidinput) {
   EXPECT_EQ(options.GetAdditionalProjects(), "projectA,projectB");
   EXPECT_EQ(options.GetQueryProperties(), "property1=value1");
   EXPECT_EQ(options.GetActivationThreshold(), "10000");
-  EXPECT_EQ(options.GetMaxThreads(), "8");
+  EXPECT_EQ(options.GetMaxThreads(), "10");
+  EXPECT_EQ(options.GetMaxRetries(), "9");
 }
 TEST_F(AdvanceOptionsTest, SetValuesMissingkeys) {
   Section attribute_map = {
@@ -110,6 +112,7 @@ TEST_F(AdvanceOptionsTest, SetValuesMissingkeys) {
 
   EXPECT_EQ(options.GetLanguageDialect(), "GoogleSQL");
   EXPECT_EQ(options.GetMaxThreads(), "8");
+  EXPECT_EQ(options.GetMaxRetries(), "6");
   EXPECT_EQ(options.GetDatasetName(), "");
   EXPECT_EQ(options.GetEncryptionKey(), "");
   EXPECT_EQ(options.GetRowsPerBlock(), "");

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

@@ -134,6 +134,10 @@ void ConnectionHandle::SetUp(Section& dsn_section,
   dsn_.max_threads =
       !max_threads.empty() ? std::stoull(max_threads) : kDefaultMaxThreads;
 
+  std::string max_retries = dsn_section["MAXRETRIES"];
+  dsn_.max_retries =
+      !max_retries.empty() ? std::stoull(max_retries) : kDefaultMaxRetries;
+
   dsn_.pem_file = dsn_section["TRUSTEDCERTS"];
   dsn_.kms_key_name = dsn_section["KMSKEYNAME"];
   dsn_.session_location = dsn_section["SESSIONLOCATION"];

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

@@ -31,6 +31,7 @@ using google::cloud::odbc_bigquery_client_interface::ODBCBQClient;
 
 std::string const kDefaultDestDatasetId = "_bqodbc_temp_tables";
 std::string const kDefaultLargeResultsTableExpiration = "3600000";
+inline int const kDefaultMaxRetries = 6;
 
 // Details of authentication provided in the odbc.ini/Windows Registry
 struct Authentication {
@@ -59,6 +60,7 @@ struct Dsn {
   // TODO(jsrinnn): Remove this if it is not being used.
   std::string refresh_token;
   std::uint32_t max_threads;
+  std::uint32_t max_retries = kDefaultMaxRetries;
   bool is_bq_legacy_sql = false;
   bool is_job_creation_required = false;
   bool sessions_enabled = false;

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "google/cloud/odbc/bq_driver/internal/odbc_internal_commons.h"
+#include "google/cloud/odbc/bq_client_interface/utils.h"
 #include "google/cloud/odbc/bq_driver/internal/trace_utils.h"
 #include "google/cloud/odbc/bq_driver/internal/utils.h"
 #include <cmath>
@@ -41,6 +42,7 @@ using ::google::cloud::bigquery_v2_minimal_internal::TableFieldSchema;
 using ::google::cloud::bigquery_v2_minimal_internal::TableReference;
 #endif  // (!defined(_WIN32) || defined(_WIN64)) && !defined(NO_ARROW)
 using ::google::cloud::bigquery_v2_minimal_internal::TableSchema;
+using google::cloud::odbc_bigquery_client_interface::MaxRetriesOption;
 using ::google::cloud::odbc_internal::SQLStates;
 using ::google::cloud::odbc_internal::StatusRecord;
 using ::google::cloud::odbc_internal::StatusRecordOr;
@@ -808,12 +810,13 @@ StatusRecordOr<Job> CancelBQJob(ConnectionHandle& conn_handle,
   }
 
   Options options;
+  options.set<MaxRetriesOption>(conn_handle.GetDsn().max_retries);
   return bq_client->CancelJob(project_id, job_id, location, options);
 }
 
 StatusRecordOr<PostQueryResults> PostQueryWithoutResults(
     std::shared_ptr<ODBCBQClient> const& bq_client,
-    PostQueryRequest const& post_query_request) {
+    PostQueryRequest const& post_query_request, Options const& options) {
   if (!bq_client) {
     LOG(ERROR)
         << "PostQueryWithoutResults:: Invalid or null BQ Client within the "
@@ -824,7 +827,6 @@ StatusRecordOr<PostQueryResults> PostQueryWithoutResults(
   }
   // For now , we use default options.
   // We can set timeout here as needed later.
-  Options options;
   auto pq_status = bq_client->PostQuery(post_query_request, options);
   if (!pq_status) {
     LOG(ERROR) << "PostQueryWithoutResults::PostQuery:: "
@@ -843,8 +845,10 @@ StatusRecordOr<PostQueryResults> PostQueryWithoutResults(
     return StatusRecord{SQLStates::k_08S01(),
                         "Connection to the data source is broken"};
   }
-  auto pq_status =
-      PostQueryWithoutResults(conn_handle.GetClient(), post_query_request);
+  Options options;
+  options.set<MaxRetriesOption>(conn_handle.GetDsn().max_retries);
+  auto pq_status = PostQueryWithoutResults(conn_handle.GetClient(),
+                                           post_query_request, options);
   if (!pq_status) {
     return pq_status.GetStatusRecord();
   }

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

@@ -542,7 +542,8 @@ odbc_internal::StatusRecordOr<
 PostQueryWithoutResults(
     std::shared_ptr<ODBCBQClient> const& bq_client,
     google::cloud::bigquery_v2_minimal_internal::PostQueryRequest const&
-        post_query_request);
+        post_query_request,
+    Options const& options);
 
 odbc_internal::StatusRecordOr<
     google::cloud::bigquery_v2_minimal_internal::PostQueryResults>