84 lines
2.5 KiB
Go
84 lines
2.5 KiB
Go
package repo
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
|
|
"github.com/ente-io/museum/ente"
|
|
"github.com/ente-io/stacktrace"
|
|
"github.com/lib/pq"
|
|
log "github.com/sirupsen/logrus"
|
|
)
|
|
|
|
// GetFilesInfo returns map of fileIDs to ente.FileInfo for a given userID.
|
|
func (repo *FileRepository) GetFilesInfo(ctx context.Context, fileIDs []int64, userID int64) (map[int64]*ente.FileInfo, error) {
|
|
rows, err := repo.DB.QueryContext(ctx, `SELECT file_id, info from files where file_id = ANY($1) and owner_id = $2`, pq.Array(fileIDs), userID)
|
|
if err != nil {
|
|
return nil, stacktrace.Propagate(err, "")
|
|
}
|
|
defer rows.Close()
|
|
result := make(map[int64]*ente.FileInfo, 0)
|
|
for rows.Next() {
|
|
var fileID int64
|
|
var info *ente.FileInfo
|
|
if err = rows.Scan(&fileID, &info); err != nil {
|
|
return nil, stacktrace.Propagate(err, "")
|
|
}
|
|
result[fileID] = info
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
// UpdateSizeInfo updates the size info for a given map of fileIDs to ente.FileInfo.
|
|
func (repo *FileRepository) UpdateSizeInfo(ctx context.Context, sizeInfo map[int64]*ente.FileInfo) error {
|
|
// Update the size info for each file using a batched transaction.
|
|
tx, err := repo.DB.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return stacktrace.Propagate(err, "")
|
|
}
|
|
defer tx.Rollback()
|
|
for fileID, info := range sizeInfo {
|
|
_, err := tx.ExecContext(ctx, `UPDATE files SET info = $1 WHERE file_id = $2 and info is NULL`, info, fileID)
|
|
if err != nil {
|
|
return stacktrace.Propagate(err, "")
|
|
}
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return stacktrace.Propagate(err, "")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetFileInfoFromObjectKeys returns the file info for a given list of fileIDs.
|
|
func (repo *FileRepository) GetFileInfoFromObjectKeys(ctx context.Context, fileIDs []int64) (map[int64]*ente.FileInfo, error) {
|
|
rows, err := repo.DB.QueryContext(ctx, `SELECT file_id, size, o_type FROM object_keys WHERE file_id = ANY($1)`, pq.Array(fileIDs))
|
|
if err != nil {
|
|
return nil, stacktrace.Propagate(err, "")
|
|
}
|
|
defer func(rows *sql.Rows) {
|
|
err := rows.Close()
|
|
if err != nil {
|
|
log.Errorf("error closing rows: %v", err)
|
|
}
|
|
}(rows)
|
|
result := make(map[int64]*ente.FileInfo, 0)
|
|
for rows.Next() {
|
|
var fileID int64
|
|
var size int64
|
|
var oType ente.ObjectType
|
|
if err = rows.Scan(&fileID, &size, &oType); err != nil {
|
|
return nil, stacktrace.Propagate(err, "")
|
|
}
|
|
if _, ok := result[fileID]; !ok {
|
|
result[fileID] = &ente.FileInfo{}
|
|
}
|
|
switch oType {
|
|
case ente.FILE:
|
|
result[fileID].FileSize = size
|
|
case ente.THUMBNAIL:
|
|
result[fileID].ThumbnailSize = size
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|