snapshot_update.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 <functional>
#include <memory>
#include <optional>
#include <span>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
 
#include "iceberg/iceberg_export.h"
#include "iceberg/result.h"
#include "iceberg/snapshot.h"
#include "iceberg/type_fwd.h"
#include "iceberg/update/pending_update.h"
#include "iceberg/util/executor.h"
 
namespace iceberg {
 
class ICEBERG_EXPORT SnapshotUpdate : public PendingUpdate {
 public:
  struct ApplyResult {
    std::shared_ptr<Snapshot> snapshot;
    std::string target_branch;
    bool stage_only = false;
  };
 
  ~SnapshotUpdate() override;
 
  Kind kind() const override { return Kind::kUpdateSnapshot; }
  bool IsRetryable() const override { return true; }
 
  auto& DeleteWith(this auto& self,
                   std::function<Status(const std::string&)> delete_func) {
    if (self.delete_func_) {
      return self.AddError(ErrorKind::kInvalidArgument,
                           "Cannot set delete callback more than once");
    }
    self.delete_func_ = std::move(delete_func);
    return self;
  }
 
  auto& StageOnly(this auto& self) {
    self.stage_only_ = true;
    return self;
  }
 
  auto& ScanManifestsWith(this auto& self, Executor& executor) {
    self.plan_executor_ = std::ref(executor);
    return self;
  }
 
  auto& ToBranch(this auto& self, const std::string& branch) {
    if (branch.empty()) [[unlikely]] {
      return self.AddError(ErrorKind::kInvalidArgument, "Branch name cannot be empty");
    }
 
    if (auto ref_it = self.base().refs.find(branch); ref_it != self.base().refs.end()) {
      if (ref_it->second->type() != SnapshotRefType::kBranch) {
        return self.AddError(ErrorKind::kInvalidArgument,
                             "{} is a tag, not a branch. Tags cannot be targets for "
                             "producing snapshots",
                             branch);
      }
    }
 
    self.target_branch_ = branch;
    return self;
  }
 
  auto& Set(this auto& self, const std::string& property, const std::string& value) {
    static_cast<SnapshotUpdate&>(self).SetSummaryProperty(property, value);
    return self;
  }
 
  auto& WriteManifestsWith(this auto& self, Executor& executor, int32_t parallelism) {
    if (parallelism <= 0) [[unlikely]] {
      return self.AddError(
          ErrorKind::kInvalidArgument,
          "Manifest write parallelism must be greater than 0, but was: {}", parallelism);
    }
 
    self.write_manifest_executor_ = std::ref(executor);
    self.write_manifest_parallelism_ = parallelism;
    return self;
  }
 
  Result<ApplyResult> Apply();
 
  Status Finalize(Result<const TableMetadata*> commit_result) override;
 
 protected:
  struct ContentFileWithSequenceNumber {
    std::shared_ptr<DataFile> file;
    std::optional<int64_t> data_sequence_number;
  };
 
  explicit SnapshotUpdate(std::shared_ptr<TransactionContext> ctx);
 
  Result<std::vector<ManifestFile>> WriteDataManifests(
      std::span<const std::shared_ptr<DataFile>> files,
      const std::shared_ptr<PartitionSpec>& spec,
      std::optional<int64_t> data_sequence_number = std::nullopt);
 
  Result<std::vector<ManifestFile>> WriteDeleteManifests(
      std::span<const ContentFileWithSequenceNumber> files,
      const std::shared_ptr<PartitionSpec>& spec);
 
  const std::string& target_branch() const { return target_branch_; }
  bool can_inherit_snapshot_id() const { return can_inherit_snapshot_id_; }
  const std::string& commit_uuid() const { return commit_uuid_; }
  int32_t manifest_count() const {
    return manifest_count_.load(std::memory_order_relaxed);
  }
  int32_t attempt() const { return attempt_.load(std::memory_order_relaxed); }
  int64_t target_manifest_size_bytes() const { return target_manifest_size_bytes_; }
 
  virtual Status CleanUncommitted(const std::unordered_set<std::string>& committed) = 0;
 
  virtual std::string operation() = 0;
 
  virtual Status Validate(const TableMetadata& current_metadata,
                          const std::shared_ptr<Snapshot>& snapshot) {
    return {};
  };
 
  virtual Result<std::vector<ManifestFile>> Apply(
      const TableMetadata& metadata_to_update,
      const std::shared_ptr<Snapshot>& snapshot) = 0;
 
  virtual std::unordered_map<std::string, std::string> Summary() = 0;
 
  virtual void SetSummaryProperty(const std::string& property, const std::string& value);
 
  virtual bool CleanupAfterCommit() const { return true; }
 
  int64_t SnapshotId();
 
  Status DeleteFile(const std::string& path);
 
  std::string ManifestPath();
  std::string ManifestListPath();
  SnapshotSummaryBuilder& summary_builder() { return summary_; }
  SnapshotSummaryBuilder BuildManifestCountSummary(
      std::span<const ManifestFile> manifests, int32_t replaced_manifests_count);
 
 private:
  Result<std::unordered_map<std::string, std::string>> ComputeSummary(
      const TableMetadata& previous);
 
  Status CleanAll();
 
 protected:
  SnapshotSummaryBuilder summary_;
 
 private:
  const bool can_inherit_snapshot_id_{true};
  const std::string commit_uuid_;
  OptionalExecutor write_manifest_executor_;
  int32_t write_manifest_parallelism_{1};
  std::atomic<int32_t> manifest_count_{0};
  std::atomic<int32_t> attempt_{0};
  std::vector<std::string> manifest_lists_;
  const int64_t target_manifest_size_bytes_;
  std::optional<int64_t> snapshot_id_;
  OptionalExecutor plan_executor_;
  bool stage_only_{false};
  std::function<Status(const std::string&)> delete_func_;
  std::string target_branch_{SnapshotRef::kMainBranch};
  std::shared_ptr<Snapshot> staged_snapshot_;
};
 
}  // namespace iceberg