internal/postgres: get DB user info
Add a method for getting information about a DB user. We plan to use
this on the worker to see if we can make better load-shedding
decisions.
For golang/go#48010
Change-Id: I80f2811f657ac47d94446a47a38e00502ae29ae8
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/348933
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Jamal Carvalho <jamal@golang.org>
TryBot-Result: kokoro <noreply+kokoro@google.com>
diff --git a/internal/postgres/postgres.go b/internal/postgres/postgres.go
index e53983c..8b48f1d 100644
--- a/internal/postgres/postgres.go
+++ b/internal/postgres/postgres.go
@@ -130,3 +130,37 @@
}
return total, new, nil
}
+
+// UserInfo holds information about a DB user.
+type UserInfo struct {
+ User string
+ NumTotal int // number of processes running as that user
+ NumWaiting int // number of that user's processes waiting for locks
+}
+
+// GetUserInfo returns information about a database user.
+func (db *DB) GetUserInfo(ctx context.Context, user string) (_ *UserInfo, err error) {
+ derrors.Wrap(&err, "GetLockInfo(%q)", user)
+
+ ui := UserInfo{User: user}
+ // Count the total number of processes running as user.
+ err = db.db.QueryRow(ctx, `
+ SELECT COUNT(DISTINCT pid) FROM pg_stat_activity WHERE usename = $1
+ `, user).Scan(&ui.NumTotal)
+ if err != nil {
+ return nil, err
+ }
+ // Count the number of processes waiting for locks. Note that we can't add
+ // the number of processes where granted = true to the number where granted
+ // = false to get the total, because a process can hold one lock while
+ // waiting for another.
+ err = db.db.QueryRow(ctx, `
+ Select COUNT(DISTINCT l.pid)
+ FROM pg_locks l INNER JOIN pg_stat_activity a USING (pid)
+ WHERE a.usename = $1 AND NOT l.granted
+ `, user).Scan(&ui.NumWaiting)
+ if err != nil {
+ return nil, err
+ }
+ return &ui, nil
+}
diff --git a/internal/postgres/postgres_test.go b/internal/postgres/postgres_test.go
index 4240ed1..5550b7d 100644
--- a/internal/postgres/postgres_test.go
+++ b/internal/postgres/postgres_test.go
@@ -111,3 +111,20 @@
})
}
}
+
+func TestGetUserInfo(t *testing.T) {
+ // We can't know what we'll get from this query, so just perform some basic
+ // sanity checks.
+ t.Parallel()
+ ctx := context.Background()
+ testDB, release := acquire(t)
+ defer release()
+
+ got, err := testDB.GetUserInfo(ctx, "postgres")
+ if err != nil {
+ t.Fatal(err)
+ }
+ if got.NumTotal < 1 {
+ t.Errorf("total = %d, wanted >= 1", got.NumTotal)
+ }
+}
