feat: Refresh method support for table (#152)

Co-authored-by: shuxu.li <[email protected]>

Author: lishuxu

src/iceberg/catalog/in_memory_catalog.h

@@ -90,8 +90,8 @@ class ICEBERG_EXPORT InMemoryCatalog
       const TableIdentifier& identifier,
       const std::string& metadata_file_location) override;
 
-  std::unique_ptr<Catalog::TableBuilder> BuildTable(const TableIdentifier& identifier,
-                                                    const Schema& schema) const override;
+  std::unique_ptr<TableBuilder> BuildTable(const TableIdentifier& identifier,
+                                           const Schema& schema) const override;
 
  private:
   std::string catalog_name_;

src/iceberg/table.cc

@@ -21,16 +21,36 @@
 
 #include <algorithm>
 
+#include "iceberg/catalog.h"
 #include "iceberg/partition_spec.h"
 #include "iceberg/schema.h"
 #include "iceberg/sort_order.h"
 #include "iceberg/table_metadata.h"
 #include "iceberg/table_scan.h"
+#include "iceberg/util/macros.h"
 
 namespace iceberg {
 
 const std::string& Table::uuid() const { return metadata_->table_uuid; }
 
+Status Table::Refresh() {
+  if (!catalog_) {
+    return NotSupported("Refresh is not supported for table without a catalog");
+  }
+
+  ICEBERG_ASSIGN_OR_RAISE(auto refreshed_table, catalog_->LoadTable(identifier_));
+  if (metadata_location_ != refreshed_table->metadata_location_) {
+    metadata_ = std::move(refreshed_table->metadata_);
+    metadata_location_ = std::move(refreshed_table->metadata_location_);
+    io_ = std::move(refreshed_table->io_);
+
+    schemas_map_.reset();
+    partition_spec_map_.reset();
+    sort_orders_map_.reset();
+  }
+  return {};
+}
+
 Result<std::shared_ptr<Schema>> Table::schema() const { return metadata_->Schema(); }
 
 const std::shared_ptr<std::unordered_map<int32_t, std::shared_ptr<Schema>>>&

src/iceberg/table.h

@@ -57,6 +57,9 @@ class ICEBERG_EXPORT Table {
   /// \brief Returns the UUID of the table
   const std::string& uuid() const;
 
+  /// \brief Refresh the current table metadata
+  Status Refresh();
+
   /// \brief Return the schema for this table, return NotFoundError if not found
   Result<std::shared_ptr<Schema>> schema() const;
 
@@ -116,7 +119,7 @@ class ICEBERG_EXPORT Table {
  private:
   const TableIdentifier identifier_;
   std::shared_ptr<TableMetadata> metadata_;
-  const std::string metadata_location_;
+  std::string metadata_location_;
   std::shared_ptr<FileIO> io_;
   std::shared_ptr<Catalog> catalog_;
 

test/in_memory_catalog_test.cc

@@ -26,9 +26,11 @@
 #include <gtest/gtest.h>
 
 #include "iceberg/arrow/arrow_fs_file_io_internal.h"
+#include "iceberg/schema.h"
 #include "iceberg/table.h"
 #include "iceberg/table_metadata.h"
 #include "matchers.h"
+#include "mock_catalog.h"
 #include "test_common.h"
 
 namespace iceberg {
@@ -115,6 +117,56 @@ TEST_F(InMemoryCatalogTest, RegisterTable) {
   ASSERT_EQ(table.value()->location(), "s3://bucket/test/location");
 }
 
+TEST_F(InMemoryCatalogTest, RefreshTable) {
+  TableIdentifier table_ident{.ns = {}, .name = "t1"};
+  auto schema = std::make_shared<Schema>(
+      std::vector<SchemaField>{SchemaField::MakeRequired(1, "x", int64())},
+      /*schema_id=*/1);
+
+  std::shared_ptr<FileIO> io;
+
+  auto catalog = std::make_shared<MockCatalog>();
+  // Mock 1st call to LoadTable
+  EXPECT_CALL(*catalog, LoadTable(::testing::_))
+      .WillOnce(::testing::Return(
+          std::make_unique<Table>(table_ident,
+                                  std::make_shared<TableMetadata>(TableMetadata{
+                                      .schemas = {schema},
+                                      .current_schema_id = 1,
+                                      .current_snapshot_id = 1,
+                                      .snapshots = {std::make_shared<Snapshot>(Snapshot{
+                                          .snapshot_id = 1,
+                                          .sequence_number = 1,
+                                      })}}),
+                                  "s3://location/1.json", io, catalog)));
+  auto load_table_result = catalog->LoadTable(table_ident);
+  ASSERT_THAT(load_table_result, IsOk());
+  auto loaded_table = std::move(load_table_result.value());
+  ASSERT_EQ(loaded_table->current_snapshot().value()->snapshot_id, 1);
+
+  // Mock 2nd call to LoadTable
+  EXPECT_CALL(*catalog, LoadTable(::testing::_))
+      .WillOnce(::testing::Return(
+          std::make_unique<Table>(table_ident,
+                                  std::make_shared<TableMetadata>(TableMetadata{
+                                      .schemas = {schema},
+                                      .current_schema_id = 1,
+                                      .current_snapshot_id = 2,
+                                      .snapshots = {std::make_shared<Snapshot>(Snapshot{
+                                                        .snapshot_id = 1,
+                                                        .sequence_number = 1,
+                                                    }),
+                                                    std::make_shared<Snapshot>(Snapshot{
+                                                        .snapshot_id = 2,
+                                                        .sequence_number = 2,
+                                                    })}}),
+                                  "s3://location/2.json", io, catalog)));
+  auto refreshed_result = loaded_table->Refresh();
+  ASSERT_THAT(refreshed_result, IsOk());
+  // check table is refreshed
+  ASSERT_EQ(loaded_table->current_snapshot().value()->snapshot_id, 2);
+}
+
 TEST_F(InMemoryCatalogTest, DropTable) {
   TableIdentifier tableIdent{.ns = {}, .name = "t1"};
   auto result = catalog_->DropTable(tableIdent, false);

test/mock_catalog.h

@@ -0,0 +1,88 @@
+/*
+ * 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 <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "iceberg/catalog.h"
+
+namespace iceberg {
+
+class MockCatalog : public Catalog {
+ public:
+  MockCatalog() = default;
+  ~MockCatalog() override = default;
+
+  MOCK_METHOD(std::string_view, name, (), (const, override));
+
+  MOCK_METHOD(Status, CreateNamespace,
+              (const Namespace&, (const std::unordered_map<std::string, std::string>&)),
+              (override));
+
+  MOCK_METHOD((Result<std::vector<Namespace>>), ListNamespaces, (const Namespace&),
+              (const, override));
+
+  MOCK_METHOD((Result<std::unordered_map<std::string, std::string>>),
+              GetNamespaceProperties, (const Namespace&), (const, override));
+
+  MOCK_METHOD(Status, UpdateNamespaceProperties,
+              (const Namespace&, (const std::unordered_map<std::string, std::string>&),
+               (const std::unordered_set<std::string>&)),
+              (override));
+
+  MOCK_METHOD(Status, DropNamespace, (const Namespace&), (override));
+
+  MOCK_METHOD(Result<bool>, NamespaceExists, (const Namespace&), (const, override));
+
+  MOCK_METHOD((Result<std::vector<TableIdentifier>>), ListTables, (const Namespace&),
+              (const, override));
+
+  MOCK_METHOD((Result<std::unique_ptr<Table>>), CreateTable,
+              (const TableIdentifier&, const Schema&, const PartitionSpec&,
+               const std::string&, (const std::unordered_map<std::string, std::string>&)),
+              (override));
+
+  MOCK_METHOD((Result<std::unique_ptr<Table>>), UpdateTable,
+              (const TableIdentifier&,
+               (const std::vector<std::unique_ptr<UpdateRequirement>>&),
+               (const std::vector<std::unique_ptr<MetadataUpdate>>&)),
+              (override));
+
+  MOCK_METHOD((Result<std::shared_ptr<Transaction>>), StageCreateTable,
+              (const TableIdentifier&, const Schema&, const PartitionSpec&,
+               const std::string&, (const std::unordered_map<std::string, std::string>&)),
+              (override));
+
+  MOCK_METHOD(Result<bool>, TableExists, (const TableIdentifier&), (const, override));
+
+  MOCK_METHOD(Status, DropTable, (const TableIdentifier&, bool), (override));
+
+  MOCK_METHOD((Result<std::unique_ptr<Table>>), LoadTable, (const TableIdentifier&),
+              (override));
+
+  MOCK_METHOD((Result<std::shared_ptr<Table>>), RegisterTable,
+              (const TableIdentifier&, const std::string&), (override));
+
+  MOCK_METHOD((std::unique_ptr<TableBuilder>), BuildTable,
+              (const TableIdentifier&, const Schema&), (const, override));
+};
+
+}  // namespace iceberg