feat: new db schema for model and template for engine

luke-nguyen990 · luke-nguyen990 · commit 9d1a9d8a2296 · 2024-11-12T17:51:45.000+07:00
diff --git a/engine/controllers/models.cc b/engine/controllers/models.cc
@@ -334,8 +334,17 @@ void Models::ImportModel(
     // Use relative path for model_yaml_path. In case of import, we use absolute path for model
     auto yaml_rel_path =
         fmu::ToRelativeCortexDataPath(fs::path(model_yaml_path));
-    cortex::db::ModelEntry model_entry{modelHandle, "local", "imported",
-                                       yaml_rel_path.string(), modelHandle};
+    cortex::db::ModelEntry model_entry {
+      modelHandle,         
+      "local",             
+      "imported",          
+      cortex::db::ModelStatus::Downloaded, 
+      "",                  
+      "",                  
+      "",                  
+      yaml_rel_path.string(), 
+      modelHandle          
+    };
 
     std::filesystem::create_directories(
         std::filesystem::path(model_yaml_path).parent_path());
diff --git a/engine/database/engines.cc b/engine/database/engines.cc
@@ -0,0 +1,19 @@
+#include "engines.h"
+#include "database.h"
+
+namespace cortex::db {
+
+Engines::Engines() : db_(cortex::db::Database::GetInstance().db()) {
+  db_.exec(
+      "CREATE TABLE IF NOT EXISTS engines ("
+      "engine_id TEXT PRIMARY KEY,"
+      "type TEXT,"
+      "api_key TEXT,"
+      "url TEXT,"
+      "version TEXT,"
+      "variant TEXT,"
+      "status TEXT,"
+      "metadata TEXT);");
+}
+Engines::~Engines() {}
+}
diff --git a/engine/database/engines.h b/engine/database/engines.h
@@ -0,0 +1,32 @@
+#pragma once
+
+#include <SQLiteCpp/Database.h>
+#include <trantor/utils/Logger.h>
+#include <string>
+#include <vector>
+#include "utils/result.hpp"
+
+namespace cortex::db {
+
+struct EngineEntry {
+    std::string engine;
+};
+
+class Engines {
+
+ private:
+  SQLite::Database& db_;
+
+  bool IsUnique(const std::vector<EngineEntry>& entries,
+                const std::string& model_id,
+                const std::string& model_alias) const;
+
+  cpp::result<std::vector<EngineEntry>, std::string> LoadModelListNoLock() const;
+
+ public:
+  Engines();
+  Engines(SQLite::Database& db);
+  ~Engines();
+};
+
+}  // namespace cortex::db
diff --git a/engine/database/models.cc b/engine/database/models.cc
@@ -12,6 +12,10 @@ Models::Models() : db_(cortex::db::Database::GetInstance().db()) {
   db_.exec(
       "CREATE TABLE IF NOT EXISTS models ("
       "model_id TEXT PRIMARY KEY,"
+      "model_format TEXT,"
+      "model_source TEXT,"
+      "status TEXT,"
+      "engine TEXT,"
       "author_repo_id TEXT,"
       "branch_name TEXT,"
       "path_to_model_yaml TEXT,"
@@ -22,14 +26,40 @@ Models::Models(SQLite::Database& db) : db_(db) {
   db_.exec(
       "CREATE TABLE IF NOT EXISTS models ("
       "model_id TEXT PRIMARY KEY,"
+      "model_format TEXT,"
+      "model_source TEXT,"
+      "status TEXT,"
+      "engine TEXT,"
       "author_repo_id TEXT,"
       "branch_name TEXT,"
       "path_to_model_yaml TEXT,"
-      "model_alias TEXT UNIQUE);");
+      "model_alias TEXT);");
 }
-
 Models::~Models() {}
 
+std::string Models::StatusToString(ModelStatus status) const {
+  switch (status) {
+    case ModelStatus::Remote:
+      return "remote";
+    case ModelStatus::Downloaded:
+      return "downloaded";
+    case ModelStatus::Undownloaded:
+      return "undownloaded";
+  }
+  return "unknown";
+}
+
+ModelStatus Models::StringToStatus(const std::string& status_str) const {
+  if (status_str == "remote") {
+    return ModelStatus::Remote;
+  } else if (status_str == "downloaded") {
+    return ModelStatus::Downloaded;
+  } else if (status_str == "undownloaded") {
+    return ModelStatus::Undownloaded;
+  }
+  throw std::invalid_argument("Invalid status string");
+}
+
 cpp::result<std::vector<ModelEntry>, std::string> Models::LoadModelList()
     const {
   try {
@@ -57,16 +87,21 @@ cpp::result<std::vector<ModelEntry>, std::string> Models::LoadModelListNoLock()
   try {
     std::vector<ModelEntry> entries;
     SQLite::Statement query(db_,
-                            "SELECT model_id, author_repo_id, branch_name, "
+                            "SELECT model_id, model_format, model_source, "
+                            "status, engine, author_repo_id, branch_name, "
                             "path_to_model_yaml, model_alias FROM models");
 
     while (query.executeStep()) {
       ModelEntry entry;
       entry.model = query.getColumn(0).getString();
-      entry.author_repo_id = query.getColumn(1).getString();
-      entry.branch_name = query.getColumn(2).getString();
-      entry.path_to_model_yaml = query.getColumn(3).getString();
-      entry.model_alias = query.getColumn(4).getString();
+      entry.model_format = query.getColumn(1).getString();
+      entry.model_source = query.getColumn(2).getString();
+      entry.status = StringToStatus(query.getColumn(3).getString());
+      entry.engine = query.getColumn(4).getString();
+      entry.author_repo_id = query.getColumn(5).getString();
+      entry.branch_name = query.getColumn(6).getString();
+      entry.path_to_model_yaml = query.getColumn(7).getString();
+      entry.model_alias = query.getColumn(8).getString();
       entries.push_back(entry);
     }
     return entries;
@@ -140,7 +175,8 @@ cpp::result<ModelEntry, std::string> Models::GetModelInfo(
     const std::string& identifier) const {
   try {
     SQLite::Statement query(db_,
-                            "SELECT model_id, author_repo_id, branch_name, "
+                            "SELECT model_id, model_format, model_source, "
+                            "status, engine, author_repo_id, branch_name, "
                             "path_to_model_yaml, model_alias FROM models "
                             "WHERE model_id = ? OR model_alias = ?");
 
@@ -149,10 +185,14 @@ cpp::result<ModelEntry, std::string> Models::GetModelInfo(
     if (query.executeStep()) {
       ModelEntry entry;
       entry.model = query.getColumn(0).getString();
-      entry.author_repo_id = query.getColumn(1).getString();
-      entry.branch_name = query.getColumn(2).getString();
-      entry.path_to_model_yaml = query.getColumn(3).getString();
-      entry.model_alias = query.getColumn(4).getString();
+      entry.model_format = query.getColumn(1).getString();
+      entry.model_source = query.getColumn(2).getString();
+      entry.status = StringToStatus(query.getColumn(3).getString());
+      entry.engine = query.getColumn(4).getString();
+      entry.author_repo_id = query.getColumn(5).getString();
+      entry.branch_name = query.getColumn(6).getString();
+      entry.path_to_model_yaml = query.getColumn(7).getString();
+      entry.model_alias = query.getColumn(8).getString();
       return entry;
     } else {
       return cpp::fail("Model not found: " + identifier);
@@ -164,6 +204,10 @@ cpp::result<ModelEntry, std::string> Models::GetModelInfo(
 
 void Models::PrintModelInfo(const ModelEntry& entry) const {
   LOG_INFO << "Model ID: " << entry.model;
+  LOG_INFO << "Model Format: " << entry.model_format;
+  LOG_INFO << "Model Source: " << entry.model_source;
+  LOG_INFO << "Status: " << StatusToString(entry.status);
+  LOG_INFO << "Engine: " << entry.engine;
   LOG_INFO << "Author/Repo ID: " << entry.author_repo_id;
   LOG_INFO << "Branch Name: " << entry.branch_name;
   LOG_INFO << "Path to model.yaml: " << entry.path_to_model_yaml;
@@ -188,14 +232,18 @@ cpp::result<bool, std::string> Models::AddModelEntry(ModelEntry new_entry,
 
       SQLite::Statement insert(
           db_,
-          "INSERT INTO models (model_id, author_repo_id, "
-          "branch_name, path_to_model_yaml, model_alias) VALUES (?, ?, "
-          "?, ?, ?)");
+          "INSERT INTO models (model_id, model_format, model_source, status, "
+          "engine, author_repo_id, branch_name, path_to_model_yaml, model_alias) "
+          "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)");
       insert.bind(1, new_entry.model);
-      insert.bind(2, new_entry.author_repo_id);
-      insert.bind(3, new_entry.branch_name);
-      insert.bind(4, new_entry.path_to_model_yaml);
-      insert.bind(5, new_entry.model_alias);
+      insert.bind(2, new_entry.model_format);
+      insert.bind(3, new_entry.model_source);
+      insert.bind(4, StatusToString(new_entry.status));
+      insert.bind(5, new_entry.engine);
+      insert.bind(6, new_entry.author_repo_id);
+      insert.bind(7, new_entry.branch_name);
+      insert.bind(8, new_entry.path_to_model_yaml);
+      insert.bind(9, new_entry.model_alias);
       insert.exec();
 
       return true;
@@ -215,14 +263,19 @@ cpp::result<bool, std::string> Models::UpdateModelEntry(
   try {
     SQLite::Statement upd(db_,
                           "UPDATE models "
-                          "SET author_repo_id = ?, branch_name = ?, "
+                          "SET model_format = ?, model_source = ?, status = ?, "
+                          "engine = ?, author_repo_id = ?, branch_name = ?, "
                           "path_to_model_yaml = ? "
                           "WHERE model_id = ? OR model_alias = ?");
-    upd.bind(1, updated_entry.author_repo_id);
-    upd.bind(2, updated_entry.branch_name);
-    upd.bind(3, updated_entry.path_to_model_yaml);
-    upd.bind(4, identifier);
-    upd.bind(5, identifier);
+    upd.bind(1, updated_entry.model_format);
+    upd.bind(2, updated_entry.model_source);
+    upd.bind(3, StatusToString(updated_entry.status));
+    upd.bind(4, updated_entry.engine);
+    upd.bind(5, updated_entry.author_repo_id);
+    upd.bind(6, updated_entry.branch_name);
+    upd.bind(7, updated_entry.path_to_model_yaml);
+    upd.bind(8, identifier);
+    upd.bind(9, identifier);
     return upd.exec() == 1;
   } catch (const std::exception& e) {
     return cpp::fail(e.what());
@@ -305,4 +358,5 @@ bool Models::HasModel(const std::string& identifier) const {
     return false;
   }
 }
-}  // namespace cortex::db
+
+}  // namespace cortex::db
diff --git a/engine/database/models.h b/engine/database/models.h
@@ -7,8 +7,19 @@
 #include "utils/result.hpp"
 
 namespace cortex::db {
+
+enum class ModelStatus {
+  Remote,
+  Downloaded,
+  Undownloaded
+};
+
 struct ModelEntry {
   std::string model;
+  std::string model_format;
+  std::string model_source;
+  ModelStatus status;
+  std::string engine;
   std::string author_repo_id;
   std::string branch_name;
   std::string path_to_model_yaml;
@@ -26,6 +37,9 @@ class Models {
 
   cpp::result<std::vector<ModelEntry>, std::string> LoadModelListNoLock() const;
 
+  std::string StatusToString(ModelStatus status) const;
+  ModelStatus StringToStatus(const std::string& status_str) const;
+
  public:
   static const std::string kModelListPath;
   cpp::result<std::vector<ModelEntry>, std::string> LoadModelList() const;
@@ -50,4 +64,5 @@ class Models {
       const std::string& identifier) const;
   bool HasModel(const std::string& identifier) const;
 };
-}  // namespace cortex::db
+
+}  // namespace cortex::db
diff --git a/engine/extensions/remote-engine/TemplateRenderer.cc b/engine/extensions/remote-engine/TemplateRenderer.cc
@@ -1,5 +1,7 @@
 #if defined(_WIN32) || defined(_WIN64)
 #define NOMINMAX
+#undef min
+#undef max
 #endif
 #include "TemplateRenderer.h"
 #include <regex>
diff --git a/engine/extensions/remote-engine/TemplateRenderer.h b/engine/extensions/remote-engine/TemplateRenderer.h
@@ -8,6 +8,8 @@
 // clang-format off
 #if defined(_WIN32) || defined(_WIN64)
 #define NOMINMAX
+#undef min
+#undef max
 #endif
 #include <nlohmann/json.hpp>
 #include <inja/inja.hpp>
diff --git a/engine/test/components/test_models_db.cc b/engine/test/components/test_models_db.cc
@@ -6,6 +6,7 @@ namespace cortex::db {
 namespace {
 constexpr const auto kTestDb = "./test.db";
 }
+
 class ModelsTestSuite : public ::testing::Test {
  public:
   ModelsTestSuite()
@@ -21,9 +22,9 @@ class ModelsTestSuite : public ::testing::Test {
   SQLite::Database db_;
   cortex::db::Models model_list_;
 
-  const cortex::db::ModelEntry kTestModel{"test_model_id", "test_author",
-                                          "main", "/path/to/model.yaml",
-                                          "test_alias"};
+  const cortex::db::ModelEntry kTestModel{
+      "test_model_id", "test_format", "test_source", cortex::db::ModelStatus::Downloaded, "test_engine",
+      "test_author", "main", "/path/to/model.yaml", "test_alias"};
 };
 
 TEST_F(ModelsTestSuite, TestAddModelEntry) {
@@ -33,8 +34,12 @@ TEST_F(ModelsTestSuite, TestAddModelEntry) {
   EXPECT_TRUE(retrieved_model);
   EXPECT_EQ(retrieved_model.value().model, kTestModel.model);
   EXPECT_EQ(retrieved_model.value().author_repo_id, kTestModel.author_repo_id);
+  EXPECT_EQ(retrieved_model.value().model_format, kTestModel.model_format);
+  EXPECT_EQ(retrieved_model.value().model_source, kTestModel.model_source);
+  EXPECT_EQ(retrieved_model.value().status, kTestModel.status);
+  EXPECT_EQ(retrieved_model.value().engine, kTestModel.engine);
 
-  // // Clean up
+  // Clean up
   EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model).value());
 }
 
@@ -59,14 +64,14 @@ TEST_F(ModelsTestSuite, TestUpdateModelEntry) {
   EXPECT_TRUE(model_list_.AddModelEntry(kTestModel).value());
 
   cortex::db::ModelEntry updated_model = kTestModel;
+  updated_model.status = cortex::db::ModelStatus::Downloaded;
 
   EXPECT_TRUE(
       model_list_.UpdateModelEntry(kTestModel.model, updated_model).value());
 
   auto retrieved_model = model_list_.GetModelInfo(kTestModel.model);
   EXPECT_TRUE(retrieved_model);
-  EXPECT_TRUE(
-      model_list_.UpdateModelEntry(kTestModel.model, updated_model).value());
+  EXPECT_EQ(retrieved_model.value().status, updated_model.status);
 
   // Clean up
   EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model).value());
@@ -162,4 +167,5 @@ TEST_F(ModelsTestSuite, TestHasModel) {
   // Clean up
   EXPECT_TRUE(model_list_.DeleteModelEntry(kTestModel.model).value());
 }
-}  // namespace cortex::db
+
+}  // namespace cortex::db
