utils.h

// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef CPP_BIGQUERY_ODBC_GOOGLE_CLOUD_ODBC_BQ_CLIENT_INTERFACE_UTILS_H
#define CPP_BIGQUERY_ODBC_GOOGLE_CLOUD_ODBC_BQ_CLIENT_INTERFACE_UTILS_H

#include "google/cloud/odbc/internal/sql_state_constants.h"
#include "google/cloud/odbc/internal/status_record_or.h"
#include "google/cloud/internal/backoff_policy.h"
#include "absl/log/log.h"
#include <cctype>
#include <cstdint>
#include <iomanip>
#include <regex>
#include <sstream>
#include <string>
#include <thread>

namespace google::cloud::odbc_bigquery_client_interface {
struct MaxRetriesOption {
  using Type = int;
};

// Page size requested from paginated metadata list APIs (projects.list,
// datasets.list, tables.list). When maxResults is unset the BigQuery REST API
// returns small pages (50), which multiplies sequential HTTP round trips
constexpr std::int32_t kMetadataPageSize = 1000;

// URL-encodes a value for safe use as a single path segment.
inline std::string UrlEncodeSegment(std::string const& value) {
  std::ostringstream os;
  os << std::uppercase << std::hex;
  for (unsigned char c : value) {
    if (std::isalnum(c) || c == '-' || c == '_' || c == '~') {
      os << static_cast<char>(c);
    } else {
      os << '%' << std::setw(2) << std::setfill('0') << static_cast<int>(c);
    }
  }
  return os.str();
}

template <typename Functor>
auto RetryLoop(Functor&& functor, std::string const& operation_name,
               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;

  google::cloud::internal::ExponentialBackoffPolicy backoff_policy(
      std::chrono::milliseconds(initial_delay_ms),
      std::chrono::milliseconds(max_delay_ms), backoff_multiplier);

  while (attempt <= max_retries) {
    response = functor();

    if (response.ok()) {
      LOG(INFO) << operation_name << " succeeded on attempt " << attempt;
      return response;
    }

    auto code = response.status().code();
    std::string message = response.status().message();

    bool is_rate_limit =
        (code == google::cloud::StatusCode::kPermissionDenied &&
         absl::StrContains(message, "Exceeded rate limits"));

    if ((code != google::cloud::StatusCode::kDeadlineExceeded &&
         !is_rate_limit)) {
      LOG(WARNING) << operation_name
                   << " failed permanently: " << response.status();
      return response;
    }

    auto delay = backoff_policy.OnCompletion();
    LOG(WARNING)
        << operation_name << " failed (attempt " << attempt
        << "): " << response.status() << " -- retrying after "
        << std::chrono::duration_cast<std::chrono::milliseconds>(delay).count()
        << "ms";

    std::this_thread::sleep_for(delay);
    ++attempt;
  }

  return response;
}

inline google::cloud::odbc_internal::StatusRecordOr<std::string>
ParsePartnerToken(std::string const& raw_token) {
  if (raw_token.empty()) {
    return std::string("");
  }
  std::regex pattern(R"(\(\s*(GPN:[^;]*?)\s*(?:;\s*([^)]*?))?\s*\))");
  std::smatch match;
  if (std::regex_search(raw_token, match, pattern)) {
    std::string gpn_part = match[1].str();
    std::string env_part = match[2].str();
    std::string partner_token = " (";
    partner_token += gpn_part;
    if (!env_part.empty()) {
      partner_token += "; ";
      partner_token += env_part;
    }
    partner_token += ")";
    return partner_token;
  }
  return google::cloud::odbc_internal::StatusRecord{
      google::cloud::odbc_internal::SQLStates::k_HY024(),
      "Invalid PartnerToken format."};
}

}  // namespace google::cloud::odbc_bigquery_client_interface
#endif  // CPP_BIGQUERY_ODBC_GOOGLE_CLOUD_ODBC_BQ_CLIENT_INTERFACE_UTILS_H