Implement Cosmos DB for the Media API (#11)

* - Update config to use CosmosDB for MediaAPI
- Implement CosmosDB for MediaAPI
- Fix bug in Public Room Creation

* - Remove file: config line

dendrite-config-cosmosdb.yaml

@@ -241,7 +241,7 @@ media_api:
   external_api:
     listen: http://[::]:8074
   database:
-    connection_string: file:mediaapi.db
+    connection_string: "cosmosdb:AccountEndpoint=https://localhost:8081/;AccountKey=C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==;DatabaseName=safezone_local;ContainerName=test.criticalarc.com;"
     max_open_conns: 5
     max_idle_conns: 2
     conn_max_lifetime: -1

mediaapi/storage/cosmosdb/media_repository_table.go

@@ -17,96 +17,191 @@ package cosmosdb
 
 import (
 	"context"
-	"database/sql"
+	"fmt"
 	"time"
 
-	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
+
+	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
+
 	"github.com/matrix-org/dendrite/mediaapi/types"
 	"github.com/matrix-org/gomatrixserverlib"
 )
 
-const mediaSchema = `
+// const mediaSchema = `
--- The media_repository table holds metadata for each media file stored and accessible to the local server,
+// -- The media_repository table holds metadata for each media file stored and accessible to the local server,
--- the actual file is stored separately.
+// -- the actual file is stored separately.
-CREATE TABLE IF NOT EXISTS mediaapi_media_repository (
+// CREATE TABLE IF NOT EXISTS mediaapi_media_repository (
-    -- The id used to refer to the media.
+//     -- The id used to refer to the media.
-    -- For uploads to this server this is a base64-encoded sha256 hash of the file data
+//     -- For uploads to this server this is a base64-encoded sha256 hash of the file data
-    -- For media from remote servers, this can be any unique identifier string
+//     -- For media from remote servers, this can be any unique identifier string
-    media_id TEXT NOT NULL,
+//     media_id TEXT NOT NULL,
-    -- The origin of the media as requested by the client. Should be a homeserver domain.
+//     -- The origin of the media as requested by the client. Should be a homeserver domain.
-    media_origin TEXT NOT NULL,
+//     media_origin TEXT NOT NULL,
-    -- The MIME-type of the media file as specified when uploading.
+//     -- The MIME-type of the media file as specified when uploading.
-    content_type TEXT NOT NULL,
+//     content_type TEXT NOT NULL,
-    -- Size of the media file in bytes.
+//     -- Size of the media file in bytes.
-    file_size_bytes INTEGER NOT NULL,
+//     file_size_bytes INTEGER NOT NULL,
-    -- When the content was uploaded in UNIX epoch ms.
+//     -- When the content was uploaded in UNIX epoch ms.
-    creation_ts INTEGER NOT NULL,
+//     creation_ts INTEGER NOT NULL,
-    -- The file name with which the media was uploaded.
+//     -- The file name with which the media was uploaded.
-    upload_name TEXT NOT NULL,
+//     upload_name TEXT NOT NULL,
-    -- Alternate RFC 4648 unpadded base64 encoding string representation of a SHA-256 hash sum of the file data.
+//     -- Alternate RFC 4648 unpadded base64 encoding string representation of a SHA-256 hash sum of the file data.
-    base64hash TEXT NOT NULL,
+//     base64hash TEXT NOT NULL,
-    -- The user who uploaded the file. Should be a Matrix user ID.
+//     -- The user who uploaded the file. Should be a Matrix user ID.
-    user_id TEXT NOT NULL
+//     user_id TEXT NOT NULL
-);
+// );
-CREATE UNIQUE INDEX IF NOT EXISTS mediaapi_media_repository_index ON mediaapi_media_repository (media_id, media_origin);
+// CREATE UNIQUE INDEX IF NOT EXISTS mediaapi_media_repository_index ON mediaapi_media_repository (media_id, media_origin);
-`
+// `
 
-const insertMediaSQL = `
+type MediaRepositoryCosmos struct {
-INSERT INTO mediaapi_media_repository (media_id, media_origin, content_type, file_size_bytes, creation_ts, upload_name, base64hash, user_id)
+	MediaID           string `json:"media_id"`
-    VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+	MediaOrigin       string `json:"media_origin"`
-`
+	ContentType       string `json:"content_type"`
-
+	FileSizeBytes     int64  `json:"file_size_bytes"`
-const selectMediaSQL = `
+	CreationTimestamp int64  `json:"creation_ts"`
-SELECT content_type, file_size_bytes, creation_ts, upload_name, base64hash, user_id FROM mediaapi_media_repository WHERE media_id = $1 AND media_origin = $2
+	UploadName        string `json:"upload_name"`
-`
+	Base64hash        string `json:"base64hash"`
-
+	UserID            string `json:"user_id"`
-const selectMediaByHashSQL = `
-SELECT content_type, file_size_bytes, creation_ts, upload_name, media_id, user_id FROM mediaapi_media_repository WHERE base64hash = $1 AND media_origin = $2
-`
-
-type mediaStatements struct {
-	db                    *sql.DB
-	writer                sqlutil.Writer
-	insertMediaStmt       *sql.Stmt
-	selectMediaStmt       *sql.Stmt
-	selectMediaByHashStmt *sql.Stmt
 }
 
-func (s *mediaStatements) prepare(db *sql.DB, writer sqlutil.Writer) (err error) {
+type MediaRepositoryCosmosData struct {
+	Id              string                `json:"id"`
+	Pk              string                `json:"_pk"`
+	Cn              string                `json:"_cn"`
+	ETag            string                `json:"_etag"`
+	Timestamp       int64                 `json:"_ts"`
+	MediaRepository MediaRepositoryCosmos `json:"mx_mediaapi_media_repository"`
+}
+
+// const insertMediaSQL = `
+// INSERT INTO mediaapi_media_repository (media_id, media_origin, content_type, file_size_bytes, creation_ts, upload_name, base64hash, user_id)
+//     VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+// `
+
+// const selectMediaSQL = `
+// SELECT content_type, file_size_bytes, creation_ts, upload_name, base64hash, user_id FROM mediaapi_media_repository WHERE media_id = $1 AND media_origin = $2
+// `
+
+// SELECT content_type, file_size_bytes, creation_ts, upload_name, media_id, user_id FROM mediaapi_media_repository WHERE base64hash = $1 AND media_origin = $2
+const selectMediaByHashSQL = "" +
+	"select * from c where c._cn = @x1 " +
+	"and c.mx_mediaapi_media_repository.base64hash = @x2 " +
+	"and c.mx_mediaapi_media_repository.media_origin = @x3 "
+
+type mediaStatements struct {
+	db     *Database
+	writer cosmosdbutil.Writer
+	// insertMediaStmt       *sql.Stmt
+	// selectMediaStmt       *sql.Stmt
+	selectMediaByHashStmt string
+	tableName             string
+}
+
+func queryMediaRepository(s *mediaStatements, ctx context.Context, qry string, params map[string]interface{}) ([]MediaRepositoryCosmosData, error) {
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
+	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.ContainerName, dbCollectionName)
+	var response []MediaRepositoryCosmosData
+
+	var optionsQry = cosmosdbapi.GetQueryDocumentsOptions(pk)
+	var query = cosmosdbapi.GetQuery(qry, params)
+	_, err := cosmosdbapi.GetClient(s.db.connection).QueryDocuments(
+		ctx,
+		s.db.cosmosConfig.DatabaseName,
+		s.db.cosmosConfig.ContainerName,
+		query,
+		&response,
+		optionsQry)
+
+	if err != nil {
+		return nil, err
+	}
+	return response, nil
+}
+
+func getMediaRepository(s *mediaStatements, ctx context.Context, pk string, docId string) (*MediaRepositoryCosmosData, error) {
+	response := MediaRepositoryCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
+func (s *mediaStatements) prepare(db *Database, writer cosmosdbutil.Writer) (err error) {
 	s.db = db
 	s.writer = writer
 
-	_, err = db.Exec(mediaSchema)
+	s.selectMediaByHashStmt = selectMediaByHashSQL
-	if err != nil {
+	s.tableName = "media_repository"
-		return
+	return
-	}
-
-	return statementList{
-		{&s.insertMediaStmt, insertMediaSQL},
-		{&s.selectMediaStmt, selectMediaSQL},
-		{&s.selectMediaByHashStmt, selectMediaByHashSQL},
-	}.prepare(db)
 }
 
 func (s *mediaStatements) insertMedia(
 	ctx context.Context, mediaMetadata *types.MediaMetadata,
 ) error {
 	mediaMetadata.CreationTimestamp = types.UnixMs(time.Now().UnixNano() / 1000000)
-	return s.writer.Do(s.db, nil, func(txn *sql.Tx) error {
+	// return s.writer.Do(s.db, nil, func(txn *sql.Tx) error {
-		stmt := sqlutil.TxStmt(txn, s.insertMediaStmt)
+
-		_, err := stmt.ExecContext(
+	// INSERT INTO mediaapi_media_repository (media_id, media_origin, content_type, file_size_bytes, creation_ts, upload_name, base64hash, user_id)
-			ctx,
+	// VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
-			mediaMetadata.MediaID,
+
-			mediaMetadata.Origin,
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
-			mediaMetadata.ContentType,
+	// CREATE UNIQUE INDEX IF NOT EXISTS mediaapi_media_repository_index ON mediaapi_media_repository (media_id, media_origin);
-			mediaMetadata.FileSizeBytes,
+	docId := fmt.Sprintf("%s_%s", mediaMetadata.MediaID, mediaMetadata.Origin)
-			mediaMetadata.CreationTimestamp,
+	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.ContainerName, dbCollectionName, docId)
-			mediaMetadata.UploadName,
+	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.ContainerName, dbCollectionName)
-			mediaMetadata.Base64Hash,
+
-			mediaMetadata.UserID,
+	data := MediaRepositoryCosmos{
-		)
+		MediaID:           string(mediaMetadata.MediaID),
-		return err
+		MediaOrigin:       string(mediaMetadata.Origin),
-	})
+		ContentType:       string(mediaMetadata.ContentType),
+		FileSizeBytes:     int64(mediaMetadata.FileSizeBytes),
+		CreationTimestamp: int64(mediaMetadata.CreationTimestamp),
+		UploadName:        string(mediaMetadata.UploadName),
+		Base64hash:        string(mediaMetadata.Base64Hash),
+		UserID:            string(mediaMetadata.UserID),
+	}
+
+	dbData := &MediaRepositoryCosmosData{
+		Id:              cosmosDocId,
+		Cn:              dbCollectionName,
+		Pk:              pk,
+		Timestamp:       time.Now().Unix(),
+		MediaRepository: data,
+	}
+
+	// stmt := sqlutil.TxStmt(txn, s.insertMediaStmt)
+	// _, err := stmt.ExecContext(
+	// 	ctx,
+	// 	mediaMetadata.MediaID,
+	// 	mediaMetadata.Origin,
+	// 	mediaMetadata.ContentType,
+	// 	mediaMetadata.FileSizeBytes,
+	// 	mediaMetadata.CreationTimestamp,
+	// 	mediaMetadata.UploadName,
+	// 	mediaMetadata.Base64Hash,
+	// 	mediaMetadata.UserID,
+	// )
+
+	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
+	// })
 }
 
 func (s *mediaStatements) selectMedia(
@@ -116,16 +211,34 @@ func (s *mediaStatements) selectMedia(
 		MediaID: mediaID,
 		Origin:  mediaOrigin,
 	}
-	err := s.selectMediaStmt.QueryRowContext(
+
-		ctx, mediaMetadata.MediaID, mediaMetadata.Origin,
+	// SELECT content_type, file_size_bytes, creation_ts, upload_name, base64hash, user_id FROM mediaapi_media_repository WHERE media_id = $1 AND media_origin = $2
-	).Scan(
+
-		&mediaMetadata.ContentType,
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
-		&mediaMetadata.FileSizeBytes,
+	// CREATE UNIQUE INDEX IF NOT EXISTS mediaapi_media_repository_index ON mediaapi_media_repository (media_id, media_origin);
-		&mediaMetadata.CreationTimestamp,
+	docId := fmt.Sprintf("%s_%s", mediaMetadata.MediaID, mediaMetadata.Origin)
-		&mediaMetadata.UploadName,
+	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.ContainerName, dbCollectionName, docId)
-		&mediaMetadata.Base64Hash,
+	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.ContainerName, dbCollectionName)
-		&mediaMetadata.UserID,
+
-	)
+	// err := s.selectMediaStmt.QueryRowContext(
+	// ctx, mediaMetadata.MediaID, mediaMetadata.Origin,
+	row, err := getMediaRepository(s, ctx, pk, cosmosDocId)
+
+	if err != nil {
+		return nil, err
+	}
+
+	if row == nil {
+		return nil, nil
+	}
+
+	mediaMetadata.ContentType = types.ContentType(row.MediaRepository.ContentType)
+	mediaMetadata.FileSizeBytes = types.FileSizeBytes(row.MediaRepository.FileSizeBytes)
+	mediaMetadata.CreationTimestamp = types.UnixMs(row.MediaRepository.CreationTimestamp)
+	mediaMetadata.UploadName = types.Filename(row.MediaRepository.UploadName)
+	mediaMetadata.Base64Hash = types.Base64Hash(row.MediaRepository.Base64hash)
+	mediaMetadata.UserID = types.MatrixUserID(row.MediaRepository.UserID)
+
 	return &mediaMetadata, err
 }
 
@@ -136,15 +249,36 @@ func (s *mediaStatements) selectMediaByHash(
 		Base64Hash: mediaHash,
 		Origin:     mediaOrigin,
 	}
-	err := s.selectMediaStmt.QueryRowContext(
+
-		ctx, mediaMetadata.Base64Hash, mediaMetadata.Origin,
+	// SELECT content_type, file_size_bytes, creation_ts, upload_name, media_id, user_id FROM mediaapi_media_repository WHERE base64hash = $1 AND media_origin = $2
-	).Scan(
+
-		&mediaMetadata.ContentType,
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
-		&mediaMetadata.FileSizeBytes,
+	params := map[string]interface{}{
-		&mediaMetadata.CreationTimestamp,
+		"@x1": dbCollectionName,
-		&mediaMetadata.UploadName,
+		"@x2": mediaHash,
-		&mediaMetadata.MediaID,
+		"@x3": mediaOrigin,
-		&mediaMetadata.UserID,
+	}
-	)
+
+	// err := s.selectMediaStmt.QueryRowContext(
+	// 	ctx, mediaMetadata.Base64Hash, mediaMetadata.Origin,
+	// ).Scan(
+	rows, err := queryMediaRepository(s, ctx, s.selectMediaByHashStmt, params)
+
+	if err != nil {
+		return nil, err
+	}
+
+	if len(rows) == 0 {
+		return nil, nil
+	}
+
+	row := rows[0]
+
+	mediaMetadata.ContentType = types.ContentType(row.MediaRepository.ContentType)
+	mediaMetadata.FileSizeBytes = types.FileSizeBytes(row.MediaRepository.FileSizeBytes)
+	mediaMetadata.CreationTimestamp = types.UnixMs(row.MediaRepository.CreationTimestamp)
+	mediaMetadata.UploadName = types.Filename(row.MediaRepository.UploadName)
+	mediaMetadata.MediaID = types.MediaID(row.MediaRepository.MediaID)
+	mediaMetadata.UserID = types.MatrixUserID(row.MediaRepository.UserID)
 	return &mediaMetadata, err
 }

mediaapi/storage/cosmosdb/prepare.go

@@ -1,38 +0,0 @@
-// Copyright 2017-2018 New Vector Ltd
-// Copyright 2019-2020 The Matrix.org Foundation C.I.C.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-// FIXME: This should be made internal!
-
-package cosmosdb
-
-import (
-	"database/sql"
-)
-
-// a statementList is a list of SQL statements to prepare and a pointer to where to store the resulting prepared statement.
-type statementList []struct {
-	statement **sql.Stmt
-	sql       string
-}
-
-// prepare the SQL for each statement in the list and assign the result to the prepared statement.
-func (s statementList) prepare(db *sql.DB) (err error) {
-	for _, statement := range s {
-		if *statement.statement, err = db.Prepare(statement.sql); err != nil {
-			return
-		}
-	}
-	return
-}

mediaapi/storage/cosmosdb/sql.go

@@ -15,18 +15,14 @@
 
 package cosmosdb
 
-import (
+import "github.com/matrix-org/dendrite/internal/cosmosdbutil"
-	"database/sql"
-
-	"github.com/matrix-org/dendrite/internal/sqlutil"
-)
 
 type statements struct {
 	media     mediaStatements
 	thumbnail thumbnailStatements
 }
 
-func (s *statements) prepare(db *sql.DB, writer sqlutil.Writer) (err error) {
+func (s *statements) prepare(db *Database, writer cosmosdbutil.Writer) (err error) {
 	if err = s.media.prepare(db, writer); err != nil {
 		return
 	}

mediaapi/storage/cosmosdb/storage.go

@@ -19,8 +19,11 @@ import (
 	"context"
 	"database/sql"
 
+	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
+	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
+
 	// Import the postgres database driver.
-	"github.com/matrix-org/dendrite/internal/sqlutil"
+
 	"github.com/matrix-org/dendrite/mediaapi/types"
 	"github.com/matrix-org/dendrite/setup/config"
 	"github.com/matrix-org/gomatrixserverlib"
@@ -29,21 +32,26 @@ import (
 
 // Database is used to store metadata about a repository of media files.
 type Database struct {
-	statements statements
+	statements   statements
-	db         *sql.DB
+	db           *sql.DB
-	writer     sqlutil.Writer
+	writer       cosmosdbutil.Writer
+	connection   cosmosdbapi.CosmosConnection
+	databaseName string
+	cosmosConfig cosmosdbapi.CosmosConfig
 }
 
 // Open opens a postgres database.
 func Open(dbProperties *config.DatabaseOptions) (*Database, error) {
+	conn := cosmosdbutil.GetCosmosConnection(&dbProperties.ConnectionString)
+	configCosmos := cosmosdbutil.GetCosmosConfig(&dbProperties.ConnectionString)
 	d := Database{
-		writer: sqlutil.NewExclusiveWriter(),
+		connection:   conn,
+		cosmosConfig: configCosmos,
+		writer:       cosmosdbutil.NewExclusiveWriterFake(),
+		databaseName: "mediaapi",
 	}
 	var err error
-	if d.db, err = sqlutil.Open(dbProperties); err != nil {
+	if err = d.statements.prepare(&d, d.writer); err != nil {
-		return nil, err
-	}
-	if err = d.statements.prepare(d.db, d.writer); err != nil {
 		return nil, err
 	}
 	return &d, nil

mediaapi/storage/cosmosdb/thumbnail_table.go

@@ -17,88 +17,187 @@ package cosmosdb
 
 import (
 	"context"
-	"database/sql"
+	"fmt"
 	"time"
 
-	"github.com/matrix-org/dendrite/internal"
+	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
-	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
+
 	"github.com/matrix-org/dendrite/mediaapi/types"
 	"github.com/matrix-org/gomatrixserverlib"
 )
 
-const thumbnailSchema = `
+// const thumbnailSchema = `
--- The mediaapi_thumbnail table holds metadata for each thumbnail file stored and accessible to the local server,
+// -- The mediaapi_thumbnail table holds metadata for each thumbnail file stored and accessible to the local server,
--- the actual file is stored separately.
+// -- the actual file is stored separately.
-CREATE TABLE IF NOT EXISTS mediaapi_thumbnail (
+// CREATE TABLE IF NOT EXISTS mediaapi_thumbnail (
-    media_id TEXT NOT NULL,
+//     media_id TEXT NOT NULL,
-    media_origin TEXT NOT NULL,
+//     media_origin TEXT NOT NULL,
-    content_type TEXT NOT NULL,
+//     content_type TEXT NOT NULL,
-    file_size_bytes INTEGER NOT NULL,
+//     file_size_bytes INTEGER NOT NULL,
-    creation_ts INTEGER NOT NULL,
+//     creation_ts INTEGER NOT NULL,
-    width INTEGER NOT NULL,
+//     width INTEGER NOT NULL,
-    height INTEGER NOT NULL,
+//     height INTEGER NOT NULL,
-    resize_method TEXT NOT NULL
+//     resize_method TEXT NOT NULL
-);
+// );
-CREATE UNIQUE INDEX IF NOT EXISTS mediaapi_thumbnail_index ON mediaapi_thumbnail (media_id, media_origin, width, height, resize_method);
+// CREATE UNIQUE INDEX IF NOT EXISTS mediaapi_thumbnail_index ON mediaapi_thumbnail (media_id, media_origin, width, height, resize_method);
-`
+// `
 
-const insertThumbnailSQL = `
+type ThumbnailCosmos struct {
-INSERT INTO mediaapi_thumbnail (media_id, media_origin, content_type, file_size_bytes, creation_ts, width, height, resize_method)
+	MediaID           string `json:"media_id"`
-    VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+	MediaOrigin       string `json:"media_origin"`
-`
+	ContentType       string `json:"content_type"`
-
+	FileSizeBytes     int64  `json:"file_size_bytes"`
-// Note: this selects one specific thumbnail
+	CreationTimestamp int64  `json:"creation_ts"`
-const selectThumbnailSQL = `
+	Width             int64  `json:"width"`
-SELECT content_type, file_size_bytes, creation_ts FROM mediaapi_thumbnail WHERE media_id = $1 AND media_origin = $2 AND width = $3 AND height = $4 AND resize_method = $5
+	Height            int64  `json:"height"`
-`
+	ResizeMethod      string `json:"resize_method"`
-
-// Note: this selects all thumbnails for a media_origin and media_id
-const selectThumbnailsSQL = `
-SELECT content_type, file_size_bytes, creation_ts, width, height, resize_method FROM mediaapi_thumbnail WHERE media_id = $1 AND media_origin = $2
-`
-
-type thumbnailStatements struct {
-	db                   *sql.DB
-	writer               sqlutil.Writer
-	insertThumbnailStmt  *sql.Stmt
-	selectThumbnailStmt  *sql.Stmt
-	selectThumbnailsStmt *sql.Stmt
 }
 
-func (s *thumbnailStatements) prepare(db *sql.DB, writer sqlutil.Writer) (err error) {
+type ThumbnailCosmosData struct {
-	_, err = db.Exec(thumbnailSchema)
+	Id        string          `json:"id"`
+	Pk        string          `json:"_pk"`
+	Cn        string          `json:"_cn"`
+	ETag      string          `json:"_etag"`
+	Timestamp int64           `json:"_ts"`
+	Thumbnail ThumbnailCosmos `json:"mx_mediaapi_thumbnail"`
+}
+
+// const insertThumbnailSQL = `
+// INSERT INTO mediaapi_thumbnail (media_id, media_origin, content_type, file_size_bytes, creation_ts, width, height, resize_method)
+//     VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+// `
+
+// Note: this selects one specific thumbnail
+// const selectThumbnailSQL = `
+// SELECT content_type, file_size_bytes, creation_ts FROM mediaapi_thumbnail WHERE media_id = $1 AND media_origin = $2 AND width = $3 AND height = $4 AND resize_method = $5
+// `
+
+// Note: this selects all thumbnails for a media_origin and media_id
+// SELECT content_type, file_size_bytes, creation_ts, width, height, resize_method FROM mediaapi_thumbnail WHERE media_id = $1 AND media_origin = $2
+const selectThumbnailsSQL = "" +
+	"select * from c where c._cn = @x1 " +
+	"and c.mx_mediaapi_thumbnail.media_id = @x2" +
+	"and c.mx_mediaapi_thumbnail.media_origin = @x3"
+
+type thumbnailStatements struct {
+	db     *Database
+	writer cosmosdbutil.Writer
+	// insertThumbnailStmt  *sql.Stmt
+	// selectThumbnailStmt  *sql.Stmt
+	selectThumbnailsStmt string
+	tableName            string
+}
+
+func queryThumbnail(s *thumbnailStatements, ctx context.Context, qry string, params map[string]interface{}) ([]ThumbnailCosmosData, error) {
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
+	var pk = cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.ContainerName, dbCollectionName)
+	var response []ThumbnailCosmosData
+
+	var optionsQry = cosmosdbapi.GetQueryDocumentsOptions(pk)
+	var query = cosmosdbapi.GetQuery(qry, params)
+	_, err := cosmosdbapi.GetClient(s.db.connection).QueryDocuments(
+		ctx,
+		s.db.cosmosConfig.DatabaseName,
+		s.db.cosmosConfig.ContainerName,
+		query,
+		&response,
+		optionsQry)
+
 	if err != nil {
-		return
+		return nil, err
 	}
+	return response, nil
+}
+
+func getThumbnail(s *thumbnailStatements, ctx context.Context, pk string, docId string) (*ThumbnailCosmosData, error) {
+	response := ThumbnailCosmosData{}
+	err := cosmosdbapi.GetDocumentOrNil(
+		s.db.connection,
+		s.db.cosmosConfig,
+		ctx,
+		pk,
+		docId,
+		&response)
+
+	if response.Id == "" {
+		return nil, nil
+	}
+
+	return &response, err
+}
+
+func (s *thumbnailStatements) prepare(db *Database, writer cosmosdbutil.Writer) (err error) {
 	s.db = db
 	s.writer = writer
-
+	s.tableName = "thumbnail"
-	return statementList{
+	s.selectThumbnailsStmt = selectThumbnailsSQL
-		{&s.insertThumbnailStmt, insertThumbnailSQL},
+	return
-		{&s.selectThumbnailStmt, selectThumbnailSQL},
-		{&s.selectThumbnailsStmt, selectThumbnailsSQL},
-	}.prepare(db)
 }
 
 func (s *thumbnailStatements) insertThumbnail(
 	ctx context.Context, thumbnailMetadata *types.ThumbnailMetadata,
 ) error {
 	thumbnailMetadata.MediaMetadata.CreationTimestamp = types.UnixMs(time.Now().UnixNano() / 1000000)
-	return s.writer.Do(s.db, nil, func(txn *sql.Tx) error {
+
-		stmt := sqlutil.TxStmt(txn, s.insertThumbnailStmt)
+	// INSERT INTO mediaapi_thumbnail (media_id, media_origin, content_type, file_size_bytes, creation_ts, width, height, resize_method)
-		_, err := stmt.ExecContext(
+	// VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
-			ctx,
+
-			thumbnailMetadata.MediaMetadata.MediaID,
+	// return s.writer.Do(s.db, nil, func(txn *sql.Tx) error {
-			thumbnailMetadata.MediaMetadata.Origin,
+	// stmt := sqlutil.TxStmt(txn, s.insertThumbnailStmt)
-			thumbnailMetadata.MediaMetadata.ContentType,
+
-			thumbnailMetadata.MediaMetadata.FileSizeBytes,
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
-			thumbnailMetadata.MediaMetadata.CreationTimestamp,
+	// CREATE UNIQUE INDEX IF NOT EXISTS mediaapi_thumbnail_index ON mediaapi_thumbnail (media_id, media_origin, width, height, resize_method);
-			thumbnailMetadata.ThumbnailSize.Width,
+	docId := fmt.Sprintf("%s_%s_%d_%d_s",
-			thumbnailMetadata.ThumbnailSize.Height,
+		thumbnailMetadata.MediaMetadata.MediaID,
-			thumbnailMetadata.ThumbnailSize.ResizeMethod,
+		thumbnailMetadata.MediaMetadata.Origin,
-		)
+		thumbnailMetadata.ThumbnailSize.Width,
-		return err
+		thumbnailMetadata.ThumbnailSize.Height,
-	})
+		thumbnailMetadata.ThumbnailSize.ResizeMethod,
+	)
+	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.ContainerName, dbCollectionName, docId)
+	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.ContainerName, dbCollectionName)
+
+	// _, err := stmt.ExecContext(
+	// 	ctx,
+	// 	thumbnailMetadata.MediaMetadata.MediaID,
+	// 	thumbnailMetadata.MediaMetadata.Origin,
+	// 	thumbnailMetadata.MediaMetadata.ContentType,
+	// 	thumbnailMetadata.MediaMetadata.FileSizeBytes,
+	// 	thumbnailMetadata.MediaMetadata.CreationTimestamp,
+	// 	thumbnailMetadata.ThumbnailSize.Width,
+	// 	thumbnailMetadata.ThumbnailSize.Height,
+	// 	thumbnailMetadata.ThumbnailSize.ResizeMethod,
+	// )
+
+	data := ThumbnailCosmos{
+		MediaID:           string(thumbnailMetadata.MediaMetadata.MediaID),
+		MediaOrigin:       string(thumbnailMetadata.MediaMetadata.Origin),
+		ContentType:       string(thumbnailMetadata.MediaMetadata.ContentType),
+		FileSizeBytes:     int64(thumbnailMetadata.MediaMetadata.FileSizeBytes),
+		CreationTimestamp: int64(thumbnailMetadata.MediaMetadata.CreationTimestamp),
+		Width:             int64(thumbnailMetadata.ThumbnailSize.Width),
+		Height:            int64(thumbnailMetadata.ThumbnailSize.Height),
+		ResizeMethod:      string(thumbnailMetadata.ThumbnailSize.ResizeMethod),
+	}
+
+	dbData := &ThumbnailCosmosData{
+		Id:        cosmosDocId,
+		Cn:        dbCollectionName,
+		Pk:        pk,
+		Timestamp: time.Now().Unix(),
+		Thumbnail: data,
+	}
+
+	var options = cosmosdbapi.GetUpsertDocumentOptions(dbData.Pk)
+	_, _, err := cosmosdbapi.GetClient(s.db.connection).CreateDocument(
+		ctx,
+		s.db.cosmosConfig.DatabaseName,
+		s.db.cosmosConfig.ContainerName,
+		&dbData,
+		options)
+
+	return err
 }
 
 func (s *thumbnailStatements) selectThumbnail(
@@ -119,53 +218,82 @@ func (s *thumbnailStatements) selectThumbnail(
 			ResizeMethod: resizeMethod,
 		},
 	}
-	err := s.selectThumbnailStmt.QueryRowContext(
+
-		ctx,
+	// SELECT content_type, file_size_bytes, creation_ts FROM mediaapi_thumbnail WHERE media_id = $1 AND media_origin = $2 AND width = $3 AND height = $4 AND resize_method = $5
-		thumbnailMetadata.MediaMetadata.MediaID,
+
-		thumbnailMetadata.MediaMetadata.Origin,
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
-		thumbnailMetadata.ThumbnailSize.Width,
+	// CREATE UNIQUE INDEX IF NOT EXISTS mediaapi_thumbnail_index ON mediaapi_thumbnail (media_id, media_origin, width, height, resize_method);
-		thumbnailMetadata.ThumbnailSize.Height,
+	docId := fmt.Sprintf("%s_%s_%d_%d_s",
-		thumbnailMetadata.ThumbnailSize.ResizeMethod,
+		mediaID,
-	).Scan(
+		mediaOrigin,
-		&thumbnailMetadata.MediaMetadata.ContentType,
+		width,
-		&thumbnailMetadata.MediaMetadata.FileSizeBytes,
+		height,
-		&thumbnailMetadata.MediaMetadata.CreationTimestamp,
+		resizeMethod,
 	)
-	return &thumbnailMetadata, err
+	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.ContainerName, dbCollectionName, docId)
+	pk := cosmosdbapi.GetPartitionKey(s.db.cosmosConfig.ContainerName, dbCollectionName)
+
+	// row := sqlutil.TxStmt(txn, s.selectOutboundPeeksStmt).QueryRowContext(ctx, roomID)
+	row, err := getThumbnail(s, ctx, pk, cosmosDocId)
+
+	if err != nil {
+		return nil, err
+	}
+
+	if row == nil {
+		return nil, nil
+	}
+
+	thumbnailMetadata.MediaMetadata.MediaID = types.MediaID(row.Thumbnail.MediaID)
+	thumbnailMetadata.MediaMetadata.Origin = gomatrixserverlib.ServerName(row.Thumbnail.MediaOrigin)
+	thumbnailMetadata.ThumbnailSize.Width = int(row.Thumbnail.Width)
+	thumbnailMetadata.ThumbnailSize.Height = int(row.Thumbnail.Height)
+	thumbnailMetadata.ThumbnailSize.ResizeMethod = row.Thumbnail.ResizeMethod
+	thumbnailMetadata.MediaMetadata.ContentType = types.ContentType(row.Thumbnail.ContentType)
+	thumbnailMetadata.MediaMetadata.FileSizeBytes = types.FileSizeBytes(row.Thumbnail.FileSizeBytes)
+	thumbnailMetadata.MediaMetadata.CreationTimestamp = types.UnixMs(row.Thumbnail.CreationTimestamp)
+	return &thumbnailMetadata, nil
 }
 
 func (s *thumbnailStatements) selectThumbnails(
 	ctx context.Context, mediaID types.MediaID, mediaOrigin gomatrixserverlib.ServerName,
 ) ([]*types.ThumbnailMetadata, error) {
-	rows, err := s.selectThumbnailsStmt.QueryContext(
+
-		ctx, mediaID, mediaOrigin,
+	// SELECT content_type, file_size_bytes, creation_ts, width, height, resize_method FROM mediaapi_thumbnail WHERE media_id = $1 AND media_origin = $2
-	)
+
+	var dbCollectionName = cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
+	params := map[string]interface{}{
+		"@x1": dbCollectionName,
+		"@x2": mediaID,
+		"@x3": mediaOrigin,
+	}
+
+	// rows, err := s.selectThumbnailsStmt.QueryContext(
+	// 	ctx, mediaID, mediaOrigin,
+	// )
+	rows, err := queryThumbnail(s, ctx, s.selectThumbnailsStmt, params)
+
 	if err != nil {
 		return nil, err
 	}
-	defer internal.CloseAndLogIfError(ctx, rows, "selectThumbnails: rows.close() failed")
 
 	var thumbnails []*types.ThumbnailMetadata
-	for rows.Next() {
+	for _, item := range rows {
 		thumbnailMetadata := types.ThumbnailMetadata{
 			MediaMetadata: &types.MediaMetadata{
 				MediaID: mediaID,
 				Origin:  mediaOrigin,
 			},
 		}
-		err = rows.Scan(
+		thumbnailMetadata.MediaMetadata.ContentType = types.ContentType(item.Thumbnail.ContentType)
-			&thumbnailMetadata.MediaMetadata.ContentType,
+		thumbnailMetadata.MediaMetadata.FileSizeBytes = types.FileSizeBytes(item.Thumbnail.FileSizeBytes)
-			&thumbnailMetadata.MediaMetadata.FileSizeBytes,
+		thumbnailMetadata.MediaMetadata.CreationTimestamp = types.UnixMs(item.Thumbnail.CreationTimestamp)
-			&thumbnailMetadata.MediaMetadata.CreationTimestamp,
+		thumbnailMetadata.ThumbnailSize.Width = int(item.Thumbnail.Width)
-			&thumbnailMetadata.ThumbnailSize.Width,
+		thumbnailMetadata.ThumbnailSize.Height = int(item.Thumbnail.Height)
-			&thumbnailMetadata.ThumbnailSize.Height,
+		thumbnailMetadata.ThumbnailSize.ResizeMethod = item.Thumbnail.ResizeMethod
-			&thumbnailMetadata.ThumbnailSize.ResizeMethod,
+
-		)
-		if err != nil {
-			return nil, err
-		}
 		thumbnails = append(thumbnails, &thumbnailMetadata)
 	}
 
-	return thumbnails, rows.Err()
+	return thumbnails, err
 }

roomserver/storage/cosmosdb/room_aliases_table.go

@@ -21,7 +21,6 @@ import (
 	"time"
 
 	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
-	"github.com/matrix-org/dendrite/internal/cosmosdbutil"
 	"github.com/matrix-org/dendrite/roomserver/storage/tables"
 )
 
@@ -113,7 +112,7 @@ func getRoomAlias(s *roomAliasesStatements, ctx context.Context, pk string, docI
 		&response)
 
 	if response.Id == "" {
-		return nil, cosmosdbutil.ErrNoRows
+		return nil, nil
 	}
 
 	return &response, err