dendrite/roomserver/storage/cosmosdb/invite_table.go

// 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.

package cosmosdb

import (
	"context"
	"database/sql"
	"fmt"

	"github.com/matrix-org/dendrite/internal/cosmosdbapi"
	"github.com/matrix-org/dendrite/internal/cosmosdbutil"

	"github.com/matrix-org/dendrite/roomserver/storage/tables"
	"github.com/matrix-org/dendrite/roomserver/types"
)

// const inviteSchema = `
// 	CREATE TABLE IF NOT EXISTS roomserver_invites (
// 		invite_event_id TEXT PRIMARY KEY,
// 		room_nid INTEGER NOT NULL,
// 		target_nid INTEGER NOT NULL,
// 		sender_nid INTEGER NOT NULL DEFAULT 0,
// 		retired BOOLEAN NOT NULL DEFAULT FALSE,
// 		invite_event_json TEXT NOT NULL
// 	);

// 	CREATE INDEX IF NOT EXISTS roomserver_invites_active_idx ON roomserver_invites (target_nid, room_nid)
// 		WHERE NOT retired;
// `

type inviteCosmos struct {
	InviteEventID   string `json:"invite_event_id"`
	RoomNID         int64  `json:"room_nid"`
	TargetNID       int64  `json:"target_nid"`
	SenderNID       int64  `json:"sender_nid"`
	Retired         bool   `json:"retired"`
	InviteEventJSON []byte `json:"invite_event_json"`
}

type inviteCosmosData struct {
	cosmosdbapi.CosmosDocument
	Invite inviteCosmos `json:"mx_roomserver_invite"`
}

// const insertInviteEventSQL = "" +
// 	"INSERT INTO roomserver_invites (invite_event_id, room_nid, target_nid," +
// 	" sender_nid, invite_event_json) VALUES ($1, $2, $3, $4, $5)" +
// 	" ON CONFLICT DO NOTHING"

// "SELECT invite_event_id, sender_nid FROM roomserver_invites" +
// " WHERE target_nid = $1 AND room_nid = $2" +
// " AND NOT retired"
const selectInviteActiveForUserInRoomSQL = "" +
	"select * from c where c._cn = @x1 " +
	" and c.mx_roomserver_invite.target_nid = @x2" +
	" and c.mx_roomserver_invite.room_nid = @x3" +
	" and c.mx_roomserver_invite.retired = false"

// Retire every active invite for a user in a room.
// Ideally we'd know which invite events were retired by a given update so we
// wouldn't need to remove every active invite.
// However the matrix protocol doesn't give us a way to reliably identify the
// invites that were retired, so we are forced to retire all of them.
// const updateInviteRetiredSQL = `
// 	UPDATE roomserver_invites SET retired = TRUE WHERE room_nid = $1 AND target_nid = $2 AND NOT retired
// `

// SELECT invite_event_id FROM roomserver_invites WHERE room_nid = $1 AND target_nid = $2 AND NOT retired
const selectInvitesAboutToRetireSQL = "" +
	"select * from c where c._cn = @x1 " +
	" and c.mx_roomserver_invite.room_nid = @x2" +
	" and c.mx_roomserver_invite.target_nid = @x3" +
	" and c.mx_roomserver_invite.retired = false"

type inviteStatements struct {
	db *Database
	// insertInviteEventStmt               *sql.Stmt
	selectInviteActiveForUserInRoomStmt string
	// updateInviteRetiredStmt             *sql.Stmt
	selectInvitesAboutToRetireStmt string
	tableName                      string
}

func (s *inviteStatements) getCollectionName() string {
	return cosmosdbapi.GetCollectionName(s.db.databaseName, s.tableName)
}

func (s *inviteStatements) getPartitionKey(roomNId int64) string {
	uniqueId := fmt.Sprintf("%d", roomNId)
	return cosmosdbapi.GetPartitionKeyByUniqueId(s.db.cosmosConfig.TenantName, s.getCollectionName(), uniqueId)
}

func getInvite(s *inviteStatements, ctx context.Context, pk string, docId string) (*inviteCosmosData, error) {
	response := inviteCosmosData{}
	err := cosmosdbapi.GetDocumentOrNil(
		s.db.connection,
		s.db.cosmosConfig,
		ctx,
		pk,
		docId,
		&response)

	if response.Id == "" {
		return nil, cosmosdbutil.ErrNoRows
	}

	return &response, err
}

func NewCosmosDBInvitesTable(db *Database) (tables.Invites, error) {
	s := &inviteStatements{
		db: db,
	}
	// return s, shared.StatementList{
	// 	{&s.insertInviteEventStmt, insertInviteEventSQL},
	s.selectInviteActiveForUserInRoomStmt = selectInviteActiveForUserInRoomSQL
	// 	{&s.updateInviteRetiredStmt, updateInviteRetiredSQL},
	s.selectInvitesAboutToRetireStmt = selectInvitesAboutToRetireSQL
	// }.Prepare(db)
	s.tableName = "invites"
	return s, nil
}

func (s *inviteStatements) InsertInviteEvent(
	ctx context.Context,
	txn *sql.Tx, inviteEventID string, roomNID types.RoomNID,
	targetUserNID, senderUserNID types.EventStateKeyNID,
	inviteEventJSON []byte,
) (bool, error) {

	// "INSERT INTO roomserver_invites (invite_event_id, room_nid, target_nid," +
	// " sender_nid, invite_event_json) VALUES ($1, $2, $3, $4, $5)" +
	// " ON CONFLICT DO NOTHING"

	// 		invite_event_id TEXT PRIMARY KEY,
	docId := inviteEventID
	cosmosDocId := cosmosdbapi.GetDocumentId(s.db.cosmosConfig.TenantName, s.getCollectionName(), docId)

	data := inviteCosmos{
		InviteEventID:   inviteEventID,
		InviteEventJSON: inviteEventJSON,
		Retired:         false,
		RoomNID:         int64(roomNID),
		SenderNID:       int64(senderUserNID),
		TargetNID:       int64(targetUserNID),
	}

	var dbData = inviteCosmosData{
		CosmosDocument: cosmosdbapi.GenerateDocument(s.getCollectionName(), s.db.cosmosConfig.TenantName, s.getPartitionKey(int64(roomNID)), cosmosDocId),
		Invite:         data,
	}

	var options = cosmosdbapi.GetCreateDocumentOptions(dbData.Pk)
	_, _, err := cosmosdbapi.GetClient(s.db.connection).CreateDocument(
		ctx,
		s.db.cosmosConfig.DatabaseName,
		s.db.cosmosConfig.ContainerName,
		&dbData,
		options)

	if err != nil {
		return false, err
	}
	// TODO: Is this important?
	// count, err = result.RowsAffected()
	// return count != 0, err
	return true, nil
}

func (s *inviteStatements) UpdateInviteRetired(
	ctx context.Context,
	txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID,
) (eventIDs []string, err error) {

	// "SELECT invite_event_id, sender_nid FROM roomserver_invites" +
	// " WHERE target_nid = $1 AND room_nid = $2" +
	// " AND NOT retired"

	// gather all the event IDs we will retire
	params := map[string]interface{}{
		"@x1": s.getCollectionName(),
		"@x2": targetUserNID,
		"@x3": roomNID,
	}

	var rows []inviteCosmosData
	err = cosmosdbapi.PerformQuery(ctx,
		s.db.connection,
		s.db.cosmosConfig.DatabaseName,
		s.db.cosmosConfig.ContainerName,
		s.getPartitionKey(int64(roomNID)), s.selectInvitesAboutToRetireStmt, params, &rows)

	if err != nil {
		return
	}

	for _, item := range rows {
		eventIDs = append(eventIDs, item.Invite.InviteEventID)
		// 	UPDATE roomserver_invites SET retired = TRUE WHERE room_nid = $1 AND target_nid = $2 AND NOT retired

		// now retire the invites
		item.SetUpdateTime()
		item.Invite.Retired = true
		_, err = cosmosdbapi.UpdateDocument(ctx, s.db.connection, s.db.cosmosConfig.DatabaseName, s.db.cosmosConfig.ContainerName, item.Pk, item.ETag, item.Id, item)
	}

	return
}

// selectInviteActiveForUserInRoom returns a list of sender state key NIDs
func (s *inviteStatements) SelectInviteActiveForUserInRoom(
	ctx context.Context,
	targetUserNID types.EventStateKeyNID, roomNID types.RoomNID,
) ([]types.EventStateKeyNID, []string, error) {

	// SELECT invite_event_id FROM roomserver_invites WHERE room_nid = $1 AND target_nid = $2 AND NOT retired

	params := map[string]interface{}{
		"@x1": s.getCollectionName(),
		"@x2": roomNID,
		"@x3": targetUserNID,
	}

	var rows []inviteCosmosData
	err := cosmosdbapi.PerformQuery(ctx,
		s.db.connection,
		s.db.cosmosConfig.DatabaseName,
		s.db.cosmosConfig.ContainerName,
		s.getPartitionKey(int64(roomNID)), s.selectInviteActiveForUserInRoomStmt, params, &rows)

	if err != nil {
		return nil, nil, err
	}

	var result []types.EventStateKeyNID
	var eventIDs []string
	for _, item := range rows {
		var eventID = item.Invite.InviteEventID
		var senderUserNID = item.Invite.SenderNID
		result = append(result, types.EventStateKeyNID(senderUserNID))
		eventIDs = append(eventIDs, eventID)
	}
	return result, eventIDs, nil
}