Upload environment.py

server/environment.py

@@ -322,10 +322,33 @@ class SQLQualityEnvironment:
                 # Find duplicate keys
                 cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
                 columns = [row[1] for row in cursor.fetchall()]
+                ground_truth["all_columns"] = columns
                 
                 if columns:
                     # Try to find duplicates on first column
                     key_col = columns[0]
+                    ground_truth["column"] = key_col
+                    ground_truth["offending_key"] = key_col
+                    
+                    # Count total duplicates
+                    cursor = conn.execute(f'''
+                        SELECT COUNT(*) FROM (
+                            SELECT "{key_col}"
+                            FROM "{target_table}" 
+                            GROUP BY "{key_col}" 
+                            HAVING COUNT(*) > 1
+                        )
+                    ''')
+                    dup_groups = cursor.fetchone()[0]
+                    
+                    # Count expected rows after dedup (unique keys)
+                    cursor = conn.execute(f'''
+                        SELECT COUNT(DISTINCT "{key_col}") FROM "{target_table}"
+                    ''')
+                    expected_clean = cursor.fetchone()[0]
+                    ground_truth["expected_clean_count"] = expected_clean
+                    
+                    # Get first duplicate example
                     cursor = conn.execute(f'''
                         SELECT "{key_col}", COUNT(*) as cnt 
                         FROM "{target_table}" 
@@ -337,6 +360,103 @@ class SQLQualityEnvironment:
                     if result:
                         ground_truth["duplicate_key"] = result[0]
                         ground_truth["duplicate_count"] = result[1]
+                        ground_truth["duplicate_groups"] = dup_groups
+                        
+            elif task_id == "type_repair" and target_table:
+                # Find columns that might have type issues (TEXT containing numbers)
+                cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
+                columns = [row[1] for row in cursor.fetchall()]
+                ground_truth["all_columns"] = columns
+                
+                # Find columns with numeric-like text values
+                for col in columns:
+                    try:
+                        # Check if column contains numeric-looking values stored as TEXT
+                        cursor = conn.execute(f'''
+                            SELECT "{col}" FROM "{target_table}" 
+                            WHERE "{col}" IS NOT NULL 
+                            LIMIT 100
+                        ''')
+                        values = [row[0] for row in cursor.fetchall()]
+                        
+                        # Check if majority look like numbers but stored as text
+                        numeric_like = 0
+                        for v in values:
+                            if v and str(v).strip():
+                                try:
+                                    float(str(v).replace(',', '').replace('$', ''))
+                                    numeric_like += 1
+                                except (ValueError, TypeError):
+                                    pass
+                        
+                        if numeric_like > len(values) * 0.5:
+                            ground_truth["column"] = col
+                            ground_truth["target_column"] = col
+                            ground_truth["issue_type"] = "TEXT_STORED_NUMERIC"
+                            break
+                    except Exception:
+                        continue
+                
+                # If no type issue found, default to first column
+                if "column" not in ground_truth and columns:
+                    ground_truth["column"] = columns[0]
+                    ground_truth["target_column"] = columns[0]
+                    
+            elif task_id == "fk_integrity" and target_table:
+                # For FK integrity, we need to simulate orphan records
+                # Since uploaded data may not have parent/child tables,
+                # we'll create a virtual parent table based on existing data
+                cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
+                columns = [row[1] for row in cursor.fetchall()]
+                ground_truth["all_columns"] = columns
+                
+                # Look for columns that might be foreign keys (ending in _id)
+                fk_candidates = [c for c in columns if c.lower().endswith('_id') and c.lower() != 'id']
+                
+                if fk_candidates:
+                    fk_col = fk_candidates[0]
+                    ground_truth["column"] = fk_col
+                    
+                    # Find values that would be orphans (using modulo to simulate missing parents)
+                    cursor = conn.execute(f'''
+                        SELECT DISTINCT "{fk_col}" FROM "{target_table}" 
+                        WHERE "{fk_col}" IS NOT NULL
+                    ''')
+                    all_ids = [row[0] for row in cursor.fetchall()]
+                    
+                    if all_ids:
+                        # Mark some IDs as invalid to simulate orphans
+                        import random
+                        random.seed(42)  # Deterministic for testing
+                        valid_ids = all_ids[:max(1, len(all_ids) * 2 // 3)]
+                        invalid_ids = [v for v in all_ids if v not in valid_ids]
+                        
+                        ground_truth["parent_table"] = f"{target_table}_parent"
+                        ground_truth["parent_key"] = fk_col.replace('_id', '')
+                        ground_truth["valid_invoice_ids"] = valid_ids
+                        ground_truth["expected_clean_count"] = len([v for v in all_ids if v in valid_ids])
+                        ground_truth["orphan_count"] = len(invalid_ids)
+                        
+            elif task_id == "schema_drift" and target_table:
+                # Schema drift compares two versions of same table
+                # For uploaded data, we'll treat it as canonical and create expected schema
+                cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
+                columns_info = cursor.fetchall()
+                columns = [row[1] for row in columns_info]
+                
+                ground_truth["table"] = target_table
+                ground_truth["canonical_table"] = target_table
+                ground_truth["expected_columns"] = columns
+                
+                # Get row count
+                cursor = conn.execute(f'SELECT COUNT(*) FROM "{target_table}"')
+                ground_truth["expected_row_count"] = cursor.fetchone()[0]
+                
+                # Get distinct IDs if available
+                if 'id' in [c.lower() for c in columns]:
+                    id_col = next(c for c in columns if c.lower() == 'id')
+                    cursor = conn.execute(f'SELECT COUNT(DISTINCT "{id_col}") FROM "{target_table}"')
+                    ground_truth["expected_distinct_ids"] = cursor.fetchone()[0]
                         
         except Exception as e:
             # Ground truth analysis failed, but we can still run the task