Line data Source code
1 : // SPDX-FileCopyrightText: 2025 PairInteraction Developers
2 : // SPDX-License-Identifier: LGPL-3.0-or-later
3 :
4 : #include "pairinteraction/database/ParquetManager.hpp"
5 :
6 : #include "pairinteraction/database/GitHubDownloader.hpp"
7 : #include "pairinteraction/version.hpp"
8 :
9 : #include <cpptrace/cpptrace.hpp>
10 : #include <ctime>
11 : #include <duckdb.hpp>
12 : #include <filesystem>
13 : #include <fmt/core.h>
14 : #include <fmt/format.h>
15 : #include <fstream>
16 : #include <future>
17 : #include <iomanip>
18 : #include <miniz.h>
19 : #include <nlohmann/json.hpp>
20 : #include <regex>
21 : #include <set>
22 : #include <spdlog/spdlog.h>
23 : #include <sstream>
24 : #include <stdexcept>
25 : #include <string>
26 : #include <system_error>
27 :
28 : namespace fs = std::filesystem;
29 : using json = nlohmann::json;
30 :
31 : namespace pairinteraction {
32 : namespace {
33 :
34 4 : std::string format_time(std::time_t time_val) {
35 4 : std::tm *ptm = std::localtime(&time_val);
36 4 : std::ostringstream oss;
37 4 : oss << std::put_time(ptm, "%Y-%m-%d %H:%M:%S");
38 8 : return oss.str();
39 4 : }
40 :
41 0 : json load_json(const fs::path &file) {
42 0 : std::ifstream in(file);
43 0 : in.exceptions(std::ifstream::failbit | std::ifstream::badbit);
44 0 : json doc;
45 0 : in >> doc;
46 0 : return doc;
47 0 : }
48 :
49 4 : void save_json(const fs::path &file, const json &doc) {
50 4 : std::ofstream out(file);
51 4 : if (!out) {
52 0 : throw std::runtime_error(fmt::format("Failed to open {} for writing", file.string()));
53 : }
54 4 : out << doc;
55 4 : out.close();
56 4 : }
57 :
58 : } // namespace
59 :
60 0 : void ParquetManager::react_on_exception(const std::string &context, const std::exception &e) {
61 0 : repo_paths_.clear();
62 0 : remote_asset_info.clear();
63 0 : SPDLOG_ERROR("{}: {}. The download of database tables is disabled.", context, e.what());
64 0 : }
65 :
66 0 : void ParquetManager::react_on_error_code(const std::string &context, int error_code) {
67 0 : repo_paths_.clear();
68 0 : remote_asset_info.clear();
69 0 : SPDLOG_ERROR("{}: status code {}. The download of database tables is disabled.", context,
70 : error_code);
71 0 : }
72 :
73 1 : void ParquetManager::react_on_invalid_response(const std::string &context) {
74 1 : repo_paths_.clear();
75 1 : remote_asset_info.clear();
76 1 : SPDLOG_ERROR("{}. The download of database tables is disabled.", context);
77 1 : }
78 :
79 0 : void ParquetManager::react_on_rate_limit_reached(std::time_t reset_time) {
80 0 : repo_paths_.clear();
81 0 : remote_asset_info.clear();
82 0 : SPDLOG_ERROR("Rate limit reached, resets at {}. The download of database tables is disabled.",
83 : format_time(reset_time));
84 0 : }
85 :
86 9 : ParquetManager::ParquetManager(std::filesystem::path directory, const GitHubDownloader &downloader,
87 : std::vector<std::string> repo_paths, duckdb::Connection &con,
88 9 : bool use_cache)
89 9 : : directory_(std::move(directory)), downloader(downloader), repo_paths_(std::move(repo_paths)),
90 9 : con(con), use_cache_(use_cache) {
91 : // Ensure the local directory exists
92 9 : if (!std::filesystem::exists(directory_ / "tables")) {
93 0 : fs::create_directories(directory_ / "tables");
94 : }
95 :
96 : // If repo paths are provided, check the GitHub rate limit
97 9 : if (!repo_paths_.empty()) {
98 4 : GitHubDownloader::Result result;
99 : try {
100 4 : result = downloader.download("/rate_limit", "", false).get();
101 0 : } catch (const std::exception &e) {
102 0 : react_on_exception("Failed obtaining the rate limit", e);
103 0 : return;
104 0 : }
105 4 : if (result.status_code != 200) {
106 0 : react_on_error_code("Failed obtaining the rate limit", result.status_code);
107 0 : return;
108 : }
109 4 : if (result.rate_limit.remaining == 0) {
110 0 : react_on_rate_limit_reached(result.rate_limit.reset_time);
111 0 : return;
112 : }
113 4 : SPDLOG_INFO("Remaining GitHub API requests: {}. Rate limit resets at {}.",
114 : result.rate_limit.remaining, format_time(result.rate_limit.reset_time));
115 4 : }
116 0 : }
117 :
118 9 : void ParquetManager::scan_remote() {
119 : // If repo_paths_ is empty, we have nothing to do
120 9 : if (repo_paths_.empty()) {
121 6 : return;
122 : }
123 :
124 4 : remote_asset_info.clear();
125 :
126 : struct RepoDownload {
127 : std::string endpoint;
128 : fs::path cache_file;
129 : json cached_doc;
130 : std::future<GitHubDownloader::Result> future;
131 : };
132 4 : std::vector<RepoDownload> downloads;
133 4 : downloads.reserve(repo_paths_.size());
134 :
135 : // For each repo path, load its cached JSON (or an empty JSON) and issue the download
136 9 : for (const auto &endpoint : repo_paths_) {
137 : // Generate a unique cache filename per endpoint
138 : std::string cache_filename =
139 5 : "homepage_cache_" + std::to_string(std::hash<std::string>{}(endpoint)) + ".json";
140 5 : auto cache_file = directory_ / cache_filename;
141 :
142 : // Load cached JSON from file if it exists
143 5 : json cached_doc;
144 5 : if (std::filesystem::exists(cache_file)) {
145 : try {
146 0 : cached_doc = load_json(cache_file);
147 0 : } catch (const std::exception &e) {
148 0 : SPDLOG_WARN("Error reading {}: {}. Discarding homepage cache.", cache_file.string(),
149 : e.what());
150 0 : cached_doc = json{};
151 0 : }
152 : }
153 :
154 : // Extract the etag header from the cached JSON if available
155 5 : std::string etag;
156 5 : if (!cached_doc.is_null() && cached_doc.contains("etag")) {
157 0 : etag = cached_doc["etag"].get<std::string>();
158 : }
159 :
160 : // Issue the asynchronous download using the cached etag.
161 : try {
162 5 : downloads.push_back(
163 5 : {endpoint, cache_file, cached_doc, downloader.download(endpoint, etag)});
164 0 : } catch (const std::exception &e) {
165 0 : react_on_exception(
166 0 : fmt::format("Failed to download overview of available tables from {}", endpoint),
167 : e);
168 0 : return;
169 0 : }
170 5 : }
171 :
172 : // Process downloads for each repo path
173 8 : for (auto &dl : downloads) {
174 5 : auto result = dl.future.get();
175 5 : if ((result.status_code == 403 || result.status_code == 429) &&
176 0 : result.rate_limit.remaining == 0) {
177 0 : react_on_rate_limit_reached(result.rate_limit.reset_time);
178 0 : return;
179 : }
180 5 : if (result.status_code != 200 && result.status_code != 304) {
181 0 : react_on_error_code(
182 0 : fmt::format("Failed to download overview of available tables from {}", dl.endpoint),
183 : result.status_code);
184 0 : return;
185 : }
186 :
187 5 : json doc;
188 5 : if (result.status_code == 304) {
189 0 : if (dl.cached_doc.is_null() || dl.cached_doc.empty()) {
190 0 : throw std::runtime_error(
191 0 : fmt::format("Received 304 Not Modified but cached response {} does not exist.",
192 0 : dl.cache_file.string()));
193 : }
194 0 : doc = dl.cached_doc;
195 0 : SPDLOG_INFO("Using cached overview of available tables from {}.", dl.endpoint);
196 : } else {
197 : // The endpoint for listing releases returns an array of releases, the endpoint for
198 : // requesting a single release returns the release itself
199 5 : json parsed = json::parse(result.body, nullptr, /*allow_exceptions=*/false);
200 5 : if (parsed.is_discarded()) {
201 1 : react_on_invalid_response(fmt::format(
202 1 : "Failed to parse the overview of available tables from {}", dl.endpoint));
203 1 : return;
204 : }
205 6 : doc["releases"] = parsed.is_array() ? parsed : json::array({parsed});
206 4 : doc["etag"] = result.etag;
207 4 : SPDLOG_INFO("Using downloaded overview of available tables from {}.", dl.endpoint);
208 5 : }
209 :
210 : // Validate the JSON response
211 4 : if (!doc["releases"].is_array() || doc["releases"].empty()) {
212 : // Discard the cache so that the next run issues an unconditional request instead of
213 : // receiving 304 Not Modified and reading the invalid cache again
214 0 : std::error_code ec;
215 0 : fs::remove(dl.cache_file, ec);
216 0 : react_on_invalid_response(fmt::format(
217 0 : "No releases contained in the overview of available tables from {}", dl.endpoint));
218 0 : return;
219 : }
220 :
221 4 : if (result.status_code != 304) {
222 4 : save_json(dl.cache_file, doc);
223 : }
224 :
225 : // Update remote_asset_info based on the asset entries of the latest release that is
226 : // compatible with this version of the software. Because the releases are ordered from the
227 : // latest to the oldest one, this is not necessarily the latest release of the repository.
228 4 : bool found_compatible_release = false;
229 8 : for (const auto &release : doc["releases"]) {
230 : // Skip malformed entries, drafts, and pre-releases because they are not meant to
231 : // be used
232 15 : if (!release.is_object() || release.value("draft", false) ||
233 15 : release.value("prerelease", false)) {
234 2 : continue;
235 : }
236 6 : if (!release.contains("assets") || !release["assets"].is_array()) {
237 0 : continue;
238 : }
239 :
240 12 : for (const auto &asset : release["assets"]) {
241 : // Skip malformed entries
242 6 : if (!asset.is_object() || !asset.contains("name") || !asset.contains("url")) {
243 2 : continue;
244 : }
245 :
246 6 : std::string name = asset["name"].get<std::string>();
247 6 : std::smatch match;
248 :
249 6 : if (std::regex_match(name, match, remote_regex) && match.size() == 4) {
250 6 : std::string key = match[1].str();
251 6 : int version_major = std::stoi(match[2].str());
252 6 : int version_minor = std::stoi(match[3].str());
253 :
254 6 : if (version_major != COMPATIBLE_DATABASE_VERSION_MAJOR) {
255 2 : continue;
256 : }
257 4 : found_compatible_release = true;
258 :
259 4 : auto it = remote_asset_info.find(key);
260 4 : if (it == remote_asset_info.end() || version_minor > it->second.version_minor) {
261 4 : std::string remote_url = asset["url"].get<std::string>();
262 4 : const std::string host = downloader.get_host();
263 8 : remote_asset_info[key] = {version_minor, remote_url.erase(0, host.size())};
264 4 : }
265 6 : }
266 8 : }
267 :
268 6 : if (found_compatible_release) {
269 4 : break;
270 : }
271 : }
272 6 : }
273 :
274 : // Ensure that scan_remote was successful
275 3 : if (!downloads.empty() && remote_asset_info.empty()) {
276 0 : throw std::runtime_error(
277 : "No compatible database tables were found in the remote repositories. Consider "
278 0 : "upgrading PairInteraction to a newer version.");
279 : }
280 18 : }
281 :
282 9 : void ParquetManager::scan_local() {
283 9 : local_asset_info.clear();
284 :
285 : // Iterate over files in the directory to update local_asset_info
286 89 : for (const auto &entry : fs::directory_iterator(directory_ / "tables")) {
287 40 : std::string name = entry.path().filename().string();
288 40 : std::smatch match;
289 :
290 80 : if (entry.is_directory() && std::regex_match(name, match, local_regex) &&
291 40 : match.size() == 4) {
292 40 : std::string key = match[1].str();
293 40 : int version_major = std::stoi(match[2].str());
294 40 : int version_minor = std::stoi(match[3].str());
295 :
296 40 : if (version_major != COMPATIBLE_DATABASE_VERSION_MAJOR) {
297 6 : continue;
298 : }
299 :
300 34 : auto it = local_asset_info.find(key);
301 34 : if (it == local_asset_info.end() || version_minor > it->second.version_minor) {
302 33 : local_asset_info[key].version_minor = version_minor;
303 345 : for (const auto &subentry : fs::directory_iterator(entry)) {
304 156 : if (subentry.is_regular_file() && subentry.path().extension() == ".parquet") {
305 312 : local_asset_info[key].paths[subentry.path().stem().string()] = {
306 468 : subentry.path().string(), false};
307 : }
308 33 : }
309 : }
310 40 : }
311 55 : }
312 165 : }
313 :
314 8524 : void ParquetManager::update_local_asset(const std::string &key) {
315 8524 : assert(remote_asset_info.empty() == repo_paths_.empty());
316 :
317 : // If remote_asset_info is empty, we have nothing to do
318 8523 : if (remote_asset_info.empty()) {
319 8529 : return;
320 : }
321 :
322 : // Get remote version if available
323 4 : int remote_version = -1;
324 4 : auto remote_it = remote_asset_info.find(key);
325 4 : if (remote_it != remote_asset_info.end()) {
326 3 : remote_version = remote_it->second.version_minor;
327 : }
328 :
329 : // Get local version if available and check if it is up-to-date
330 : {
331 4 : int local_version = -1;
332 4 : std::shared_lock<std::shared_mutex> lock(mtx_local);
333 4 : auto local_it = local_asset_info.find(key);
334 4 : if (local_it != local_asset_info.end()) {
335 3 : local_version = local_it->second.version_minor;
336 : }
337 4 : if (local_version >= remote_version) {
338 1 : return;
339 : }
340 4 : }
341 :
342 : // If it is not up-to-date, acquire a unique lock for updating the table
343 3 : std::unique_lock<std::shared_mutex> lock(mtx_local);
344 :
345 : // Re-check if the table is up to date because another thread might have updated it
346 3 : int local_version = -1;
347 3 : auto local_it = local_asset_info.find(key);
348 3 : if (local_it != local_asset_info.end()) {
349 3 : local_version = local_it->second.version_minor;
350 : }
351 3 : if (local_version >= remote_version) {
352 0 : return;
353 : }
354 :
355 : // Download the remote file
356 3 : std::string endpoint = remote_it->second.endpoint;
357 3 : SPDLOG_INFO("Downloading {}_v{}.{} from {}", key, COMPATIBLE_DATABASE_VERSION_MAJOR,
358 : remote_version, endpoint);
359 :
360 3 : GitHubDownloader::Result result;
361 : try {
362 3 : result = downloader.download(endpoint, "", true).get();
363 0 : } catch (const std::exception &e) {
364 0 : react_on_exception(fmt::format("Failed to download table {}", endpoint), e);
365 0 : return;
366 0 : }
367 3 : if ((result.status_code == 403 || result.status_code == 429) &&
368 0 : result.rate_limit.remaining == 0) {
369 0 : react_on_rate_limit_reached(result.rate_limit.reset_time);
370 0 : return;
371 : }
372 3 : if (result.status_code != 200) {
373 0 : react_on_error_code(fmt::format("Failed to download table {}", endpoint),
374 : result.status_code);
375 0 : return;
376 : }
377 :
378 : // Unzip the downloaded file
379 3 : mz_zip_archive zip_archive{};
380 3 : if (mz_zip_reader_init_mem(&zip_archive, result.body.data(), result.body.size(), 0) == 0) {
381 0 : throw std::runtime_error("Failed to initialize zip archive.");
382 : }
383 :
384 6 : for (mz_uint i = 0; i < mz_zip_reader_get_num_files(&zip_archive); i++) {
385 : mz_zip_archive_file_stat file_stat;
386 3 : if (mz_zip_reader_file_stat(&zip_archive, i, &file_stat) == 0) {
387 0 : throw std::runtime_error("Failed to get file stat from zip archive.");
388 : }
389 :
390 : // Skip directories
391 3 : const char *filename = static_cast<const char *>(file_stat.m_filename);
392 3 : size_t len = std::strlen(filename);
393 3 : if (len > 0 && filename[len - 1] == '/') {
394 0 : continue;
395 : }
396 :
397 : // Ensure that the filename matches the expectations
398 3 : std::string dir = fs::path(filename).parent_path().string();
399 3 : std::string stem = fs::path(filename).stem().string();
400 3 : std::string suffix = fs::path(filename).extension().string();
401 3 : std::smatch match;
402 6 : if (!std::regex_match(dir, match, local_regex) || match.size() != 4 ||
403 3 : suffix != ".parquet") {
404 0 : throw std::runtime_error(
405 0 : fmt::format("Unexpected filename {} in zip archive.", filename));
406 : }
407 :
408 : // Ensure that the directory belongs to the asset that was requested
409 6 : if (match[1].str() != key ||
410 12 : std::stoi(match[2].str()) != COMPATIBLE_DATABASE_VERSION_MAJOR ||
411 6 : std::stoi(match[3].str()) != remote_version) {
412 0 : throw std::runtime_error(
413 0 : fmt::format("Expected the zip archive to contain the directory {}_v{}.{}, but it "
414 : "contains {}.",
415 0 : key, COMPATIBLE_DATABASE_VERSION_MAJOR, remote_version, dir));
416 : }
417 :
418 : // Construct the path to store the table
419 3 : auto path = directory_ / "tables" / dir / (stem + suffix);
420 3 : SPDLOG_INFO("Storing table to {}", path.string());
421 :
422 : // Extract the file to memory
423 3 : std::vector<char> buffer(file_stat.m_uncomp_size);
424 3 : if (mz_zip_reader_extract_to_mem(&zip_archive, i, buffer.data(), buffer.size(), 0) == 0) {
425 0 : throw std::runtime_error(fmt::format("Failed to extract {}.", filename));
426 : }
427 :
428 : // Ensure the parent directory exists
429 3 : fs::create_directories(path.parent_path());
430 :
431 : // Save the extracted file
432 3 : std::ofstream out(path.string(), std::ios::binary);
433 3 : if (!out) {
434 0 : throw std::runtime_error(fmt::format("Failed to open {} for writing", path.string()));
435 : }
436 3 : out.write(buffer.data(), static_cast<std::streamsize>(buffer.size()));
437 3 : out.close();
438 :
439 : // Update the local asset/table info
440 3 : local_asset_info[key].version_minor = remote_version;
441 3 : local_asset_info[key].paths[path.stem().string()] = {path.string(), false};
442 3 : }
443 :
444 3 : mz_zip_reader_end(&zip_archive);
445 6 : }
446 :
447 7266 : void ParquetManager::cache_table(std::unordered_map<std::string, PathInfo>::iterator table_it) {
448 : // Check if the table is already cached
449 : {
450 7266 : std::shared_lock<std::shared_mutex> lock(mtx_local);
451 7265 : if (table_it->second.cached) {
452 7228 : return;
453 : }
454 7263 : }
455 :
456 : // Acquire a unique lock for caching the table
457 35 : std::unique_lock<std::shared_mutex> lock(mtx_local);
458 :
459 : // Re-check if the table is already cached because another thread might have cached it
460 35 : if (table_it->second.cached) {
461 0 : return;
462 : }
463 :
464 : // Cache the table in memory
465 35 : std::string table_name;
466 : {
467 35 : auto result = con.Query(R"(SELECT UUID()::varchar)");
468 35 : if (result->HasError()) {
469 0 : throw cpptrace::runtime_error("Error selecting a unique table name: " +
470 0 : result->GetError());
471 : }
472 : table_name =
473 35 : duckdb::FlatVector::GetData<duckdb::string_t>(result->Fetch()->data[0])[0].GetString();
474 35 : }
475 :
476 : {
477 35 : auto result = con.Query(fmt::format(R"(CREATE TEMP TABLE '{}' AS SELECT * FROM '{}')",
478 70 : table_name, table_it->second.path));
479 35 : if (result->HasError()) {
480 0 : throw cpptrace::runtime_error("Error creating table: " + result->GetError());
481 : }
482 35 : }
483 :
484 35 : table_it->second.path = table_name;
485 35 : table_it->second.cached = true;
486 35 : }
487 :
488 : std::unordered_map<std::string, ParquetManager::LocalAssetInfo>::iterator
489 8543 : ParquetManager::get_local_asset(const std::string &key) {
490 : // Update the local table if a newer version is available remotely
491 8543 : this->update_local_asset(key);
492 :
493 : // Ensure availability of the local asset
494 8549 : auto asset_it = local_asset_info.find(key);
495 8534 : if (asset_it == local_asset_info.end()) {
496 : // If we do not know about any table that can be downloaded, downloading might be blocked.
497 : // Otherwise, the species might be misspelled.
498 1 : if (remote_asset_info.empty()) {
499 0 : throw std::runtime_error(
500 0 : "No tables found for species '" + key +
501 : "'. Check whether you have allowed downloading missing tables, or download the "
502 0 : "tables manually via `pairinteraction database download " +
503 0 : key + "`.");
504 : }
505 1 : throw std::runtime_error("No tables found for species '" + key +
506 2 : "'. Check the spelling of the species.");
507 : }
508 8543 : return asset_it;
509 : }
510 :
511 1275 : bool ParquetManager::has_table(const std::string &key, const std::string &table) {
512 1275 : auto asset_it = this->get_local_asset(key);
513 2510 : return asset_it->second.paths.contains(table);
514 : }
515 :
516 7276 : std::string ParquetManager::get_path(const std::string &key, const std::string &table) {
517 : // Ensure availability of the local table file
518 7276 : auto asset_it = this->get_local_asset(key);
519 7272 : auto table_it = asset_it->second.paths.find(table);
520 7270 : if (table_it == asset_it->second.paths.end()) {
521 1 : throw std::runtime_error("No table '" + table + ".parquet' found for species '" + key +
522 2 : "'. The tables for the species are incomplete.");
523 : }
524 :
525 : // Cache the local table in memory if requested
526 7269 : if (use_cache_) {
527 7266 : this->cache_table(table_it);
528 : }
529 :
530 : // Return the path to the local table file
531 14542 : return table_it->second.path;
532 : }
533 :
534 3 : std::string ParquetManager::get_versions_info() const {
535 : // Helper lambda returns the version string if available
536 54 : auto get_version = [](const auto &map, const std::string &table) -> int {
537 54 : if (auto it = map.find(table); it != map.end()) {
538 27 : return it->second.version_minor;
539 : }
540 27 : return -1;
541 : };
542 :
543 : // Gather all unique table names
544 3 : std::set<std::string> tables;
545 30 : for (const auto &entry : local_asset_info) {
546 27 : tables.insert(entry.first);
547 : }
548 3 : for (const auto &entry : remote_asset_info) {
549 0 : tables.insert(entry.first);
550 : }
551 :
552 : // Output versions info
553 3 : std::ostringstream oss;
554 :
555 3 : oss << " ";
556 3 : oss << std::left << std::setw(17) << "Asset";
557 3 : oss << std::left << std::setw(6 + 4) << "Local";
558 3 : oss << std::left << std::setw(7) << "Remote\n";
559 3 : oss << std::string(35, '-') << "\n";
560 :
561 30 : for (const auto &table : tables) {
562 27 : int local_version = get_version(local_asset_info, table);
563 27 : int remote_version = get_version(remote_asset_info, table);
564 :
565 : std::string comparator = (local_version < remote_version)
566 : ? "<"
567 27 : : ((local_version > remote_version) ? ">" : "==");
568 : std::string local_version_str = local_version == -1
569 27 : ? "N/A"
570 54 : : "v" + std::to_string(COMPATIBLE_DATABASE_VERSION_MAJOR) + "." +
571 81 : std::to_string(local_version);
572 : std::string remote_version_str = remote_version == -1
573 27 : ? "N/A"
574 27 : : "v" + std::to_string(COMPATIBLE_DATABASE_VERSION_MAJOR) + "." +
575 54 : std::to_string(remote_version);
576 :
577 27 : oss << " ";
578 27 : oss << std::left << std::setw(17) << table;
579 27 : oss << std::left << std::setw(6) << local_version_str;
580 27 : oss << std::left << std::setw(4) << comparator;
581 27 : oss << std::left << std::setw(7) << remote_version_str << "\n";
582 27 : }
583 6 : return oss.str();
584 3 : }
585 :
586 : } // namespace pairinteraction
|
