Skip to content

Commit

Permalink
feat(GODT-2522): New Database Layout
Browse files Browse the repository at this point in the history
New database includes the following changes

* Indices and composite primary keys for message and mailbox flags.
* Every mailbox gets a dedicated table.
* UIDs are auto generated via incrementing integer
* Each mailbox messages entry also stores the remote ID of the message
  for faster queries
* New table to track which mailboxes a message is inserted in
  • Loading branch information
LBeernaertProton committed Jun 27, 2023
1 parent cefbb2d commit 8258a0f
Show file tree
Hide file tree
Showing 28 changed files with 1,777 additions and 486 deletions.
3 changes: 2 additions & 1 deletion db/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,14 @@ package db

import (
"context"
"github.com/ProtonMail/gluon/imap"
"path/filepath"
)

const ChunkLimit = 1000

type Client interface {
Init(ctx context.Context) error
Init(ctx context.Context, generator imap.UIDValidityGenerator) error
Read(ctx context.Context, op func(context.Context, ReadOnly) error) error
Write(ctx context.Context, op func(context.Context, Transaction) error) error
Close() error
Expand Down
1 change: 1 addition & 0 deletions db/errors.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import "errors"

var ErrNotFound = errors.New("value not found")
var ErrTransactionFailed = errors.New("transaction failed")
var ErrMigrationFailed = errors.New("database migration failed")

func IsErrNotFound(err error) bool {
if err == nil {
Expand Down
12 changes: 1 addition & 11 deletions db/ops_mailbox.go
Original file line number Diff line number Diff line change
Expand Up @@ -54,13 +54,7 @@ type MailboxReadOps interface {

MailboxFilterContains(ctx context.Context, mboxID imap.InternalMailboxID, messageIDs []MessageIDPair) ([]imap.InternalMessageID, error)

MailboxFilterContainsInternalID(ctx context.Context, mboxID imap.InternalMailboxID, messageIDs []imap.InternalMessageID) ([]imap.InternalMessageID, error)

GetMailboxCount(ctx context.Context) (int, error)

// GetMessageUIDsWithFlagsAfterAddOrUIDBump exploits a property of adding a message to or bumping the UIDs of existing message in mailbox. It can only be
// used if you can guarantee that the messageID list contains only IDs that have recently added or bumped in the mailbox.
GetMailboxMessageUIDsWithFlagsAfterAddOrUIDBump(ctx context.Context, mboxID imap.InternalMailboxID, messageIDs []imap.InternalMessageID) ([]UIDWithFlags, error)
}

type MailboxWriteOps interface {
Expand All @@ -87,11 +81,7 @@ type MailboxWriteOps interface {

DeleteMailboxWithRemoteID(ctx context.Context, mboxID imap.MailboxID) error

BumpMailboxUIDNext(ctx context.Context, mboxID imap.InternalMailboxID, count int) error

AddMessagesToMailbox(ctx context.Context, mboxID imap.InternalMailboxID, messageIDs []imap.InternalMessageID) ([]UIDWithFlags, error)

BumpMailboxUIDsForMessage(ctx context.Context, mboxID imap.InternalMailboxID, messageIDs []imap.InternalMessageID) ([]UIDWithFlags, error)
AddMessagesToMailbox(ctx context.Context, mboxID imap.InternalMailboxID, messageIDs []MessageIDPair) ([]UIDWithFlags, error)

RemoveMessagesFromMailbox(ctx context.Context, mboxID imap.InternalMailboxID, messageIDs []imap.InternalMessageID) error

Expand Down
1 change: 0 additions & 1 deletion db/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,6 @@ type Mailbox struct {
ID imap.InternalMailboxID
RemoteID imap.MailboxID
Name string
UIDNext imap.UID
UIDValidity imap.UID
Subscribed bool
Flags []*MailboxFlag
Expand Down
43 changes: 28 additions & 15 deletions internal/backend/connector_updates.go
Original file line number Diff line number Diff line change
Expand Up @@ -208,7 +208,7 @@ func (user *user) applyMessagesCreated(ctx context.Context, update *imap.Message
// collect all unique messages to create
messagesToCreate := make([]*DBRequestWithLiteral, 0, len(update.Messages))
messagesToCreateFilter := make(map[imap.MessageID]imap.InternalMessageID, len(update.Messages)/2)
messageForMBox := make(map[imap.InternalMailboxID][]imap.InternalMessageID)
messageForMBox := make(map[imap.InternalMailboxID][]db.MessageIDPair)
mboxInternalIDMap := make(map[imap.MailboxID]imap.InternalMailboxID)

err := user.db.Write(ctx, func(ctx context.Context, tx db.Transaction) error {
Expand Down Expand Up @@ -271,12 +271,12 @@ func (user *user) applyMessagesCreated(ctx context.Context, update *imap.Message

messageList, ok := messageForMBox[v]
if !ok {
messageList = []imap.InternalMessageID{}
messageList = []db.MessageIDPair{}
messageForMBox[v] = messageList
}

if !slices.Contains(messageList, internalID) {
messageList = append(messageList, internalID)
if !slices.ContainsFunc(messageList, func(id db.MessageIDPair) bool { return id.InternalID == internalID }) {
messageList = append(messageList, db.MessageIDPair{InternalID: internalID, RemoteID: message.Message.ID})
messageForMBox[v] = messageList
}
}
Expand Down Expand Up @@ -313,13 +313,13 @@ func (user *user) applyMessagesCreated(ctx context.Context, update *imap.Message

// Assign all the messages to the mailbox
for mboxID, msgList := range messageForMBox {
inMailbox, err := tx.MailboxFilterContainsInternalID(ctx, mboxID, msgList)
inMailbox, err := tx.MailboxFilterContains(ctx, mboxID, msgList)
if err != nil {
return err
}

toAdd := xslices.Filter(msgList, func(id imap.InternalMessageID) bool {
return !slices.Contains(inMailbox, id)
toAdd := xslices.Filter(msgList, func(id db.MessageIDPair) bool {
return !slices.Contains(inMailbox, id.InternalID)
})

if len(toAdd) != 0 {
Expand Down Expand Up @@ -374,7 +374,10 @@ func (user *user) applyMessageMailboxesUpdated(ctx context.Context, update *imap
return err
}

if err := user.setMessageMailboxes(ctx, tx, internalMsgID, internalMBoxIDs); err != nil {
if err := user.setMessageMailboxes(ctx, tx, db.MessageIDPair{
InternalID: internalMsgID,
RemoteID: update.MessageID,
}, internalMBoxIDs); err != nil {
return err
}

Expand Down Expand Up @@ -431,20 +434,20 @@ func (user *user) applyMessageIDChanged(ctx context.Context, update *imap.Messag
return nil
}

func (user *user) setMessageMailboxes(ctx context.Context, tx db.Transaction, messageID imap.InternalMessageID, mboxIDs []imap.InternalMailboxID) error {
curMailboxIDs, err := tx.GetMessageMailboxIDs(ctx, messageID)
func (user *user) setMessageMailboxes(ctx context.Context, tx db.Transaction, messageID db.MessageIDPair, mboxIDs []imap.InternalMailboxID) error {
curMailboxIDs, err := tx.GetMessageMailboxIDs(ctx, messageID.InternalID)
if err != nil {
return err
}

for _, mboxID := range xslices.Filter(mboxIDs, func(mboxID imap.InternalMailboxID) bool { return !slices.Contains(curMailboxIDs, mboxID) }) {
if _, err := user.applyMessagesAddedToMailbox(ctx, tx, mboxID, []imap.InternalMessageID{messageID}); err != nil {
if _, err := user.applyMessagesAddedToMailbox(ctx, tx, mboxID, []db.MessageIDPair{messageID}); err != nil {
return err
}
}

for _, mboxID := range xslices.Filter(curMailboxIDs, func(mboxID imap.InternalMailboxID) bool { return !slices.Contains(mboxIDs, mboxID) }) {
if err := user.applyMessagesRemovedFromMailbox(ctx, tx, mboxID, []imap.InternalMessageID{messageID}); err != nil {
if err := user.applyMessagesRemovedFromMailbox(ctx, tx, mboxID, []imap.InternalMessageID{messageID.InternalID}); err != nil {
return err
}
}
Expand All @@ -453,7 +456,7 @@ func (user *user) setMessageMailboxes(ctx context.Context, tx db.Transaction, me
}

// applyMessagesAddedToMailbox adds the messages to the given mailbox.
func (user *user) applyMessagesAddedToMailbox(ctx context.Context, tx db.Transaction, mboxID imap.InternalMailboxID, messageIDs []imap.InternalMessageID) ([]db.UIDWithFlags, error) {
func (user *user) applyMessagesAddedToMailbox(ctx context.Context, tx db.Transaction, mboxID imap.InternalMailboxID, messageIDs []db.MessageIDPair) ([]db.UIDWithFlags, error) {
messageUIDs, update, err := state.AddMessagesToMailbox(ctx, tx, mboxID, messageIDs, nil, user.imapLimits)
if err != nil {
return nil, err
Expand Down Expand Up @@ -633,7 +636,10 @@ func (user *user) applyMessageUpdated(ctx context.Context, update *imap.MessageU
targetMailboxes = append(targetMailboxes, internalMBoxID)
}

return user.setMessageMailboxes(ctx, tx, internalMessageID, targetMailboxes)
return user.setMessageMailboxes(ctx, tx, db.MessageIDPair{
InternalID: internalMessageID,
RemoteID: update.Message.ID,
}, targetMailboxes)
} else {
log.Debug("Message has new literal, applying update")

Expand Down Expand Up @@ -696,7 +702,14 @@ func (user *user) applyMessageUpdated(ctx context.Context, update *imap.MessageU
return err
}

_, update, err := state.AddMessagesToMailbox(ctx, tx, internalMBoxID, []imap.InternalMessageID{newInternalID}, nil, user.imapLimits)
_, update, err := state.AddMessagesToMailbox(
ctx,
tx,
internalMBoxID,
[]db.MessageIDPair{{InternalID: newInternalID, RemoteID: update.Message.ID}},
nil,
user.imapLimits,
)
if err != nil {
return err
}
Expand Down
2 changes: 1 addition & 1 deletion internal/backend/user.go
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ func newUser(
uidValidityGenerator imap.UIDValidityGenerator,
panicHandler async.PanicHandler,
) (*user, error) {
if err := database.Init(ctx); err != nil {
if err := database.Init(ctx, uidValidityGenerator); err != nil {
return nil, err
}

Expand Down
62 changes: 39 additions & 23 deletions internal/db_impl/sqlite3/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"database/sql"
"errors"
"fmt"
"github.com/ProtonMail/gluon/imap"
"io/fs"
"os"
"path/filepath"
Expand All @@ -26,12 +27,29 @@ type Client struct {
trace bool
}

func (c *Client) Init(ctx context.Context) error {
if _, err := c.db.ExecContext(ctx, "PRAGMA foreign_keys = ON"); err != nil {
return fmt.Errorf("failed to enable db pragma: %w", err)
func NewClient(dir string, userID string, debug, trace bool) (*Client, bool, error) {
if err := os.MkdirAll(dir, 0o700); err != nil {
return nil, false, err
}

if _, err := c.db.ExecContext(ctx, "PRAGMA journal_mode = WAL"); err != nil {
path := getDatabasePath(dir, userID)

// Check if the database already exists.
exists, err := pathExists(path)
if err != nil {
return nil, false, err
}

client, err := sql.Open("sqlite3", getDatabaseConn(dir, userID, path))
if err != nil {
return nil, false, err
}

return &Client{db: client, debug: debug, trace: trace}, !exists, nil
}

func (c *Client) Init(ctx context.Context, generator imap.UIDValidityGenerator) error {
if _, err := c.db.ExecContext(ctx, "PRAGMA foreign_keys = ON"); err != nil {
return fmt.Errorf("failed to enable db pragma: %w", err)
}

Expand All @@ -41,7 +59,22 @@ func (c *Client) Init(ctx context.Context) error {

return c.wrapTx(ctx, func(ctx context.Context, tx *sql.Tx, entry *logrus.Entry) error {
entry.Debugf("Running database migrations")
return RunMigrations(ctx, utils.TXWrapper{TX: tx})
var qw utils.QueryWrapper = &utils.TXWrapper{
TX: tx,
}

if c.debug {
qw = &utils.DebugQueryWrapper{
QW: qw,
Entry: entry,
}
}

if err := RunMigrations(ctx, qw, generator); err != nil {
return fmt.Errorf("%w: %v", db.ErrMigrationFailed, err)
}

return nil
})
}

Expand Down Expand Up @@ -232,24 +265,7 @@ func NewBuilder(options ...Option) db.ClientInterface {
}

func (b Builder) New(dir string, userID string) (db.Client, bool, error) {
if err := os.MkdirAll(dir, 0o700); err != nil {
return nil, false, err
}

path := getDatabasePath(dir, userID)

// Check if the database already exists.
exists, err := pathExists(path)
if err != nil {
return nil, false, err
}

client, err := sql.Open("sqlite3", getDatabaseConn(dir, userID, path))
if err != nil {
return nil, false, err
}

return &Client{db: client, debug: b.debug, trace: b.trace}, !exists, nil
return NewClient(dir, userID, b.debug, b.trace)
}

func (Builder) Delete(dir string, userID string) error {
Expand Down
Loading

0 comments on commit 8258a0f

Please sign in to comment.