Add CT and UT to all documents and refactor (#17)

* - Create CosmosDocument as a base class
- Add CT and UT
- Refactor all tables to use the CosmosDocument

* - Add UpsertDocument method to perform updates in a generic way
- Add SetUpdateTime() to update the UT for updates
- Refactor it all

Co-authored-by: alexf@example.com <alexf@example.com>

appservice/storage/cosmosdb/appservice_events_table.go

@@ -57,12 +57,7 @@ type EventNumberCosmosData struct {
 }
 
 type EventCosmosData struct {
-	Id        string      `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string      `json:"_pk"`
-	Tn        string      `json:"_sid"`
-	Cn        string      `json:"_cn"`
-	ETag      string      `json:"_etag"`
-	Timestamp int64       `json:"_ts"`
 	Event EventCosmos `json:"mx_appservice_event"`
 }
 
@@ -166,6 +161,23 @@ func queryEventEventNumber(s *eventsStatements, ctx context.Context, qry string,
 	return response, nil
 }
 
+func getEvent(s *eventsStatements, ctx context.Context, pk string, docId string) (*EventCosmosData, error) {
+	response := EventCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func setEvent(s *eventsStatements, ctx context.Context, event EventCosmosData) (*EventCosmosData, error) {
 	var optionsReplace = cosmosdbapi.GetReplaceDocumentOptions(event.Pk, event.ETag)
 	var _, _, ex = cosmosdbapi.GetClient(s.db.connection).ReplaceDocument(
@@ -338,17 +350,22 @@ func (s *eventsStatements) insertEvent(
 	// "INSERT INTO appservice_events(as_id, headered_event_json, txn_id) " +
 	// "VALUES ($1, $2, $3)"
 
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
+	docId := fmt.Sprintf("%s", appServiceID)
+	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
+	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
+
+	dbData, err := getEvent(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.Event.HeaderedEventJSON = eventJSON
+	} else {
 		// 	id INTEGER PRIMARY KEY AUTOINCREMENT,
 		idSeq, seqErr := GetNextEventID(s, ctx)
 		if seqErr != nil {
 			return seqErr
 		}
 
-	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
-	docId := fmt.Sprintf("%d", idSeq)
-	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
-	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
-
 		// appServiceID,
 		// eventJSON,
 		// -1, // No transaction ID yet
@@ -359,25 +376,19 @@ func (s *eventsStatements) insertEvent(
 			TXNID:             -1,
 		}
 
-	dbData := &EventCosmosData{
+		dbData = &EventCosmosData{
-		Id:        cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 			Event:          data,
 		}
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	}
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+
-		ctx,
+	return cosmosdbapi.UpsertDocument(ctx,
+		s.db.connection,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		&dbData)
-
-	return err
-
 }
 
 // updateTxnIDForEvents sets the transactionID for a collection of events. Done

federationsender/storage/cosmosdb/blacklist_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 
@@ -38,12 +37,7 @@ type BlacklistCosmos struct {
 }
 
 type BlacklistCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	Blacklist BlacklistCosmos `json:"mx_federationsender_blacklist"`
 }
 
@@ -149,30 +143,28 @@ func (s *blacklistStatements) InsertBlacklist(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	dbData, _ := getBlacklist(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+	} else {
 		data := BlacklistCosmos{
 			ServerName: string(serverName),
 		}
 
-	dbData := &BlacklistCosmosData{
+		dbData = &BlacklistCosmosData{
-		Id:        cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 			Blacklist:      data,
 		}
+	}
 
 	// _, err := stmt.ExecContext(ctx, serverName)
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	_, _, err := cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		&dbData)
-
-	return err
 }
 
 // selectRoomForUpdate locks the row for the room and returns the last_event_id.

federationsender/storage/cosmosdb/inbound_peeks_table.go

@@ -48,12 +48,7 @@ type InboundPeekCosmos struct {
 }
 
 type InboundPeekCosmosData struct {
-	Id          string             `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk          string             `json:"_pk"`
-	Tn          string             `json:"_sid"`
-	Cn          string             `json:"_cn"`
-	ETag        string             `json:"_etag"`
-	Timestamp   int64              `json:"_ts"`
 	InboundPeek InboundPeekCosmos `json:"mx_federationsender_inbound_peek"`
 }
 
@@ -183,6 +178,12 @@ func (s *inboundPeeksStatements) InsertInboundPeek(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	dbData, _ := getInboundPeek(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.InboundPeek.RenewedTimestamp = nowMilli
+		dbData.InboundPeek.RenewalInterval = renewalInterval
+	} else {
 		data := InboundPeekCosmos{
 			RoomID:            roomID,
 			ServerName:        string(serverName),
@@ -192,24 +193,20 @@ func (s *inboundPeeksStatements) InsertInboundPeek(
 			RenewalInterval:   renewalInterval,
 		}
 
-	dbData := &InboundPeekCosmosData{
+		dbData = &InboundPeekCosmosData{
-		Id:          cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:          s.db.cosmosConfig.TenantName,
-		Cn:          dbCollectionName,
-		Pk:          pk,
-		Timestamp:   time.Now().Unix(),
 			InboundPeek:    data,
 		}
+	}
 
 	// _, err = stmt.ExecContext(ctx, roomID, serverName, peekID, nowMilli, nowMilli, renewalInterval)
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	err = cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		&dbData)
 
 	return
 }

federationsender/storage/cosmosdb/joined_hosts_table.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/federationsender/types"
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
@@ -53,12 +52,7 @@ type JoinedHostCosmos struct {
 }
 
 type JoinedHostCosmosData struct {
-	Id         string           `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk         string           `json:"_pk"`
-	Tn         string           `json:"_sid"`
-	Cn         string           `json:"_cn"`
-	ETag       string           `json:"_etag"`
-	Timestamp  int64            `json:"_ts"`
 	JoinedHost JoinedHostCosmos `json:"mx_federationsender_joined_host"`
 }
 
@@ -102,6 +96,23 @@ type joinedHostsStatements struct {
 	tableName string
 }
 
+func getJoinedHost(s *joinedHostsStatements, ctx context.Context, pk string, docId string) (*JoinedHostCosmosData, error) {
+	response := JoinedHostCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryJoinedHostDistinct(s *joinedHostsStatements, ctx context.Context, qry string, params map[string]interface{}) ([]JoinedHostCosmos, error) {
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
@@ -190,32 +201,28 @@ func (s *joinedHostsStatements) InsertJoinedHosts(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	dbData, _ := getJoinedHost(s, ctx, pk, cosmosDocId)
+	if dbData == nil {
 		data := JoinedHostCosmos{
 			EventID:    eventID,
 			RoomID:     roomID,
 			ServerName: string(serverName),
 		}
 
-	dbData := &JoinedHostCosmosData{
+		dbData = &JoinedHostCosmosData{
-		Id: cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn: s.db.cosmosConfig.TenantName,
-		Cn: dbCollectionName,
-		Pk: pk,
-		Timestamp:  time.Now().Unix(),
 			JoinedHost:     data,
 		}
-
 		// _, err := stmt.ExecContext(ctx, roomID, eventID, serverName)
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+		return cosmosdbapi.UpsertDocument(ctx,
-	_, _, err := cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+			s.db.connection,
-		ctx,
 			s.db.cosmosConfig.DatabaseName,
 			s.db.cosmosConfig.ContainerName,
-		&dbData,
+			dbData.Pk,
-		options)
+			&dbData)
-
+	}
-	return err
+	return nil
 }
 
 func (s *joinedHostsStatements) DeleteJoinedHosts(

federationsender/storage/cosmosdb/outbound_peeks_table.go

@@ -48,12 +48,7 @@ type OutboundPeekCosmos struct {
 }
 
 type OutboundPeekCosmosData struct {
-	Id           string             `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk           string             `json:"_pk"`
-	Tn           string             `json:"_sid"`
-	Cn           string             `json:"_cn"`
-	ETag         string             `json:"_etag"`
-	Timestamp    int64              `json:"_ts"`
 	OutboundPeek OutboundPeekCosmos `json:"mx_federationsender_outbound_peek"`
 }
 
@@ -179,6 +174,13 @@ func (s *outboundPeeksStatements) InsertOutboundPeek(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	dbData, _ := getOutboundPeek(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.OutboundPeek.RenewalInterval = renewalInterval
+		dbData.OutboundPeek.RenewedTimestamp = nowMilli
+
+	} else {
 		data := OutboundPeekCosmos{
 			RoomID:            roomID,
 			ServerName:        string(serverName),
@@ -188,24 +190,21 @@ func (s *outboundPeeksStatements) InsertOutboundPeek(
 			RenewalInterval:   renewalInterval,
 		}
 
-	dbData := &OutboundPeekCosmosData{
+		dbData = &OutboundPeekCosmosData{
-		Id:           cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:           s.db.cosmosConfig.TenantName,
-		Cn:           dbCollectionName,
-		Pk:           pk,
-		Timestamp:    time.Now().Unix(),
 			OutboundPeek:   data,
 		}
 
+	}
+
 	// _, err = stmt.ExecContext(ctx, roomID, serverName, peekID, nowMilli, nowMilli, renewalInterval)
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	err = cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		&dbData)
 
 	return
 }

federationsender/storage/cosmosdb/queue_edus_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 
@@ -50,12 +49,7 @@ type QueueEDUCosmosNumber struct {
 }
 
 type QueueEDUCosmosData struct {
-	Id        string         `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string         `json:"_pk"`
-	Tn        string         `json:"_sid"`
-	Cn        string         `json:"_cn"`
-	ETag      string         `json:"_etag"`
-	Timestamp int64          `json:"_ts"`
 	QueueEDU QueueEDUCosmos `json:"mx_federationsender_queue_edu"`
 }
 
@@ -218,11 +212,7 @@ func (s *queueEDUsStatements) InsertQueueEDU(
 	}
 
 	dbData := &QueueEDUCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		QueueEDU:       data,
 	}
 

federationsender/storage/cosmosdb/queue_json_table.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 )
@@ -42,12 +41,7 @@ type QueueJSONCosmos struct {
 }
 
 type QueueJSONCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	QueueJSON QueueJSONCosmos `json:"mx_federationsender_queue_json"`
 }
 
@@ -143,11 +137,7 @@ func (s *queueJSONStatements) InsertQueueJSON(
 	}
 
 	dbData := &QueueJSONCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		QueueJSON:      data,
 	}
 

federationsender/storage/cosmosdb/queue_pdus_table.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 
@@ -51,12 +50,7 @@ type QueuePDUCosmosNumber struct {
 }
 
 type QueuePDUCosmosData struct {
-	Id        string         `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string         `json:"_pk"`
-	Tn        string         `json:"_sid"`
-	Cn        string         `json:"_cn"`
-	ETag      string         `json:"_etag"`
-	Timestamp int64          `json:"_ts"`
 	QueuePDU QueuePDUCosmos `json:"mx_federationsender_queue_pdu"`
 }
 
@@ -230,11 +224,7 @@ func (s *queuePDUsStatements) InsertQueuePDU(
 	}
 
 	dbData := &QueuePDUCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		QueuePDU:       data,
 	}
 

internal/cosmosdbapi/client.go

@@ -1,6 +1,9 @@
 package cosmosdbapi
 
 import (
+	"context"
+	"time"
+
 	cosmosapi "github.com/vippsas/go-cosmosdb/cosmosapi"
 )
 
@@ -22,3 +25,63 @@ func GetClient(conn CosmosConnection) *cosmosapi.Client {
 	}
 	return cosmosapi.New(conn.Url, cfg, nil, nil)
 }
+
+func UpsertDocument(ctx context.Context,
+	conn CosmosConnection,
+	databaseName string,
+	containerName string,
+	partitonKey string,
+	dbData interface{}) error {
+	var options = getUpsertDocumentOptions(partitonKey)
+	_, _, err := GetClient(conn).CreateDocument(
+		ctx,
+		databaseName,
+		containerName,
+		&dbData,
+		options)
+	return err
+}
+
+func (doc *CosmosDocument) SetUpdateTime() {
+	now := time.Now().UTC()
+	doc.Ut = now.Format(time.RFC3339)
+}
+
+func GenerateDocument(
+	collection string,
+	tenantName string,
+	partitionKey string,
+	docId string,
+) CosmosDocument {
+	doc := CosmosDocument{}
+	now := time.Now().UTC()
+	doc.Timestamp = now.Unix()
+	doc.Ct = now.Format(time.RFC3339)
+	doc.Ut = now.Format(time.RFC3339)
+	doc.Cn = collection
+	doc.Tn = tenantName
+	doc.Pk = partitionKey
+	doc.Id = docId
+	return doc
+}
+
+func GetDocumentOrNil(connection CosmosConnection, config CosmosConfig, ctx context.Context, partitionKey string, cosmosDocId string, dbData interface{}) error {
+	var _, err = GetClient(connection).GetDocument(
+		ctx,
+		config.DatabaseName,
+		config.ContainerName,
+		cosmosDocId,
+		GetGetDocumentOptions(partitionKey),
+		&dbData,
+	)
+
+	if err != nil {
+		if err.Error() == "Resource that no longer exists" {
+			dbData = nil
+			return nil
+		}
+		return err
+	}
+
+	return nil
+}

internal/cosmosdbapi/document.go

@@ -1,11 +1,21 @@
 package cosmosdbapi
 
 import (
-	"context"
 	"fmt"
 	"strings"
 )
 
+type CosmosDocument struct {
+	Id        string `json:"id"`
+	Pk        string `json:"_pk"`
+	Tn        string `json:"_sid"`
+	Cn        string `json:"_cn"`
+	Ct        string `json:"_ct"`
+	Ut        string `json:"_ut"`
+	ETag      string `json:"_etag"`
+	Timestamp int64  `json:"_ts"`
+}
+
 func removeSpecialChars(docId string) string {
 	// The following characters are restricted and cannot be used in the Id property: '/', '\', '?', '#'
 	invalidChars := [4]string{"/", "\\", "?", "#"}
@@ -25,24 +35,3 @@ func GetDocumentId(tenantName string, collectionName string, id string) string {
 func GetPartitionKey(tenantName string, collectionName string) string {
 	return fmt.Sprintf("%s,%s", collectionName, tenantName)
 }
-
-func GetDocumentOrNil(connection CosmosConnection, config CosmosConfig, ctx context.Context, partitionKey string, cosmosDocId string, dbData interface{}) error {
-	var _, err = GetClient(connection).GetDocument(
-		ctx,
-		config.DatabaseName,
-		config.ContainerName,
-		cosmosDocId,
-		GetGetDocumentOptions(partitionKey),
-		&dbData,
-	)
-
-	if err != nil {
-		if err.Error() == "Resource that no longer exists" {
-			dbData = nil
-			return nil
-		}
-		return err
-	}
-
-	return nil
-}

internal/cosmosdbapi/documentoperations.go

@@ -11,7 +11,7 @@ func GetCreateDocumentOptions(pk string) cosmosapi.CreateDocumentOptions {
 	}
 }
 
-func GetUpsertDocumentOptions(pk string) cosmosapi.CreateDocumentOptions {
+func getUpsertDocumentOptions(pk string) cosmosapi.CreateDocumentOptions {
 	return cosmosapi.CreateDocumentOptions{
 		IsUpsert:          true,
 		PartitionKeyValue: pk,

internal/cosmosdbutil/document_seq.go

@@ -8,12 +8,7 @@ import (
 )
 
 type SequenceCosmosData struct {
-	Id        string `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string `json:"_pk"`
-	Tn        string `json:"_sid"`
-	Cn        string `json:"_cn"`
-	ETag      string `json:"_etag"`
-	Timestamp int64  `json:"_ts"`
 	Value int64 `json:"_value"`
 }
 
@@ -43,10 +38,7 @@ func GetNextSequence(
 
 	if dbData.Id == "" {
 		dbData = SequenceCosmosData{}
-		dbData.Id = cosmosDocId
+		dbData.CosmosDocument = cosmosdbapi.GenerateDocument(dbCollectionName, config.TenantName, pk, cosmosDocId)
-		dbData.Pk = pk
-		dbData.Tn = config.TenantName
-		dbData.Cn = dbCollectionName
 		dbData.Value = initial
 		var optionsCreate = cosmosdbapi.GetCreateDocumentOptions(dbData.Pk)
 		var _, _, err = cosmosdbapi.GetClient(connection).CreateDocument(

internal/cosmosdbutil/partition_offset_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 
@@ -54,12 +53,7 @@ type PartitionOffsetCosmos struct {
 }
 
 type PartitionOffsetCosmosData struct {
-	Id              string                `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk              string                `json:"_pk"`
-	Tn              string                `json:"_sid"`
-	Cn              string                `json:"_cn"`
-	ETag            string                `json:"_etag"`
-	Timestamp       int64                 `json:"_ts"`
 	PartitionOffset PartitionOffsetCosmos `json:"mx_partition_offset"`
 }
 
@@ -90,6 +84,23 @@ type PartitionOffsetStatements struct {
 	tableName string
 }
 
+func getPartitionOffset(s *PartitionOffsetStatements, ctx context.Context, pk string, docId string) (*PartitionOffsetCosmosData, error) {
+	response := PartitionOffsetCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.Connection,
+		s.db.CosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryPartitionOffset(s *PartitionOffsetStatements, ctx context.Context, qry string, params map[string]interface{}) ([]PartitionOffsetCosmosData, error) {
 	var dbCollectionName = getCollectionName(*s)
 	var pk = cosmosdbapi.GetPartitionKey(s.db.CosmosConfig.ContainerName, dbCollectionName)
@@ -192,33 +203,32 @@ func (s *PartitionOffsetStatements) upsertPartitionOffset(
 		cosmosDocId := cosmosdbapi.GetDocumentId(s.db.CosmosConfig.TenantName, dbCollectionName, docId)
 		pk := cosmosdbapi.GetPartitionKey(s.db.CosmosConfig.ContainerName, dbCollectionName)
 
+		dbData, _ := getPartitionOffset(s, ctx, pk, cosmosDocId)
+		if dbData != nil {
+			dbData.SetUpdateTime()
+			dbData.PartitionOffset.PartitionOffset = offset
+		} else {
 			data := PartitionOffsetCosmos{
 				Partition:       partition,
 				PartitionOffset: offset,
 				Topic:           topic,
 			}
 
-		dbData := &PartitionOffsetCosmosData{
+			dbData = &PartitionOffsetCosmosData{
-			Id: cosmosDocId,
+				CosmosDocument:  cosmosdbapi.GenerateDocument(dbCollectionName, s.db.CosmosConfig.TenantName, pk, cosmosDocId),
-			Tn: s.db.CosmosConfig.TenantName,
-			Cn: dbCollectionName,
-			Pk: pk,
-			// nowMilli := time.Now().UnixNano() / int64(time.Millisecond)
-			Timestamp:       time.Now().Unix(),
 				PartitionOffset: data,
 			}
 
+		}
+
 		// _, err := stmt.ExecContext(ctx, topic, partition, offset)
 
-		var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+		return cosmosdbapi.UpsertDocument(ctx,
-		_, _, err := cosmosdbapi.GetClient(s.db.Connection).CreateDocument(
+			s.db.Connection,
-			ctx,
 			s.db.CosmosConfig.DatabaseName,
 			s.db.CosmosConfig.ContainerName,
-			&dbData,
+			dbData.Pk,
-			options)
+			&dbData)
-
-		return err
 	})
 }
 

internal/naffka/naffkacosmosdb/naffka_topics_table.go

@@ -38,12 +38,7 @@ type TopicCosmosNumber struct {
 }
 
 type TopicCosmosData struct {
-	Id        string      `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string      `json:"_pk"`
-	Tn        string      `json:"_sid"`
-	Cn        string      `json:"_cn"`
-	ETag      string      `json:"_etag"`
-	Timestamp int64       `json:"_ts"`
 	Topic TopicCosmos `json:"mx_naffka_topic"`
 }
 
@@ -56,12 +51,7 @@ type MessageCosmos struct {
 }
 
 type MessageCosmosData struct {
-	Id        string        `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string        `json:"_pk"`
-	Tn        string        `json:"_sid"`
-	Cn        string        `json:"_cn"`
-	ETag      string        `json:"_etag"`
-	Timestamp int64         `json:"_ts"`
 	Message MessageCosmos `json:"mx_naffka_message"`
 }
 
@@ -228,31 +218,30 @@ func (t *topicsStatements) InsertTopic(
 	cosmosDocId := cosmosdbapi.GetDocumentId(t.DB.cosmosConfig.ContainerName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(t.DB.cosmosConfig.ContainerName, dbCollectionName)
 
+	dbData, _ := getTopic(t, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.Topic.TopicName = topicName
+	} else {
 		data := TopicCosmos{
 			TopicNID:  topicNID,
 			TopicName: topicName,
 		}
 
-	dbData := &TopicCosmosData{
+		dbData = &TopicCosmosData{
-		Id:        cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, t.DB.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        t.DB.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 			Topic:          data,
 		}
+	}
 
 	// _, err := stmt.ExecContext(ctx, topicName, topicNID)
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	_, _, err := cosmosdbapi.GetClient(t.DB.connection).CreateDocument(
+		t.DB.connection,
-		ctx,
 		t.DB.cosmosConfig.DatabaseName,
 		t.DB.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
-
-	return err
 }
 
 func (t *topicsStatements) SelectNextTopicNID(
@@ -366,11 +355,7 @@ func (t *topicsStatements) InsertTopics(
 	}
 
 	dbData := &MessageCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, t.DB.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        t.DB.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		Message:        data,
 	}
 
@@ -383,7 +368,6 @@ func (t *topicsStatements) InsertTopics(
 		t.DB.cosmosConfig.ContainerName,
 		&dbData,
 		options)
-
 	return err
 }
 

keyserver/storage/cosmosdb/cross_signing_keys_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/keyserver/storage/tables"
@@ -42,12 +41,7 @@ type CrossSigningKeysCosmos struct {
 }
 
 type CrossSigningKeysCosmosData struct {
-	Id               string                 `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk               string                 `json:"_pk"`
-	Tn               string                 `json:"_sid"`
-	Cn               string                 `json:"_cn"`
-	ETag             string                 `json:"_etag"`
-	Timestamp        int64                  `json:"_ts"`
 	CrossSigningKeys CrossSigningKeysCosmos `json:"mx_keyserver_cross_signing_keys"`
 }
 
@@ -68,6 +62,23 @@ type crossSigningKeysStatements struct {
 	tableName string
 }
 
+func getCrossSigningKeys(s *crossSigningKeysStatements, ctx context.Context, pk string, docId string) (*CrossSigningKeysCosmosData, error) {
+	response := CrossSigningKeysCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryCrossSigningKeys(s *crossSigningKeysStatements, ctx context.Context, qry string, params map[string]interface{}) ([]CrossSigningKeysCosmosData, error) {
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
@@ -149,31 +160,29 @@ func (s *crossSigningKeysStatements) UpsertCrossSigningKeysForUser(
 	docId := fmt.Sprintf("%s_%s", userID, keyType)
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 
+	dbData, _ := getCrossSigningKeys(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.CrossSigningKeys.KeyData = keyData
+	} else {
 		data := CrossSigningKeysCosmos{
 			UserID:  userID,
 			KeyType: int64(keyTypeInt),
 			KeyData: keyData,
 		}
 
-	dbData := CrossSigningKeysCosmosData{
+		dbData = &CrossSigningKeysCosmosData{
-		Id:               cosmosDocId,
+			CosmosDocument:   cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:               s.db.cosmosConfig.TenantName,
-		Cn:               dbCollectionName,
-		Pk:               pk,
-		Timestamp:        time.Now().Unix(),
 			CrossSigningKeys: data,
 		}
-
+	}
 	// if _, err := sqlutil.TxStmt(txn, s.upsertCrossSigningKeysForUserStmt).ExecContext(ctx, userID, keyTypeInt, keyData); err != nil {
 	// 	return fmt.Errorf("s.upsertCrossSigningKeysForUserStmt: %w", err)
 	// }
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	var _, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		dbData,
+		dbData.Pk,
-		options)
+		dbData)
-
-	return err
 }

keyserver/storage/cosmosdb/cross_signing_sigs_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/keyserver/storage/tables"
@@ -46,12 +45,7 @@ type CrossSigningSigsCosmos struct {
 }
 
 type CrossSigningSigsCosmosData struct {
-	Id               string                 `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk               string                 `json:"_pk"`
-	Tn               string                 `json:"_sid"`
-	Cn               string                 `json:"_cn"`
-	ETag             string                 `json:"_etag"`
-	Timestamp        int64                  `json:"_ts"`
 	CrossSigningSigs CrossSigningSigsCosmos `json:"mx_keyserver_cross_signing_sigs"`
 }
 
@@ -80,6 +74,23 @@ type crossSigningSigsStatements struct {
 	tableName                           string
 }
 
+func getCrossSigningSigs(s *crossSigningSigsStatements, ctx context.Context, pk string, docId string) (*CrossSigningSigsCosmosData, error) {
+	response := CrossSigningSigsCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryCrossSigningSigs(s *crossSigningSigsStatements, ctx context.Context, qry string, params map[string]interface{}) ([]CrossSigningSigsCosmosData, error) {
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
@@ -182,6 +193,12 @@ func (s *crossSigningSigsStatements) UpsertCrossSigningSigsForTarget(
 	docId := fmt.Sprintf("%s_%s_%s", originUserID, targetUserID, targetKeyID)
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 
+	dbData, _ := getCrossSigningSigs(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.CrossSigningSigs.OriginKeyId = string(originKeyID)
+		dbData.CrossSigningSigs.Signature = signature
+	} else {
 		data := CrossSigningSigsCosmos{
 			TargetUserId: targetUserID,
 			TargetKeyId:  string(targetKeyID),
@@ -190,26 +207,20 @@ func (s *crossSigningSigsStatements) UpsertCrossSigningSigsForTarget(
 			Signature:    signature,
 		}
 
-	dbData := CrossSigningSigsCosmosData{
+		dbData = &CrossSigningSigsCosmosData{
-		Id:               cosmosDocId,
+			CosmosDocument:   cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:               s.db.cosmosConfig.TenantName,
-		Cn:               dbCollectionName,
-		Pk:               pk,
-		Timestamp:        time.Now().Unix(),
 			CrossSigningSigs: data,
 		}
-
+	}
 	// if _, err := sqlutil.TxStmt(txn, s.upsertCrossSigningSigsForTargetStmt).ExecContext(ctx, originUserID, originKeyID, targetUserID, targetKeyID, signature); err != nil {
 	// 	return fmt.Errorf("s.upsertCrossSigningSigsForTargetStmt: %w", err)
 	// }
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	var _, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		dbData,
+		dbData.Pk,
-		options)
+		dbData)
-	return err
 }
 
 func (s *crossSigningSigsStatements) DeleteCrossSigningSigsForTarget(

keyserver/storage/cosmosdb/device_keys_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -55,12 +54,7 @@ type DeviceKeyCosmosNumber struct {
 }
 
 type DeviceKeyCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	DeviceKey DeviceKeyCosmos `json:"mx_keyserver_device_key"`
 }
 
@@ -167,13 +161,26 @@ func getDeviceKey(s *deviceKeysStatements, ctx context.Context, pk string, docId
 }
 
 func insertDeviceKeyCore(s *deviceKeysStatements, ctx context.Context, dbData DeviceKeyCosmosData) error {
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	// "INSERT INTO keyserver_device_keys (user_id, device_id, ts_added_secs, key_json, stream_id, display_name)" +
-	var _, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+	// " VALUES ($1, $2, $3, $4, $5, $6)" +
-		ctx,
+	// " ON CONFLICT (user_id, device_id)" +
+	// " DO UPDATE SET key_json = $4, stream_id = $5, display_name = $6"
+	existing, _ := getDeviceKey(s, ctx, dbData.Pk, dbData.Id)
+	if existing != nil {
+		existing.SetUpdateTime()
+		existing.DeviceKey.KeyJSON = dbData.DeviceKey.KeyJSON
+		existing.DeviceKey.StreamID = dbData.DeviceKey.StreamID
+		existing.DeviceKey.DisplayName = dbData.DeviceKey.DisplayName
+	} else {
+		existing = &dbData
+	}
+
+	err := cosmosdbapi.UpsertDocument(ctx,
+		s.db.connection,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		dbData,
+		existing.Pk,
-		options)
+		existing)
 
 	if err != nil {
 		return err
@@ -445,17 +452,12 @@ func (s *deviceKeysStatements) InsertDeviceKeys(ctx context.Context, txn *sql.Tx
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
 	for _, key := range keys {
-		now := time.Now().Unix()
 		//     UNIQUE (user_id, device_id)
 		docId := fmt.Sprintf("%s_%s", key.UserID, key.DeviceID)
 		cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 
 		dbData := &DeviceKeyCosmosData{
-			Id:        cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-			Tn:        s.db.cosmosConfig.TenantName,
-			Cn:        dbCollectionName,
-			Pk:        pk,
-			Timestamp: now,
 			DeviceKey:      mapFromDeviceKeyMessage(key),
 		}
 

keyserver/storage/cosmosdb/key_changes_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"fmt"
 	"math"
-	"time"
 
 	"github.com/Shopify/sarama"
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
@@ -48,12 +47,7 @@ type KeyChangeUserMaxCosmosData struct {
 }
 
 type KeyChangeCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	KeyChange KeyChangeCosmos `json:"mx_keyserver_key_change"`
 }
 
@@ -84,6 +78,23 @@ type keyChangesStatements struct {
 	tableName            string
 }
 
+func getKeyChangeUser(s *keyChangesStatements, ctx context.Context, pk string, docId string) (*KeyChangeCosmosData, error) {
+	response := KeyChangeCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryKeyChangeUserMax(s *keyChangesStatements, ctx context.Context, qry string, params map[string]interface{}) ([]KeyChangeUserMaxCosmosData, error) {
 	var response []KeyChangeUserMaxCosmosData
 
@@ -127,31 +138,30 @@ func (s *keyChangesStatements) InsertKeyChange(ctx context.Context, partition in
 	docId := fmt.Sprintf("%d_%d", partition, offset)
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 
+	dbData, _ := getKeyChangeUser(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.KeyChange.UserID = userID
+	} else {
 		data := KeyChangeCosmos{
 			Offset:    offset,
 			Partition: partition,
 			UserID:    userID,
 		}
 
-	dbData := KeyChangeCosmosData{
+		dbData = &KeyChangeCosmosData{
-		Id:        cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 			KeyChange:      data,
 		}
+	}
 
 	// _, err := s.upsertKeyChangeStmt.ExecContext(ctx, partition, offset, userID)
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	var _, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		dbData,
+		dbData.Pk,
-		options)
+		dbData)
-
-	return err
 }
 
 func (s *keyChangesStatements) SelectKeyChanges(

keyserver/storage/cosmosdb/one_time_keys_table.go

@@ -19,7 +19,6 @@ import (
 	"database/sql"
 	"encoding/json"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 
@@ -59,12 +58,7 @@ type OneTimeKeyAlgoNumberCosmosData struct {
 }
 
 type OneTimeKeyCosmosData struct {
-	Id         string           `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk         string           `json:"_pk"`
-	Tn         string           `json:"_sid"`
-	Cn         string           `json:"_cn"`
-	ETag       string           `json:"_etag"`
-	Timestamp  int64            `json:"_ts"`
 	OneTimeKey OneTimeKeyCosmos `json:"mx_keyserver_one_time_key"`
 }
 
@@ -108,6 +102,23 @@ type oneTimeKeysStatements struct {
 	tableName string
 }
 
+func getOneTimeKey(s *oneTimeKeysStatements, ctx context.Context, pk string, docId string) (*OneTimeKeyCosmosData, error) {
+	response := OneTimeKeyCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryOneTimeKey(s *oneTimeKeysStatements, ctx context.Context, qry string, params map[string]interface{}) ([]OneTimeKeyCosmosData, error) {
 	var response []OneTimeKeyCosmosData
 
@@ -155,13 +166,23 @@ func queryOneTimeKeyAlgoCount(s *oneTimeKeysStatements, ctx context.Context, qry
 }
 
 func insertOneTimeKeyCore(s *oneTimeKeysStatements, ctx context.Context, dbData OneTimeKeyCosmosData) error {
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	// "INSERT INTO keyserver_one_time_keys (user_id, device_id, key_id, algorithm, ts_added_secs, key_json)" +
-	var _, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+	// " VALUES ($1, $2, $3, $4, $5, $6)" +
-		ctx,
+	// " ON CONFLICT (user_id, device_id, key_id, algorithm)" +
+	// " DO UPDATE SET key_json = $6"
+	existing, _ := getOneTimeKey(s, ctx, dbData.Pk, dbData.Id)
+	if existing != nil {
+		existing.SetUpdateTime()
+		existing.OneTimeKey.KeyJSON = dbData.OneTimeKey.KeyJSON
+	} else {
+		existing = &dbData
+	}
+	err := cosmosdbapi.UpsertDocument(ctx,
+		s.db.connection,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		dbData,
+		existing.Pk,
-		options)
+		existing)
 
 	if err != nil {
 		return err
@@ -266,7 +287,6 @@ func (s *oneTimeKeysStatements) CountOneTimeKeys(ctx context.Context, userID, de
 func (s *oneTimeKeysStatements) InsertOneTimeKeys(
 	ctx context.Context, txn *sql.Tx, keys api.OneTimeKeys,
 ) (*api.OneTimeKeysCount, error) {
-	now := time.Now().Unix()
 	counts := &api.OneTimeKeysCount{
 		DeviceID: keys.DeviceID,
 		UserID:   keys.UserID,
@@ -298,11 +318,7 @@ func (s *oneTimeKeysStatements) InsertOneTimeKeys(
 		}
 
 		dbData := &OneTimeKeyCosmosData{
-			Id:         cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-			Tn:         s.db.cosmosConfig.TenantName,
-			Cn:         dbCollectionName,
-			Pk:         pk,
-			Timestamp:  now,
 			OneTimeKey:     data,
 		}
 

keyserver/storage/cosmosdb/stale_device_lists.go

@@ -39,17 +39,11 @@ type StaleDeviceListCosmos struct {
 	UserID    string `json:"user_id"`
 	Domain    string `json:"domain"`
 	IsStale   bool   `json:"is_stale"`
-	// Use the CosmosDB.Timestamp for this one
+	AddedSecs int64  `json:"ts_added_secs"`
-	// ts_added_secs    int64  `json:"ts_added_secs"`
 }
 
 type StaleDeviceListCosmosData struct {
-	Id              string                `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk              string                `json:"_pk"`
-	Tn              string                `json:"_sid"`
-	Cn              string                `json:"_cn"`
-	ETag            string                `json:"_etag"`
-	Timestamp       int64                 `json:"_ts"`
 	StaleDeviceList StaleDeviceListCosmos `json:"mx_keyserver_stale_device_list"`
 }
 
@@ -78,6 +72,23 @@ type staleDeviceListsStatements struct {
 	tableName                             string
 }
 
+func getStaleDeviceList(s *staleDeviceListsStatements, ctx context.Context, pk string, docId string) (*StaleDeviceListCosmosData, error) {
+	response := StaleDeviceListCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryStaleDeviceList(s *staleDeviceListsStatements, ctx context.Context, qry string, params map[string]interface{}) ([]StaleDeviceListCosmosData, error) {
 	var response []StaleDeviceListCosmosData
 
@@ -126,31 +137,30 @@ func (s *staleDeviceListsStatements) InsertStaleDeviceList(ctx context.Context,
 	docId := userID
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 
+	dbData, _ := getStaleDeviceList(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.StaleDeviceList.IsStale = isStale
+		dbData.StaleDeviceList.AddedSecs = time.Now().Unix()
+	} else {
 		data := StaleDeviceListCosmos{
 			Domain:  string(domain),
 			IsStale: isStale,
 			UserID:  userID,
 		}
 
-	dbData := StaleDeviceListCosmosData{
+		dbData = &StaleDeviceListCosmosData{
-		Id:              cosmosDocId,
+			CosmosDocument:  cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:              s.db.cosmosConfig.TenantName,
-		Cn:              dbCollectionName,
-		Pk:              pk,
-		Timestamp:       time.Now().Unix(),
 			StaleDeviceList: data,
 		}
-
+	}
 	// _, err := s.upsertKeyChangeStmt.ExecContext(ctx, partition, offset, userID)
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		dbData,
+		dbData.Pk,
-		options)
+		dbData)
-
-	return err
 }
 
 func (s *staleDeviceListsStatements) SelectUserIDsWithStaleDeviceLists(ctx context.Context, domains []gomatrixserverlib.ServerName) ([]string, error) {

mediaapi/storage/cosmosdb/media_repository_table.go

@@ -66,12 +66,7 @@ type MediaRepositoryCosmos struct {
 }
 
 type MediaRepositoryCosmosData struct {
-	Id              string                `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk              string                `json:"_pk"`
-	Tn              string                `json:"_sid"`
-	Cn              string                `json:"_cn"`
-	ETag            string                `json:"_etag"`
-	Timestamp       int64                 `json:"_ts"`
 	MediaRepository MediaRepositoryCosmos `json:"mx_mediaapi_media_repository"`
 }
 
@@ -173,11 +168,7 @@ func (s *mediaStatements) insertMedia(
 	}
 
 	dbData := &MediaRepositoryCosmosData{
-		Id:              cosmosDocId,
+		CosmosDocument:  cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:              s.db.cosmosConfig.TenantName,
-		Cn:              dbCollectionName,
-		Pk:              pk,
-		Timestamp:       time.Now().Unix(),
 		MediaRepository: data,
 	}
 

mediaapi/storage/cosmosdb/thumbnail_table.go

@@ -55,12 +55,7 @@ type ThumbnailCosmos struct {
 }
 
 type ThumbnailCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	Thumbnail ThumbnailCosmos `json:"mx_mediaapi_thumbnail"`
 }
 
@@ -149,7 +144,7 @@ func (s *thumbnailStatements) insertThumbnail(
 
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	// CREATE UNIQUE INDEX IF NOT EXISTS mediaapi_thumbnail_index ON mediaapi_thumbnail (media_id, media_origin, width, height, resize_method);
-	docId := fmt.Sprintf("%s_%s_%d_%d_s",
+	docId := fmt.Sprintf("%s_%s_%d_%d_%s",
 		thumbnailMetadata.MediaMetadata.MediaID,
 		thumbnailMetadata.MediaMetadata.Origin,
 		thumbnailMetadata.ThumbnailSize.Width,
@@ -183,22 +178,17 @@ func (s *thumbnailStatements) insertThumbnail(
 	}
 
 	dbData := &ThumbnailCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		Thumbnail:      data,
 	}
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	var options = cosmosdbapi.GetCreateDocumentOptions(dbData.Pk)
 	_, _, err := cosmosdbapi.GetClient(s.db.connection).CreateDocument(
 		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
 		&dbData,
 		options)
-
 	return err
 }
 
@@ -225,7 +215,7 @@ func (s *thumbnailStatements) selectThumbnail(
 
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	// CREATE UNIQUE INDEX IF NOT EXISTS mediaapi_thumbnail_index ON mediaapi_thumbnail (media_id, media_origin, width, height, resize_method);
-	docId := fmt.Sprintf("%s_%s_%d_%d_s",
+	docId := fmt.Sprintf("%s_%s_%d_%d_%s",
 		mediaID,
 		mediaOrigin,
 		width,

roomserver/storage/cosmosdb/event_json_table.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 
@@ -40,12 +39,7 @@ type EventJSONCosmos struct {
 }
 
 type EventJSONCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	EventJSON EventJSONCosmos `json:"mx_roomserver_event_json"`
 }
 
@@ -71,6 +65,23 @@ type eventJSONStatements struct {
 	tableName               string
 }
 
+func getEventJSON(s *eventJSONStatements, ctx context.Context, pk string, docId string) (*EventJSONCosmosData, error) {
+	response := EventJSONCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryEventJSON(s *eventJSONStatements, ctx context.Context, qry string, params map[string]interface{}) ([]EventJSONCosmosData, error) {
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
@@ -121,30 +132,29 @@ func (s *eventJSONStatements) InsertEventJSON(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	dbData, _ := getEventJSON(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.EventJSON.EventJSON = eventJSON
+	} else {
 		data := EventJSONCosmos{
 			EventNID:  int64(eventNID),
 			EventJSON: eventJSON,
 		}
 
-	var dbData = EventJSONCosmosData{
+		dbData = &EventJSONCosmosData{
-		Id:        cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 			EventJSON:      data,
 		}
+	}
 
 	//Insert OR Replace
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	var _, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
-
-	return err
 }
 
 func (s *eventJSONStatements) BulkSelectEventJSON(

roomserver/storage/cosmosdb/event_state_keys_table.go

@@ -18,7 +18,6 @@ package cosmosdb
 import (
 	"context"
 	"database/sql"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 
@@ -44,12 +43,7 @@ type EventStateKeysCosmos struct {
 }
 
 type EventStateKeysCosmosData struct {
-	Id             string               `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk             string               `json:"_pk"`
-	Tn             string               `json:"_sid"`
-	Cn             string               `json:"_cn"`
-	ETag           string               `json:"_etag"`
-	Timestamp      int64                `json:"_ts"`
 	EventStateKeys EventStateKeysCosmos `json:"mx_roomserver_event_state_keys"`
 }
 
@@ -163,11 +157,7 @@ func ensureEventStateKeys(s *eventStateKeyStatements, ctx context.Context) {
 
 	//     event_state_key_nid INTEGER PRIMARY KEY AUTOINCREMENT,
 	dbData := EventStateKeysCosmosData{
-		Id:             cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:             s.db.cosmosConfig.TenantName,
-		Cn:             dbCollectionName,
-		Pk:             pk,
-		Timestamp:      time.Now().Unix(),
 		EventStateKeys: data,
 	}
 
@@ -175,13 +165,12 @@ func ensureEventStateKeys(s *eventStateKeyStatements, ctx context.Context) {
 }
 
 func insertEventStateKeyCore(s *eventStateKeyStatements, ctx context.Context, dbData EventStateKeysCosmosData) error {
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	err := cosmosdbapi.UpsertDocument(ctx,
-	var _, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		dbData,
+		dbData.Pk,
-		options)
+		dbData)
 
 	if err != nil {
 		return err
@@ -223,14 +212,11 @@ func (s *eventStateKeyStatements) InsertEventStateKeyNID(
 
 		//     event_state_key_nid INTEGER PRIMARY KEY AUTOINCREMENT,
 		dbData = EventStateKeysCosmosData{
-			Id:             cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-			Tn:             s.db.cosmosConfig.TenantName,
-			Cn:             dbCollectionName,
-			Pk:             pk,
-			Timestamp:      time.Now().Unix(),
 			EventStateKeys: data,
 		}
 	} else {
+		dbData.SetUpdateTime()
 		dbData.EventStateKeys = existing.EventStateKeys
 	}
 

roomserver/storage/cosmosdb/event_types_table.go

@@ -18,7 +18,6 @@ package cosmosdb
 import (
 	"context"
 	"database/sql"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 
@@ -44,12 +43,7 @@ import (
 // `
 
 type EventTypeCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	EventType EventTypeCosmos `json:"mx_roomserver_event_type"`
 }
 
@@ -172,11 +166,7 @@ func insertEventTypeCore(s *eventTypeStatements, ctx context.Context, eventType
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
 	var dbData = EventTypeCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		EventType:      eventType,
 	}
 

roomserver/storage/cosmosdb/events_table.go

@@ -20,7 +20,6 @@ import (
 	"database/sql"
 	"fmt"
 	"sort"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 
@@ -66,12 +65,7 @@ type EventCosmosMaxDepth struct {
 }
 
 type EventCosmosData struct {
-	Id        string      `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string      `json:"_pk"`
-	Tn        string      `json:"_sid"`
-	Cn        string      `json:"_cn"`
-	ETag      string      `json:"_etag"`
-	Timestamp int64       `json:"_ts"`
 	Event EventCosmos `json:"mx_roomserver_event"`
 }
 
@@ -377,11 +371,7 @@ func (s *eventStatements) InsertEvent(
 		}
 
 		dbData = &EventCosmosData{
-			Id:        cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-			Tn:        s.db.cosmosConfig.TenantName,
-			Cn:        dbCollectionName,
-			Pk:        pk,
-			Timestamp: time.Now().Unix(),
 			Event:          data,
 		}
 	} else {
@@ -410,17 +400,16 @@ func (s *eventStatements) InsertEvent(
 		dbData.Event.ReferenceSha256 = referenceSHA256
 		dbData.Event.RoomNID = int64(roomNID)
 
-		dbData.Timestamp = time.Now().Unix()
+		dbData.SetUpdateTime()
 	}
 
 	// ON CONFLICT DO NOTHING; - Do Upsert
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	err := cosmosdbapi.UpsertDocument(ctx,
-	_, _, err := cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
 
 	if err != nil {
 		return 0, 0, err

roomserver/storage/cosmosdb/invite_table.go

@@ -18,7 +18,6 @@ package cosmosdb
 import (
 	"context"
 	"database/sql"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -51,12 +50,7 @@ type InviteCosmos struct {
 }
 
 type InviteCosmosData struct {
-	Id        string       `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string       `json:"_pk"`
-	Tn        string       `json:"_sid"`
-	Cn        string       `json:"_cn"`
-	ETag      string       `json:"_etag"`
-	Timestamp int64        `json:"_ts"`
 	Invite InviteCosmos `json:"mx_roomserver_invite"`
 }
 
@@ -191,11 +185,7 @@ func (s *inviteStatements) InsertInviteEvent(
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
 	var dbData = InviteCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		Invite:         data,
 	}
 

roomserver/storage/cosmosdb/membership_table.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -53,12 +52,7 @@ type MembershipCosmos struct {
 }
 
 type MembershipCosmosData struct {
-	Id         string           `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk         string           `json:"_pk"`
-	Tn         string           `json:"_sid"`
-	Cn         string           `json:"_cn"`
-	ETag       string           `json:"_etag"`
-	Timestamp  int64            `json:"_ts"`
 	Membership MembershipCosmos `json:"mx_roomserver_membership"`
 }
 
@@ -303,6 +297,7 @@ func (s *membershipStatements) InsertMembership(
 		exists.Membership.RoomNID = int64(roomNID)
 		exists.Membership.TargetNID = int64(targetUserNID)
 		exists.Membership.TargetLocal = localTarget
+		exists.SetUpdateTime()
 		_, errSet := setMembership(s, ctx, *exists)
 		return errSet
 	}
@@ -318,23 +313,18 @@ func (s *membershipStatements) InsertMembership(
 	}
 
 	var dbData = MembershipCosmosData{
-		Id:         cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:         s.db.cosmosConfig.TenantName,
-		Cn:         dbCollectionName,
-		Pk:         pk,
-		Timestamp:  time.Now().Unix(),
 		Membership:     data,
 	}
 
 	// " ON CONFLICT DO NOTHING"
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	var options = cosmosdbapi.GetCreateDocumentOptions(dbData.Pk)
 	_, _, err := cosmosdbapi.GetClient(s.db.connection).CreateDocument(
 		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
 		&dbData,
 		options)
-
 	return err
 }
 

roomserver/storage/cosmosdb/previous_events_table.go

@@ -20,7 +20,6 @@ import (
 	"database/sql"
 	"fmt"
 	"strings"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 
@@ -51,12 +50,7 @@ type PreviousEventCosmos struct {
 }
 
 type PreviousEventCosmosData struct {
-	Id            string              `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk            string              `json:"_pk"`
-	Tn            string              `json:"_sid"`
-	Cn            string              `json:"_cn"`
-	ETag          string              `json:"_etag"`
-	Timestamp     int64               `json:"_ts"`
 	PreviousEvent PreviousEventCosmos `json:"mx_roomserver_previous_event"`
 }
 
@@ -160,15 +154,12 @@ func (s *previousEventStatements) InsertPreviousEvent(
 		}
 
 		dbData = PreviousEventCosmosData{
-			Id:            cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-			Tn:            s.db.cosmosConfig.TenantName,
-			Cn:            dbCollectionName,
-			Pk:            pk,
-			Timestamp:     time.Now().Unix(),
 			PreviousEvent:  data,
 		}
 	} else {
 		dbData = *existing
+		dbData.SetUpdateTime()
 	}
 
 	var nids []string
@@ -188,15 +179,12 @@ func (s *previousEventStatements) InsertPreviousEvent(
 	//   (previous_event_id, previous_reference_sha256, event_nids)
 	//   VALUES ($1, $2, $3)
 
-	var optionsReplace = cosmosdbapi.GetUpsertDocumentOptions(pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		optionsReplace,
+		dbData)
-	)
-	return err
 }
 
 // Check if the event reference exists

roomserver/storage/cosmosdb/published_table.go

@@ -17,7 +17,6 @@ package cosmosdb
 import (
 	"context"
 	"database/sql"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -41,12 +40,7 @@ type PublishCosmos struct {
 }
 
 type PublishCosmosData struct {
-	Id        string        `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string        `json:"_pk"`
-	Tn        string        `json:"_sid"`
-	Cn        string        `json:"_cn"`
-	ETag      string        `json:"_etag"`
-	Timestamp int64         `json:"_ts"`
 	Publish PublishCosmos `json:"mx_roomserver_publish"`
 }
 
@@ -137,30 +131,29 @@ func (s *publishedStatements) UpsertRoomPublished(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	dbData, _ := getPublish(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.Publish.Published = published
+	} else {
 		data := PublishCosmos{
 			RoomID:    roomID,
 			Published: false,
 		}
 
-	var dbData = PublishCosmosData{
+		dbData = &PublishCosmosData{
-		Id:        cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 			Publish:        data,
 		}
+	}
 
 	// 	"INSERT OR REPLACE INTO roomserver_published (room_id, published) VALUES ($1, $2)"
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	_, _, err := cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
-
-	return err
 }
 
 func (s *publishedStatements) SelectPublishedFromRoomID(

roomserver/storage/cosmosdb/redactions_table.go

@@ -17,7 +17,6 @@ package cosmosdb
 import (
 	"context"
 	"database/sql"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -45,12 +44,7 @@ type RedactionCosmos struct {
 }
 
 type RedactionCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	Redaction RedactionCosmos `json:"mx_roomserver_redaction"`
 }
 
@@ -165,11 +159,7 @@ func (s *redactionStatements) InsertRedaction(
 	}
 
 	var dbData = RedactionCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		Redaction:      data,
 	}
 

roomserver/storage/cosmosdb/room_aliases_table.go

@@ -18,7 +18,6 @@ package cosmosdb
 import (
 	"context"
 	"database/sql"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/roomserver/storage/tables"
@@ -41,12 +40,7 @@ type RoomAliasCosmos struct {
 }
 
 type RoomAliasCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	RoomAlias RoomAliasCosmos `json:"mx_roomserver_room_alias"`
 }
 
@@ -157,11 +151,7 @@ func (s *roomAliasesStatements) InsertRoomAlias(
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
 	var dbData = RoomAliasCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		RoomAlias:      data,
 	}
 

roomserver/storage/cosmosdb/rooms_table.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 
@@ -42,12 +41,7 @@ import (
 // `
 
 type RoomCosmosData struct {
-	Id        string     `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string     `json:"_pk"`
-	Tn        string     `json:"_sid"`
-	Cn        string     `json:"_cn"`
-	ETag      string     `json:"_etag"`
-	Timestamp int64      `json:"_ts"`
 	Room RoomCosmos `json:"mx_roomserver_room"`
 }
 
@@ -273,23 +267,21 @@ func (s *roomStatements) InsertRoomNID(
 		}
 
 		dbData = &RoomCosmosData{
-			Id:        cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-			Tn:        s.db.cosmosConfig.TenantName,
-			Cn:        dbCollectionName,
-			Pk:        pk,
-			Timestamp: time.Now().Unix(),
 			Room:           data,
 		}
+	} else {
+		dbData.SetUpdateTime()
+		dbData.Room.RoomVersion = string(roomVersion)
 	}
 
 	// ON CONFLICT DO NOTHING; - Do Upsert
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	err = cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
 
 	if err != nil {
 		return 0, fmt.Errorf("s.SelectRoomNID: %w", err)

roomserver/storage/cosmosdb/state_block_table.go

@@ -21,7 +21,6 @@ import (
 	"encoding/hex"
 	"fmt"
 	"sort"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -54,12 +53,7 @@ type StateBlockCosmosMaxNID struct {
 }
 
 type StateBlockCosmosData struct {
-	Id         string           `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk         string           `json:"_pk"`
-	Tn         string           `json:"_sid"`
-	Cn         string           `json:"_cn"`
-	ETag       string           `json:"_etag"`
-	Timestamp  int64            `json:"_ts"`
 	StateBlock StateBlockCosmos `json:"mx_roomserver_state_block"`
 }
 
@@ -192,6 +186,7 @@ func (s *stateBlockStatements) BulkInsertStateData(
 	if existing != nil {
 		//if exists, just update and dont create a new seq
 		existing.StateBlock.EventNIDs = ids
+		existing.SetUpdateTime()
 		_, err = setStateBlock(s, ctx, *existing)
 		if err != nil {
 			return 0, err
@@ -211,21 +206,16 @@ func (s *stateBlockStatements) BulkInsertStateData(
 	}
 
 	var dbData = StateBlockCosmosData{
-		Id:         cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:         s.db.cosmosConfig.TenantName,
-		Cn:         dbCollectionName,
-		Pk:         pk,
-		Timestamp:  time.Now().Unix(),
 		StateBlock:     data,
 	}
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	err = cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
 
 	return
 }

roomserver/storage/cosmosdb/state_snapshot_table.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/roomserver/storage/tables"
@@ -56,12 +55,7 @@ type StateSnapshotCosmos struct {
 }
 
 type StateSnapshotCosmosData struct {
-	Id            string              `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk            string              `json:"_pk"`
-	Tn            string              `json:"_sid"`
-	Cn            string              `json:"_cn"`
-	ETag          string              `json:"_etag"`
-	Timestamp     int64               `json:"_ts"`
 	StateSnapshot StateSnapshotCosmos `json:"mx_roomserver_state_snapshot"`
 }
 
@@ -154,11 +148,7 @@ func (s *stateSnapshotStatements) InsertState(
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
 	var dbData = StateSnapshotCosmosData{
-		Id:            cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:            s.db.cosmosConfig.TenantName,
-		Cn:            dbCollectionName,
-		Pk:            pk,
-		Timestamp:     time.Now().Unix(),
 		StateSnapshot:  data,
 	}
 

roomserver/storage/cosmosdb/transactions_table.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -44,12 +43,7 @@ type TransactionCosmos struct {
 }
 
 type TransactionCosmosData struct {
-	Id          string            `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk          string            `json:"_pk"`
-	Tn          string            `json:"_sid"`
-	Cn          string            `json:"_cn"`
-	ETag        string            `json:"_etag"`
-	Timestamp   int64             `json:"_ts"`
 	Transaction TransactionCosmos `json:"mx_roomserver_transaction"`
 }
 
@@ -124,11 +118,7 @@ func (s *transactionStatements) InsertTransaction(
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
 	var dbData = TransactionCosmosData{
-		Id:          cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:          s.db.cosmosConfig.TenantName,
-		Cn:          dbCollectionName,
-		Pk:          pk,
-		Timestamp:   time.Now().Unix(),
 		Transaction:    data,
 	}
 

signingkeyserver/storage/cosmosdb/server_key_table.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -62,12 +61,7 @@ type ServerKeyCosmos struct {
 }
 
 type ServerKeyCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	ServerKey ServerKeyCosmos `json:"mx_keydb_server_key"`
 }
 
@@ -93,6 +87,23 @@ type serverKeyStatements struct {
 	tableName string
 }
 
+func getServerKey(s *serverKeyStatements, ctx context.Context, pk string, docId string) (*ServerKeyCosmosData, error) {
+	response := ServerKeyCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryServerKey(s *serverKeyStatements, ctx context.Context, qry string, params map[string]interface{}) ([]ServerKeyCosmosData, error) {
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
@@ -208,6 +219,13 @@ func (s *serverKeyStatements) upsertServerKeys(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	dbData, _ := getServerKey(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.ServerKey.ValidUntilTimestamp = int64(key.ValidUntilTS)
+		dbData.ServerKey.ExpiredTimestamp = int64(key.ExpiredTS)
+		dbData.ServerKey.ServerKey = key.Key.Encode()
+	} else {
 		data := ServerKeyCosmos{
 			ServerName:          string(request.ServerName),
 			ServerKeyID:         string(request.KeyID),
@@ -217,15 +235,11 @@ func (s *serverKeyStatements) upsertServerKeys(
 			ServerKey:           key.Key.Encode(),
 		}
 
-	dbData := &ServerKeyCosmosData{
+		dbData = &ServerKeyCosmosData{
-		Id:        cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 			ServerKey:      data,
 		}
-
+	}
 	// _, err := stmt.ExecContext(
 	// 	ctx,
 	// 	string(request.ServerName),
@@ -236,15 +250,12 @@ func (s *serverKeyStatements) upsertServerKeys(
 	// 	key.Key.Encode(),
 	// )
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	_, _, err := cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
-
-	return err
 }
 
 func nameAndKeyID(request gomatrixserverlib.PublicKeyLookupRequest) string {

syncapi/storage/cosmosdb/account_data_table.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 
@@ -52,12 +51,7 @@ type AccountDataTypeNumberCosmosData struct {
 }
 
 type AccountDataTypeCosmosData struct {
-	Id              string                `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk              string                `json:"_pk"`
-	Tn              string                `json:"_sid"`
-	Cn              string                `json:"_cn"`
-	ETag            string                `json:"_etag"`
-	Timestamp       int64                 `json:"_ts"`
 	AccountDataType AccountDataTypeCosmos `json:"mx_syncapi_account_data_type"`
 }
 
@@ -89,6 +83,23 @@ type accountDataStatements struct {
 	tableName                    string
 }
 
+func getAccountDataType(s *accountDataStatements, ctx context.Context, pk string, docId string) (*AccountDataTypeCosmosData, error) {
+	response := AccountDataTypeCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryAccountDataType(s *accountDataStatements, ctx context.Context, qry string, params map[string]interface{}) ([]AccountDataTypeCosmosData, error) {
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
@@ -163,6 +174,11 @@ func (s *accountDataStatements) InsertAccountData(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	dbData, _ := getAccountDataType(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+		dbData.AccountDataType.ID = int64(pos)
+	} else {
 		data := AccountDataTypeCosmos{
 			ID:       int64(pos),
 			UserID:   userID,
@@ -170,23 +186,19 @@ func (s *accountDataStatements) InsertAccountData(
 			DataType: dataType,
 		}
 
-	dbData := &AccountDataTypeCosmosData{
+		dbData = &AccountDataTypeCosmosData{
-		Id:              cosmosDocId,
+			CosmosDocument:  cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:              s.db.cosmosConfig.TenantName,
-		Cn:              dbCollectionName,
-		Pk:              pk,
-		Timestamp:       time.Now().Unix(),
 			AccountDataType: data,
 		}
+	}
 
 	// _, err = sqlutil.TxStmt(txn, s.insertAccountDataStmt).ExecContext(ctx, pos, userID, roomID, dataType, pos)
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	err = cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
 
 	return
 }

syncapi/storage/cosmosdb/backwards_extremities_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 
@@ -45,12 +44,7 @@ type BackwardExtremityCosmos struct {
 }
 
 type BackwardExtremityCosmosData struct {
-	Id                string                  `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk                string                  `json:"_pk"`
-	Tn                string                  `json:"_sid"`
-	Cn                string                  `json:"_cn"`
-	ETag              string                  `json:"_etag"`
-	Timestamp         int64                   `json:"_ts"`
 	BackwardExtremity BackwardExtremityCosmos `json:"mx_syncapi_backward_extremity"`
 }
 
@@ -84,6 +78,23 @@ type backwardExtremitiesStatements struct {
 	tableName                            string
 }
 
+func getBackwardExtremity(s *backwardExtremitiesStatements, ctx context.Context, pk string, docId string) (*BackwardExtremityCosmosData, error) {
+	response := BackwardExtremityCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryBackwardExtremity(s *backwardExtremitiesStatements, ctx context.Context, qry string, params map[string]interface{}) ([]BackwardExtremityCosmosData, error) {
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
@@ -147,28 +158,28 @@ func (s *backwardExtremitiesStatements) InsertsBackwardExtremity(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	dbData, _ := getBackwardExtremity(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		dbData.SetUpdateTime()
+	} else {
 		data := BackwardExtremityCosmos{
 			EventID:     eventID,
 			PrevEventID: prevEventID,
 			RoomID:      roomID,
 		}
 
-	dbData := &BackwardExtremityCosmosData{
+		dbData = &BackwardExtremityCosmosData{
-		Id:                cosmosDocId,
+			CosmosDocument:    cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:                s.db.cosmosConfig.TenantName,
-		Cn:                dbCollectionName,
-		Pk:                pk,
-		Timestamp:         time.Now().Unix(),
 			BackwardExtremity: data,
 		}
+	}
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	err = cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
 
 	return
 }

syncapi/storage/cosmosdb/current_room_state_table.go

@@ -20,7 +20,6 @@ import (
 	"database/sql"
 	"encoding/json"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 
@@ -66,12 +65,7 @@ type CurrentRoomStateCosmos struct {
 }
 
 type CurrentRoomStateCosmosData struct {
-	Id               string                 `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk               string                 `json:"_pk"`
-	Tn               string                 `json:"_sid"`
-	Cn               string                 `json:"_cn"`
-	ETag             string                 `json:"_etag"`
-	Timestamp        int64                  `json:"_ts"`
 	CurrentRoomState CurrentRoomStateCosmos `json:"mx_syncapi_current_room_state"`
 }
 
@@ -424,6 +418,17 @@ func (s *currentRoomStateStatements) UpsertRoomState(
 		membershipData = *membership
 	}
 
+	dbData, _ := getEvent(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		// " DO UPDATE SET event_id = $2, sender=$4, contains_url=$5, headered_event_json = $7, membership = $8, added_at = $9"
+		dbData.SetUpdateTime()
+		dbData.CurrentRoomState.EventID = event.EventID()
+		dbData.CurrentRoomState.Sender = event.Sender()
+		dbData.CurrentRoomState.ContainsUrl = containsURL
+		dbData.CurrentRoomState.HeaderedEventJSON = headeredJSON
+		dbData.CurrentRoomState.Membership = membershipData
+		dbData.CurrentRoomState.AddedAt = int64(addedAt)
+	} else {
 		data := CurrentRoomStateCosmos{
 			RoomID:            event.RoomID(),
 			EventID:           event.EventID(),
@@ -436,25 +441,19 @@ func (s *currentRoomStateStatements) UpsertRoomState(
 			AddedAt:           int64(addedAt),
 		}
 
-	dbData := &CurrentRoomStateCosmosData{
+		dbData = &CurrentRoomStateCosmosData{
-		Id:               cosmosDocId,
+			CosmosDocument:   cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:               s.db.cosmosConfig.TenantName,
-		Cn:               dbCollectionName,
-		Pk:               pk,
-		Timestamp:        time.Now().Unix(),
 			CurrentRoomState: data,
 		}
+	}
 
 	// _, err = sqlutil.TxStmt(txn, s.insertAccountDataStmt).ExecContext(ctx, pos, userID, roomID, dataType, pos)
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
-
-	return err
 }
 
 func minOfInts(a, b int) int {

syncapi/storage/cosmosdb/filter_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -49,12 +48,7 @@ type FilterCosmos struct {
 }
 
 type FilterCosmosData struct {
-	Id        string       `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string       `json:"_pk"`
-	Tn        string       `json:"_sid"`
-	Cn        string       `json:"_cn"`
-	ETag      string       `json:"_etag"`
-	Timestamp int64        `json:"_ts"`
 	Filter FilterCosmos `json:"mx_syncapi_filter"`
 }
 
@@ -235,11 +229,7 @@ func (s *filterStatements) InsertFilter(
 	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
 	var dbData = FilterCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		Filter:         data,
 	}
 

syncapi/storage/cosmosdb/invites_table.go

@@ -20,7 +20,6 @@ import (
 	"database/sql"
 	"encoding/json"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -58,12 +57,7 @@ type InviteEventCosmosMaxNumber struct {
 }
 
 type InviteEventCosmosData struct {
-	Id          string            `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk          string            `json:"_pk"`
-	Tn          string            `json:"_sid"`
-	Cn          string            `json:"_cn"`
-	ETag        string            `json:"_etag"`
-	Timestamp   int64             `json:"_ts"`
 	InviteEvent InviteEventCosmos `json:"mx_syncapi_invite_event"`
 }
 
@@ -228,11 +222,7 @@ func (s *inviteEventsStatements) InsertInviteEvent(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 
 	var dbData = InviteEventCosmosData{
-		Id:          cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:          s.db.cosmosConfig.TenantName,
-		Cn:          dbCollectionName,
-		Pk:          pk,
-		Timestamp:   time.Now().Unix(),
 		InviteEvent:    data,
 	}
 

syncapi/storage/cosmosdb/memberships_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 
@@ -62,12 +61,7 @@ type MembershipCosmos struct {
 }
 
 type MembershipCosmosData struct {
-	Id         string           `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk         string           `json:"_pk"`
-	Tn         string           `json:"_sid"`
-	Cn         string           `json:"_cn"`
-	ETag       string           `json:"_etag"`
-	Timestamp  int64            `json:"_ts"`
 	Membership MembershipCosmos `json:"mx_syncapi_membership"`
 }
 
@@ -94,6 +88,23 @@ type membershipsStatements struct {
 	tableName string
 }
 
+func getMembership(s *membershipsStatements, ctx context.Context, pk string, docId string) (*MembershipCosmosData, error) {
+	response := MembershipCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryMembership(s *membershipsStatements, ctx context.Context, qry string, params map[string]interface{}) ([]MembershipCosmosData, error) {
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
@@ -147,6 +158,20 @@ func (s *membershipsStatements) UpsertMembership(
 	// 	topologicalPos,
 	// )
 
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
+	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
+	// 	UNIQUE (room_id, user_id, membership)
+	docId := fmt.Sprintf("%s_%s_%s", event.RoomID(), *event.StateKey(), membership)
+	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
+
+	dbData, _ := getMembership(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		// " DO UPDATE SET event_id = $4, stream_pos = $5, topological_pos = $6"
+		dbData.SetUpdateTime()
+		dbData.Membership.EventID = event.EventID()
+		dbData.Membership.StreamPos = int64(streamPos)
+		dbData.Membership.TopologicalPos = int64(topologicalPos)
+	} else {
 		data := MembershipCosmos{
 			RoomID:         event.RoomID(),
 			UserID:         *event.StateKey(),
@@ -156,30 +181,18 @@ func (s *membershipsStatements) UpsertMembership(
 			TopologicalPos: int64(topologicalPos),
 		}
 
-	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
+		dbData = &MembershipCosmosData{
-	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-	// 	UNIQUE (room_id, user_id, membership)
-	docId := fmt.Sprintf("%s_%s_%s", event.RoomID(), *event.StateKey(), membership)
-	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
-
-	var dbData = MembershipCosmosData{
-		Id:         cosmosDocId,
-		Tn:         s.db.cosmosConfig.TenantName,
-		Cn:         dbCollectionName,
-		Pk:         pk,
-		Timestamp:  time.Now().Unix(),
 			Membership:     data,
 		}
+	}
 
-	var optionsCreate = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		dbData,
+		dbData.Pk,
-		optionsCreate)
+		dbData)
-
-	return err
 }
 
 func (s *membershipsStatements) SelectMembership(

syncapi/storage/cosmosdb/output_room_events_table.go

@@ -21,7 +21,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"sort"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 
@@ -73,12 +72,7 @@ type OutputRoomEventCosmosMaxNumber struct {
 }
 
 type OutputRoomEventCosmosData struct {
-	Id              string                `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk              string                `json:"_pk"`
-	Tn              string                `json:"_sid"`
-	Cn              string                `json:"_cn"`
-	ETag            string                `json:"_etag"`
-	Timestamp       int64                 `json:"_ts"`
 	OutputRoomEvent OutputRoomEventCosmos `json:"mx_syncapi_output_room_event"`
 }
 
@@ -495,11 +489,7 @@ func (s *outputRoomEventsStatements) InsertEvent(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 
 	var dbData = OutputRoomEventCosmosData{
-		Id:              cosmosDocId,
+		CosmosDocument:  cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:              s.db.cosmosConfig.TenantName,
-		Cn:              dbCollectionName,
-		Pk:              pk,
-		Timestamp:       time.Now().Unix(),
 		OutputRoomEvent: data,
 	}
 

syncapi/storage/cosmosdb/output_room_events_topology_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/syncapi/storage/tables"
@@ -48,12 +47,7 @@ type OutputRoomEventTopologyCosmos struct {
 }
 
 type OutputRoomEventTopologyCosmosData struct {
-	Id                      string                        `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk                      string                        `json:"_pk"`
-	Tn                      string                        `json:"_sid"`
-	Cn                      string                        `json:"_cn"`
-	ETag                    string                        `json:"_etag"`
-	Timestamp               int64                         `json:"_ts"`
 	OutputRoomEventTopology OutputRoomEventTopologyCosmos `json:"mx_syncapi_output_room_event_topology"`
 }
 
@@ -153,6 +147,23 @@ func queryOutputRoomEventTopology(s *outputRoomEventsTopologyStatements, ctx con
 	return response, nil
 }
 
+func getOutputRoomEventTopology(s *outputRoomEventsTopologyStatements, ctx context.Context, pk string, docId string) (*OutputRoomEventTopologyCosmosData, error) {
+	response := OutputRoomEventTopologyCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func deleteOutputRoomEventTopology(s *outputRoomEventsTopologyStatements, ctx context.Context, dbData OutputRoomEventTopologyCosmosData) error {
 	var options = cosmosdbapi.GetDeleteDocumentOptions(dbData.Pk)
 	var _, err = cosmosdbapi.GetClient(s.db.connection).DeleteDocument(
@@ -198,6 +209,11 @@ func (s *outputRoomEventsTopologyStatements) InsertEventInTopology(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	var err error
+	dbData, _ := getOutputRoomEventTopology(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		// 	" ON CONFLICT DO NOTHING"
+	} else {
 		data := OutputRoomEventTopologyCosmos{
 			EventID:             event.EventID(),
 			TopologicalPosition: event.Depth(),
@@ -205,28 +221,23 @@ func (s *outputRoomEventsTopologyStatements) InsertEventInTopology(
 			StreamPosition:      int64(pos),
 		}
 
-	dbData := &OutputRoomEventTopologyCosmosData{
+		dbData = &OutputRoomEventTopologyCosmosData{
-		Id:                      cosmosDocId,
+			CosmosDocument:          cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:                      s.db.cosmosConfig.TenantName,
-		Cn:                      dbCollectionName,
-		Pk:                      pk,
-		Timestamp:               time.Now().Unix(),
 			OutputRoomEventTopology: data,
 		}
-
 		// _, err := sqlutil.TxStmt(txn, s.insertEventInTopologyStmt).ExecContext(
 		// 	ctx, event.EventID(), event.Depth(), event.RoomID(), pos,
 		// )
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+		err = cosmosdbapi.UpsertDocument(ctx,
-	_, _, err := cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+			s.db.connection,
-		ctx,
 			s.db.cosmosConfig.DatabaseName,
 			s.db.cosmosConfig.ContainerName,
-		&dbData,
+			dbData.Pk,
-		options)
+			dbData)
+	}
 
-	return types.StreamPosition(event.Depth()), err
+	return types.StreamPosition(dbData.OutputRoomEventTopology.TopologicalPosition), err
 }
 
 func (s *outputRoomEventsTopologyStatements) SelectEventIDsInRange(

syncapi/storage/cosmosdb/peeks_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 
@@ -58,12 +57,7 @@ type PeekCosmosMaxNumber struct {
 }
 
 type PeekCosmosData struct {
-	Id        string     `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string     `json:"_pk"`
-	Tn        string     `json:"_sid"`
-	Cn        string     `json:"_cn"`
-	ETag      string     `json:"_etag"`
-	Timestamp int64      `json:"_ts"`
 	Peek PeekCosmos `json:"mx_syncapi_peek"`
 }
 
@@ -163,6 +157,23 @@ func queryPeekMaxNumber(s *peekStatements, ctx context.Context, qry string, para
 	return response, nil
 }
 
+func getPeek(s *peekStatements, ctx context.Context, pk string, docId string) (*PeekCosmosData, error) {
+	response := PeekCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func setPeek(s *peekStatements, ctx context.Context, peek PeekCosmosData) (*PeekCosmosData, error) {
 	var optionsReplace = cosmosdbapi.GetReplaceDocumentOptions(peek.Pk, peek.ETag)
 	var _, _, ex = cosmosdbapi.GetClient(s.db.connection).ReplaceDocument(
@@ -208,6 +219,13 @@ func (s *peekStatements) InsertPeek(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
+	dbData, _ := getPeek(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		// " (id, room_id, user_id, device_id, creation_ts, deleted)" +
+		// " VALUES ($1, $2, $3, $4, $5, false)"
+		dbData.SetUpdateTime()
+		dbData.Peek.Deleted = false
+	} else {
 		data := PeekCosmos{
 			ID:       int64(streamPos),
 			RoomID:   roomID,
@@ -215,25 +233,20 @@ func (s *peekStatements) InsertPeek(
 			DeviceID: deviceID,
 		}
 
-	dbData := &PeekCosmosData{
+		dbData = &PeekCosmosData{
-		Id: cosmosDocId,
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn: s.db.cosmosConfig.TenantName,
-		Cn: dbCollectionName,
-		Pk: pk,
-		// nowMilli := time.Now().UnixNano() / int64(time.Millisecond)
-		Timestamp: time.Now().Unix(),
 			Peek:           data,
 		}
+	}
 
 	// _, err = sqlutil.TxStmt(txn, s.insertPeekStmt).ExecContext(ctx, streamPos, roomID, userID, deviceID, nowMilli)
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	err = cosmosdbapi.UpsertDocument(ctx,
-	_, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		&dbData,
+		dbData.Pk,
-		options)
+		dbData)
 
 	return
 }

syncapi/storage/cosmosdb/receipt_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/eduserver/api"
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
@@ -56,12 +55,7 @@ type ReceiptCosmosMaxNumber struct {
 }
 
 type ReceiptCosmosData struct {
-	Id        string        `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string        `json:"_pk"`
-	Tn        string        `json:"_sid"`
-	Cn        string        `json:"_cn"`
-	ETag      string        `json:"_etag"`
-	Timestamp int64         `json:"_ts"`
 	Receipt ReceiptCosmos `json:"mx_syncapi_receipt"`
 }
 
@@ -174,11 +168,7 @@ func (r *receiptStatements) UpsertReceipt(ctx context.Context, txn *sql.Tx, room
 	cosmosDocId := cosmosdbapi.GetDocumentId(r.db.cosmosConfig.ContainerName, dbCollectionName, docId)
 
 	var dbData = ReceiptCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, r.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        r.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		Receipt:        data,
 	}
 

syncapi/storage/cosmosdb/send_to_device_table.go

@@ -19,7 +19,6 @@ import (
 	"database/sql"
 	"encoding/json"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/syncapi/storage/tables"
@@ -53,12 +52,7 @@ type SendToDeviceCosmosMaxNumber struct {
 }
 
 type SendToDeviceCosmosData struct {
-	Id           string             `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk           string             `json:"_pk"`
-	Tn           string             `json:"_sid"`
-	Cn           string             `json:"_cn"`
-	ETag         string             `json:"_etag"`
-	Timestamp    int64              `json:"_ts"`
 	SendToDevice SendToDeviceCosmos `json:"mx_syncapi_send_to_device"`
 }
 
@@ -200,11 +194,7 @@ func (s *sendToDeviceStatements) InsertSendToDeviceMessage(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 
 	var dbData = SendToDeviceCosmosData{
-		Id:           cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:           s.db.cosmosConfig.TenantName,
-		Cn:           dbCollectionName,
-		Pk:           pk,
-		Timestamp:    time.Now().Unix(),
 		SendToDevice:   data,
 	}
 

userapi/storage/accounts/cosmosdb/account_data_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 )
@@ -40,12 +39,7 @@ import (
 // `
 
 type AccountDataCosmosData struct {
-	Id          string            `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk          string            `json:"_pk"`
-	Tn          string            `json:"_sid"`
-	Cn          string            `json:"_cn"`
-	ETag        string            `json:"_etag"`
-	Timestamp   int64             `json:"_ts"`
 	AccountData AccountDataCosmos `json:"mx_userapi_accountdata"`
 }
 
@@ -72,6 +66,23 @@ func (s *accountDataStatements) prepare(db *Database) (err error) {
 	return
 }
 
+func getAccountData(s *accountDataStatements, ctx context.Context, pk string, docId string) (*AccountDataCosmosData, error) {
+	response := AccountDataCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
 func queryAccountData(s *accountDataStatements, ctx context.Context, qry string, params map[string]interface{}) ([]AccountDataCosmosData, error) {
 	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
 	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
@@ -99,6 +110,24 @@ func (s *accountDataStatements) insertAccountData(
 
 	// 	INSERT INTO account_data(localpart, room_id, type, content) VALUES($1, $2, $3, $4)
 	// 	ON CONFLICT (localpart, room_id, type) DO UPDATE SET content = $4
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.db.accountDatas.tableName)
+	id := ""
+	if roomID == "" {
+		id = fmt.Sprintf("%s_%s", localpart, dataType)
+	} else {
+		id = fmt.Sprintf("%s_%s_%s", localpart, roomID, dataType)
+	}
+
+	docId := id
+	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
+	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
+
+	dbData, _ := getAccountData(s, ctx, pk, cosmosDocId)
+	if dbData != nil {
+		// 	ON CONFLICT (localpart, room_id, type) DO UPDATE SET content = $4
+		dbData.SetUpdateTime()
+		dbData.AccountData.Content = content
+	} else {
 		var result = AccountDataCosmos{
 			LocalPart: localpart,
 			RoomId:    roomID,
@@ -106,36 +135,18 @@ func (s *accountDataStatements) insertAccountData(
 			Content:   content,
 		}
 
-	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.db.accountDatas.tableName)
+		dbData = &AccountDataCosmosData{
-	id := ""
+			CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-	if roomID == "" {
-		id = fmt.Sprintf("%s_%s", result.LocalPart, result.Type)
-	} else {
-		id = fmt.Sprintf("%s_%s_%s", result.LocalPart, result.RoomId, result.Type)
-	}
-
-	docId := id
-	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
-	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
-
-	var dbData = AccountDataCosmosData{
-		Id:          cosmosDocId,
-		Tn:          s.db.cosmosConfig.TenantName,
-		Cn:          dbCollectionName,
-		Pk:          pk,
-		Timestamp:   time.Now().Unix(),
 			AccountData:    result,
 		}
+	}
 
-	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	return cosmosdbapi.UpsertDocument(ctx,
-	var _, _, err = cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		s.db.connection,
-		ctx,
 		s.db.cosmosConfig.DatabaseName,
 		s.db.cosmosConfig.ContainerName,
-		dbData,
+		dbData.Pk,
-		options)
+		dbData)
-
-	return err
 }
 
 func (s *accountDataStatements) selectAccountData(

userapi/storage/accounts/cosmosdb/accounts_table.go

@@ -57,12 +57,7 @@ type AccountCosmos struct {
 }
 
 type AccountCosmosData struct {
-	Id        string        `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string        `json:"_pk"`
-	Tn        string        `json:"_sid"`
-	Cn        string        `json:"_cn"`
-	ETag      string        `json:"_etag"`
-	Timestamp int64         `json:"_ts"`
 	Account AccountCosmos `json:"mx_userapi_account"`
 }
 
@@ -187,11 +182,7 @@ func (s *accountsStatements) insertAccount(
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
 	var dbData = AccountCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		Account:        data,
 	}
 

userapi/storage/accounts/cosmosdb/key_backup_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/userapi/api"
@@ -42,12 +41,7 @@ import (
 // `
 
 type KeyBackupCosmosData struct {
-	Id        string          `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string          `json:"_pk"`
-	Tn        string          `json:"_sid"`
-	Cn        string          `json:"_cn"`
-	ETag      string          `json:"_etag"`
-	Timestamp int64           `json:"_ts"`
 	KeyBackup KeyBackupCosmos `json:"mx_userapi_account_e2e_room_keys"`
 }
 
@@ -252,11 +246,7 @@ func (s *keyBackupStatements) insertBackupKey(
 	}
 
 	dbData := &KeyBackupCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		KeyBackup:      data,
 	}
 

userapi/storage/accounts/cosmosdb/key_backup_version_table.go

@@ -19,7 +19,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"strconv"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -43,12 +42,7 @@ import (
 // `
 
 type KeyBackupVersionCosmosData struct {
-	Id               string                 `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk               string                 `json:"_pk"`
-	Tn               string                 `json:"_sid"`
-	Cn               string                 `json:"_cn"`
-	ETag             string                 `json:"_etag"`
-	Timestamp        int64                  `json:"_ts"`
 	KeyBackupVersion KeyBackupVersionCosmos `json:"mx_userapi_account_e2e_room_keys_versions"`
 }
 
@@ -194,11 +188,7 @@ func (s *keyBackupVersionStatements) insertKeyBackup(
 	}
 
 	dbData := &KeyBackupVersionCosmosData{
-		Id:               cosmosDocId,
+		CosmosDocument:   cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:               s.db.cosmosConfig.TenantName,
-		Cn:               dbCollectionName,
-		Pk:               pk,
-		Timestamp:        time.Now().Unix(),
 		KeyBackupVersion: data,
 	}
 

userapi/storage/accounts/cosmosdb/openid_table.go

@@ -2,7 +2,6 @@ package cosmosdb
 
 import (
 	"context"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 
@@ -30,12 +29,7 @@ type OpenIDTokenCosmos struct {
 }
 
 type OpenIdTokenCosmosData struct {
-	Id          string            `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk          string            `json:"_pk"`
-	Tn          string            `json:"_sid"`
-	Cn          string            `json:"_cn"`
-	ETag        string            `json:"_etag"`
-	Timestamp   int64             `json:"_ts"`
 	OpenIdToken OpenIDTokenCosmos `json:"mx_userapi_openidtoken"`
 }
 
@@ -114,11 +108,7 @@ func (s *tokenStatements) insertToken(
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
 	var dbData = OpenIdTokenCosmosData{
-		Id:          cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:          s.db.cosmosConfig.TenantName,
-		Cn:          dbCollectionName,
-		Pk:          pk,
-		Timestamp:   time.Now().Unix(),
 		OpenIdToken:    mapToToken(*result),
 	}
 

userapi/storage/accounts/cosmosdb/profile_table.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
@@ -46,12 +45,7 @@ type ProfileCosmos struct {
 }
 
 type ProfileCosmosData struct {
-	Id        string        `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string        `json:"_pk"`
-	Tn        string        `json:"_sid"`
-	Cn        string        `json:"_cn"`
-	ETag      string        `json:"_etag"`
-	Timestamp int64         `json:"_ts"`
 	Profile ProfileCosmos `json:"mx_userapi_profile"`
 }
 
@@ -155,11 +149,7 @@ func (s *profilesStatements) insertProfile(
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 
 	var dbData = ProfileCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		Profile:        mapToProfile(*result),
 	}
 

userapi/storage/accounts/cosmosdb/threepid_table.go

@@ -17,7 +17,6 @@ package cosmosdb
 import (
 	"context"
 	"fmt"
-	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
 
@@ -44,12 +43,7 @@ type ThreePIDCosmos struct {
 }
 
 type ThreePIDCosmosData struct {
-	Id        string         `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string         `json:"_pk"`
-	Tn        string         `json:"_sid"`
-	Cn        string         `json:"_cn"`
-	ETag      string         `json:"_etag"`
-	Timestamp int64          `json:"_ts"`
 	ThreePID ThreePIDCosmos `json:"mx_userapi_threepid"`
 }
 
@@ -163,11 +157,7 @@ func (s *threepidStatements) insertThreePID(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.TenantName, dbCollectionName)
 	var dbData = ThreePIDCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		ThreePID:       result,
 	}
 

userapi/storage/devices/cosmosdb/devices_table.go

@@ -70,12 +70,7 @@ type DeviceCosmos struct {
 }
 
 type DeviceCosmosData struct {
-	Id        string       `json:"id"`
+	cosmosdbapi.CosmosDocument
-	Pk        string       `json:"_pk"`
-	Tn        string       `json:"_sid"`
-	Cn        string       `json:"_cn"`
-	ETag      string       `json:"_etag"`
-	Timestamp int64        `json:"_ts"`
 	Device DeviceCosmos `json:"mx_userapi_device"`
 }
 
@@ -238,11 +233,7 @@ func (s *devicesStatements) insertDevice(
 	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, dbCollectionName, docId)
 
 	var dbData = DeviceCosmosData{
-		Id:        cosmosDocId,
+		CosmosDocument: cosmosdbapi.GenerateDocument(dbCollectionName, s.db.cosmosConfig.TenantName, pk, cosmosDocId),
-		Tn:        s.db.cosmosConfig.TenantName,
-		Cn:        dbCollectionName,
-		Pk:        pk,
-		Timestamp: time.Now().Unix(),
 		Device:         data,
 	}