db schema update follow up

Author: luke-nguyen990

engine/database/models.cc

@@ -22,6 +22,19 @@ Models::Models() : db_(cortex::db::Database::GetInstance().db()) {
       "model_alias TEXT);");
 }
 
+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);");
+}
 Models::~Models() {}
 
 cpp::result<std::vector<ModelEntry>, std::string> Models::LoadModelList()
@@ -51,16 +64,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 = 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;
@@ -134,7 +152,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 = ?");
 
@@ -143,10 +162,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 = 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);
@@ -158,6 +181,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: " << 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;
@@ -182,14 +209,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, 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;
@@ -209,14 +240,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, 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());
@@ -299,4 +335,5 @@ bool Models::HasModel(const std::string& identifier) const {
     return false;
   }
 }
-}  // namespace cortex::db
+
+}  // namespace cortex::db

engine/database/models.h

@@ -7,8 +7,13 @@
 #include "utils/result.hpp"
 
 namespace cortex::db {
+
 struct ModelEntry {
   std::string model;
+  std::string model_format;  // Ensure this is present
+  std::string model_source;
+  std::string status;
+  std::string engine;
   std::string author_repo_id;
   std::string branch_name;
   std::string path_to_model_yaml;
@@ -50,4 +55,5 @@ class Models {
       const std::string& identifier) const;
   bool HasModel(const std::string& identifier) const;
 };
-}  // namespace cortex::db
+
+}  // namespace cortex::db