Optimize marshalling for jsonb types (#19898)

We check for the presence of binary_parameters in the DSN and add the 0x01 byte accordingly. This helps us avoid casting to string and efficiently use the database. ```release-note NONE ``` Co-authored-by: Mattermod <mattermod@users.noreply.github.com>
2022-04-06 15:01:32 +05:30
--- a/model/utils.go
+++ b/model/utils.go
@@ -33,6 +33,7 @@ const (
 	UppercaseLetters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
 	NUMBERS          = "0123456789"
 	SYMBOLS          = " !\"\\#$%&'()*+,-./:;<=>?@[]^_`|~"
+	BinaryParamKey   = "MM_BINARY_PARAMETERS"
 )

 type StringInterface map[string]interface{}
@@ -124,12 +125,19 @@ func (m *StringMap) Scan(value interface{}) error {

 // Value converts StringMap to database value
 func (m StringMap) Value() (driver.Value, error) {
-	j, err := json.Marshal(m)
+	ok := m[BinaryParamKey]
+	delete(m, BinaryParamKey)
+	buf, err := json.Marshal(m)
 	if err != nil {
 		return nil, err
 	}
-	// non utf8 characters are not supported https://mattermost.atlassian.net/browse/MM-41066
-	return string(j), err
+	if ok == "true" {
+		return append([]byte{0x01}, buf...), nil
+	} else if ok == "false" {
+		return buf, nil
+	}
+	// Key wasn't found. We fall back to the default case.
+	return string(buf), nil
 }

 func (StringMap) ImplementsGraphQLType(name string) bool {
--- a/store/sqlstore/job_store.go
+++ b/store/sqlstore/job_store.go
@@ -34,10 +34,13 @@ func (jss SqlJobStore) Save(job *model.Job) (*model.Job, error) {
 	if err != nil {
 		return nil, errors.Wrap(err, "failed marshalling job data")
 	}
+	if jss.IsBinaryParamEnabled() {
+		jsonData = jss.AppendBinaryFlag(jsonData)
+	}
 	query := jss.getQueryBuilder().
 		Insert("Jobs").
 		Columns("Id", "Type", "Priority", "CreateAt", "StartAt", "LastActivityAt", "Status", "Progress", "Data").
-		Values(job.Id, job.Type, job.Priority, job.CreateAt, job.StartAt, job.LastActivityAt, job.Status, job.Progress, string(jsonData))
+		Values(job.Id, job.Type, job.Priority, job.CreateAt, job.StartAt, job.LastActivityAt, job.Status, job.Progress, jsonData)

 	queryString, args, err := query.ToSql()
 	if err != nil {
@@ -56,11 +59,14 @@ func (jss SqlJobStore) UpdateOptimistically(job *model.Job, currentStatus string
 	if jsonErr != nil {
 		return false, errors.Wrap(jsonErr, "failed to encode job's data to JSON")
 	}
+	if jss.IsBinaryParamEnabled() {
+		dataJSON = jss.AppendBinaryFlag(dataJSON)
+	}
 	query, args, err := jss.getQueryBuilder().
 		Update("Jobs").
 		Set("LastActivityAt", model.GetMillis()).
 		Set("Status", job.Status).
-		Set("Data", string(dataJSON)).
+		Set("Data", dataJSON).
 		Set("Progress", job.Progress).
 		Where(sq.Eq{"Id": job.Id, "Status": currentStatus}).ToSql()
 	if err != nil {
--- a/store/sqlstore/link_metadata_store.go
+++ b/store/sqlstore/link_metadata_store.go
@@ -32,16 +32,19 @@ func (s SqlLinkMetadataStore) Save(metadata *model.LinkMetadata) (*model.LinkMet
 	if err != nil {
 		return nil, errors.Wrap(err, "could not serialize metadataBytes to JSON")
 	}
+	if s.IsBinaryParamEnabled() {
+		metadataBytes = s.AppendBinaryFlag(metadataBytes)
+	}

 	query := s.getQueryBuilder().
 		Insert("LinkMetadata").
 		Columns("Hash", "URL", "Timestamp", "Type", "Data").
-		Values(metadata.Hash, metadata.URL, metadata.Timestamp, metadata.Type, string(metadataBytes))
+		Values(metadata.Hash, metadata.URL, metadata.Timestamp, metadata.Type, metadataBytes)

 	if s.DriverName() == model.DatabaseDriverMysql {
-		query = query.SuffixExpr(sq.Expr("ON DUPLICATE KEY UPDATE URL = ?, Timestamp = ?, Type = ?, Data = ?", metadata.URL, metadata.Timestamp, metadata.Type, string(metadataBytes)))
+		query = query.SuffixExpr(sq.Expr("ON DUPLICATE KEY UPDATE URL = ?, Timestamp = ?, Type = ?, Data = ?", metadata.URL, metadata.Timestamp, metadata.Type, metadataBytes))
 	} else {
-		query = query.SuffixExpr(sq.Expr("ON CONFLICT (hash) DO UPDATE SET URL = ?, Timestamp = ?, Type = ?, Data = ?", metadata.URL, metadata.Timestamp, metadata.Type, string(metadataBytes)))
+		query = query.SuffixExpr(sq.Expr("ON CONFLICT (hash) DO UPDATE SET URL = ?, Timestamp = ?, Type = ?, Data = ?", metadata.URL, metadata.Timestamp, metadata.Type, metadataBytes))
 	}

 	q, args, err := query.ToSql()
--- a/store/sqlstore/session_store.go
+++ b/store/sqlstore/session_store.go
@@ -43,10 +43,14 @@ func (me SqlSessionStore) Save(session *model.Session) (*model.Session, error) {
 		return nil, errors.Wrap(err, "failed marshalling session props")
 	}

+	if me.IsBinaryParamEnabled() {
+		jsonProps = me.AppendBinaryFlag(jsonProps)
+	}
+
 	query, args, err := me.getQueryBuilder().
 		Insert("Sessions").
 		Columns("Id", "Token", "CreateAt", "ExpiresAt", "LastActivityAt", "UserId", "DeviceId", "Roles", "IsOAuth", "ExpiredNotify", "Props").
-		Values(session.Id, session.Token, session.CreateAt, session.ExpiresAt, session.LastActivityAt, session.UserId, session.DeviceId, session.Roles, session.IsOAuth, session.ExpiredNotify, string(jsonProps)).
+		Values(session.Id, session.Token, session.CreateAt, session.ExpiresAt, session.LastActivityAt, session.UserId, session.DeviceId, session.Roles, session.IsOAuth, session.ExpiredNotify, jsonProps).
 		ToSql()
 	if err != nil {
 		return nil, errors.Wrap(err, "sessions_tosql")
@@ -263,9 +267,12 @@ func (me SqlSessionStore) UpdateProps(session *model.Session) error {
 	if err != nil {
 		return errors.Wrap(err, "failed marshalling session props")
 	}
+	if me.IsBinaryParamEnabled() {
+		jsonProps = me.AppendBinaryFlag(jsonProps)
+	}
 	query, args, err := me.getQueryBuilder().
 		Update("Sessions").
-		Set("Props", string(jsonProps)).
+		Set("Props", jsonProps).
 		Where(sq.Eq{"Id": session.Id}).
 		ToSql()
 	if err != nil {
--- a/store/sqlstore/store.go
+++ b/store/sqlstore/store.go
@@ -9,6 +9,7 @@ import (
 	dbsql "database/sql"
 	"fmt"
 	"log"
+	"net/url"
 	"path/filepath"
 	"strconv"
 	"strings"
@@ -128,6 +129,8 @@ type SqlStore struct {
 	license           *model.License
 	licenseMutex      sync.RWMutex
 	metrics           einterfaces.MetricsInterface
+
+	isBinaryParam bool
 }

 func New(settings model.SqlSettings, metrics einterfaces.MetricsInterface) *SqlStore {
@@ -160,6 +163,11 @@ func New(settings model.SqlSettings, metrics einterfaces.MetricsInterface) *SqlS
 		mlog.Fatal("Failed to apply database migrations.", mlog.Err(err))
 	}

+	store.isBinaryParam, err = store.computeBinaryParam()
+	if err != nil {
+		mlog.Fatal("Failed to compute binary param", mlog.Err(err))
+	}
+
 	store.stores.team = newSqlTeamStore(store)
 	store.stores.channel = newSqlChannelStore(store, metrics)
 	store.stores.post = newSqlPostStore(store, metrics)
@@ -323,6 +331,29 @@ func (ss *SqlStore) DriverName() string {
 	return *ss.settings.DriverName
 }

+// computeBinaryParam returns whether the data source uses binary_parameters
+// when using Postgres
+func (ss *SqlStore) computeBinaryParam() (bool, error) {
+	if ss.DriverName() != model.DatabaseDriverPostgres {
+		return false, nil
+	}
+
+	url, err := url.Parse(*ss.settings.DataSource)
+	if err != nil {
+		return false, err
+	}
+	return url.Query().Get("binary_parameters") == "yes", nil
+}
+
+func (ss *SqlStore) IsBinaryParamEnabled() bool {
+	return ss.isBinaryParam
+}
+
+// AppendBinaryFlag updates the byte slice to work using binary_parameters=yes.
+func (ss *SqlStore) AppendBinaryFlag(buf []byte) []byte {
+	return append([]byte{0x01}, buf...)
+}
+
 func (ss *SqlStore) getCurrentSchemaVersion() (string, error) {
 	var version string
 	err := ss.GetMasterX().Get(&version, "SELECT Value FROM Systems WHERE Name='Version'")
--- a/store/sqlstore/store_test.go
+++ b/store/sqlstore/store_test.go
@@ -485,6 +485,57 @@ func TestEnsureMinimumDBVersion(t *testing.T) {
 	}
 }

+func TestIsBinaryParamEnabled(t *testing.T) {
+	tests := []struct {
+		store    SqlStore
+		expected bool
+	}{
+		{
+			store: SqlStore{
+				settings: &model.SqlSettings{
+					DriverName: model.NewString(model.DatabaseDriverPostgres),
+					DataSource: model.NewString("postgres://mmuser:mostest@localhost/loadtest?sslmode=disable\u0026binary_parameters=yes"),
+				},
+			},
+			expected: true,
+		},
+		{
+			store: SqlStore{
+				settings: &model.SqlSettings{
+					DriverName: model.NewString(model.DatabaseDriverMysql),
+					DataSource: model.NewString("postgres://mmuser:mostest@localhost/loadtest?sslmode=disable\u0026binary_parameters=yes"),
+				},
+			},
+			expected: false,
+		},
+		{
+			store: SqlStore{
+				settings: &model.SqlSettings{
+					DriverName: model.NewString(model.DatabaseDriverPostgres),
+					DataSource: model.NewString("postgres://mmuser:mostest@localhost/loadtest?sslmode=disable&binary_parameters=yes"),
+				},
+			},
+			expected: true,
+		},
+		{
+			store: SqlStore{
+				settings: &model.SqlSettings{
+					DriverName: model.NewString(model.DatabaseDriverPostgres),
+					DataSource: model.NewString("postgres://mmuser:mostest@localhost/loadtest?sslmode=disable"),
+				},
+			},
+			expected: false,
+		},
+	}
+
+	for i := range tests {
+		ok, err := tests[i].store.computeBinaryParam()
+		require.NoError(t, err)
+		assert.Equal(t, tests[i].expected, ok)
+	}
+
+}
+
 func TestUpAndDownMigrations(t *testing.T) {
 	testDrivers := []string{
 		model.DatabaseDriverPostgres,
--- a/store/sqlstore/user_store.go
+++ b/store/sqlstore/user_store.go
@@ -71,6 +71,7 @@ func (us SqlUserStore) insert(user *model.User) (sql.Result, error) {
 			:Props, :NotifyProps, :LastPasswordUpdate, :LastPictureUpdate, :FailedAttempts,
 			:Locale, :Timezone, :MfaActive, :MfaSecret, :RemoteId)`

+	user.Props = wrapBinaryParamStringMap(us.IsBinaryParamEnabled(), user.Props)
 	return us.GetMasterX().NamedExec(query, user)
 }

@@ -201,6 +202,7 @@ func (us SqlUserStore) Update(user *model.User, trustedUpdateData bool) (*model.
 				MfaSecret=:MfaSecret, RemoteId=:RemoteId
 			WHERE Id=:Id`

+	user.Props = wrapBinaryParamStringMap(us.IsBinaryParamEnabled(), user.Props)
 	res, err := us.GetMasterX().NamedExec(query, user)
 	if err != nil {
 		if IsUniqueConstraintError(err, []string{"Email", "users_email_key", "idx_users_email_unique"}) {
@@ -226,9 +228,17 @@ func (us SqlUserStore) Update(user *model.User, trustedUpdateData bool) (*model.
 }

 func (us SqlUserStore) UpdateNotifyProps(userID string, props map[string]string) error {
+	buf, err := json.Marshal(props)
+	if err != nil {
+		return errors.Wrap(err, "failed marshalling session props")
+	}
+	if us.IsBinaryParamEnabled() {
+		buf = us.AppendBinaryFlag(buf)
+	}
+
 	if _, err := us.GetMasterX().Exec(`UPDATE Users
 		SET NotifyProps = ?
-		WHERE Id = ?`, model.MapToJSON(props), userID); err != nil {
+		WHERE Id = ?`, buf, userID); err != nil {
 		return errors.Wrapf(err, "failed to update User with userId=%s", userID)
 	}

--- a/store/sqlstore/utils.go
+++ b/store/sqlstore/utils.go
@@ -9,6 +9,7 @@ import (
 	"strings"
 	"unicode"

+	"github.com/mattermost/mattermost-server/v6/model"
 	"github.com/mattermost/mattermost-server/v6/shared/mlog"
 )

@@ -137,6 +138,14 @@ func constructArrayArgs(ids []string) (string, []interface{}) {
 	return "(" + placeholder.String() + ")", values
 }

+func wrapBinaryParamStringMap(ok bool, props model.StringMap) model.StringMap {
+	if props == nil {
+		props = make(model.StringMap)
+	}
+	props[model.BinaryParamKey] = strconv.FormatBool(ok)
+	return props
+}
+
 // morphWriter is a target to pass to the logger instance of morph.
 // For now, everything is just logged at a debug level. If we need to log
 // errors/warnings from the library also, that needs to be seen later.