From 71ba303a9465ecb668685bb5f9a4a6bf9d9ea610 Mon Sep 17 00:00:00 2001
From: Vincent Varoquaux <vincent.varoquaux@airbus.com>
Date: Fri, 11 Mar 2022 16:31:15 +0100
Subject: [PATCH] add worker.Count to return the number of pending jobs

---
 queries.go | 14 ++++++++++++++
 worker.go  |  5 +++++
 2 files changed, 19 insertions(+)

diff --git a/queries.go b/queries.go
index 89b414e..101d8df 100644
--- a/queries.go
+++ b/queries.go
@@ -8,6 +8,7 @@ import (
 	"github.com/guregu/null"
 	"github.com/jmoiron/sqlx"
 	"github.com/joomcode/errorx"
+	"github.com/lib/pq"
 )
 
 func getNextJob(tx *sqlx.Tx, queueNames []string) (*Job, error) {
@@ -40,6 +41,19 @@ func getNextJob(tx *sqlx.Tx, queueNames []string) (*Job, error) {
 	}
 }
 
+func count(execer DB, queueNames []string) (int64, error) {
+	var count int64
+	if err := execer.QueryRow(`
+		SELECT count(*) FROM pgq_jobs
+		WHERE
+			queue_name = ANY($1)
+			AND run_after < $2
+			AND ran_at IS NULL;`, pq.Array(queueNames), time.Now()).Scan(&count); err != nil {
+		return -1, errorx.Decorate(err, "could not count jobs")
+	}
+	return count, nil
+}
+
 // DB is anything with the DB methods on it that we need. (like a DB or a Tx)
 type DB interface {
 	Exec(string, ...interface{}) (sql.Result, error)
diff --git a/worker.go b/worker.go
index 9424b2e..9eda331 100644
--- a/worker.go
+++ b/worker.go
@@ -122,6 +122,11 @@ func (worker *Worker) Run() error {
 	}
 }
 
+// Count returns the number of pending jobs for the given queues
+func (worker *Worker) Count(queueNames ...string) (int64, error) {
+	return count(worker.db, queueNames)
+}
+
 func (worker *Worker) getQueueNames() []string {
 	names := []string{}
 	now := time.Now()