fix(userstore): batch allowed-library and profile lookups when listing (#380)
* fix(userstore): batch allowed-library and profile lookups when listing Listing profiles cost 1 + P queries: one allowed-libraries lookup fired per profile row while the cursor was still open, which on Postgres also checks out a second pooled connection mid-scan. The admin sessions dashboard makes it worse, calling ListProfiles once per streaming user just to resolve names. The SQLite store had the same pattern for both profiles and collections, even though the Postgres collections path was already written with array_agg to avoid exactly this. Collect the rows first, then fetch the child lists in one batched query and stitch them together in Go. No behaviour change, just fewer round trips. * fix(userstore): avoid sqlite batch variable limits --------- Co-authored-by: Quick104 <31828688+Quick104@users.noreply.github.com>
This commit is contained in:
@@ -114,13 +114,53 @@ func ListCollections(db *sql.DB, profileID string) ([]Collection, error) {
|
||||
if err := rows.Scan(&c.ID, &c.ProfileID, &c.CreatorProfileID, &c.Name, &c.CollectionType, &c.IsShared, &c.QueryDefinition, &c.SortConfig, &c.CreatedAt, &c.UpdatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.AllowedProfileIDs, err = listCollectionProfiles(db, c.ID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
collections = append(collections, c)
|
||||
}
|
||||
return collections, rows.Err()
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := attachCollectionProfiles(db, profileID, collections); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return collections, nil
|
||||
}
|
||||
|
||||
// attachCollectionProfiles fills AllowedProfileIDs for every collection with a
|
||||
// single batched query, avoiding the N+1 round trip a per-collection lookup
|
||||
// would create.
|
||||
func attachCollectionProfiles(db *sql.DB, profileID string, collections []Collection) error {
|
||||
if len(collections) == 0 {
|
||||
return nil
|
||||
}
|
||||
rows, err := db.Query(
|
||||
`SELECT allowed.collection_id, allowed.profile_id
|
||||
FROM personal_collection_profiles AS visible
|
||||
JOIN personal_collection_profiles AS allowed
|
||||
ON allowed.collection_id = visible.collection_id
|
||||
WHERE visible.profile_id = ?
|
||||
ORDER BY allowed.profile_id ASC`,
|
||||
profileID,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
byCollection := make(map[string][]string, len(collections))
|
||||
for rows.Next() {
|
||||
var collectionID, profileID string
|
||||
if err := rows.Scan(&collectionID, &profileID); err != nil {
|
||||
return err
|
||||
}
|
||||
byCollection[collectionID] = append(byCollection[collectionID], profileID)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
for i := range collections {
|
||||
collections[i].AllowedProfileIDs = byCollection[collections[i].ID]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateCollection renames a collection and updates its updated_at timestamp.
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
package userdb
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"github.com/Silo-Server/silo-server/internal/userstore"
|
||||
)
|
||||
|
||||
func newListingTestDB(t *testing.T) *sql.DB {
|
||||
t.Helper()
|
||||
db, err := sql.Open("sqlite3", ":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("open sqlite: %v", err)
|
||||
}
|
||||
db.SetMaxOpenConns(1)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
if err := InitSchema(db); err != nil {
|
||||
t.Fatalf("InitSchema: %v", err)
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
func TestListProfilesAttachesAllowedLibrariesInOrder(t *testing.T) {
|
||||
db := newListingTestDB(t)
|
||||
for _, profile := range []Profile{
|
||||
{ID: "p1", Name: "One", AllowedLibraryIDs: []int{5, 2}, CreatedAt: "2026-01-01T00:00:00Z"},
|
||||
{ID: "p2", Name: "Two", AllowedLibraryIDs: []int{3}, CreatedAt: "2026-01-02T00:00:00Z"},
|
||||
{ID: "p3", Name: "Three", CreatedAt: "2026-01-03T00:00:00Z"},
|
||||
} {
|
||||
if err := CreateProfile(db, profile); err != nil {
|
||||
t.Fatalf("CreateProfile(%s): %v", profile.ID, err)
|
||||
}
|
||||
}
|
||||
|
||||
profiles, err := ListProfiles(db)
|
||||
if err != nil {
|
||||
t.Fatalf("ListProfiles: %v", err)
|
||||
}
|
||||
if len(profiles) != 3 {
|
||||
t.Fatalf("ListProfiles returned %d profiles, want 3", len(profiles))
|
||||
}
|
||||
if got, want := profiles[0].AllowedLibraryIDs, []int{2, 5}; !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("p1 AllowedLibraryIDs = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := profiles[1].AllowedLibraryIDs, []int{3}; !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("p2 AllowedLibraryIDs = %v, want %v", got, want)
|
||||
}
|
||||
if profiles[2].AllowedLibraryIDs != nil {
|
||||
t.Fatalf("p3 AllowedLibraryIDs = %v, want nil", profiles[2].AllowedLibraryIDs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListCollectionsAttachesAllowedProfilesInOrder(t *testing.T) {
|
||||
db := newListingTestDB(t)
|
||||
collection, err := CreateCollection(db, userstore.CreateCollectionInput{
|
||||
CreatorProfileID: "p1",
|
||||
Name: "Shared",
|
||||
IsShared: true,
|
||||
AllowedProfileIDs: []string{"p3", "p2"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateCollection: %v", err)
|
||||
}
|
||||
|
||||
collections, err := ListCollections(db, "p2")
|
||||
if err != nil {
|
||||
t.Fatalf("ListCollections: %v", err)
|
||||
}
|
||||
if len(collections) != 1 || collections[0].ID != collection.ID {
|
||||
t.Fatalf("ListCollections = %+v, want collection %s", collections, collection.ID)
|
||||
}
|
||||
if got, want := collections[0].AllowedProfileIDs, []string{"p1", "p2", "p3"}; !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("AllowedProfileIDs = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchAttachDoesNotScaleBindParametersWithRows(t *testing.T) {
|
||||
db := newListingTestDB(t)
|
||||
const aboveSQLiteVariableLimit = 32767
|
||||
|
||||
profiles := make([]Profile, aboveSQLiteVariableLimit)
|
||||
if err := attachAllowedLibraries(db, profiles); err != nil {
|
||||
t.Fatalf("attachAllowedLibraries: %v", err)
|
||||
}
|
||||
|
||||
collections := make([]Collection, aboveSQLiteVariableLimit)
|
||||
if err := attachCollectionProfiles(db, "viewer", collections); err != nil {
|
||||
t.Fatalf("attachCollectionProfiles: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,8 @@ package userdb
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
|
||||
"github.com/Silo-Server/silo-server/internal/userstore"
|
||||
)
|
||||
|
||||
func listProfileAllowedLibraries(db *sql.DB, profileID string) ([]int, error) {
|
||||
@@ -29,6 +31,38 @@ func listProfileAllowedLibraries(db *sql.DB, profileID string) ([]int, error) {
|
||||
return libraryIDs, nil
|
||||
}
|
||||
|
||||
// attachAllowedLibraries fills AllowedLibraryIDs for every profile with a
|
||||
// single batched query, avoiding the N+1 round trip a per-profile lookup
|
||||
// would create.
|
||||
func attachAllowedLibraries(db *sql.DB, profiles []Profile) error {
|
||||
if len(profiles) == 0 {
|
||||
return nil
|
||||
}
|
||||
rows, err := db.Query(
|
||||
`SELECT profile_id, library_id
|
||||
FROM profile_allowed_libraries
|
||||
ORDER BY library_id ASC`,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("listing allowed libraries: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var allowedLibraries []userstore.ProfileAllowedLibrary
|
||||
for rows.Next() {
|
||||
var allowedLibrary userstore.ProfileAllowedLibrary
|
||||
if err := rows.Scan(&allowedLibrary.ProfileID, &allowedLibrary.LibraryID); err != nil {
|
||||
return fmt.Errorf("scanning allowed library: %w", err)
|
||||
}
|
||||
allowedLibraries = append(allowedLibraries, allowedLibrary)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return fmt.Errorf("iterating allowed libraries: %w", err)
|
||||
}
|
||||
userstore.AttachAllowedLibraries(profiles, allowedLibraries)
|
||||
return nil
|
||||
}
|
||||
|
||||
func replaceProfileAllowedLibrariesTx(tx *sql.Tx, profileID string, libraryIDs []int) error {
|
||||
if _, err := tx.Exec("DELETE FROM profile_allowed_libraries WHERE profile_id = ?", profileID); err != nil {
|
||||
return fmt.Errorf("clearing allowed libraries for profile %s: %w", profileID, err)
|
||||
|
||||
@@ -128,15 +128,14 @@ func ListProfiles(db *sql.DB) ([]Profile, error) {
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("scanning profile row: %w", err)
|
||||
}
|
||||
p.AllowedLibraryIDs, err = listProfileAllowedLibraries(db, p.ID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
profiles = append(profiles, p)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("iterating profile rows: %w", err)
|
||||
}
|
||||
if err := attachAllowedLibraries(db, profiles); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return profiles, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -6,6 +6,8 @@ import (
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
|
||||
"github.com/Silo-Server/silo-server/internal/userstore"
|
||||
)
|
||||
|
||||
type profileLibraryQuerier interface {
|
||||
@@ -41,6 +43,45 @@ func listProfileAllowedLibraries(ctx context.Context, q profileLibraryQuerier, u
|
||||
return libraryIDs, nil
|
||||
}
|
||||
|
||||
// attachAllowedLibraries fills AllowedLibraryIDs for every profile with a
|
||||
// single batched query, avoiding the N+1 round trip a per-profile lookup
|
||||
// would create.
|
||||
func (s *PostgresUserStore) attachAllowedLibraries(ctx context.Context, profiles []userstore.Profile) error {
|
||||
if len(profiles) == 0 {
|
||||
return nil
|
||||
}
|
||||
ids := make([]string, len(profiles))
|
||||
for i := range profiles {
|
||||
ids[i] = profiles[i].ID
|
||||
}
|
||||
rows, err := s.pool.Query(ctx,
|
||||
`SELECT profile_id, library_id
|
||||
FROM user_profile_allowed_libraries
|
||||
WHERE user_id = $1 AND profile_id = ANY($2)
|
||||
ORDER BY library_id ASC`,
|
||||
s.userID,
|
||||
ids,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("listing allowed libraries for user %d: %w", s.userID, err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var allowedLibraries []userstore.ProfileAllowedLibrary
|
||||
for rows.Next() {
|
||||
var allowedLibrary userstore.ProfileAllowedLibrary
|
||||
if err := rows.Scan(&allowedLibrary.ProfileID, &allowedLibrary.LibraryID); err != nil {
|
||||
return fmt.Errorf("scanning allowed library for user %d: %w", s.userID, err)
|
||||
}
|
||||
allowedLibraries = append(allowedLibraries, allowedLibrary)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return fmt.Errorf("iterating allowed libraries for user %d: %w", s.userID, err)
|
||||
}
|
||||
userstore.AttachAllowedLibraries(profiles, allowedLibraries)
|
||||
return nil
|
||||
}
|
||||
|
||||
func replaceProfileAllowedLibraries(ctx context.Context, q profileLibraryQuerier, userID int, profileID string, libraryIDs []int) error {
|
||||
if _, err := q.Exec(ctx,
|
||||
"DELETE FROM user_profile_allowed_libraries WHERE user_id = $1 AND profile_id = $2",
|
||||
|
||||
@@ -126,15 +126,14 @@ func (s *PostgresUserStore) ListProfiles(ctx context.Context) ([]userstore.Profi
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("scanning profile row: %w", err)
|
||||
}
|
||||
p.AllowedLibraryIDs, err = listProfileAllowedLibraries(ctx, s.pool, s.userID, p.ID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
profiles = append(profiles, *p)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("iterating profile rows: %w", err)
|
||||
}
|
||||
if err := s.attachAllowedLibraries(ctx, profiles); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return profiles, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
package userstore
|
||||
|
||||
// ProfileAllowedLibrary associates one profile with one allowed library.
|
||||
type ProfileAllowedLibrary struct {
|
||||
ProfileID string
|
||||
LibraryID int
|
||||
}
|
||||
|
||||
// AttachAllowedLibraries replaces each profile's AllowedLibraryIDs with the
|
||||
// matching associations, preserving their input order.
|
||||
func AttachAllowedLibraries(profiles []Profile, allowedLibraries []ProfileAllowedLibrary) {
|
||||
byProfile := make(map[string][]int, len(profiles))
|
||||
for _, allowedLibrary := range allowedLibraries {
|
||||
byProfile[allowedLibrary.ProfileID] = append(
|
||||
byProfile[allowedLibrary.ProfileID],
|
||||
allowedLibrary.LibraryID,
|
||||
)
|
||||
}
|
||||
for i := range profiles {
|
||||
profiles[i].AllowedLibraryIDs = byProfile[profiles[i].ID]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package userstore
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAttachAllowedLibraries(t *testing.T) {
|
||||
profiles := []Profile{
|
||||
{ID: "p1", AllowedLibraryIDs: []int{99}},
|
||||
{ID: "p2", AllowedLibraryIDs: []int{99}},
|
||||
{ID: "p3", AllowedLibraryIDs: []int{99}},
|
||||
}
|
||||
allowedLibraries := []ProfileAllowedLibrary{
|
||||
{ProfileID: "p2", LibraryID: 1},
|
||||
{ProfileID: "p1", LibraryID: 2},
|
||||
{ProfileID: "p1", LibraryID: 5},
|
||||
}
|
||||
|
||||
AttachAllowedLibraries(profiles, allowedLibraries)
|
||||
|
||||
if got, want := profiles[0].AllowedLibraryIDs, []int{2, 5}; !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("p1 AllowedLibraryIDs = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := profiles[1].AllowedLibraryIDs, []int{1}; !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("p2 AllowedLibraryIDs = %v, want %v", got, want)
|
||||
}
|
||||
if profiles[2].AllowedLibraryIDs != nil {
|
||||
t.Fatalf("p3 AllowedLibraryIDs = %v, want nil", profiles[2].AllowedLibraryIDs)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user