Skip to content

Commit

Permalink
Check if id column is present and is autoincrement
Browse files Browse the repository at this point in the history
  • Loading branch information
sytherax committed Dec 1, 2024
1 parent 76fcf48 commit 87b2f03
Showing 1 changed file with 113 additions and 0 deletions.
113 changes: 113 additions & 0 deletions src/db.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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!(
Expand Down

0 comments on commit 87b2f03

Please sign in to comment.