Add bulk import optimization: track_lookup cache, batch inserts, BulkSubmitter

This commit is contained in:
safierinx-a 2026-03-27 03:29:33 +05:30
parent 0ec7b458cc
commit ae373a7090
21 changed files with 1296 additions and 125 deletions

View file

@ -101,6 +101,16 @@ type DB interface {
MergeAlbums(ctx context.Context, fromId, toId int32, replaceImage bool) error
MergeArtists(ctx context.Context, fromId, toId int32, replaceImage bool) error
// Track Lookup Cache
GetTrackLookup(ctx context.Context, key string) (*TrackLookupResult, error)
SaveTrackLookup(ctx context.Context, opts SaveTrackLookupOpts) error
InvalidateTrackLookup(ctx context.Context, opts InvalidateTrackLookupOpts) error
// Batch
SaveListensBatch(ctx context.Context, opts []SaveListenOpts) (int64, error)
// Etc
ImageHasAssociation(ctx context.Context, image uuid.UUID) (bool, error)

View file

@ -160,3 +160,22 @@ type GetInterestOpts struct {
ArtistID int32
TrackID int32
}
type TrackLookupResult struct {
ArtistID int32
AlbumID int32
TrackID int32
}
type SaveTrackLookupOpts struct {
Key string
ArtistID int32
AlbumID int32
TrackID int32
}
type InvalidateTrackLookupOpts struct {
ArtistID int32
AlbumID int32
TrackID int32
}

View file

@ -338,6 +338,7 @@ func (d *Psql) SaveAlbumAliases(ctx context.Context, id int32, aliases []string,
}
func (d *Psql) DeleteAlbum(ctx context.Context, id int32) error {
d.q.DeleteTrackLookupByAlbum(ctx, id)
return d.q.DeleteRelease(ctx, id)
}
func (d *Psql) DeleteAlbumAlias(ctx context.Context, id int32, alias string) error {

View file

@ -119,6 +119,7 @@ func (d *Psql) SaveArtistAliases(ctx context.Context, id int32, aliases []string
}
func (d *Psql) DeleteArtist(ctx context.Context, id int32) error {
d.q.DeleteTrackLookupByArtist(ctx, id)
return d.q.DeleteArtist(ctx, id)
}

View file

@ -11,6 +11,7 @@ import (
"github.com/gabehf/koito/internal/logger"
"github.com/gabehf/koito/internal/models"
"github.com/gabehf/koito/internal/repository"
"github.com/jackc/pgx/v5"
)
func (d *Psql) GetListensPaginated(ctx context.Context, opts db.GetItemsOpts) (*db.PaginatedResponse[*models.Listen], error) {
@ -197,6 +198,67 @@ func (d *Psql) SaveListen(ctx context.Context, opts db.SaveListenOpts) error {
})
}
func (d *Psql) SaveListensBatch(ctx context.Context, opts []db.SaveListenOpts) (int64, error) {
if len(opts) == 0 {
return 0, nil
}
tx, err := d.conn.BeginTx(ctx, pgx.TxOptions{})
if err != nil {
return 0, fmt.Errorf("SaveListensBatch: BeginTx: %w", err)
}
defer tx.Rollback(ctx)
_, err = tx.Exec(ctx, `
CREATE TEMP TABLE tmp_listens (
track_id INT,
listened_at TIMESTAMPTZ,
user_id INT,
client TEXT
) ON COMMIT DROP
`)
if err != nil {
return 0, fmt.Errorf("SaveListensBatch: create temp table: %w", err)
}
rows := make([][]interface{}, len(opts))
for i, o := range opts {
var client interface{}
if o.Client != "" {
client = o.Client
}
t := o.Time
if t.IsZero() {
t = time.Now()
}
rows[i] = []interface{}{o.TrackID, t, o.UserID, client}
}
_, err = tx.CopyFrom(ctx,
pgx.Identifier{"tmp_listens"},
[]string{"track_id", "listened_at", "user_id", "client"},
pgx.CopyFromRows(rows),
)
if err != nil {
return 0, fmt.Errorf("SaveListensBatch: CopyFrom: %w", err)
}
tag, err := tx.Exec(ctx, `
INSERT INTO listens (track_id, listened_at, user_id, client)
SELECT track_id, listened_at, user_id, client FROM tmp_listens
ON CONFLICT DO NOTHING
`)
if err != nil {
return 0, fmt.Errorf("SaveListensBatch: insert: %w", err)
}
if err := tx.Commit(ctx); err != nil {
return 0, fmt.Errorf("SaveListensBatch: Commit: %w", err)
}
return tag.RowsAffected(), nil
}
func (d *Psql) DeleteListen(ctx context.Context, trackId int32, listenedAt time.Time) error {
l := logger.FromContext(ctx)
if trackId == 0 {

View file

@ -12,6 +12,8 @@ import (
func (d *Psql) MergeTracks(ctx context.Context, fromId, toId int32) error {
l := logger.FromContext(ctx)
l.Info().Msgf("Merging track %d into track %d", fromId, toId)
d.q.DeleteTrackLookupByTrack(ctx, fromId)
d.q.DeleteTrackLookupByTrack(ctx, toId)
tx, err := d.conn.BeginTx(ctx, pgx.TxOptions{})
if err != nil {
l.Err(err).Msg("Failed to begin transaction")
@ -61,6 +63,8 @@ func (d *Psql) MergeTracks(ctx context.Context, fromId, toId int32) error {
func (d *Psql) MergeAlbums(ctx context.Context, fromId, toId int32, replaceImage bool) error {
l := logger.FromContext(ctx)
l.Info().Msgf("Merging album %d into album %d", fromId, toId)
d.q.DeleteTrackLookupByAlbum(ctx, fromId)
d.q.DeleteTrackLookupByAlbum(ctx, toId)
tx, err := d.conn.BeginTx(ctx, pgx.TxOptions{})
if err != nil {
l.Err(err).Msg("Failed to begin transaction")
@ -117,6 +121,8 @@ func (d *Psql) MergeAlbums(ctx context.Context, fromId, toId int32, replaceImage
func (d *Psql) MergeArtists(ctx context.Context, fromId, toId int32, replaceImage bool) error {
l := logger.FromContext(ctx)
l.Info().Msgf("Merging artist %d into artist %d", fromId, toId)
d.q.DeleteTrackLookupByArtist(ctx, fromId)
d.q.DeleteTrackLookupByArtist(ctx, toId)
tx, err := d.conn.BeginTx(ctx, pgx.TxOptions{})
if err != nil {
l.Err(err).Msg("Failed to begin transaction")

View file

@ -241,6 +241,9 @@ func (d *Psql) SaveTrackAliases(ctx context.Context, id int32, aliases []string,
func (d *Psql) DeleteTrack(ctx context.Context, id int32) error {
l := logger.FromContext(ctx)
d.q.DeleteTrackLookupByTrack(ctx, id)
tx, err := d.conn.BeginTx(ctx, pgx.TxOptions{})
if err != nil {
l.Err(err).Msg("Failed to begin transaction")

View file

@ -0,0 +1,52 @@
package psql
import (
"context"
"github.com/gabehf/koito/internal/db"
"github.com/gabehf/koito/internal/repository"
"github.com/jackc/pgx/v5"
)
func (d *Psql) GetTrackLookup(ctx context.Context, key string) (*db.TrackLookupResult, error) {
row, err := d.q.GetTrackLookup(ctx, key)
if err != nil {
if err == pgx.ErrNoRows {
return nil, err
}
return nil, err
}
return &db.TrackLookupResult{
ArtistID: row.ArtistID,
AlbumID: row.AlbumID,
TrackID: row.TrackID,
}, nil
}
func (d *Psql) SaveTrackLookup(ctx context.Context, opts db.SaveTrackLookupOpts) error {
return d.q.InsertTrackLookup(ctx, repository.InsertTrackLookupParams{
LookupKey: opts.Key,
ArtistID: opts.ArtistID,
AlbumID: opts.AlbumID,
TrackID: opts.TrackID,
})
}
func (d *Psql) InvalidateTrackLookup(ctx context.Context, opts db.InvalidateTrackLookupOpts) error {
if opts.ArtistID != 0 {
if err := d.q.DeleteTrackLookupByArtist(ctx, opts.ArtistID); err != nil {
return err
}
}
if opts.AlbumID != 0 {
if err := d.q.DeleteTrackLookupByAlbum(ctx, opts.AlbumID); err != nil {
return err
}
}
if opts.TrackID != 0 {
if err := d.q.DeleteTrackLookupByTrack(ctx, opts.TrackID); err != nil {
return err
}
}
return nil
}