From 87b2f03ce8b3e997014f5eff346c489889489815 Mon Sep 17 00:00:00 2001
From: sytherax <swooruptrader@gmail.com>
Date: Sun, 1 Dec 2024 22:41:46 +0545
Subject: [PATCH] Check if id column is present and is autoincrement

---
 src/db.rs | 113 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 113 insertions(+)

diff --git a/src/db.rs b/src/db.rs
index ecbbe9b9..7f7d4ce8 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -276,6 +276,108 @@ where
     Ok(())
 }
 
+/// Checks if the specified table has an 'id' column.
+///
+/// This function checks if the specified table has an 'id' column, which is a common primary key column.
+/// It supports Postgres, SQLite, and MySQL database backends.
+///
+/// # Arguments
+///
+/// - `db`: A reference to the `DatabaseConnection`.
+/// - `db_backend`: A reference to the `DatabaseBackend`.
+/// - `table_name`: The name of the table to check.
+///
+/// # Returns
+///
+/// A `Result` containing a `bool` indicating whether the table has an 'id' column.
+async fn has_id_column(
+    db: &DatabaseConnection,
+    db_backend: &DatabaseBackend,
+    table_name: &str,
+) -> crate::Result<bool> {
+    // First check if 'id' column exists
+    let result = match db_backend {
+        DatabaseBackend::Postgres => {
+            let query = format!(
+                "SELECT EXISTS (
+              SELECT 1 
+              FROM information_schema.columns 
+              WHERE table_name = '{table_name}' 
+              AND column_name = 'id'
+          )"
+            );
+            let result = db
+                .query_one(Statement::from_string(DatabaseBackend::Postgres, query))
+                .await?;
+            result.map_or(false, |row| {
+                row.try_get::<bool>("", "exists").unwrap_or(false)
+            })
+        }
+        DatabaseBackend::Sqlite => {
+            let query = format!(
+                "SELECT COUNT(*) as count 
+          FROM pragma_table_info('{table_name}') 
+          WHERE name = 'id'"
+            );
+            let result = db
+                .query_one(Statement::from_string(DatabaseBackend::Sqlite, query))
+                .await?;
+            result.map_or(false, |row| {
+                row.try_get::<i32>("", "count").unwrap_or(0) > 0
+            })
+        }
+        DatabaseBackend::MySql => {
+            return Err(Error::Message(
+                "Unsupported database backend: MySQL".to_string(),
+            ))
+        }
+    };
+
+    Ok(result)
+}
+
+/// Checks whether the specified table has an auto-increment 'id' column.
+///
+/// # Returns
+///
+/// A `Result` containing a `bool` indicating whether the table has an auto-increment 'id' column.
+async fn is_auto_increment(
+    db: &DatabaseConnection,
+    db_backend: &DatabaseBackend,
+    table_name: &str,
+) -> crate::Result<bool> {
+    let result = match db_backend {
+        DatabaseBackend::Postgres => {
+            let query = format!(
+                "SELECT pg_get_serial_sequence('{table_name}', 'id') IS NOT NULL as is_serial"
+            );
+            let result = db
+                .query_one(Statement::from_string(DatabaseBackend::Postgres, query))
+                .await?;
+            result.map_or(false, |row| {
+                row.try_get::<bool>("", "is_serial").unwrap_or(false)
+            })
+        }
+        DatabaseBackend::Sqlite => {
+            let query =
+                format!("SELECT sql FROM sqlite_master WHERE type='table' AND name='{table_name}'");
+            let result = db
+                .query_one(Statement::from_string(DatabaseBackend::Sqlite, query))
+                .await?;
+            result.map_or(false, |row| {
+                row.try_get::<String>("", "sql")
+                    .map_or(false, |sql| sql.to_lowercase().contains("autoincrement"))
+            })
+        }
+        DatabaseBackend::MySql => {
+            return Err(Error::Message(
+                "Unsupported database backend: MySQL".to_string(),
+            ))
+        }
+    };
+    Ok(result)
+}
+
 /// Function to reset auto-increment
 /// # Errors
 /// Returns error if it fails
@@ -284,6 +386,17 @@ pub async fn reset_autoincrement(
     table_name: &str,
     db: &DatabaseConnection,
 ) -> crate::Result<()> {
+    // Check if 'id' column exists
+    let has_id_column = has_id_column(db, &db_backend, table_name).await?;
+    if !has_id_column {
+        return Ok(());
+    }
+    // Check if 'id' column is auto-increment
+    let is_auto_increment = is_auto_increment(db, &db_backend, table_name).await?;
+    if !is_auto_increment {
+        return Ok(());
+    }
+
     match db_backend {
         DatabaseBackend::Postgres => {
             let query_str = format!(