From 9359bd54978fa1236343d50141bd8e25c45daf81 Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Thu, 3 Aug 2023 17:38:34 +0200
Subject: [PATCH] Only update required entries if we already got some

---
 internal/device_data.go    |  2 +-
 internal/device_lists.go   | 20 ++++++++++++++++++--
 state/device_data_table.go | 17 +++++++++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/internal/device_data.go b/internal/device_data.go
index 09cee9f..0c5af57 100644
--- a/internal/device_data.go
+++ b/internal/device_data.go
@@ -23,7 +23,7 @@ func isBitSet(n int, bit int) bool {
 type DeviceData struct {
 	// Contains the latest device_one_time_keys_count values.
 	// Set whenever this field arrives down the v2 poller, and it replaces what was previously there.
-	OTKCounts map[string]int `json:"otk"`
+	OTKCounts MapStringInt `json:"otk"`
 	// Contains the latest device_unused_fallback_key_types value
 	// Set whenever this field arrives down the v2 poller, and it replaces what was previously there.
 	FallbackKeyTypes []string `json:"fallback"`
diff --git a/internal/device_lists.go b/internal/device_lists.go
index 1595068..260e052 100644
--- a/internal/device_lists.go
+++ b/internal/device_lists.go
@@ -1,5 +1,10 @@
 package internal
 
+import (
+	"database/sql/driver"
+	"encoding/json"
+)
+
 const (
 	DeviceListChanged = 1
 	DeviceListLeft    = 2
@@ -7,8 +12,19 @@ const (
 
 type DeviceLists struct {
 	// map user_id -> DeviceList enum
-	New  map[string]int `json:"n"`
-	Sent map[string]int `json:"s"`
+	New  MapStringInt `json:"n"`
+	Sent MapStringInt `json:"s"`
+}
+
+type MapStringInt map[string]int
+
+// Value implements driver.Valuer
+func (dl MapStringInt) Value() (driver.Value, error) {
+	if len(dl) == 0 {
+		return "{}", nil
+	}
+	v, err := json.Marshal(dl)
+	return v, err
 }
 
 func (dl DeviceLists) Combine(newer DeviceLists) DeviceLists {
diff --git a/state/device_data_table.go b/state/device_data_table.go
index 0f5f718..7b9ec63 100644
--- a/state/device_data_table.go
+++ b/state/device_data_table.go
@@ -6,6 +6,7 @@ import (
 	"reflect"
 
 	"github.com/jmoiron/sqlx"
+	"github.com/lib/pq"
 	"github.com/matrix-org/sliding-sync/internal"
 	"github.com/matrix-org/sliding-sync/sqlutil"
 )
@@ -119,6 +120,22 @@ func (t *DeviceDataTable) Upsert(dd *internal.DeviceData) (err error) {
 		}
 		tempDD.DeviceLists = tempDD.DeviceLists.Combine(dd.DeviceLists)
 
+		// we already got something in the database - update by just sending the new data
+		if len(row.Data) > 0 {
+			if tempDD.FallbackKeyTypes == nil {
+				// If fallback is null, it would, for some reason, nuke the whole
+				// column, so make sure we have something set.
+				tempDD.FallbackKeyTypes = []string{}
+			}
+			_, err = txn.Exec(`UPDATE syncv3_device_data SET data = 
+				jsonb_set(jsonb_set(jsonb_set(jsonb_set(jsonb_set(data, '{otk}', $3, false), '{fallback}', to_jsonb($4::text[]), false), '{c}', $5, false), '{dl,s}', $6, false),'{dl,n}', $7, false)
+				WHERE user_id = $1 AND device_id = $2
+				`,
+				dd.UserID, dd.DeviceID, tempDD.OTKCounts, pq.StringArray(tempDD.FallbackKeyTypes), tempDD.ChangedBits, tempDD.DeviceLists.Sent, tempDD.DeviceLists.New,
+			)
+			return err
+		}
+
 		data, err := json.Marshal(tempDD)
 		if err != nil {
 			return err