annotateAll

Signed-off-by: Shlomi Noach <[email protected]>

go/vt/schemadiff/annotations.go

@@ -167,6 +167,16 @@ func annotatedStatement(stmt string, annotationType TextualAnnotationType, annot
 	return result
 }
 
+// annotateAll blindly annotates all lines of the given statement with the given annotation type.
+func annotateAll(stmt string, annotationType TextualAnnotationType) *TextualAnnotations {
+	stmtLines := strings.Split(stmt, "\n")
+	result := NewTextualAnnotations()
+	for _, line := range stmtLines {
+		result.mark(line, annotationType)
+	}
+	return result
+}
+
 // unifiedAnnotated takes two annotations of from, to statements and returns a unified annotation.
 func unifiedAnnotated(from *TextualAnnotations, to *TextualAnnotations) *TextualAnnotations {
 	unified := NewTextualAnnotations()

go/vt/schemadiff/annotations_test.go

@@ -25,6 +25,22 @@ import (
 	"vitess.io/vitess/go/vt/sqlparser"
 )
 
+func TestAnnotateAll(t *testing.T) {
+	stmt := `create table t(
+	id int,
+	name varchar(100),
+	primary key(id)
+) engine=innodb`
+	annotations := annotateAll(stmt, RemovedTextualAnnotationType)
+	assert.Equal(t, 5, annotations.Len())
+	expect := `-create table t(
+-	id int,
+-	name varchar(100),
+-	primary key(id)
+-) engine=innodb`
+	assert.Equal(t, expect, annotations.Export())
+}
+
 func TestUnifiedAnnotated(t *testing.T) {
 	tcases := []struct {
 		name            string
@@ -67,3 +83,20 @@ func TestUnifiedAnnotated(t *testing.T) {
 		})
 	}
 }
+
+func TestUnifiedAnnotatedAll(t *testing.T) {
+	stmt := `create table t(
+	id int,
+	name varchar(100),
+	primary key(id)
+) engine=innodb`
+	annotatedTo := annotateAll(stmt, AddedTextualAnnotationType)
+	annotatedFrom := NewTextualAnnotations()
+	unified := unifiedAnnotated(annotatedFrom, annotatedTo)
+	expect := `+create table t(
++	id int,
++	name varchar(100),
++	primary key(id)
++) engine=innodb`
+	assert.Equal(t, expect, unified.Export())
+}