Add new loading mechanism for GeoIP context

libtenzir/builtins/contexts/geoip.cpp

@@ -6,6 +6,9 @@
 // SPDX-FileCopyrightText: (c) 2023 The VAST Contributors
 // SPDX-License-Identifier: BSD-3-Clause
 
+#include "tenzir/detail/posix.hpp"
+#include "tenzir/logger.hpp"
+
 #include <tenzir/arrow_table_slice.hpp>
 #include <tenzir/data.hpp>
 #include <tenzir/error.hpp>
@@ -22,9 +25,11 @@
 
 #include <chrono>
 #include <cstdint>
+#include <cstdio>
 #include <maxminddb.h>
 #include <memory>
 #include <string>
+#include <utility>
 
 namespace tenzir::plugins::geoip {
 
@@ -44,6 +49,11 @@ struct mmdb_deleter final {
 using mmdb_ptr = std::unique_ptr<MMDB_s, mmdb_deleter>;
 
 auto make_mmdb(const std::string& path) -> caf::expected<mmdb_ptr> {
+  if (!std::filesystem::exists(path)) {
+    return diagnostic::error("")
+      .note("failed to find path `{}`", path)
+      .to_error();
+  }
   auto ptr = new MMDB_s;
   const auto status = MMDB_open(path.c_str(), MMDB_MODE_MMAP, ptr);
   if (status != MMDB_SUCCESS) {
@@ -302,6 +312,11 @@ class ctx final : public virtual context {
   /// Emits context information for every event in `slice` in order.
   auto apply(series array, bool replace) const
     -> caf::expected<std::vector<series>> override {
+    if (!mmdb_) {
+      return caf::make_error(ec::lookup_error,
+                             fmt::format("no GeoIP data currently exists for "
+                                         "this context"));
+    }
     auto status = 0;
     MMDB_entry_data_list_s* entry_data_list = nullptr;
     auto builder = series_builder{};
@@ -475,6 +490,11 @@ class ctx final : public virtual context {
   }
 
   auto reset() -> caf::expected<void> override {
+    if (!mmdb_) {
+      return caf::make_error(ec::lookup_error,
+                             fmt::format("no GeoIP data currently exists for "
+                                         "this context"));
+    }
     auto mmdb = make_mmdb(db_path_);
     if (not mmdb) {
       return mmdb.error();
@@ -490,18 +510,19 @@ class ctx final : public virtual context {
   }
 
   auto save() const -> caf::expected<save_result> override {
-    auto builder = flatbuffers::FlatBufferBuilder{};
-    auto path = builder.CreateString(db_path_);
-    fbs::context::geoip::GeoIPDataBuilder geoip_builder(builder);
-    geoip_builder.add_url(path);
-    auto geoip_data = geoip_builder.Finish();
-    fbs::context::geoip::FinishGeoIPDataBuffer(builder, geoip_data);
-    return save_result{.data = tenzir::fbs::release(builder), .version = 1};
+    if (!mmdb_) {
+      return caf::make_error(ec::lookup_error,
+                             fmt::format("no GeoIP data currently exists for "
+                                         "this context"));
+    }
+    return save_result{.data = std::move(chunk::mmap(db_path_).value()),
+                       .version = latest_version};
   }
 
 private:
   std::string db_path_;
   mmdb_ptr mmdb_;
+  int latest_version = 2;
 };
 
 struct v1_loader : public context_loader {
@@ -531,9 +552,59 @@ struct v1_loader : public context_loader {
   }
 };
 
+struct v2_loader : public context_loader {
+  explicit v2_loader(record global_config)
+    : global_config_{std::move(global_config)} {
+  }
+
+  auto version() const -> int {
+    return 2;
+  }
+
+  auto load(chunk_ptr serialized) const
+    -> caf::expected<std::unique_ptr<context>> {
+    const auto* cache_dir
+      = get_if<std::string>(&global_config_, "tenzir.cache-directory");
+    std::filesystem::create_directories(*cache_dir);
+    const auto current_time = std::time(nullptr);
+    std::string temp_file_name = *cache_dir + std::to_string(current_time);
+    auto temp_file = std::fstream(temp_file_name, std::ios_base::out);
+    if (!temp_file) {
+      return caf::make_error(ec::filesystem_error,
+                             fmt::format("failed to open temp file on "
+                                         "data load: {}",
+                                         detail::describe_errno()));
+    }
+    temp_file.write(reinterpret_cast<const char*>(serialized->data()),
+                    static_cast<std::streamsize>(serialized->size()));
+    if (!temp_file) {
+      return caf::make_error(ec::filesystem_error,
+                             fmt::format("failed write the temp file "
+                                         "on data load: {}",
+                                         detail::describe_errno()));
+    }
+    temp_file.close();
+    if (!temp_file) {
+      return caf::make_error(ec::filesystem_error,
+                             fmt::format("failed close the temp file: {}",
+                                         detail::describe_errno()));
+    }
+    auto mmdb = make_mmdb(temp_file_name);
+    if (not mmdb) {
+      return mmdb.error();
+    }
+    return std::make_unique<ctx>(std::move(temp_file_name), std::move(*mmdb));
+  }
+
+private:
+  const record global_config_;
+};
+
 class plugin : public virtual context_plugin {
-  auto initialize(const record&, const record&) -> caf::error override {
+  auto initialize(const record&, const record& global_config)
+    -> caf::error override {
     register_loader(std::make_unique<v1_loader>());
+    register_loader(std::make_unique<v2_loader>(global_config));
     return caf::none;
   }
 
@@ -559,9 +630,7 @@ class plugin : public virtual context_plugin {
         .to_error();
     }
     if (db_path.empty()) {
-      return diagnostic::error("missing required db-path option")
-        .usage("context create <name> geoip --db-path <path>")
-        .to_error();
+      return std::make_unique<ctx>(std::move(db_path), nullptr);
     }
     auto mmdb = make_mmdb(db_path);
     if (not mmdb) {

nix/tenzir/plugins/source.json

@@ -2,7 +2,8 @@
   "name": "tenzir-plugins",
   "url": "git@github.com:tenzir/tenzir-plugins",
   "ref": "main",
-  "rev": "135363991f09bb9bb3682bbd81d145f7a2b2e60c",
+  "rev": "14ceb6f3f509c6ef99928004e3e6b51e6379757d",
   "submodules": true,
-  "shallow": true
+  "shallow": true,
+  "allRefs": true
 }