Skip to content

Commit c03f043

Browse files
zhjwpkuzhjwpku
committed
feat: add UpdateSnapshotReference
1 parent d07b756 commit c03f043

File tree

11 files changed

+459
-3
lines changed

11 files changed

+459
-3
lines changed

‎src/iceberg/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ set(ICEBERG_SOURCES
8585
update/update_properties.cc
8686
update/update_schema.cc
8787
update/update_sort_order.cc
88+
update/update_snapshot_reference.cc
8889
util/bucket_util.cc
8990
util/content_file_util.cc
9091
util/conversions.cc

‎src/iceberg/meson.build‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ iceberg_sources = files(
105105
'update/update_partition_spec.cc',
106106
'update/update_properties.cc',
107107
'update/update_schema.cc',
108+
'update/update_snapshot_reference.cc',
108109
'update/update_sort_order.cc',
109110
'util/bucket_util.cc',
110111
'util/content_file_util.cc',

‎src/iceberg/transaction.cc‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
#include "iceberg/update/update_partition_spec.h"
3333
#include "iceberg/update/update_properties.h"
3434
#include "iceberg/update/update_schema.h"
35+
#include "iceberg/update/update_snapshot_reference.h"
3536
#include "iceberg/update/update_sort_order.h"
3637
#include "iceberg/util/checked_cast.h"
3738
#include "iceberg/util/macros.h"
@@ -113,6 +114,24 @@ Status Transaction::Apply(PendingUpdate& update) {
113114
metadata_builder_->SetCurrentSchema(std::move(result.schema),
114115
result.new_last_column_id);
115116
} break;
117+
case PendingUpdate::Kind::kUpdateSnapshotReference: {
118+
auto& update_ref = internal::checked_cast<UpdateSnapshotReference&>(update);
119+
ICEBERG_ASSIGN_OR_RAISE(auto updated_refs, update_ref.Apply());
120+
const auto& current_refs = current().refs;
121+
// Identify references which have been removed
122+
for (const auto& [name, ref] : current_refs) {
123+
if (updated_refs.find(name) == updated_refs.end()) {
124+
metadata_builder_->RemoveRef(name);
125+
}
126+
}
127+
// Identify references which have been created or updated.
128+
for (const auto& [name, ref] : updated_refs) {
129+
auto current_it = current_refs.find(name);
130+
if (current_it == current_refs.end() || *current_it->second != *ref) {
131+
metadata_builder_->SetRef(name, ref);
132+
}
133+
}
134+
} break;
116135
default:
117136
return NotSupported("Unsupported pending update: {}",
118137
static_cast<int32_t>(update.kind()));
@@ -193,4 +212,12 @@ Result<std::shared_ptr<UpdateSchema>> Transaction::NewUpdateSchema() {
193212
return update_schema;
194213
}
195214

215+
Result<std::shared_ptr<UpdateSnapshotReference>>
216+
Transaction::NewUpdateSnapshotReference() {
217+
ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<UpdateSnapshotReference> update_ref,
218+
UpdateSnapshotReference::Make(shared_from_this()));
219+
ICEBERG_RETURN_UNEXPECTED(AddUpdate(update_ref));
220+
return update_ref;
221+
}
222+
196223
} // namespace iceberg

‎src/iceberg/transaction.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,10 @@ class ICEBERG_EXPORT Transaction : public std::enable_shared_from_this<Transacti
7272
/// changes.
7373
Result<std::shared_ptr<UpdateSchema>> NewUpdateSchema();
7474

75+
/// \brief Create a new UpdateSnapshotReference to update snapshot references (branches
76+
/// and tags) and commit the changes.
77+
Result<std::shared_ptr<UpdateSnapshotReference>> NewUpdateSnapshotReference();
78+
7579
private:
7680
Transaction(std::shared_ptr<Table> table, Kind kind, bool auto_commit,
7781
std::unique_ptr<TableMetadataBuilder> metadata_builder);

‎src/iceberg/type_fwd.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,7 @@ class UpdatePartitionSpec;
192192
class UpdateProperties;
193193
class UpdateSchema;
194194
class UpdateSortOrder;
195+
class UpdateSnapshotReference;
195196

196197
/// ----------------------------------------------------------------------------
197198
/// TODO: Forward declarations below are not added yet.

‎src/iceberg/update/meson.build‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ install_headers(
2020
'pending_update.h',
2121
'update_partition_spec.h',
2222
'update_schema.h',
23+
'update_snapshot_reference.h',
2324
'update_sort_order.h',
2425
'update_properties.h',
2526
],

‎src/iceberg/update/pending_update.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ class ICEBERG_EXPORT PendingUpdate : public ErrorCollector {
4545
kUpdatePartitionSpec,
4646
kUpdateProperties,
4747
kUpdateSchema,
48+
kUpdateSnapshotReference,
4849
kUpdateSortOrder,
4950
};
5051

‎src/iceberg/update/update_snapshot_reference.cc‎

Lines changed: 230 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,230 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
#include "iceberg/update/update_snapshot_reference.h"
21+
22+
#include <memory>
23+
#include <optional>
24+
#include <string>
25+
#include <unordered_map>
26+
27+
#include "iceberg/result.h"
28+
#include "iceberg/snapshot.h"
29+
#include "iceberg/table_metadata.h"
30+
#include "iceberg/transaction.h"
31+
#include "iceberg/util/error_collector.h"
32+
#include "iceberg/util/macros.h"
33+
#include "iceberg/util/snapshot_util_internal.h"
34+
35+
namespace iceberg {
36+
37+
Result<std::shared_ptr<UpdateSnapshotReference>> UpdateSnapshotReference::Make(
38+
std::shared_ptr<Transaction> transaction) {
39+
ICEBERG_PRECHECK(transaction != nullptr,
40+
"Cannot create UpdateSnapshotReference without a transaction");
41+
return std::shared_ptr<UpdateSnapshotReference>(
42+
new UpdateSnapshotReference(std::move(transaction)));
43+
}
44+
45+
UpdateSnapshotReference::UpdateSnapshotReference(std::shared_ptr<Transaction> transaction)
46+
: PendingUpdate(std::move(transaction)) {
47+
// Initialize updated_refs_ with current refs from base metadata
48+
const auto& base_refs = transaction_->current().refs;
49+
for (const auto& [name, ref] : base_refs) {
50+
updated_refs_[name] = ref;
51+
}
52+
}
53+
54+
UpdateSnapshotReference::~UpdateSnapshotReference() = default;
55+
56+
UpdateSnapshotReference& UpdateSnapshotReference::CreateBranch(const std::string& name,
57+
int64_t snapshot_id) {
58+
ICEBERG_BUILDER_CHECK(!name.empty(), "Branch name cannot be empty");
59+
// TODO(anyone): use MakeBranch after #408 got merged
60+
auto branch = std::make_shared<SnapshotRef>(
61+
SnapshotRef{.snapshot_id = snapshot_id, .retention = SnapshotRef::Branch{}});
62+
auto [_, inserted] = updated_refs_.emplace(name, std::move(branch));
63+
ICEBERG_BUILDER_CHECK(inserted, "Ref '{}' already exists", name);
64+
return *this;
65+
}
66+
67+
UpdateSnapshotReference& UpdateSnapshotReference::CreateTag(const std::string& name,
68+
int64_t snapshot_id) {
69+
ICEBERG_BUILDER_CHECK(!name.empty(), "Tag name cannot be empty");
70+
auto tag = std::make_shared<SnapshotRef>(
71+
SnapshotRef{.snapshot_id = snapshot_id, .retention = SnapshotRef::Tag{}});
72+
auto [_, inserted] = updated_refs_.emplace(name, std::move(tag));
73+
ICEBERG_BUILDER_CHECK(inserted, "Ref '{}' already exists", name);
74+
return *this;
75+
}
76+
77+
UpdateSnapshotReference& UpdateSnapshotReference::RemoveBranch(const std::string& name) {
78+
ICEBERG_BUILDER_CHECK(!name.empty(), "Branch name cannot be empty");
79+
ICEBERG_BUILDER_CHECK(name != SnapshotRef::kMainBranch, "Cannot remove main branch");
80+
auto it = updated_refs_.find(name);
81+
ICEBERG_BUILDER_CHECK(it != updated_refs_.end(), "Branch does not exist: {}", name);
82+
ICEBERG_BUILDER_CHECK(it->second->type() == SnapshotRefType::kBranch,
83+
"Ref '{}' is a tag not a branch", name);
84+
updated_refs_.erase(it);
85+
return *this;
86+
}
87+
88+
UpdateSnapshotReference& UpdateSnapshotReference::RemoveTag(const std::string& name) {
89+
ICEBERG_BUILDER_CHECK(!name.empty(), "Tag name cannot be empty");
90+
auto it = updated_refs_.find(name);
91+
ICEBERG_BUILDER_CHECK(it != updated_refs_.end(), "Tag does not exist: {}", name);
92+
ICEBERG_BUILDER_CHECK(it->second->type() == SnapshotRefType::kTag,
93+
"Ref '{}' is a branch not a tag", name);
94+
updated_refs_.erase(it);
95+
return *this;
96+
}
97+
98+
UpdateSnapshotReference& UpdateSnapshotReference::RenameBranch(
99+
const std::string& name, const std::string& new_name) {
100+
ICEBERG_BUILDER_CHECK(!name.empty(), "Branch to rename cannot be empty");
101+
ICEBERG_BUILDER_CHECK(!new_name.empty(), "New branch name cannot be empty");
102+
ICEBERG_BUILDER_CHECK(name != SnapshotRef::kMainBranch, "Cannot rename main branch");
103+
auto it = updated_refs_.find(name);
104+
ICEBERG_BUILDER_CHECK(it != updated_refs_.end(), "Branch does not exist: {}", name);
105+
ICEBERG_BUILDER_CHECK(it->second->type() == SnapshotRefType::kBranch,
106+
"Ref '{}' is a tag not a branch", name);
107+
auto [_, inserted] = updated_refs_.emplace(new_name, it->second);
108+
ICEBERG_BUILDER_CHECK(inserted, "Ref '{}' already exists", new_name);
109+
updated_refs_.erase(it);
110+
return *this;
111+
}
112+
113+
UpdateSnapshotReference& UpdateSnapshotReference::ReplaceBranch(const std::string& name,
114+
int64_t snapshot_id) {
115+
ICEBERG_BUILDER_CHECK(!name.empty(), "Branch name cannot be empty");
116+
auto it = updated_refs_.find(name);
117+
ICEBERG_BUILDER_CHECK(it != updated_refs_.end(), "Branch does not exist: {}", name);
118+
ICEBERG_BUILDER_CHECK(it->second->type() == SnapshotRefType::kBranch,
119+
"Ref '{}' is a tag not a branch", name);
120+
it->second->snapshot_id = snapshot_id;
121+
return *this;
122+
}
123+
124+
UpdateSnapshotReference& UpdateSnapshotReference::ReplaceBranch(const std::string& from,
125+
const std::string& to) {
126+
return ReplaceBranchInternal(from, to, false);
127+
}
128+
129+
UpdateSnapshotReference& UpdateSnapshotReference::FastForward(const std::string& from,
130+
const std::string& to) {
131+
return ReplaceBranchInternal(from, to, true);
132+
}
133+
134+
UpdateSnapshotReference& UpdateSnapshotReference::ReplaceBranchInternal(
135+
const std::string& from, const std::string& to, bool fast_forward) {
136+
ICEBERG_BUILDER_CHECK(!from.empty(), "Branch to update cannot be empty");
137+
ICEBERG_BUILDER_CHECK(!to.empty(), "Destination ref cannot be empty");
138+
auto to_it = updated_refs_.find(to);
139+
ICEBERG_BUILDER_CHECK(to_it != updated_refs_.end(), "Ref does not exist: {}", to);
140+
141+
auto from_it = updated_refs_.find(from);
142+
if (from_it == updated_refs_.end()) {
143+
// Create branch if it doesn't exist
144+
return CreateBranch(from, to_it->second->snapshot_id);
145+
}
146+
147+
ICEBERG_BUILDER_CHECK(from_it->second->type() == SnapshotRefType::kBranch,
148+
"Ref '{}' is a tag not a branch", from);
149+
150+
// Nothing to replace if snapshot IDs are the same
151+
if (to_it->second->snapshot_id == from_it->second->snapshot_id) {
152+
return *this;
153+
}
154+
155+
if (fast_forward) {
156+
// Validate that current branch snapshot (from) is an ancestor of target snapshot (to)
157+
// This ensures we can fast-forward from 'from' to 'to'
158+
// TODO(anyone): use base() after #408
159+
const auto& base_metadata = transaction_->current();
160+
auto lookup = [&base_metadata](int64_t id) -> Result<std::shared_ptr<Snapshot>> {
161+
return base_metadata.SnapshotById(id);
162+
};
163+
ICEBERG_BUILDER_ASSIGN_OR_RETURN(
164+
bool target_is_ancestor,
165+
SnapshotUtil::IsAncestorOf(to_it->second->snapshot_id,
166+
from_it->second->snapshot_id, lookup));
167+
ICEBERG_BUILDER_CHECK(target_is_ancestor,
168+
"Cannot fast-forward: {} is not an ancestor of {}", from, to);
169+
}
170+
171+
from_it->second->snapshot_id = to_it->second->snapshot_id;
172+
return *this;
173+
}
174+
175+
UpdateSnapshotReference& UpdateSnapshotReference::ReplaceTag(const std::string& name,
176+
int64_t snapshot_id) {
177+
ICEBERG_BUILDER_CHECK(!name.empty(), "Tag name cannot be empty");
178+
auto it = updated_refs_.find(name);
179+
ICEBERG_BUILDER_CHECK(it != updated_refs_.end(), "Tag does not exist: {}", name);
180+
ICEBERG_BUILDER_CHECK(it->second->type() == SnapshotRefType::kTag,
181+
"Ref '{}' is a branch not a tag", name);
182+
it->second->snapshot_id = snapshot_id;
183+
return *this;
184+
}
185+
186+
UpdateSnapshotReference& UpdateSnapshotReference::SetMinSnapshotsToKeep(
187+
const std::string& name, int32_t min_snapshots_to_keep) {
188+
ICEBERG_BUILDER_CHECK(!name.empty(), "Branch name cannot be empty");
189+
auto it = updated_refs_.find(name);
190+
ICEBERG_BUILDER_CHECK(it != updated_refs_.end(), "Branch does not exist: {}", name);
191+
ICEBERG_BUILDER_CHECK(it->second->type() == SnapshotRefType::kBranch,
192+
"Ref '{}' is a tag not a branch", name);
193+
// TODO(anyone): use interface added in #408
194+
std::get<SnapshotRef::Branch>(it->second->retention).min_snapshots_to_keep =
195+
min_snapshots_to_keep;
196+
return *this;
197+
}
198+
199+
UpdateSnapshotReference& UpdateSnapshotReference::SetMaxSnapshotAgeMs(
200+
const std::string& name, int64_t max_snapshot_age_ms) {
201+
ICEBERG_BUILDER_CHECK(!name.empty(), "Branch name cannot be empty");
202+
auto it = updated_refs_.find(name);
203+
ICEBERG_BUILDER_CHECK(it != updated_refs_.end(), "Branch does not exist: {}", name);
204+
ICEBERG_BUILDER_CHECK(it->second->type() == SnapshotRefType::kBranch,
205+
"Ref '{}' is a tag not a branch", name);
206+
std::get<SnapshotRef::Branch>(it->second->retention).max_snapshot_age_ms =
207+
max_snapshot_age_ms;
208+
return *this;
209+
}
210+
211+
UpdateSnapshotReference& UpdateSnapshotReference::SetMaxRefAgeMs(const std::string& name,
212+
int64_t max_ref_age_ms) {
213+
ICEBERG_BUILDER_CHECK(!name.empty(), "Reference name cannot be empty");
214+
auto it = updated_refs_.find(name);
215+
ICEBERG_BUILDER_CHECK(it != updated_refs_.end(), "Ref does not exist: {}", name);
216+
if (it->second->type() == SnapshotRefType::kBranch) {
217+
std::get<SnapshotRef::Branch>(it->second->retention).max_ref_age_ms = max_ref_age_ms;
218+
} else {
219+
std::get<SnapshotRef::Tag>(it->second->retention).max_ref_age_ms = max_ref_age_ms;
220+
}
221+
return *this;
222+
}
223+
224+
Result<std::unordered_map<std::string, std::shared_ptr<SnapshotRef>>>
225+
UpdateSnapshotReference::Apply() {
226+
ICEBERG_RETURN_UNEXPECTED(CheckErrors());
227+
return updated_refs_;
228+
}
229+
230+
} // namespace iceberg

0 commit comments

Comments
 (0)