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:
Jeremy Smith
2026-07-20 11:25:24 -04:00
committed by GitHub
co-authored by Quick104
parent e56a1b3e03
commit f637404c46
8 changed files with 271 additions and 13 deletions
+45 -5
View File
@@ -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.
+92
View File
@@ -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)
}
}
+34
View File
@@ -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)
+3 -4
View File
@@ -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",
+3 -4
View File
@@ -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
}
+22
View File
@@ -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)
}
}