Line data Source code
1 : #include <filesystem>
2 : #include <memory>
3 : #include <string>
4 : #include <vector>
5 :
6 : #include "absl/status/status.h"
7 : #include "absl/strings/str_cat.h"
8 : #include "absl/types/span.h"
9 : #include "backend/storage/duckdb/duckdb_storage.h"
10 : #include "backend/storage/duckdb/duckdb_storage_internal.h"
11 : #include "backend/storage/duckdb/duckdb_storage_test_fixture.h"
12 : #include "backend/storage/duckdb/duckdb_storage_version_log.h"
13 : #include "backend/storage/storage.h"
14 : #include "gtest/gtest.h"
15 :
16 : namespace bigquery_emulator {
17 : namespace backend {
18 : namespace storage {
19 : namespace duckdb {
20 : namespace {
21 :
22 1 : TEST_F(DuckDBStorageTest, DatasetTombstoneRoundTripRestoresTables) {
23 1 : auto store_or = DuckDBStorage::Open(data_dir_.string());
24 2 : ASSERT_TRUE(store_or.ok()) << store_or.status();
25 1 : auto& store = **store_or;
26 :
27 1 : const DatasetId ds{"proj-1", "ds_undrop"};
28 1 : const TableId table{"proj-1", "ds_undrop", "people"};
29 1 : ASSERT_TRUE(store.CreateDataset(ds, "EU").ok());
30 1 : ASSERT_TRUE(store.CreateTable(table, PeopleSchema()).ok());
31 1 : std::vector<Row> rows;
32 1 : rows.push_back(MakePerson(1, "ada"));
33 1 : ASSERT_TRUE(store.AppendRows(table, absl::MakeConstSpan(rows)).ok());
34 :
35 1 : ASSERT_TRUE(store.DropDataset(ds, /*delete_contents=*/true).ok());
36 1 : auto list_after_drop = store.ListDatasets("proj-1");
37 1 : ASSERT_TRUE(list_after_drop.ok());
38 1 : EXPECT_TRUE(list_after_drop->empty());
39 :
40 1 : ASSERT_TRUE(store.RestoreDataset(ds).ok());
41 1 : auto list_after_restore = store.ListDatasets("proj-1");
42 1 : ASSERT_TRUE(list_after_restore.ok());
43 1 : ASSERT_EQ(list_after_restore->size(), 1u);
44 :
45 1 : auto iter_or = store.ScanRows(table);
46 1 : ASSERT_TRUE(iter_or.ok());
47 1 : Row r;
48 1 : auto has = (*iter_or)->Next(&r);
49 1 : ASSERT_TRUE(has.ok());
50 1 : ASSERT_TRUE(*has);
51 1 : EXPECT_EQ(r.cells[0].int64_value(), 1);
52 1 : EXPECT_EQ(r.cells[1].string_value(), "ada");
53 1 : }
54 :
55 1 : TEST_F(DuckDBStorageTest, DatasetTombstoneRoundTripRestoresViewsAndRoutines) {
56 1 : auto store_or = DuckDBStorage::Open(data_dir_.string());
57 2 : ASSERT_TRUE(store_or.ok()) << store_or.status();
58 1 : auto& store = **store_or;
59 :
60 1 : const DatasetId ds{"proj-1", "ds_registry"};
61 1 : ASSERT_TRUE(store.CreateDataset(ds, "US").ok());
62 :
63 1 : ViewRecord view;
64 1 : view.id = ViewId{"proj-1", "ds_registry", "v_items"};
65 1 : view.ddl_sql = "CREATE VIEW ds_registry.v_items AS SELECT 1 AS id";
66 1 : ASSERT_TRUE(store.UpsertView(view).ok());
67 :
68 1 : RoutineRecord routine;
69 1 : routine.id = RoutineId{"proj-1", "ds_registry", "add_one"};
70 1 : routine.kind = RoutineKind::kScalarFunction;
71 1 : routine.language = "SQL";
72 1 : routine.ddl_sql =
73 1 : "CREATE FUNCTION ds_registry.add_one(x INT64) RETURNS INT64 AS (x + 1)";
74 1 : ASSERT_TRUE(store.UpsertRoutine(routine).ok());
75 :
76 1 : ASSERT_TRUE(store.DropDataset(ds, /*delete_contents=*/true).ok());
77 :
78 1 : auto routines_after_drop = store.ListRoutines(ds);
79 1 : ASSERT_TRUE(routines_after_drop.ok());
80 1 : EXPECT_TRUE(routines_after_drop->empty());
81 :
82 1 : ASSERT_TRUE(store.RestoreDataset(ds).ok());
83 :
84 1 : auto routines_after_restore = store.ListRoutines(ds);
85 1 : ASSERT_TRUE(routines_after_restore.ok());
86 1 : ASSERT_EQ(routines_after_restore->size(), 1u);
87 1 : EXPECT_EQ((*routines_after_restore)[0].id.routine_id, "add_one");
88 :
89 1 : auto views_after_restore = store.ListAllViews();
90 1 : ASSERT_TRUE(views_after_restore.ok());
91 1 : bool found_view = false;
92 1 : for (const ViewRecord& rec : *views_after_restore) {
93 1 : if (rec.id.view_id == "v_items") {
94 1 : found_view = true;
95 1 : break;
96 1 : }
97 1 : }
98 1 : EXPECT_TRUE(found_view);
99 1 : }
100 :
101 1 : TEST_F(DuckDBStorageTest, RestoreDatasetOnLiveDatasetIsAlreadyExists) {
102 1 : auto store_or = DuckDBStorage::Open(data_dir_.string());
103 2 : ASSERT_TRUE(store_or.ok()) << store_or.status();
104 1 : auto& store = **store_or;
105 :
106 1 : const DatasetId ds{"proj-1", "ds_live"};
107 1 : ASSERT_TRUE(store.CreateDataset(ds, "US").ok());
108 1 : absl::Status restored = store.RestoreDataset(ds);
109 1 : ASSERT_FALSE(restored.ok());
110 1 : EXPECT_EQ(restored.code(), absl::StatusCode::kAlreadyExists);
111 1 : }
112 :
113 1 : TEST_F(DuckDBStorageTest, TableTombstonePathIncludesDataset) {
114 1 : auto store_or = DuckDBStorage::Open(data_dir_.string());
115 2 : ASSERT_TRUE(store_or.ok()) << store_or.status();
116 1 : auto& store = **store_or;
117 :
118 1 : const TableId table{"proj-1", "ds_path", "tbl"};
119 1 : const std::string dir =
120 1 : internal::TableTombstoneDir(store, table, /*deleted_ms=*/123);
121 1 : EXPECT_TRUE(dir.find("proj-1.ds_path") != std::string::npos);
122 1 : EXPECT_TRUE(dir.find("/tbl/123") != std::string::npos);
123 1 : }
124 :
125 1 : TEST_F(DuckDBStorageTest, DropRecreateSameIdUndropIsAlreadyExists) {
126 1 : auto store_or = DuckDBStorage::Open(data_dir_.string());
127 2 : ASSERT_TRUE(store_or.ok()) << store_or.status();
128 1 : auto& store = **store_or;
129 :
130 1 : const DatasetId ds{"proj-1", "ds_recreate"};
131 1 : ASSERT_TRUE(store.CreateDataset(ds, "US").ok());
132 1 : ASSERT_TRUE(store.DropDataset(ds, /*delete_contents=*/true).ok());
133 1 : ASSERT_TRUE(store.CreateDataset(ds, "US").ok());
134 :
135 1 : absl::Status restored = store.RestoreDataset(ds);
136 1 : ASSERT_FALSE(restored.ok());
137 1 : EXPECT_EQ(restored.code(), absl::StatusCode::kAlreadyExists);
138 1 : }
139 :
140 1 : TEST_F(DuckDBStorageTest, DoubleUndropSecondIsNotFound) {
141 1 : auto store_or = DuckDBStorage::Open(data_dir_.string());
142 2 : ASSERT_TRUE(store_or.ok()) << store_or.status();
143 1 : auto& store = **store_or;
144 :
145 1 : const DatasetId ds{"proj-1", "ds_double"};
146 1 : ASSERT_TRUE(store.CreateDataset(ds, "US").ok());
147 1 : ASSERT_TRUE(store.DropDataset(ds, /*delete_contents=*/true).ok());
148 1 : ASSERT_TRUE(store.RestoreDataset(ds).ok());
149 :
150 1 : absl::Status second = store.RestoreDataset(ds);
151 1 : ASSERT_FALSE(second.ok());
152 1 : EXPECT_EQ(second.code(), absl::StatusCode::kAlreadyExists);
153 1 : }
154 :
155 1 : TEST_F(DuckDBStorageTest, RestartDurabilityRestoresDatasetAndMetadata) {
156 1 : const DatasetId ds{"proj-1", "ds_restart"};
157 1 : const TableId table{"proj-1", "ds_restart", "items"};
158 1 : const std::string rest_metadata = R"({"labels":{"env":"test"}})";
159 :
160 1 : {
161 1 : auto store_or = DuckDBStorage::Open(data_dir_.string());
162 2 : ASSERT_TRUE(store_or.ok()) << store_or.status();
163 1 : auto& store = **store_or;
164 1 : ASSERT_TRUE(store.CreateDataset(ds, "US").ok());
165 1 : ASSERT_TRUE(store.CreateTable(table, PeopleSchema()).ok());
166 1 : ASSERT_TRUE(
167 1 : store.DropDataset(ds, /*delete_contents=*/true, rest_metadata).ok());
168 1 : ASSERT_TRUE(store.RestoreDataset(ds).ok());
169 1 : auto meta_or = store.GetDatasetRestMetadataJson(ds);
170 1 : ASSERT_TRUE(meta_or.ok());
171 1 : EXPECT_NE(meta_or->find("env"), std::string::npos);
172 1 : }
173 :
174 1 : auto reopened_or = DuckDBStorage::Open(data_dir_.string());
175 2 : ASSERT_TRUE(reopened_or.ok()) << reopened_or.status();
176 1 : auto& reopened = **reopened_or;
177 1 : auto datasets_or = reopened.ListDatasets("proj-1");
178 1 : ASSERT_TRUE(datasets_or.ok());
179 1 : ASSERT_EQ(datasets_or->size(), 1u);
180 1 : auto meta_or = reopened.GetDatasetRestMetadataJson(ds);
181 1 : ASSERT_TRUE(meta_or.ok());
182 1 : EXPECT_NE(meta_or->find("env"), std::string::npos);
183 1 : }
184 :
185 : } // namespace
186 : } // namespace duckdb
187 : } // namespace storage
188 : } // namespace backend
189 : } // namespace bigquery_emulator
|
