1 /*
2 * Copyright 2019 Google LLC.
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 * https://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15
16 // Protocol buffer related static utility functions.
17
18 #ifndef PRIVATE_JOIN_AND_COMPUTE_INTERNAL_UTIL_PROTO_UTIL_H_
19 #define PRIVATE_JOIN_AND_COMPUTE_INTERNAL_UTIL_PROTO_UTIL_H_
20
21 #include <memory>
22 #include <sstream>
23 #include <string>
24
25 #include "absl/strings/string_view.h"
26 #include "private_join_and_compute/util/recordio.h"
27 #include "private_join_and_compute/util/status.inc"
28 #include "src/google/protobuf/message_lite.h"
29
30 namespace private_join_and_compute {
31
32 class ProtoUtils {
33 public:
34 template <typename ProtoType>
35 static ProtoType FromString(absl::string_view raw_data);
36
37 static std::string ToString(const google::protobuf::MessageLite& record);
38
39 template <typename ProtoType>
40 static StatusOr<ProtoType> ReadProtoFromFile(absl::string_view filename);
41
42 template <typename ProtoType>
43 static StatusOr<std::vector<ProtoType>> ReadProtosFromFile(
44 absl::string_view filename);
45
46 static Status WriteProtoToFile(const google::protobuf::MessageLite& record,
47 absl::string_view filename);
48 template <typename ProtoType>
49 static Status WriteRecordsToFile(absl::string_view file,
50 const std::vector<ProtoType>& records);
51 };
52
53 template <typename ProtoType>
FromString(absl::string_view raw_data)54 inline ProtoType ProtoUtils::FromString(absl::string_view raw_data) {
55 ProtoType record;
56 record.ParseFromArray(raw_data.data(), raw_data.size());
57 return record;
58 }
59
ToString(const google::protobuf::MessageLite & record)60 inline std::string ProtoUtils::ToString(
61 const google::protobuf::MessageLite& record) {
62 std::ostringstream record_str_stream;
63 record.SerializeToOstream(&record_str_stream);
64 return record_str_stream.str();
65 }
66
67 template <typename ProtoType>
ReadProtoFromFile(absl::string_view filename)68 inline StatusOr<ProtoType> ProtoUtils::ReadProtoFromFile(
69 absl::string_view filename) {
70 std::unique_ptr<RecordReader> reader(RecordReader::GetRecordReader());
71 RETURN_IF_ERROR(reader->Open(filename));
72 std::string raw_record;
73 RETURN_IF_ERROR(reader->Read(&raw_record));
74 RETURN_IF_ERROR(reader->Close());
75 return ProtoUtils::FromString<ProtoType>(raw_record);
76 }
77
78 template <typename ProtoType>
ReadProtosFromFile(absl::string_view filename)79 inline StatusOr<std::vector<ProtoType>> ProtoUtils::ReadProtosFromFile(
80 absl::string_view filename) {
81 std::vector<ProtoType> result;
82 std::unique_ptr<RecordReader> reader(RecordReader::GetRecordReader());
83 RETURN_IF_ERROR(reader->Open(filename));
84 std::string raw_record;
85 ASSIGN_OR_RETURN(bool has_more, reader->HasMore());
86 while (has_more) {
87 RETURN_IF_ERROR(reader->Read(&raw_record));
88 result.push_back(ProtoUtils::FromString<ProtoType>(raw_record));
89 ASSIGN_OR_RETURN(has_more, reader->HasMore());
90 }
91 RETURN_IF_ERROR(reader->Close());
92 return std::move(result);
93 }
94
WriteProtoToFile(const google::protobuf::MessageLite & record,absl::string_view filename)95 inline Status ProtoUtils::WriteProtoToFile(
96 const google::protobuf::MessageLite& record, absl::string_view filename) {
97 std::unique_ptr<RecordWriter> writer(RecordWriter::Get());
98 RETURN_IF_ERROR(writer->Open(filename));
99 RETURN_IF_ERROR(writer->Write(ProtoUtils::ToString(record)));
100 return writer->Close();
101 }
102
103 template <typename ProtoType>
WriteRecordsToFile(absl::string_view file,const std::vector<ProtoType> & records)104 inline Status ProtoUtils::WriteRecordsToFile(
105 absl::string_view file, const std::vector<ProtoType>& records) {
106 std::unique_ptr<RecordWriter> writer(RecordWriter::Get());
107 RETURN_IF_ERROR(writer->Open(file));
108 for (const auto& record : records) {
109 RETURN_IF_ERROR(writer->Write(ProtoUtils::ToString(record)));
110 }
111 return writer->Close();
112 }
113 } // namespace private_join_and_compute
114
115 #endif // PRIVATE_JOIN_AND_COMPUTE_INTERNAL_UTIL_PROTO_UTIL_H_
116