added id to db lookup

1 files changed, 9 insertions, 2 deletions

diff --git a/src/db.rs b/src/db.rs
index abfe1dd..6c7e4f8 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -50,8 +50,6 @@ pub fn mods_get_all_ids(config: Cfg) -> Result<Vec<String>, Box<dyn std::error::
 ///
 ///Will return `MLError` when no mod id is found
 pub fn mods_get_id(data: &str, slug: &str) -> MLE<String> {
-    //TODO check if "slug" is id
-
     let data = format!("{}/data.db", data);
     let connection = Connection::open(data)?;
 
@@ -64,6 +62,15 @@ pub fn mods_get_id(data: &str, slug: &str) -> MLE<String> {
     for id in id_iter {
         mod_id = id?;
     }
+    //get from id if no slug found
+    if mod_id.is_empty() {
+        let mut stmt = connection.prepare("SELECT id FROM mods WHERE id = ?")?;
+        let id_iter = stmt.query_map([slug], |row| row.get::<usize, String>(0))?;
+
+        for id in id_iter {
+            mod_id = id?;
+        }
+    }
     //get from title if no id found from slug
     if mod_id.is_empty() {
         let mut stmt = connection.prepare("SELECT id FROM mods WHERE title = ?")?;