atomicity when creating tables

This commit is contained in:
filipriec
2025-03-14 11:21:37 +01:00
parent ec323d24be
commit 8a7188d320

View File

@@ -1,6 +1,6 @@
// src/table_definition/handlers/post_table_definition.rs
use tonic::Status;
use sqlx::PgPool;
use sqlx::{PgPool, Transaction, Postgres};
use serde_json::json;
use time::OffsetDateTime;
use common::proto::multieko2::table_definition::{PostTableDefinitionRequest, TableDefinitionResponse};
@@ -55,6 +55,35 @@ pub async fn post_table_definition(
return Err(Status::invalid_argument("Invalid table name"));
}
// Start a transaction to ensure atomicity
let mut tx = db_pool.begin().await
.map_err(|e| Status::internal(format!("Failed to start transaction: {}", e)))?;
// Execute all database operations within the transaction
let result = execute_table_definition(&mut tx, request, table_name).await;
// Commit or rollback based on the result
match result {
Ok(response) => {
// Commit the transaction
tx.commit().await
.map_err(|e| Status::internal(format!("Failed to commit transaction: {}", e)))?;
Ok(response)
},
Err(e) => {
// Transaction will be automatically rolled back when dropped
// But we can explicitly roll it back too
let _ = tx.rollback().await;
Err(e)
}
}
}
async fn execute_table_definition(
tx: &mut Transaction<'_, Postgres>,
mut request: PostTableDefinitionRequest,
table_name: String,
) -> Result<TableDefinitionResponse, Status> {
// Lookup or create profile
let profile = sqlx::query!(
"INSERT INTO profiles (name) VALUES ($1)
@@ -62,7 +91,7 @@ pub async fn post_table_definition(
RETURNING id",
request.profile_name
)
.fetch_one(db_pool)
.fetch_one(&mut **tx)
.await
.map_err(|e| Status::internal(format!("Profile error: {}", e)))?;
@@ -75,7 +104,7 @@ pub async fn post_table_definition(
profile.id,
link.linked_table_name
)
.fetch_optional(db_pool)
.fetch_optional(&mut **tx)
.await
.map_err(|e| Status::internal(format!("Linked table lookup failed: {}", e)))?;
@@ -108,7 +137,7 @@ pub async fn post_table_definition(
}
// Generate SQL with multiple links
let (create_sql, index_sql) = generate_table_sql(db_pool, &table_name, &columns, &indexes, &links).await?;
let (create_sql, index_sql) = generate_table_sql(tx, &table_name, &columns, &indexes, &links).await?;
// Store main table definition
let table_def = sqlx::query!(
@@ -121,7 +150,7 @@ pub async fn post_table_definition(
json!(columns),
json!(indexes)
)
.fetch_one(db_pool)
.fetch_one(&mut **tx)
.await
.map_err(|e| {
if let Some(db_err) = e.as_database_error() {
@@ -142,20 +171,20 @@ pub async fn post_table_definition(
linked_id,
is_required
)
.execute(db_pool)
.execute(&mut **tx)
.await
.map_err(|e| Status::internal(format!("Failed to save link: {}", e)))?;
}
// Execute generated SQL
sqlx::query(&create_sql)
.execute(db_pool)
.execute(&mut **tx)
.await
.map_err(|e| Status::internal(format!("Table creation failed: {}", e)))?;
for sql in &index_sql {
sqlx::query(sql)
.execute(db_pool)
.execute(&mut **tx)
.await
.map_err(|e| Status::internal(format!("Index creation failed: {}", e)))?;
}
@@ -167,7 +196,7 @@ pub async fn post_table_definition(
}
async fn generate_table_sql(
db_pool: &PgPool,
tx: &mut Transaction<'_, Postgres>,
table_name: &str,
columns: &[String],
indexes: &[String],
@@ -181,7 +210,7 @@ async fn generate_table_sql(
// Add foreign key columns
let mut link_info = Vec::new();
for (linked_id, required) in links {
let linked_table = get_table_name_by_id(db_pool, *linked_id).await?;
let linked_table = get_table_name_by_id(tx, *linked_id).await?;
// Extract base name after year prefix
let base_name = linked_table.split_once('_')
@@ -190,7 +219,7 @@ async fn generate_table_sql(
.to_string();
let null_clause = if *required { "NOT NULL" } else { "" };
system_columns.push(
format!("\"{0}_id\" BIGINT {1} REFERENCES \"{2}\"(id)",
format!("\"{0}_id\" BIGINT {1} REFERENCES \"{2}\"(id)",
base_name, null_clause, linked_table
)
);
@@ -231,12 +260,15 @@ async fn generate_table_sql(
Ok((create_sql, all_indexes))
}
async fn get_table_name_by_id(db_pool: &PgPool, table_id: i64) -> Result<String, Status> {
async fn get_table_name_by_id(
tx: &mut Transaction<'_, Postgres>,
table_id: i64
) -> Result<String, Status> {
let record = sqlx::query!(
"SELECT table_name FROM table_definitions WHERE id = $1",
table_id
)
.fetch_one(db_pool)
.fetch_one(&mut **tx)
.await
.map_err(|e| Status::internal(format!("Table lookup failed: {}", e)))?;