chore: Remove references to records_utils in tools and components

Bug: 385338000 Change-Id: Id36db3ab1a45c67b28693229150453f6e32f27cc GitOrigin-RevId: 633cbc2759697fda0ee6c4ddb65ba3c6c121e92d
privacysandbox · Dec 27, 2024 · f766282 · f766282
1 parent 6d8fd09
commit f766282
Show file tree

Hide file tree

Showing 10 changed files with 59 additions and 40 deletions.
diff --git a/components/data_server/data_loading/BUILD.bazel b/components/data_server/data_loading/BUILD.bazel
@@ -39,7 +39,7 @@ cc_library(
         "//public:constants",
         "//public/data_loading:data_loading_fbs",
         "//public/data_loading:filename_utils",
-        "//public/data_loading:records_utils",
+        "//public/data_loading:record_utils",
         "//public/data_loading/readers:riegeli_stream_io",
         "//public/data_loading/readers:stream_record_reader_factory",
         "//public/sharding:key_sharder",
@@ -68,7 +68,7 @@ cc_test(
         "//components/udf:code_config",
         "//components/udf:mocks",
         "//public/data_loading:filename_utils",
-        "//public/data_loading:records_utils",
+        "//public/data_loading:record_utils",
         "//public/test_util:data_record",
         "//public/test_util:mocks",
         "//public/test_util:proto_matcher",

diff --git a/components/data_server/data_loading/data_orchestrator.cc b/components/data_server/data_loading/data_orchestrator.cc
@@ -28,7 +28,7 @@
 #include "public/constants.h"
 #include "public/data_loading/data_loading_generated.h"
 #include "public/data_loading/filename_utils.h"
-#include "public/data_loading/records_utils.h"
+#include "public/data_loading/record_utils.h"
 #include "public/sharding/sharding_function.h"
 #include "src/errors/retry.h"
 #include "src/telemetry/tracing.h"

diff --git a/components/data_server/data_loading/data_orchestrator_test.cc b/components/data_server/data_loading/data_orchestrator_test.cc
@@ -31,7 +31,7 @@
 #include "gtest/gtest.h"
 #include "public/constants.h"
 #include "public/data_loading/filename_utils.h"
-#include "public/data_loading/records_utils.h"
+#include "public/data_loading/record_utils.h"
 #include "public/sharding/key_sharder.h"
 #include "public/sharding/sharding_function.h"
 #include "public/test_util/data_record.h"

diff --git a/tools/bidding_auction_data_generator/BUILD.bazel b/tools/bidding_auction_data_generator/BUILD.bazel
@@ -123,6 +123,7 @@ cc_test(
         "//public/data_loading/readers:delta_record_stream_reader",
         "//public/data_loading/readers:riegeli_stream_record_reader_factory",
         "//public/data_loading/writers:delta_record_stream_writer",
+        "//public/test_util:data_record",
         "@com_google_googletest//:gtest_main",
     ],
 )

diff --git a/tools/bidding_auction_data_generator/delta_key_value_writer_test.cc b/tools/bidding_auction_data_generator/delta_key_value_writer_test.cc
@@ -28,12 +28,13 @@ constexpr int64_t kTestLogicalCommitTime = 1234567890;
 constexpr KeyValueMutationType kTestDeltaMutationType =
     KeyValueMutationType::Update;
 
-KeyValueMutationRecordStruct GetKVMutationRecord() {
-  KeyValueMutationRecordStruct record;
-  record.key = "key1";
-  record.value = R"({"field": "test"})";
-  record.logical_commit_time = kTestLogicalCommitTime;
-  record.mutation_type = kTestDeltaMutationType;
+KeyValueMutationRecordT GetKVMutationRecord() {
+  KeyValueMutationRecordT record = {
+      .mutation_type = kTestDeltaMutationType,
+      .logical_commit_time = kTestLogicalCommitTime,
+      .key = "key1",
+  };
+  record.value.Set(StringValueT{.value = R"({"field": "test"})"});
   return record;
 }
 
@@ -56,21 +57,23 @@ TEST(DeltaKeyValueWriterTest, ValidateDeltaDataTest) {
   auto stream_reader_factory =
       std::make_unique<RiegeliStreamRecordReaderFactory>();
   auto stream_reader = stream_reader_factory->CreateReader(delta_stream);
-  EXPECT_TRUE(stream_reader
-                  ->ReadStreamRecords(
-                      [](std::string_view record_string) -> absl::Status {
-                        const auto* data_record =
-                            flatbuffers::GetRoot<DataRecord>(
-                                record_string.data());
-                        EXPECT_EQ(data_record->record_type(),
-                                  Record::KeyValueMutationRecord);
-                        const auto kv_record =
-                            GetTypedRecordStruct<KeyValueMutationRecordStruct>(
-                                *data_record);
-                        EXPECT_EQ(kv_record, GetKVMutationRecord());
-                        return absl::OkStatus();
-                      })
-                  .ok());
+  EXPECT_TRUE(
+      stream_reader
+          ->ReadStreamRecords(
+              [](std::string_view record_string) -> absl::Status {
+                return DeserializeRecord(
+                    record_string, [](const DataRecord& data_record) {
+                      DataRecordT data_record_struct;
+                      data_record.UnPackTo(&data_record_struct);
+                      EXPECT_EQ(data_record_struct.record.type,
+                                Record::KeyValueMutationRecord);
+                      const auto kv_record =
+                          *data_record_struct.record.AsKeyValueMutationRecord();
+                      EXPECT_EQ(kv_record, GetKVMutationRecord());
+                      return absl::OkStatus();
+                    });
+              })
+          .ok());
 }
 }  // namespace
 }  // namespace kv_server
diff --git a/tools/data_cli/BUILD.bazel b/tools/data_cli/BUILD.bazel
@@ -28,7 +28,7 @@ cc_binary(
     deps = [
         "//components/util:platform_initializer",
         "//public/data_loading:filename_utils",
-        "//public/data_loading:records_utils",
+        "//public/data_loading:record_utils",
         "//tools/data_cli/commands:command",
         "//tools/data_cli/commands:format_data_command",
         "//tools/data_cli/commands:generate_snapshot_command",

diff --git a/tools/data_cli/commands/format_data_command_test.cc b/tools/data_cli/commands/format_data_command_test.cc
@@ -46,23 +46,37 @@ KVFileMetadata GetMetadata() {
   return metadata;
 }
 
-class FormatDataCommandTest
-    : public testing::TestWithParam<KeyValueMutationRecordValueT> {
+template <class T>
+T GetSampleValue();
+
+template <>
+StringValueT GetSampleValue<StringValueT>() {
+  return StringValueT{.value = "value"};
+}
+
+template <>
+StringSetT GetSampleValue<StringSetT>() {
+  return StringSetT{.value =
+                        std::vector<std::string>{"value1", "value2", "value3"}};
+}
+
+template <class T>
+class FormatDataCommandTest : public testing::TestWithParam<T> {
  protected:
-  KeyValueMutationRecordValueT GetValue() { return GetParam(); }
+  T GetSampleRecordValue() { return GetSampleValue<T>(); }
 };
 
-INSTANTIATE_TEST_SUITE_P(KVMutationValue, FormatDataCommandTest,
-                         testing::Values("value",
-                                         std::vector<std::string_view>{
-                                             "value1", "value2", "value3"}));
+using testing::Types;
+typedef Types<StringValueT, StringSetT> ValueTypes;
+TYPED_TEST_SUITE(FormatDataCommandTest, ValueTypes);
 
-TEST_P(FormatDataCommandTest, ValidateGeneratingCsvToDeltaData_KVMutations) {
+TYPED_TEST(FormatDataCommandTest,
+           ValidateGeneratingCsvToDeltaData_KVMutations) {
   std::stringstream csv_stream;
   std::stringstream delta_stream;
   CsvDeltaRecordStreamWriter csv_writer(csv_stream);
   const auto& record =
-      GetNativeDataRecord(GetKVMutationRecord(GetSimpleStringValue()));
+      GetNativeDataRecord(GetKVMutationRecord(this->GetSampleRecordValue()));
   EXPECT_TRUE(csv_writer.WriteRecord(record).ok());
   EXPECT_TRUE(csv_writer.WriteRecord(record).ok());
   EXPECT_TRUE(csv_writer.WriteRecord(record).ok());
@@ -83,13 +97,14 @@ TEST_P(FormatDataCommandTest, ValidateGeneratingCsvToDeltaData_KVMutations) {
   EXPECT_TRUE(delta_reader.ReadRecords(record_callback.AsStdFunction()).ok());
 }
 
-TEST_P(FormatDataCommandTest, ValidateGeneratingDeltaToCsvData_KvMutations) {
+TYPED_TEST(FormatDataCommandTest,
+           ValidateGeneratingDeltaToCsvData_KvMutations) {
   std::stringstream delta_stream;
   std::stringstream csv_stream;
   auto delta_writer = DeltaRecordStreamWriter<std::stringstream>::Create(
       delta_stream, DeltaRecordWriter::Options{.metadata = GetMetadata()});
   const auto& record =
-      GetNativeDataRecord(GetKVMutationRecord(GetSimpleStringValue()));
+      GetNativeDataRecord(GetKVMutationRecord(this->GetSampleRecordValue()));
   EXPECT_TRUE(delta_writer.ok()) << delta_writer.status();
   EXPECT_TRUE((*delta_writer)->WriteRecord(record).ok());
   EXPECT_TRUE((*delta_writer)->WriteRecord(record).ok());

diff --git a/tools/request_simulation/BUILD.bazel b/tools/request_simulation/BUILD.bazel
@@ -364,7 +364,7 @@ cc_library(
         "//components/data/common:change_notifier",
         "//components/data/common:thread_manager",
         "//components/data/realtime:realtime_notifier",
-        "//public/data_loading:records_utils",
+        "//public/data_loading:record_utils",
         "//public/data_loading/readers:riegeli_stream_record_reader_factory",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",

diff --git a/tools/udf/udf_generator/BUILD.bazel b/tools/udf/udf_generator/BUILD.bazel
@@ -31,7 +31,7 @@ cc_binary(
         "//public:constants",
         "//public/data_loading:data_loading_fbs",
         "//public/data_loading:filename_utils",
-        "//public/data_loading:records_utils",
+        "//public/data_loading:record_utils",
         "//public/data_loading:riegeli_metadata_cc_proto",
         "//public/data_loading/writers:avro_delta_record_stream_writer",
         "//public/data_loading/writers:delta_record_stream_writer",

diff --git a/tools/udf/udf_generator/udf_delta_file_generator.cc b/tools/udf/udf_generator/udf_delta_file_generator.cc
@@ -26,7 +26,7 @@
 #include "public/constants.h"
 #include "public/data_loading/data_loading_generated.h"
 #include "public/data_loading/filename_utils.h"
-#include "public/data_loading/records_utils.h"
+#include "public/data_loading/record_utils.h"
 #include "public/data_loading/writers/avro_delta_record_stream_writer.h"
 #include "public/data_loading/writers/delta_record_stream_writer.h"
 #include "public/data_loading/writers/delta_record_writer.h"