pgq.unregister_consumer: don't unregister all subconsumers in subscription when subconsumer was given on input

3 files changed, 56 insertions, 14 deletions

diff --git a/sql/pgq/functions/pgq.unregister_consumer.sql b/sql/pgq/functions/pgq.unregister_consumer.sql
index c4150859..d495c88e 100644
--- a/sql/pgq/functions/pgq.unregister_consumer.sql
+++ b/sql/pgq/functions/pgq.unregister_consumer.sql
@@ -24,14 +24,19 @@ returns integer as $$
 declare
     x_sub_id integer;
     _sub_id_cnt integer;
+    _consumer_id integer;
+    _is_subconsumer boolean;
 begin
-    select s.sub_id into x_sub_id
-        from pgq.subscription s, pgq.consumer c, pgq.queue q
-        where s.sub_queue = q.queue_id
-          and s.sub_consumer = c.co_id
-          and q.queue_name = x_queue_name
-          and c.co_name = x_consumer_name
-        for update of s;
+    select s.sub_id, c.co_id,
+           -- subconsumers can only have both null or both not null - main consumer for subconsumers has only one not null
+           (s.sub_last_tick IS NULL AND s.sub_next_tick IS NULL) OR (s.sub_last_tick IS NOT NULL AND s.sub_next_tick IS NOT NULL)
+      into x_sub_id, _consumer_id, _is_subconsumer
+      from pgq.subscription s, pgq.consumer c, pgq.queue q
+     where s.sub_queue = q.queue_id
+       and s.sub_consumer = c.co_id
+       and q.queue_name = x_queue_name
+       and c.co_name = x_consumer_name
+       for update of s;
     if not found then
         return 0;
     end if;
@@ -41,15 +46,27 @@ begin
         from pgq.subscription
        where sub_id = x_sub_id;
 
-    -- retry events
-    delete from pgq.retry_queue
-        where ev_owner = x_sub_id;
+    -- delete only one subconsumer
+    if _sub_id_cnt > 1 and _is_subconsumer then
+        delete from pgq.subscription
+              where sub_id = x_sub_id
+                and sub_consumer = _consumer_id;
 
-    -- this will drop subconsumers too
-    delete from pgq.subscription
-        where sub_id = x_sub_id;
+        return 1;
+    else
+        -- delete main consumer (including possible subconsumers)
+
+        -- retry events
+        delete from pgq.retry_queue
+            where ev_owner = x_sub_id;
+
+        -- this will drop subconsumers too
+        delete from pgq.subscription
+            where sub_id = x_sub_id;
+    
+        return _sub_id_cnt;
+    end if;
 
-    return _sub_id_cnt;
 end;
 $$ language plpgsql security definer;
 
diff --git a/sql/pgq_coop/expected/pgq_coop_test.out b/sql/pgq_coop/expected/pgq_coop_test.out
index 11e8424e..dbb8fe56 100644
--- a/sql/pgq_coop/expected/pgq_coop_test.out
+++ b/sql/pgq_coop/expected/pgq_coop_test.out
@@ -171,3 +171,22 @@ select consumer_name, last_tick from pgq.get_consumer_info();
  maincons      |         3
 (4 rows)
 
+-- test unregistering with pure pgq api
+select pgq.unregister_consumer('testqueue', 'cmain.sc2');
+ unregister_consumer 
+---------------------
+                   1
+(1 row)
+
+select pgq.unregister_consumer('testqueue', 'cmain');
+ unregister_consumer 
+---------------------
+                   2
+(1 row)
+
+select consumer_name, last_tick from pgq.get_consumer_info();
+ consumer_name | last_tick 
+---------------+-----------
+ maincons      |         3
+(1 row)
+
diff --git a/sql/pgq_coop/sql/pgq_coop_test.sql b/sql/pgq_coop/sql/pgq_coop_test.sql
index 0801a9d9..279c0dd4 100644
--- a/sql/pgq_coop/sql/pgq_coop_test.sql
+++ b/sql/pgq_coop/sql/pgq_coop_test.sql
@@ -42,3 +42,9 @@ select pgq_coop.unregister_subconsumer('testqueue', 'maincons', 'subcons2', 0);
 select pgq_coop.next_batch('testqueue', 'cmain', 'sc1');
 select pgq_coop.next_batch('testqueue', 'cmain', 'sc2');
 select consumer_name, last_tick from pgq.get_consumer_info();
+
+-- test unregistering with pure pgq api
+select pgq.unregister_consumer('testqueue', 'cmain.sc2');
+select pgq.unregister_consumer('testqueue', 'cmain');
+select consumer_name, last_tick from pgq.get_consumer_info();
+