diff --git a/server/src/table_definition/handlers/post_table_definition.rs b/server/src/table_definition/handlers/post_table_definition.rs index 0c4d3d1..b5730bb 100644 --- a/server/src/table_definition/handlers/post_table_definition.rs +++ b/server/src/table_definition/handlers/post_table_definition.rs @@ -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 { // 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 { +async fn get_table_name_by_id( + tx: &mut Transaction<'_, Postgres>, + table_id: i64 +) -> Result { 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)))?;