consumer-pgmq 2.0.1 → 3.1.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "consumer-pgmq",
-  "version": "2.0.1",
+  "version": "3.1.0",
   "description": "The consumer of Supabase pgmq",
   "main": "dist/index.js",
   "type": "commonjs",
@@ -22,7 +22,9 @@
   "packageManager": "pnpm@10.7.1",
   "devDependencies": {
     "@types/jest": "^30.0.0",
+    "@types/pg": "^8.15.5",
     "ts-jest": "^29.4.1",
+    "ts-node": "^10.9.2",
     "ts-node-dev": "^2.0.0",
     "typescript": "^5.9.2"
   },
@@ -30,7 +32,6 @@
     "@supabase/supabase-js": "^2.56.0",
     "dotenv": "^17.2.1",
     "jest": "^30.1.1",
-    "knex": "^3.1.0",
     "pg": "^8.16.3"
   },
   "engines": {

package/schema.sql

@@ -0,0 +1,50 @@
+CREATE TABLE workers (
+    id SERIAL PRIMARY KEY,
+    status worker_status NOT NULL DEFAULT 'idle', -- Tracks the job lifecycle
+    created_at TIMESTAMP DEFAULT now(),
+    updated_at TIMESTAMP DEFAULT now()
+);
+
+CREATE TABLE jobs (
+    id SERIAL PRIMARY KEY,
+    status job_status NOT NULL DEFAULT 'pending', -- Tracks the job lifecycle
+    payload JSONB, -- Chose payload as `jsonb` but could have gone with `bytea` instead.
+    visible_at TIMESTAMP DEFAULT now(), -- SQS visibility timeout (will come back to this later)
+    retry_count INT DEFAULT 0, -- Tracks retry attempts
+    created_at TIMESTAMP DEFAULT now(),
+    updated_at TIMESTAMP DEFAULT now()
+);
+
+CREATE TABLE jobs_dlq (
+    id SERIAL PRIMARY KEY,
+    status job_status NOT NULL DEFAULT 'pending', -- Tracks the job lifecycle
+    payload JSONB, -- Chose payload as `jsonb` but could have gone with `bytea` instead.
+    visible_at TIMESTAMP DEFAULT now(), -- SQS visibility timeout (will come back to this later)
+    retry_count INT DEFAULT 0, -- Tracks retry attempts
+    created_at TIMESTAMP DEFAULT now(),
+    updated_at TIMESTAMP DEFAULT now()
+);
+
+
+WITH next_workers AS (
+            SELECT id
+            FROM workers as w
+            WHERE
+                status = 'idle'
+            LIMIT 1
+            FOR UPDATE SKIP LOCKED
+        )
+        UPDATE workers
+        SET status = 'working',
+            updated_at = now()
+        FROM next_workers
+        WHERE workers.id = next_workers.id   
+        RETURNING workers.*;
+
+CREATE INDEX ON public.workers (status);
+
+insert into workers(status) 
+    values ('idle'),
+    ('idle');
+
+    

package/src/consumer.ts

@@ -22,6 +22,10 @@ class Consumer extends EventEmitter {
      */
     private client: QueueDriver
 
+    private setTimeoutId: NodeJS.Timeout | null = null;
+
+    private id: string | null = null;
+
     constructor(
         options: Options,
         callback: HandlerCallback,
@@ -31,6 +35,14 @@ class Consumer extends EventEmitter {
         this.options = options;
         this.callback = callback;
         this.client = client;
+        this.valideOptions();
+        this.setTimeoutId = null;
+    }
+
+    private valideOptions() {
+        if (this.options.queueNameDlq && !this.options.totalRetriesBeforeSendToDlq) {
+            throw new Error("The option totalRetriesBeforeSendToDlq is required when queueNameDlq is set");
+        }
     }
 
     /**
@@ -84,6 +96,9 @@ class Consumer extends EventEmitter {
      * @private
      */
     private async pollMessage() {
+        if (this.setTimeoutId) {
+            clearTimeout(this.setTimeoutId);
+        }
         let promises: Promise<any>[] = [];
 
         try {
@@ -93,18 +108,39 @@ class Consumer extends EventEmitter {
             }
 
             if (data.length === 0 && this.options.enabledPolling) {
-                setTimeout(
+                this.setTimeoutId = setTimeout(
                     () => this.pollMessage(),
                     (this.options.timeMsWaitBeforeNextPolling || 1000) * 10
                 );
                 return;
             }
 
-
             const controller = new AbortController();
             const signal = controller.signal;
 
-            for (let i = 0; i < (data.length || 1); i++) {
+            for (let i = 0; i < data.length; i++) {
+                const hasSendToDlq = data[i] &&
+                    this.options.queueNameDlq &&
+                    this.options.totalRetriesBeforeSendToDlq &&
+                    data[i].read_ct > this.options.totalRetriesBeforeSendToDlq
+                if (
+                    hasSendToDlq
+                ) {
+                    promises.push(
+                        this.client.send(
+                            // @ts-ignore
+                            this.options.queueNameDlq,
+                            data[i].message,
+                            signal
+                        ).then(async () => {
+                            await this.deleteMessage(data[i], signal);
+                            this.emit('send-to-dlq', data[i]);
+                        })
+                    )
+                    continue;
+                }
+
+
                 promises.push(
                     this.callback(data[i].message, signal).then(async () => {
                         await this.deleteMessage(data[i], signal);
@@ -113,8 +149,12 @@ class Consumer extends EventEmitter {
                 );
             }
 
-            setTimeout(() => controller.abort(), (this.options.visibilityTime || 1) * 1000);
-            await Promise.allSettled(promises);
+            const timeoutId = setTimeout(() => controller.abort(), (this.options.visibilityTime || 1) * 1000);
+            if (promises.length > 0) {
+                await Promise.allSettled(promises);
+            }
+
+            clearTimeout(timeoutId);
             promises = [];
         } catch (err: any) {
             if (err.name === "AbortError") {
@@ -126,7 +166,7 @@ class Consumer extends EventEmitter {
             if (!this.options.enabledPolling) {
                 return;
             }
-            setTimeout(() => this.pollMessage(), this.options.timeMsWaitBeforeNextPolling || 1000);
+            this.setTimeoutId = setTimeout(() => this.pollMessage(), this.options.timeMsWaitBeforeNextPolling || 1000);
         }
     }
 
@@ -136,8 +176,21 @@ class Consumer extends EventEmitter {
      * @public
      */
     async start() {
+        if (this.options.enableControlConsumer) {
+            const { id } = await this.client.allocateConsumer()
+            this.id = id
+        }
+
+
         await this.pollMessage();
     }
+
+    async freeConsumer() {
+        if (this.id) {
+            console.log("passed on here")
+            await this.client.freeConsumer(this.id)
+        }
+    }
 }
 
 

package/src/queueDriver/PostgresQueueDriver.ts

@@ -1,14 +1,152 @@
-import { Knex } from "knex";
 import { Message, QueueDriver } from "../type";
+import { Client } from "pg";
 
 class PostgresQueueDriver implements QueueDriver {
 
 
     constructor(
-        private connection: Knex,
-        private schema: string = "public"
+        private connection: Client,
+        private schema: string = "public",
+        private isCustomQueueImplementation = true
     ) { }
 
+    /**
+     * Allocate the consumer
+     * @returns Promise<{ id: string; }>
+     */
+    async allocateConsumer(): Promise<{ id: string; }> {
+        const register = await this.connection.query(`
+            WITH next_workers AS (
+            SELECT id
+            FROM workers as w
+            WHERE
+                status = 'idle'
+            LIMIT 1
+            FOR UPDATE SKIP LOCKED
+        )
+        UPDATE workers
+        SET status = 'working',
+            updated_at = now()
+        FROM next_workers
+        WHERE workers.id = next_workers.id   
+        RETURNING workers.*;    
+        `)
+
+        if (register.rows.length == 0) {
+            throw new Error("No available consumer(worker) to allocate");
+        }
+
+        return { id: register.rows[0].id };
+    }
+
+
+    async freeConsumer(id: string): Promise<void> {
+        try {
+            await this.connection.query(`UPDATE workers set status = 'idle' WHERE id = $1`, [id])
+        } catch (error) {
+            throw error
+        }
+    }
+
+
+    private async sendCustomQueue(
+        queueName: string,
+        message: { [key: string]: any; }
+    ) {
+        try {
+            const query = `INSERT INTO ${this.schema}.${queueName}(payload) VALUES($1)`
+            await this.connection.query(query, [JSON.stringify(message)])
+            return { error: null };
+        } catch (error) {
+            return { error };
+        }
+    }
+
+    /**
+     * Send the message
+     * @param queueName The name of the queue
+     * @param message The message
+     * @returns Promise<{ error: any }>
+     */
+    async send(
+        queueName: string,
+        message: { [key: string]: any; },
+    ): Promise<{ error: any; }> {
+        if (this.isCustomQueueImplementation) {
+            return this.sendCustomQueue(queueName, message)
+        }
+
+        try {
+            await this.connection.query(`
+            SELECT * FROM ${this.schema}.send(
+                queue_name => $1,
+                msg        => $2,
+                delay      => $3
+            );
+            `, [queueName, message, 1]
+            )
+
+            return { error: null };
+        } catch (error) {
+            return { error };
+        }
+    }
+
+    private async getCustomQueue(
+        queueName: string,
+        visibilityTime: number,
+        totalMessages: number
+    ): Promise<{ data: Message[]; error: any; }> {
+        try {
+            const register = await this.connection.query(
+                `
+                WITH next_job AS (
+                SELECT id
+                FROM ${queueName} as jobs
+                WHERE
+                    (
+                        status = 'pending'
+                        OR (status = 'in_progress' AND visible_at <= now())
+                    )
+                ORDER BY created_at
+                LIMIT $1
+                FOR UPDATE SKIP LOCKED
+            )
+            UPDATE jobs
+            SET status = 'in_progress',
+                updated_at = now(),
+                visible_at = now() + interval '${visibilityTime} seconds',
+                retry_count = retry_count + 1
+            FROM next_job
+            WHERE jobs.id = next_job.id   
+            RETURNING jobs.*;
+        `,
+                [
+                    totalMessages,
+                ]
+            );
+
+            if (!register.rows) {
+                return { data: [], error: null };
+            }
+
+            const items: Message[] = [];
+            for (const row of register.rows) {
+                items.push({
+                    msg_id: row.id,
+                    read_ct: row.retry_count,
+                    enqueued_at: row.created_at,
+                    vt: row.visible_at,
+                    message: row.payload, // Assuming the message content is stored in a column named 'payload' or '
+                });
+            }
+
+            return { data: items, error: null };
+        } catch (error) {
+            return { data: [], error };
+        }
+    }
+
     /**
      * Get the message
      * @param queueName The name of the queue
@@ -17,12 +155,16 @@ class PostgresQueueDriver implements QueueDriver {
      * @returns Promise<{ data: Message[], error: any }>
      */
     async get(queueName: string, visibilityTime: number, totalMessages: number): Promise<{ data: Message[]; error: any; }> {
+        if (this.isCustomQueueImplementation) {
+            return this.getCustomQueue(queueName, visibilityTime, totalMessages);
+        }
+
         try {
-            const register = await this.connection.raw(`
+            const register = await this.connection.query(`
                 SELECT * FROM ${this.schema}.read(
-                    queue_name => ?,
-                    vt         => ?,
-                    qty        => ?
+                    queue_name => $1,
+                    vt         => $2,
+                    qty        => $3
                 );
                 `, [queueName, visibilityTime, totalMessages]
             )
@@ -44,11 +186,18 @@ class PostgresQueueDriver implements QueueDriver {
      * @returns Promise<{ data: Message[], error: any }>
      */
     async pop(queueName: string): Promise<{ data: Message[]; error: any; }> {
+        if (this.isCustomQueueImplementation) {
+            const result = await this.getCustomQueue(queueName, 30, 1);
+            if (result.data && result.data[0]) {
+                await this.delete(queueName, result.data[0].msg_id);
+            }
+            return result
+        }
         try {
-            const register = await this.connection.raw(`
+            const register = await this.connection.query(`
                 SELECT * FROM ${this.schema}.pop(
-                    queue_name => ?
-                );
+                    queue_name => $1
+                )
                 `, [queueName]
             )
 
@@ -63,6 +212,14 @@ class PostgresQueueDriver implements QueueDriver {
 
     }
 
+    private async deleteCustomQueue(queueName: string, messageID: number) {
+        await this.connection.query(`
+            DELETE FROM     ${queueName}
+            WHERE id = $1;`, [messageID]
+        )
+        return { error: null };
+    }
+
     /**
      * Delete the message
      * @param queueName The name of the queue
@@ -71,12 +228,16 @@ class PostgresQueueDriver implements QueueDriver {
      */
     async delete(queueName: string, messageID: number): Promise<{ error: any; }> {
         try {
-            await this.connection.raw(`
+            if (this.isCustomQueueImplementation) {
+                return this.deleteCustomQueue(queueName, messageID)
+            }
+            await this.connection.query(`
             SELECT * FROM ${this.schema}.delete(
-                queue_name => ?,
-                msg_id     => ?
+                queue_name => $1,
+                msg_id     => $2
             );
-            `, [queueName, messageID]
+            `, [queueName, messageID],
+
             )
 
             return { error: null };

package/src/queueDriver/SupabaseQueueDriver.ts

@@ -7,6 +7,21 @@ class SupabaseQueueDriver implements QueueDriver {
         private supabase: SupabaseClient
     ) { }
 
+    /**
+     * Send the message
+     * @param queueName The name of the queue
+     * @param message The message
+     * @returns Promise<{ error: any }>
+     */
+    async send(queueName: string, message: { [key: string]: any; }): Promise<{ error: any; }> {
+        const { error } = await this.supabase.rpc("send", {
+            queue_name: queueName,
+            message: message
+        });
+
+        return { error };
+    }
+
     /**
      * Get the message
      * @param queueName The name of the queue
@@ -52,6 +67,14 @@ class SupabaseQueueDriver implements QueueDriver {
         return { error };
     }
 
+    allocateConsumer(): Promise<{ id: string; }> {
+        throw new Error("method logic no implemented")
+    }
+
+    freeConsumer(id: string): Promise<void> {
+        throw new Error("method logic no implemented")
+    }
+
 }
 
 export default SupabaseQueueDriver;

package/src/type.ts

@@ -30,6 +30,21 @@ interface Options {
      * The enabled polling. PS: if true, the consumer will poll the message
      */
     enabledPolling: boolean;
+
+    /**
+     * The name of the queue DLQ
+     */
+    queueNameDlq?: string;
+
+    /**
+     * The total retries before send the message to DLQ. PS: if set queueNameDlq, this option is required
+     */
+    totalRetriesBeforeSendToDlq?: number;
+
+    /**
+     * The enable control consumer. PS: if true, set the allocate a consumer from table workers, because that way you can control the number of consumers(workers)
+     */
+    enableControlConsumer?: boolean;
 }
 
 /**
@@ -65,6 +80,18 @@ interface Message {
 
 interface QueueDriver {
 
+    /**
+     * Send the message
+     * @param queueName The name of the queue
+     * @param message The message
+     * @returns Promise<{ error: any }>
+     */
+    send(
+        queueName: string,
+        message: { [key: string]: any },
+        signal: AbortSignal
+    ): Promise<{ error: any }>;
+
     /**
      * Get the message
      * @param queueName The name of the queue
@@ -97,6 +124,10 @@ interface QueueDriver {
         queueName: string,
         messageID: number
     ): Promise<{ error: any }>;
+
+    allocateConsumer(): Promise<{ id: string }>;
+
+    freeConsumer(id: string): Promise<void>;
 }
 
 export type { Options, HandlerCallback, Message, QueueDriver }