From 4ddf683bb5515406cf269cffc1d2d9627306f77d Mon Sep 17 00:00:00 2001
From: Neil Macneale IV <neil@dolthub.com>
Date: Thu, 12 Dec 2024 16:51:08 -0800
Subject: [PATCH] Still produce warning when the schema diff is suspect

---
 .../sqle/dtables/commit_diff_table.go         |  2 +-
 .../doltcore/sqle/dtables/diff_table.go       | 30 ++++++++++++-------
 .../sqle/enginetest/dolt_queries_diff.go      |  1 +
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/go/libraries/doltcore/sqle/dtables/commit_diff_table.go b/go/libraries/doltcore/sqle/dtables/commit_diff_table.go
index 4eeee2a812..dcc50a12d5 100644
--- a/go/libraries/doltcore/sqle/dtables/commit_diff_table.go
+++ b/go/libraries/doltcore/sqle/dtables/commit_diff_table.go
@@ -216,7 +216,7 @@ func (dt *CommitDiffTable) LookupPartitions(ctx *sql.Context, i sql.IndexLookup)
 		fromSch:  dt.targetSchema,
 	}
 
-	isDiffable, err := dp.isDiffablePartition(ctx)
+	isDiffable, _, err := dp.isDiffablePartition(ctx)
 	if err != nil {
 		return nil, err
 	}
diff --git a/go/libraries/doltcore/sqle/dtables/diff_table.go b/go/libraries/doltcore/sqle/dtables/diff_table.go
index a41ef9c5de..450f3327ec 100644
--- a/go/libraries/doltcore/sqle/dtables/diff_table.go
+++ b/go/libraries/doltcore/sqle/dtables/diff_table.go
@@ -685,39 +685,39 @@ func (dp DiffPartition) GetRowIter(ctx *sql.Context, ddb *doltdb.DoltDB, joiner
 // isDiffablePartition checks if the commit pair for this partition is "diffable".
 // If the primary key sets changed between the two commits, it may not be
 // possible to diff them.
-func (dp *DiffPartition) isDiffablePartition(ctx *sql.Context) (bool, error) {
+func (dp *DiffPartition) isDiffablePartition(ctx *sql.Context) (simpleDiff bool, fuzzyDiff bool, err error) {
 	// dp.to is nil when a table has been deleted previously. In this case, we return
 	// false, to stop processing diffs, since that previously deleted table is considered
 	// a logically different table and we don't want to mix the diffs together.
 	if dp.to == nil {
-		return false, nil
+		return false, false, nil
 	}
 
 	// dp.from is nil when the to commit created a new table
 	if dp.from == nil {
-		return true, nil
+		return true, false, nil
 	}
 
 	fromSch, err := dp.from.GetSchema(ctx)
 	if err != nil {
-		return false, err
+		return false, false, err
 	}
 
 	toSch, err := dp.to.GetSchema(ctx)
 	if err != nil {
-		return false, err
+		return false, false, err
 	}
 
 	easyDiff := schema.ArePrimaryKeySetsDiffable(dp.from.Format(), fromSch, toSch)
 	if easyDiff {
-		return true, nil
+		return true, false, nil
 	}
 
 	_, _, err = schema.MapSchemaBasedOnTagAndName(fromSch, toSch)
 	if err == nil {
-		return true, nil
+		return false, true, nil
 	}
-	return false, nil
+	return false, false, nil
 }
 
 type partitionSelectFunc func(*sql.Context, DiffPartition) (bool, error)
@@ -771,6 +771,7 @@ type DiffPartitions struct {
 	selectFunc      partitionSelectFunc
 	toSch           schema.Schema
 	fromSch         schema.Schema
+	stopNext        bool
 }
 
 // processCommit is called in a commit iteration loop. Adds partitions when it finds a commit and its parent that have
@@ -830,6 +831,10 @@ func (dps *DiffPartitions) processCommit(ctx *sql.Context, cmHash hash.Hash, cm
 }
 
 func (dps *DiffPartitions) Next(ctx *sql.Context) (sql.Partition, error) {
+	if dps.stopNext {
+		return nil, io.EOF
+	}
+
 	for {
 		cmHash, optCmt, err := dps.cmItr.Next(ctx)
 		if err != nil {
@@ -861,16 +866,21 @@ func (dps *DiffPartitions) Next(ctx *sql.Context) (sql.Partition, error) {
 
 		if next != nil {
 			// If we can't diff this commit with its parent, don't traverse any lower
-			canDiff, err := next.isDiffablePartition(ctx)
+			simpleDiff, fuzzyDiff, err := next.isDiffablePartition(ctx)
 			if err != nil {
 				return nil, err
 			}
 
-			if !canDiff {
+			if !simpleDiff && !fuzzyDiff {
 				ctx.Warn(PrimaryKeyChangeWarningCode, fmt.Sprintf(PrimaryKeyChangeWarning, next.fromName, next.toName))
 				return nil, io.EOF
 			}
 
+			if !simpleDiff && fuzzyDiff {
+				ctx.Warn(PrimaryKeyChangeWarningCode, fmt.Sprintf(PrimaryKeyChangeWarning, next.fromName, next.toName))
+				dps.stopNext = true
+			}
+
 			return *next, nil
 		}
 	}
diff --git a/go/libraries/doltcore/sqle/enginetest/dolt_queries_diff.go b/go/libraries/doltcore/sqle/enginetest/dolt_queries_diff.go
index d44b9a961e..a05fecc8a1 100644
--- a/go/libraries/doltcore/sqle/enginetest/dolt_queries_diff.go
+++ b/go/libraries/doltcore/sqle/enginetest/dolt_queries_diff.go
@@ -561,6 +561,7 @@ var DiffSystemTableScriptTests = []queries.ScriptTest{
 					{nil, nil, nil, 23, "removed"},
 					{nil, 23, nil, nil, "added"},
 				},
+				ExpectedWarningsCount:           1,
 				ExpectedWarning:                 1105,
 				ExpectedWarningMessageSubstring: "due to primary key set change",
 			},