put endpoint is now stronq boiii

server/src/tables_data/handlers/put_table_data.rs

@@ -1,8 +1,9 @@
 // src/tables_data/handlers/put_table_data.rs
 
 // TODO HEAVY REWORK IS NEEDED, SCRIPT FIXES ARE NEEDED FIRST
+
 use tonic::Status;
-use sqlx::{PgPool, Arguments};
+use sqlx::{PgPool, Arguments, Row};
 use sqlx::postgres::PgArguments;
 use chrono::{DateTime, Utc};
 use common::proto::multieko2::tables_data::{PutTableDataRequest, PutTableDataResponse};
@@ -15,6 +16,7 @@ use std::str::FromStr;
 use crate::steel::server::execution::{self, Value};
 use crate::steel::server::functions::SteelContext;
 use crate::indexer::{IndexCommand, IndexCommandData};
+use crate::table_script::handlers::dependency_analyzer::DependencyAnalyzer;
 use tokio::sync::mpsc;
 use tracing::error;
 
@@ -27,7 +29,6 @@ pub async fn put_table_data(
     let table_name = request.table_name;
     let record_id = request.id;
 
-    // An update with no fields is a no-op; we can return success early.
     if request.data.is_empty() {
         return Ok(PutTableDataResponse {
             success: true,
@@ -36,34 +37,26 @@ pub async fn put_table_data(
         });
     }
 
-    // --- Start of logic copied and adapted from post_table_data ---
+    // --- Metadata Fetching ---
-
+    let schema = sqlx::query!("SELECT id FROM schemas WHERE name = $1", profile_name)
-    let schema = sqlx::query!(
+        .fetch_optional(db_pool).await
-        "SELECT id FROM schemas WHERE name = $1",
+        .map_err(|e| Status::internal(format!("Profile lookup error: {}", e)))?
-        profile_name
+        .ok_or_else(|| Status::not_found("Profile not found"))?;
-    )
+    let schema_id = schema.id;
-    .fetch_optional(db_pool)
-    .await
-    .map_err(|e| Status::internal(format!("Profile lookup error: {}", e)))?;
-
-    let schema_id = schema.ok_or_else(|| Status::not_found("Profile not found"))?.id;
 
     let table_def = sqlx::query!(
-        r#"SELECT id, columns FROM table_definitions
+        "SELECT id, columns FROM table_definitions WHERE schema_id = $1 AND table_name = $2",
-        WHERE schema_id = $1 AND table_name = $2"#,
+        schema_id, table_name
-        schema_id,
-        table_name
     )
-    .fetch_optional(db_pool)
+    .fetch_optional(db_pool).await
-    .await
+    .map_err(|e| Status::internal(format!("Table lookup error: {}", e)))?
-    .map_err(|e| Status::internal(format!("Table lookup error: {}", e)))?;
+    .ok_or_else(|| Status::not_found("Table not found"))?;
-
-    let table_def = table_def.ok_or_else(|| Status::not_found("Table not found"))?;
 
     let columns_json: Vec<String> = serde_json::from_value(table_def.columns.clone())
         .map_err(|e| Status::internal(format!("Column parsing error: {}", e)))?;
 
     let mut columns = Vec::new();
+    let mut user_column_names = Vec::new();
     for col_def in columns_json {
         let parts: Vec<&str> = col_def.splitn(2, ' ').collect();
         if parts.len() != 2 {
@@ -71,9 +64,11 @@ pub async fn put_table_data(
         }
         let name = parts[0].trim_matches('"').to_string();
         let sql_type = parts[1].to_string();
+        user_column_names.push(name.clone());
         columns.push((name, sql_type));
     }
 
+    // --- Validate Column Permissions ---
     let fk_columns = sqlx::query!(
         r#"SELECT ltd.table_name
         FROM table_definition_links tdl
@@ -86,13 +81,14 @@ pub async fn put_table_data(
     .map_err(|e| Status::internal(format!("Foreign key lookup error: {}", e)))?;
 
     let mut system_columns = vec!["deleted".to_string()];
-    for fk in fk_columns {
+    // FIX 1: Change from `fk_columns` to `&fk_columns` to avoid move
+    for fk in &fk_columns {
         system_columns.push(format!("{}_id", fk.table_name));
     }
 
     let system_columns_set: std::collections::HashSet<_> = system_columns.iter().map(|s| s.as_str()).collect();
-
     let user_columns: Vec<&String> = columns.iter().map(|(name, _)| name).collect();
+
     for key in request.data.keys() {
         if !system_columns_set.contains(key.as_str()) &&
            !user_columns.contains(&&key.to_string()) {
@@ -100,72 +96,182 @@ pub async fn put_table_data(
         }
     }
 
-    let mut string_data_for_scripts = HashMap::new();
+    // --- [OPTIMIZATION] Smart Data Fetching: Only fetch what scripts need ---
-    for (key, proto_value) in &request.data {
+    let scripts = sqlx::query!("SELECT target_column, script FROM table_scripts WHERE table_definitions_id = $1", table_def.id)
-        let str_val = match &proto_value.kind {
+        .fetch_all(db_pool).await
-            Some(Kind::StringValue(s)) => {
-                let trimmed = s.trim();
-                if trimmed.is_empty() {
-                    continue;
-                }
-                trimmed.to_string()
-            },
-            Some(Kind::NumberValue(n)) => n.to_string(),
-            Some(Kind::BoolValue(b)) => b.to_string(),
-            Some(Kind::NullValue(_)) | None => continue,
-            Some(Kind::StructValue(_)) | Some(Kind::ListValue(_)) => {
-                return Err(Status::invalid_argument(format!("Unsupported type for script validation in column '{}'", key)));
-            }
-        };
-        string_data_for_scripts.insert(key.clone(), str_val);
-    }
-
-    let scripts = sqlx::query!(
-        "SELECT target_column, script FROM table_scripts WHERE table_definitions_id = $1",
-        table_def.id
-    )
-    .fetch_all(db_pool)
-    .await
         .map_err(|e| Status::internal(format!("Failed to fetch scripts: {}", e)))?;
 
+    let mut required_columns = std::collections::HashSet::new();
+
+    // Always need: id, target columns of scripts, and columns being updated
+    required_columns.insert("id".to_string());
+    for script_record in &scripts {
+        required_columns.insert(script_record.target_column.clone());
+    }
+    for key in request.data.keys() {
+        required_columns.insert(key.clone());
+    }
+
+    // Analyze script dependencies to find what columns scripts actually access
+    if !scripts.is_empty() {
+        let analyzer = DependencyAnalyzer::new(schema_id, db_pool.clone());
+
+        for script_record in &scripts {
+            let dependencies = analyzer
+                .analyze_script_dependencies(&script_record.script)
+                .map_err(|e| Status::internal(format!("Failed to analyze script dependencies: {:?}", e)))?;
+
+            for dep in dependencies {
+                // If it references this table, add the columns it uses
+                if dep.target_table == table_name {
+                    match dep.dependency_type {
+                        crate::table_script::handlers::dependency_analyzer::DependencyType::ColumnAccess { column } |
+                        crate::table_script::handlers::dependency_analyzer::DependencyType::IndexedAccess { column, .. } => {
+                            required_columns.insert(column);
+                        }
+                        _ => {} // SQL queries handled differently
+                    }
+                }
+                // If it references linked tables, add their foreign key columns
+                else {
+                    let fk_column = format!("{}_id", dep.target_table.split('_').last().unwrap_or(&dep.target_table));
+                    required_columns.insert(fk_column);
+                }
+            }
+        }
+    }
+
+    // Build optimized SELECT query with only required columns
+    let qualified_table = crate::shared::schema_qualifier::qualify_table_name_for_data(db_pool, &profile_name, &table_name).await?;
+
+    let columns_clause = required_columns
+        .iter()
+        .map(|name| format!("COALESCE(\"{0}\"::TEXT, '') AS \"{0}\"", name))
+        .collect::<Vec<_>>()
+        .join(", ");
+
+    let select_sql = format!("SELECT {} FROM {} WHERE id = $1", columns_clause, qualified_table);
+
+    let current_row = sqlx::query(&select_sql).bind(record_id).fetch_optional(db_pool).await
+        .map_err(|e| Status::internal(format!("Failed to fetch current row state: {}", e)))?
+        .ok_or_else(|| Status::not_found("Record not found"))?;
+
+    let mut current_row_data = HashMap::new();
+    for col_name in &required_columns {
+        let value: String = current_row.try_get(col_name.as_str()).unwrap_or_default();
+        current_row_data.insert(col_name.clone(), value);
+    }
+
+    // --- Data Merging Logic ---
+    let mut update_data = HashMap::new();
+    for (key, proto_value) in &request.data {
+        let str_val = match &proto_value.kind {
+            Some(Kind::StringValue(s)) => s.trim().to_string(),
+            Some(Kind::NumberValue(n)) => n.to_string(),
+            Some(Kind::BoolValue(b)) => b.to_string(),
+            Some(Kind::NullValue(_)) | None => String::new(),
+            _ => return Err(Status::invalid_argument(format!("Unsupported type for column '{}'", key))),
+        };
+        if !str_val.is_empty() {
+            update_data.insert(key.clone(), str_val);
+        }
+    }
+
+    let mut final_context_data = current_row_data.clone();
+    final_context_data.extend(update_data.clone());
+
+    // FIX 2: Type-aware script validation
     for script_record in scripts {
         let target_column = script_record.target_column;
 
-        if let Some(user_value) = string_data_for_scripts.get(&target_column) {
+        // Find the SQL type for this target column
+        let target_sql_type = if let Some((_, stype)) = columns.iter().find(|(name, _)| name == &target_column) {
+            stype.as_str()
+        } else {
+            "TEXT" // Default fallback for system columns
+        };
+
         let context = SteelContext {
             current_table: table_name.clone(),
             schema_id,
             schema_name: profile_name.clone(),
-                row_data: string_data_for_scripts.clone(),
+            row_data: final_context_data.clone(),
             db_pool: Arc::new(db_pool.clone()),
         };
 
-            let script_result = execution::execute_script(
+        let script_result = execution::execute_script(script_record.script, "STRINGS", Arc::new(db_pool.clone()), context)
-                script_record.script,
+            .map_err(|e| Status::invalid_argument(format!("Script execution failed for '{}': {}", target_column, e)))?;
-                "STRINGS",
-                Arc::new(db_pool.clone()),
-                context,
-            )
-            .map_err(|e| Status::invalid_argument(
-                format!("Script execution failed for '{}': {}", target_column, e)
-            ))?;
 
-            let Value::Strings(mut script_output) = script_result else {
+        let Value::Strings(mut script_output_vec) = script_result else {
             return Err(Status::internal("Script must return string values"));
         };
+        let script_output = script_output_vec.pop().ok_or_else(|| Status::internal("Script returned no values"))?;
 
-            let expected_value = script_output.pop()
+        if update_data.contains_key(&target_column) {
-                .ok_or_else(|| Status::internal("Script returned no values"))?;
+            // Case A: Column is being updated. Validate user input against script.
+            let user_value = update_data.get(&target_column).unwrap();
             
-            if user_value != &expected_value {
+            // TYPE-AWARE COMPARISON based on SQL type
-                return Err(Status::invalid_argument(format!(
+            let values_match = match target_sql_type {
-                    "Validation failed for column '{}': Expected '{}', Got '{}'",
+                s if s.starts_with("NUMERIC") => {
-                    target_column, expected_value, user_value
+                    // For NUMERIC columns, compare as decimals
-                )));
+                    let user_decimal = Decimal::from_str(user_value).map_err(|_| Status::invalid_argument(format!("Invalid decimal format for column '{}'", target_column)))?;
+                    let script_decimal = Decimal::from_str(&script_output).map_err(|_| Status::internal(format!("Script for '{}' produced invalid decimal", target_column)))?;
+                    user_decimal == script_decimal
+                },
+                "INTEGER" | "BIGINT" => {
+                    // For integer columns, compare as integers
+                    let user_int: i64 = user_value.parse().map_err(|_| Status::invalid_argument(format!("Invalid integer format for column '{}'", target_column)))?;
+                    let script_int: i64 = script_output.parse().map_err(|_| Status::internal(format!("Script for '{}' produced invalid integer", target_column)))?;
+                    user_int == script_int
+                },
+                "BOOLEAN" => {
+                    // For boolean columns, compare as booleans
+                    let user_bool: bool = user_value.parse().map_err(|_| Status::invalid_argument(format!("Invalid boolean format for column '{}'", target_column)))?;
+                    let script_bool: bool = script_output.parse().map_err(|_| Status::internal(format!("Script for '{}' produced invalid boolean", target_column)))?;
+                    user_bool == script_bool
+                },
+                _ => {
+                    // For TEXT, TIMESTAMPTZ, DATE, etc. - compare as strings
+                    user_value == &script_output
+                }
+            };
+
+            if !values_match {
+                return Err(Status::invalid_argument(format!("Validation failed for column '{}': Script calculated '{}', but user provided '{}'", target_column, script_output, user_value)));
+            }
+        } else {
+            // Case B: Column is NOT being updated. Prevent unauthorized changes.
+            let current_value = current_row_data.get(&target_column).cloned().unwrap_or_default();
+            
+            let values_match = match target_sql_type {
+                s if s.starts_with("NUMERIC") => {
+                    let current_decimal = Decimal::from_str(&current_value).unwrap_or_default();
+                    let script_decimal = Decimal::from_str(&script_output).unwrap_or_default();
+                    current_decimal == script_decimal
+                },
+                "INTEGER" | "BIGINT" => {
+                    let current_int: i64 = current_value.parse().unwrap_or_default();
+                    let script_int: i64 = script_output.parse().unwrap_or_default();
+                    current_int == script_int
+                },
+                "BOOLEAN" => {
+                    let current_bool: bool = current_value.parse().unwrap_or(false);
+                    let script_bool: bool = script_output.parse().unwrap_or(false);
+                    current_bool == script_bool
+                },
+                _ => {
+                    current_value == script_output
+                }
+            };
+
+            if !values_match {
+                return Err(Status::failed_precondition(format!("Script for column '{}' was triggered and would change its value from '{}' to '{}'. To apply this change, please include '{}' in your update request.", target_column, current_value, script_output, target_column)));
             }
         }
     }
 
+    // --- Database Update with Full Validation ---
     let mut params = PgArguments::default();
     let mut set_clauses = Vec::new();
     let mut param_idx = 1;
@@ -293,8 +399,6 @@ pub async fn put_table_data(
         param_idx += 1;
     }
 
-    // --- End of copied logic ---
-
     if set_clauses.is_empty() {
         return Ok(PutTableDataResponse {
             success: true,
@@ -303,13 +407,6 @@ pub async fn put_table_data(
         });
     }
 
-    let qualified_table = crate::shared::schema_qualifier::qualify_table_name_for_data(
-        db_pool,
-        &profile_name,
-        &table_name,
-    )
-    .await?;
-
     let set_clause = set_clauses.join(", ");
     let sql = format!(
         "UPDATE {} SET {} WHERE id = ${} RETURNING id",
@@ -340,6 +437,7 @@ pub async fn put_table_data(
         }
     };
 
+    // --- Indexer Command ---
     let command = IndexCommand::AddOrUpdate(IndexCommandData {
         table_name: table_name.clone(),
         row_id: updated_id,