src/kudu/tools/table_scanner.h

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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
//
//   http://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.

#pragma once

#include <atomic>
#include <cstdint>
#include <functional>
#include <iosfwd>
#include <optional>
#include <memory>
#include <string>
#include <vector>

#include "kudu/client/client.h"
#include "kudu/client/scan_batch.h"
#include "kudu/client/shared_ptr.h" // IWYU pragma: keep
#include "kudu/util/mutex.h"
#include "kudu/util/status.h"
#include "kudu/util/threadpool.h"

namespace kudu {

namespace client {
class KuduSchema;
} // namespace client

namespace tools {

// This class is not thread-safe.
class TableScanner {
 public:
  TableScanner(client::sp::shared_ptr<client::KuduClient> client,
               std::string table_name,
               std::optional<client::sp::shared_ptr<client::KuduClient>> dst_client =
                   std::nullopt,
               std::optional<std::string> dst_table_name = std::nullopt);

  // Set output stream of this tool, or disable output if not set.
  // 'out' must remain valid for the lifetime of this class.
  void SetOutput(std::ostream* out);

  // Set read mode, see KuduScanner::SetReadMode().
  void SetReadMode(client::KuduScanner::ReadMode mode);

  // Set replica selection for scan operations.
  Status SetReplicaSelection(const std::string& selection);

  // Set the size for scan result batch size, in bytes. A negative value has
  // the semantics of relying on the server-side default: see the
  // --scanner_default_batch_size_bytes flag.
  void SetScanBatchSize(int32_t scan_batch_size);

  Status StartScan();
  Status StartCopy();

  uint64_t TotalScannedCount() const {
    return total_count_;
  }

 private:
  enum class WorkType {
    kScan,
    kCopy
  };

  static Status AddRow(
      const client::sp::shared_ptr<client::KuduTable>& table,
      const client::KuduSchema& table_schema,
      const client::KuduScanBatch::RowPtr& src_row,
      const client::sp::shared_ptr<client::KuduSession>& session);

  // Convert replica selection from string into the KuduClient::ReplicaSelection
  // enumerator.
  static Status ParseReplicaSelection(
      const std::string& selection_str,
      client::KuduClient::ReplicaSelection* selection);

  Status StartWork(WorkType type);
  Status ScanData(const std::vector<client::KuduScanToken*>& tokens,
                  const std::function<void(const client::KuduScanBatch& batch)>& cb);
  void ScanTask(const std::vector<client::KuduScanToken*>& tokens,
                Status* thread_status);
  void CopyTask(const std::vector<client::KuduScanToken*>& tokens,
                Status* thread_status);

  std::atomic<uint64_t> total_count_;
  std::optional<client::KuduScanner::ReadMode> mode_;
  client::sp::shared_ptr<client::KuduClient> client_;
  std::string table_name_;
  client::KuduClient::ReplicaSelection replica_selection_;
  std::optional<client::sp::shared_ptr<client::KuduClient>> dst_client_;
  std::optional<std::string> dst_table_name_;
  int32_t scan_batch_size_;
  std::unique_ptr<ThreadPool> thread_pool_;

  // Protects output to 'out_' so that rows don't get interleaved.
  Mutex output_lock_;
  std::ostream* out_;
};

} // namespace tools
} // namespace kudu