cascaded unregister-location

Instead of running drop_node function on every node in cascade,
unregister-location event added to queue by master

4 files changed, 30 insertions, 7 deletions

diff --git a/python/pgq/cascade/admin.py b/python/pgq/cascade/admin.py
index 75fea6d6..f5693bc7 100644
--- a/python/pgq/cascade/admin.py
+++ b/python/pgq/cascade/admin.py
@@ -145,7 +145,7 @@ class CascadeAdmin(skytools.AdminScript):
         if info['node_type'] is not None:
             self.log.info("Node is already initialized as %s" % info['node_type'])
             return
-        
+
         self.log.info("Initializing node")
 
         worker_name = self.options.worker
@@ -286,7 +286,7 @@ class CascadeAdmin(skytools.AdminScript):
                 state = self.get_node_info(node)
                 consumer = state.worker_name
             return (node, consumer)
-        
+
         # global consumer search
         if self.find_consumer_check(self.local_node, consumer):
             return (self.local_node, consumer)
@@ -298,7 +298,7 @@ class CascadeAdmin(skytools.AdminScript):
                 continue
             if self.find_consumer_check(node, consumer):
                 return (node, consumer)
-        
+
         raise Exception('Consumer not found')
 
     def install_code(self, db):
@@ -478,9 +478,19 @@ class CascadeAdmin(skytools.AdminScript):
             raise Exception('node still has subscribers')
 
         # remove the node from all the databases
-        for n in self.queue_info.member_map.values():
-            q = "select * from pgq_node.drop_node(%s, %s)"
-            self.node_cmd(n.name, q, [self.queue_name, node_name])
+        # for n in self.queue_info.member_map.values():
+        #    q = "select * from pgq_node.drop_node(%s, %s)"
+        #    self.node_cmd(n.name, q, [self.queue_name, node_name])
+
+        # drop node info
+        db = self.get_node_database(node_name)
+        q = "select * from pgq_node.drop_node(%s, %s)"
+        self.exec_cmd(db, q, [self.queue_name, node_name])
+
+        # unregister node location from root node (event will be added to queue)
+        root_db = self.find_root_db()
+        q = "select * from pgq_node.unregister_location(%s, %s)"
+        self.exec_cmd(root_db, q, [self.queue_name, node_name])
 
     def node_depends(self, sub_node, top_node):
         cur_node = sub_node
diff --git a/python/pgq/cascade/worker.py b/python/pgq/cascade/worker.py
index 68f9e388..1d3f8325 100644
--- a/python/pgq/cascade/worker.py
+++ b/python/pgq/cascade/worker.py
@@ -240,6 +240,10 @@ class CascadedWorker(CascadedConsumer):
             dead = ev.ev_extra3
             q = "select * from pgq_node.register_location(%s, %s, %s, %s)"
             dst_curs.execute(q, [self.pgq_queue_name, node, loc, dead])
+        elif t == "pgq.unregister-location":
+            node = ev.ev_data
+            q = "select * from pgq_node.unregister_location(%s, %s)"
+            dst_curs.execute(q, [self.pgq_queue_name, node])
         elif t == "pgq.global-watermark":
             if st.process_global_wm:
                 tick_id = int(ev.ev_data)
diff --git a/sql/pgq_node/functions/pgq_node.drop_node.sql b/sql/pgq_node/functions/pgq_node.drop_node.sql
index fbf13ae7..ab48f32d 100644
--- a/sql/pgq_node/functions/pgq_node.drop_node.sql
+++ b/sql/pgq_node/functions/pgq_node.drop_node.sql
@@ -22,7 +22,7 @@ returns record as $$
 -- Return Codes:
 --      200 - Ok
 --      404 - No such queue
--- ----------------------------------------------------------------------
+------------------------------------------------------------------------
 declare
     _is_local   boolean;
 begin
diff --git a/sql/pgq_node/functions/pgq_node.unregister_location.sql b/sql/pgq_node/functions/pgq_node.unregister_location.sql
index 1d3bf586..7314c5ae 100644
--- a/sql/pgq_node/functions/pgq_node.unregister_location.sql
+++ b/sql/pgq_node/functions/pgq_node.unregister_location.sql
@@ -45,12 +45,21 @@ begin
     delete from pgq_node.node_location
      where queue_name = i_queue_name
        and node_name = i_node_name;
+
     if found then
         select 200, 'Ok' into ret_code, ret_note;
     else
         select 301, 'Location not found: ' || i_queue_name || '/' || i_node_name
           into ret_code, ret_note;
     end if;
+
+    if node.node_type = 'root' then
+        perform pgq.insert_event(i_queue_name, 'pgq.unregister-location',
+                                 i_node_name, i_queue_name, null, null, null)
+           from pgq_node.node_info n
+         where n.queue_name = i_queue_name;
+    end if;
+
     return;
 end;
 $$ language plpgsql security definer;