Move truncation to textutil function

Signed-off-by: Matt Lord <[email protected]>

go/textutil/strings.go

@@ -17,6 +17,7 @@ limitations under the License.
 package textutil
 
 import (
+	"fmt"
 	"net/url"
 	"regexp"
 	"strings"
@@ -28,6 +29,13 @@ import (
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 )
 
+type TruncationLocation int
+
+const (
+	TruncationLocationMiddle TruncationLocation = iota
+	TruncationLocationEnd
+)
+
 var (
 	delimitedListRegexp      = regexp.MustCompile(`[ ,;]+`)
 	SimulatedNullString      = sqltypes.NULL.String()
@@ -133,3 +141,23 @@ func Title(s string) string {
 		},
 		s)
 }
+
+func TruncateText(text string, maxLen int, location TruncationLocation, truncationIndicator string) (string, error) {
+	if len(truncationIndicator)+2 >= maxLen {
+		return "", fmt.Errorf("the truncation indicator is too long for the provided text")
+	}
+	ol := len(text)
+	if ol <= maxLen {
+		return text, nil
+	}
+	switch location {
+	case TruncationLocationMiddle:
+		prefix := int((float64(maxLen) * 0.5) - float64(len(truncationIndicator)))
+		suffix := int(ol - int(prefix+len(truncationIndicator)) + 1)
+		return fmt.Sprintf("%s%s%s", text[:prefix], truncationIndicator, text[suffix:]), nil
+	case TruncationLocationEnd:
+		return text[:maxLen-len(truncationIndicator)] + truncationIndicator, nil
+	default:
+		return "", fmt.Errorf("invalid truncation location: %d", location)
+	}
+}

go/textutil/strings_test.go

@@ -17,9 +17,11 @@ limitations under the License.
 package textutil
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -208,3 +210,77 @@ func TestTitle(t *testing.T) {
 		})
 	}
 }
+
+func TestTruncateText(t *testing.T) {
+	defaultLocation := TruncationLocationMiddle
+	defaultMaxLen := 100
+	defaultTruncationIndicator := "..."
+
+	tests := []struct {
+		name                string
+		text                string
+		maxLen              int
+		location            TruncationLocation
+		truncationIndicator string
+		want                string
+		wantErr             string
+	}{
+		{
+			name:     "no truncation",
+			text:     "hello world",
+			maxLen:   defaultMaxLen,
+			location: defaultLocation,
+			want:     "hello world",
+		},
+		{
+			name:     "no truncation - exact",
+			text:     strings.Repeat("a", defaultMaxLen),
+			maxLen:   defaultMaxLen,
+			location: defaultLocation,
+			want:     strings.Repeat("a", defaultMaxLen),
+		},
+		{
+			name:                "barely too long - mid",
+			text:                strings.Repeat("a", defaultMaxLen+1),
+			truncationIndicator: defaultTruncationIndicator,
+			maxLen:              defaultMaxLen,
+			location:            defaultLocation,
+			want:                "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa...aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+		},
+		{
+			name:                "barely too long - end",
+			text:                strings.Repeat("a", defaultMaxLen+1),
+			truncationIndicator: defaultTruncationIndicator,
+			maxLen:              defaultMaxLen,
+			location:            TruncationLocationEnd,
+			want:                "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa...",
+		},
+		{
+			name:                "too small",
+			text:                strings.Repeat("a", defaultMaxLen),
+			truncationIndicator: defaultTruncationIndicator,
+			maxLen:              4,
+			location:            defaultLocation,
+			wantErr:             "the truncation indicator is too long for the provided text",
+		},
+		{
+			name:                "bad location",
+			text:                strings.Repeat("a", defaultMaxLen+1),
+			truncationIndicator: defaultTruncationIndicator,
+			maxLen:              defaultMaxLen,
+			location:            100,
+			wantErr:             "invalid truncation location: 100",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			val, err := TruncateText(tt.text, tt.maxLen, tt.location, tt.truncationIndicator)
+			if tt.wantErr != "" {
+				require.EqualError(t, err, tt.wantErr)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.want, val)
+			}
+		})
+	}
+}

go/vt/binlog/binlogplayer/mock_dbclient.go

@@ -245,6 +245,13 @@ func (dc *MockDBClient) ExecuteFetchMulti(query string, maxrows int) ([]*sqltype
 	return results, nil
 }
 
+// AddInvariant can be used to customize the behavior of the mock client.
+func (dc *MockDBClient) AddInvariant(query string, result *sqltypes.Result) {
+	dc.expectMu.Lock()
+	defer dc.expectMu.Unlock()
+	dc.invariants[query] = result
+}
+
 // RemoveInvariant can be used to customize the behavior of the mock client.
 func (dc *MockDBClient) RemoveInvariant(query string) {
 	dc.expectMu.Lock()

go/vt/vttablet/tabletmanager/vreplication/utils.go

@@ -24,6 +24,7 @@ import (
 	"vitess.io/vitess/go/constants/sidecar"
 	"vitess.io/vitess/go/mysql/sqlerror"
 	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/textutil"
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
@@ -42,7 +43,7 @@ const (
 
 // vrepliationLogTruncationStr is the string that is used to indicate that a message has been
 // truncated, in the middle, before being inserted into the vreplication_log table.
-var vrepliationLogTruncationStr = fmt.Sprintf(" ... %s ... ", sqlparser.TruncationText)
+var vreplicationLogTruncationStr = fmt.Sprintf(" ... %s ... ", sqlparser.TruncationText)
 
 const (
 	// Enum values for type column in the vreplication_log table.
@@ -111,12 +112,11 @@ func insertLog(dbClient *vdbClient, typ string, vreplID int32, state, message st
 		// We perform the truncation, if needed, in the middle of the message as the end of the message is likely to
 		// be the most important part as it often explains WHY we e.g. failed to execute an INSERT in the workflow.
 		if len(message) > maxVReplicationLogMessageLen {
-			mid := (len(message) / 2) - len(vrepliationLogTruncationStr)
-			for mid > (maxVReplicationLogMessageLen / 2) {
-				mid = mid / 2
+			message, err = textutil.TruncateText(message, maxVReplicationLogMessageLen, textutil.TruncationLocationMiddle, vreplicationLogTruncationStr)
+			if err != nil {
+				log.Errorf("Could not insert vreplication_log record because we failed to truncate the message: %v", err)
+				return
 			}
-			tail := (len(message) - (mid + len(vrepliationLogTruncationStr))) + 1
-			message = fmt.Sprintf("%s%s%s", message[:mid], vrepliationLogTruncationStr, message[tail:])
 		}
 		buf.Myprintf("insert into %s.vreplication_log(vrepl_id, type, state, message) values(%s, %s, %s, %s)",
 			sidecar.GetIdentifier(), strconv.Itoa(int(vreplID)), encodeString(typ), encodeString(state), encodeString(message))

go/vt/vttablet/tabletmanager/vreplication/utils_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/textutil"
 	"vitess.io/vitess/go/vt/binlog/binlogplayer"
 
 	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
@@ -77,15 +78,11 @@ func TestInsertLogTruncation(t *testing.T) {
 		t.Run("insertLog", func(t *testing.T) {
 			var messageOut string
 			if tc.expectTruncation {
-				mid := (len(tc.message) / 2) - len(vrepliationLogTruncationStr)
-				for mid > (maxVReplicationLogMessageLen / 2) {
-					mid = mid / 2
-				}
-				tail := (len(tc.message) - (mid + len(vrepliationLogTruncationStr))) + 1
-				messageOut = fmt.Sprintf("%s%s%s", tc.message[:mid], vrepliationLogTruncationStr, tc.message[tail:])
+				messageOut, err := textutil.TruncateText(tc.message, maxVReplicationLogMessageLen, textutil.TruncationLocationMiddle, vreplicationLogTruncationStr)
+				require.NoError(t, err)
 				require.True(t, strings.HasPrefix(messageOut, tc.message[:1024]))                 // Confirm we still have the same beginning
 				require.True(t, strings.HasSuffix(messageOut, tc.message[len(tc.message)-1024:])) // Confirm we still have the same end
-				require.True(t, strings.Contains(messageOut, vrepliationLogTruncationStr))        // Confirm we have the truncation text
+				require.True(t, strings.Contains(messageOut, vreplicationLogTruncationStr))       // Confirm we have the truncation text
 				t.Logf("Original message length: %d, truncated message length: %d", len(tc.message), len(messageOut))
 			} else {
 				messageOut = tc.message