updates proto_util

3 files changed, 42 insertions, 1 deletions

diff --git a/README.md b/README.md
index d0bc900..2212840 100644
--- a/README.md
+++ b/README.md
@@ -56,7 +56,7 @@ the Private Join and Compute library and dependencies using Bazel:
 
 ```bash
 cd private-join-and-compute
-bazel build --cxxopt='-std=c++17' :all
+bazel build --cxxopt='-std=c++17' //private_join_and_compute/:all
 ```
 
 (All the following instructions must be run from inside the
diff --git a/private_join_and_compute/util/proto_util.h b/private_join_and_compute/util/proto_util.h
index ea28d87..ee0af75 100644
--- a/private_join_and_compute/util/proto_util.h
+++ b/private_join_and_compute/util/proto_util.h
@@ -38,6 +38,10 @@ class ProtoUtils {
   template <typename ProtoType>
   static StatusOr<ProtoType> ReadProtoFromFile(absl::string_view filename);
 
+  template <typename ProtoType>
+  static StatusOr<std::vector<ProtoType>> ReadProtosFromFile(
+      absl::string_view filename);
+
   static Status WriteProtoToFile(const google::protobuf::MessageLite& record,
                                  absl::string_view filename);
   template <typename ProtoType>
@@ -70,6 +74,23 @@ inline StatusOr<ProtoType> ProtoUtils::ReadProtoFromFile(
   return ProtoUtils::FromString<ProtoType>(raw_record);
 }
 
+template <typename ProtoType>
+inline StatusOr<std::vector<ProtoType>> ProtoUtils::ReadProtosFromFile(
+    absl::string_view filename) {
+  std::vector<ProtoType> result;
+  std::unique_ptr<RecordReader> reader(RecordReader::GetRecordReader());
+  RETURN_IF_ERROR(reader->Open(filename));
+  std::string raw_record;
+  ASSIGN_OR_RETURN(bool has_more, reader->HasMore());
+  while (has_more) {
+    RETURN_IF_ERROR(reader->Read(&raw_record));
+    result.push_back(ProtoUtils::FromString<ProtoType>(raw_record));
+    ASSIGN_OR_RETURN(has_more, reader->HasMore());
+  }
+  RETURN_IF_ERROR(reader->Close());
+  return std::move(result);
+}
+
 inline Status ProtoUtils::WriteProtoToFile(
     const google::protobuf::MessageLite& record, absl::string_view filename) {
   std::unique_ptr<RecordWriter> writer(RecordWriter::Get());
diff --git a/private_join_and_compute/util/proto_util_test.cc b/private_join_and_compute/util/proto_util_test.cc
index b3bdffe..814683b 100644
--- a/private_join_and_compute/util/proto_util_test.cc
+++ b/private_join_and_compute/util/proto_util_test.cc
@@ -52,6 +52,26 @@ TEST(ProtoUtilsTest, ReadWriteToFile) {
   EXPECT_EQ(actual_test_proto.dummy(), expected_test_proto.dummy());
 }
 
+TEST(ProtoUtilsTest, ReadWriteManyToFile) {
+  std::string filename = ::testing::TempDir() + "/proto_file";
+
+  TestProto expected_test_proto;
+  expected_test_proto.set_record("data");
+  expected_test_proto.set_dummy("dummy");
+
+  std::vector<TestProto> test_vector = {
+      expected_test_proto, expected_test_proto, expected_test_proto};
+
+  ASSERT_TRUE(ProtoUtils::WriteRecordsToFile(filename, test_vector).ok());
+  ASSERT_OK_AND_ASSIGN(std::vector<TestProto> result,
+                       ProtoUtils::ReadProtosFromFile<TestProto>(filename));
+  EXPECT_EQ(result.size(), test_vector.size());
+  for (const TestProto& result_element : result) {
+    EXPECT_EQ(result_element.record(), expected_test_proto.record());
+    EXPECT_EQ(result_element.dummy(), expected_test_proto.dummy());
+  }
+}
+
 }  // namespace
 
 }  // namespace private_join_and_compute