Review

Author: attdona

src/rembus/admin.py

@@ -13,6 +13,38 @@ def upsert_twin(lst, twin):
     else:
         lst.append(twin)
 
+
+def add_exposer(router, twin, topic):
+    logger.debug("[%s] adding [%s] exposer for topic [%s]", router, twin, topic)
+    if topic not in router.exposers:
+        router.exposers[topic] = []
+    upsert_twin(router.exposers[topic], twin)
+
+def remove_exposer(router, twin, topic):
+    if topic in router.exposers:
+        if twin in router.exposers[topic]:
+            router.exposers[topic].remove(twin)
+    logger.debug(
+        "[%s] removed [%s] exposer for topic [%s]", router, twin, topic
+    )
+
+
+def add_subscriber(router, twin, topic):
+    logger.debug(
+        "[%s] adding [%s] subscriber for topic [%s]", router, twin, topic
+    )
+    if topic not in router.subscribers:
+        router.subscribers[topic] = []
+    upsert_twin(router.subscribers[topic], twin)
+
+def remove_subscriber(router, twin, topic):
+    if topic in router.subscribers:
+        if twin in router.subscribers[topic]:
+            router.subscribers[topic].remove(twin)
+    logger.debug(
+        "[%s] removed [%s] subscriber for topic [%s]", router, twin, topic
+    )
+
 async def admin_command(msg: rp.AdminMsg):
     """Handle admin commands"""
     twin = msg.twin
@@ -25,32 +57,12 @@ async def admin_command(msg: rp.AdminMsg):
     router = twin.router
     cmd = msg.data[rp.COMMAND]
     if cmd == rp.ADD_IMPL:
-        logger.debug(
-            "[%s] adding [%s] exposer for topic [%s]", router, twin, topic
-        )
-        if topic not in router.exposers:
-            router.exposers[topic] = []
-
-        upsert_twin(router.exposers[topic], twin)
-
+        add_exposer(router, twin, topic)
     elif cmd == rp.REMOVE_IMPL:
-        if topic in router.exposers:
-            if twin in router.exposers[topic]:
-                router.exposers[topic].remove(twin)
-        logger.info("[%s] removed [%s] exposer for topic [%s]", router, twin, topic)
-        logger.info("exposers: %s", router.exposers)
+        remove_exposer(router, twin, topic)
     elif cmd == rp.ADD_INTEREST:
-        logger.debug(
-            "[%s] adding [%s] subscriber for topic [%s]", router, twin, topic
-        )
-        if topic not in router.subscribers:
-            router.subscribers[topic] = []
-
-        upsert_twin(router.subscribers[topic], twin)
-
+        add_subscriber(router, twin, topic)
     elif cmd == rp.REMOVE_INTEREST:
-        if topic in router.subscribers:
-            if twin in router.subscribers[topic]:
-                router.subscribers[topic].remove(twin)
+        remove_subscriber(router, twin, topic)
 
     await twin.response(rp.STS_OK, msg)

src/rembus/core.py

@@ -407,6 +407,23 @@ async def _shutdown(self):
         if self.db is not None:
             self.db.close()
 
+    async def _broadcast(self, msg):
+        twin = msg.twin
+        data = rp.tag2df(msg.data)
+        try:
+            if msg.topic in self.handler:
+                await self.evaluate(msg.twin, msg.topic, data)
+
+            subs = self.subscribers.get(msg.topic, [])
+            for t in subs:
+                if t != twin:
+                    # Do not send back to publisher.
+                    await t.send(msg)
+
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            logger.warning("[%s] error in method invocation: %s", self, e)
+            traceback.print_exc()
+
     async def _pubsub_msg(self, msg: rp.PubSubMsg):
         msg.recvts = int(time.time())
         twin = msg.twin
@@ -423,26 +440,11 @@ async def _pubsub_msg(self, msg: rp.PubSubMsg):
                     # Save the message id to guarantee exactly one delivery.
                     twin.ackdf[msg.id] = int(time.time())
 
-        data = rp.tag2df(msg.data)
-
         if self.db is not None:
             # save the message into msg_cache
             self.append_message(msg)
 
-        try:
-            if msg.topic in self.handler:
-                await self.evaluate(msg.twin, msg.topic, data)
-
-            subs = self.subscribers.get(msg.topic, [])
-            for t in subs:
-                if t != twin:
-                    # Do not send back to publisher.
-                    await t.send(msg)
-
-        except Exception as e:  # pylint: disable=broad-exception-caught
-            logger.warning("[%s] error in method invocation: %s", self, e)
-            traceback.print_exc()
-        return
+        await self._broadcast(msg)
 
     def append_message(self, msg: rp.PubSubMsg):
         """Append a message to the message cache."""

src/rembus/db.py

@@ -342,7 +342,7 @@ def append(con: duckdb.DuckDBPyConnection, tabledef, msgs):
     for msg in msgs:
         try:
             values = msg.data
-            # key_value format
+            # key_value
             if fmt == "key_value":
                 obj = getobj(topic, values)
                 set_default(msg, tabledef, obj, add_nullable=True)
@@ -356,7 +356,7 @@ def append(con: duckdb.DuckDBPyConnection, tabledef, msgs):
 
                 fields = [obj[f] for f in tblfields]
 
-            # dataframe format
+            # dataframe
             elif fmt == "dataframe":
                 df = tag2df(values[0])
                 if list(df.columns) != tblfields:
@@ -424,90 +424,66 @@ def schema_to_polars(tabledef: Table):
     return schema
 
 
-def upsert(con: duckdb.DuckDBPyConnection, tabledef, messages):
-    tname = tabledef.table
-    fmt = tabledef.format
-    col_names = columns(tabledef) + list(tabledef.extras.values())
-    indexes = list(tabledef.keys)
-    records = []
-    dataframes = []
+def add_extras(table, msg, obj):
+    if "recv_ts" in table.extras:
+        obj[table.extras["recv_ts"]] = msg.recvts
+    if "slot" in table.extras:
+        obj[table.extras["slot"]] = msg.slot
 
-    # Build final DataFrame (tdf)
-    for msg in messages:
-        try:
-            values = msg.data
 
-            # key_value
-            if fmt == "key_value":
-                obj = getobj(tname, values)
-                set_default(msg, tabledef, obj, add_nullable=True)
+def handle_key_value(msg, table, col_names, records, tname):
+    obj = getobj(tname, msg.data)
+    set_default(msg, table, obj, add_nullable=True)
+    add_extras(table, msg, obj)
 
-                if "recv_ts" in tabledef.extras:
-                    obj[tabledef.extras["recv_ts"]] = msg.recvts
-                if "slot" in tabledef.extras:
-                    obj[tabledef.extras["slot"]] = msg.slot
+    if all(k in obj for k in col_names):
+        records.append(obj)
+    else:
+        logger.warning(
+            "[%s] unsaved %s missing required fields %s", tname, obj, col_names
+        )
 
-                # Ensure all fields exist
-                if all(k in obj for k in col_names):
-                    records.append(obj)
-                else:
-                    logger.warning(
-                        "[%s] unsaved %s missing required fields %s",
-                        tname,
-                        obj,
-                        col_names,
-                    )
-                    continue
 
-            # dataframe
-            elif fmt == "dataframe":
-                df = tag2df(values[0])
-                df = df_extras(tabledef, df, msg)
-                if list(df.columns) == col_names:
-                    dataframes.append(df)
-                else:
-                    logger.warning("[%s] unsaved df (mismatched fields)", tname)
-                    continue
+def handle_dataframe(msg, table, col_names, dataframes, tname):
+    df = tag2df(msg.data[0])
+    df = df_extras(table, df, msg)
 
-            # default format
-            else:
-                vals = values
-                if len(vals) == len(tabledef.columns):
-                    extra_vals = extras(tabledef, msg)
-                    all_vals = vals + extra_vals
-                    records.append(all_vals)
-                else:
-                    logger.warning(
-                        "[%s] unsaved %s with mismatched fields", tname, vals
-                    )
-                    continue
-        except Exception as e:
-            logger.error("[upsert] %s: %s", tname, e)
+    if list(df.columns) == col_names:
+        dataframes.append(df)
+    else:
+        logger.warning("[%s] unsaved df (mismatched fields)", tname)
+
+
+def handle_default(msg, table, col_names, records, tname):
+    vals = msg.data
+    if len(vals) != len(table.columns):
+        logger.warning("[%s] unsaved %s with mismatched fields", tname, vals)
+        return
+
+    extra_vals = extras(table, msg)
+    records.append(vals + extra_vals)
+
+
+def execute_upsert(con, table, col_names, indexes, records, dataframes):
+    tname = table.table
 
-    tdf = pl.DataFrame(records, schema=schema_to_polars(tabledef), orient="row")
+    tdf = pl.DataFrame(records, schema=schema_to_polars(table), orient="row")
     tdf = pl.concat([tdf, *dataframes], how="vertical")
+
     if tdf.is_empty():
         return
 
     if indexes:
-        tdf = (
-            tdf.sort(indexes)  # Sort so last row is last
-            .group_by(indexes)
-            .agg([pl.all().last()])  # Take last row of each group
-        )
-
-    # logger.debug("[%s] upserting dataframe:\n%s", tname, tdf)
+        tdf = tdf.sort(indexes).group_by(indexes).agg([pl.all().last()])
 
     con.register("df_view", tdf)
 
-    conds = [f"df_view.{k} = {tname}.{k}" for k in indexes]
-    cond_str = " AND ".join(conds)
-
+    cond_str = " AND ".join(f"df_view.{k} = {tname}.{k}" for k in indexes)
     col_list = ", ".join(col_names)
     val_list = ", ".join(f"df_view.{c}" for c in col_names)
 
-    update_columns = [c for c in col_names if c not in indexes]
-    update_list = ", ".join(f"{c} = df_view.{c}" for c in update_columns)
+    update_cols = [c for c in col_names if c not in indexes]
+    update_list = ", ".join(f"{c} = df_view.{c}" for c in update_cols)
 
     sql = f"""
         MERGE INTO {tname}
@@ -521,6 +497,33 @@ def upsert(con: duckdb.DuckDBPyConnection, tabledef, messages):
     con.unregister("df_view")
 
 
+def upsert(con: duckdb.DuckDBPyConnection, table, messages):
+    """Insert/update a record in a table with keywords."""
+    tname = table.table
+    fmt = table.format
+    col_names = columns(table) + list(table.extras.values())
+    indexes = list(table.keys)
+
+    records = []
+    dataframes = []
+
+    for msg in messages:
+        try:
+            if fmt == "key_value":
+                handle_key_value(msg, table, col_names, records, tname)
+
+            elif fmt == "dataframe":
+                handle_dataframe(msg, table, col_names, dataframes, tname)
+
+            else:
+                handle_default(msg, table, col_names, records, tname)
+
+        except Exception as e:  # pylint: disable=broad-except
+            logger.error("[upsert] %s: %s", tname, e)
+
+    execute_upsert(con, table, col_names, indexes, records, dataframes)
+
+
 def expand(msg: PubSubMsg):
     """
     Given a message `msg` containing a topic string

src/rembus/keyspace.py

@@ -74,6 +74,28 @@ async def unsubscribe_handler(self, msg):
                 if el in self.spaces[retopic]:
                     self.spaces[retopic].remove(el)
 
+    async def broadcast(self, topic, msg, space_twins):
+        """Brodcast the message to all subscribed spaces."""
+        to_remove = []
+        for space_twin in space_twins:
+            pattern = space_twin.space
+            logger.debug(
+                "[keyspace] publish to %s: %s, %s",
+                pattern,
+                topic,
+                msg.data,
+            )
+            twid = space_twin.twid
+            if twid in self.broker.id_twin:
+                tw = self.broker.id_twin[twid]
+                if tw.isopen():
+                    await tw.publish(pattern, topic, *msg.data)
+            else:
+                # cleanup, the twin has gone
+                to_remove.append(space_twin)
+        for el in to_remove:
+            space_twins.remove(el)
+
     async def publish_interceptor(self, msg):
         """Parse the topic and dispatch to all twins subscribed to spaces
         with regex matching the topic"""
@@ -84,35 +106,15 @@ async def publish_interceptor(self, msg):
                 if m is not None:
                     unsealed = True
 
-                    # m.groups() == captures in Julia
                     for capture in m.groups():
                         if "@" in capture:
                             # a regex chunk matched a verbatim chunk → reject
                             unsealed = False
                             break
 
                     if unsealed:
-                        to_remove = []
-                        # for pattern, twid in space_twins:
-                        for space_twin in space_twins:
-                            pattern = space_twin.space
-                            logger.debug(
-                                "[keyspace] publish to %s: %s, %s",
-                                pattern,
-                                topic,
-                                msg.data,
-                            )
-                            twid = space_twin.twid
-                            if twid in self.broker.id_twin:
-                                tw = self.broker.id_twin[twid]
-                                if tw.isopen():
-                                    await tw.publish(pattern, topic, *msg.data)
-                            else:
-                                # cleanup, the twin has gone
-                                to_remove.append(space_twin)
-
-                        for el in to_remove:
-                            space_twins.remove(el)
+                        await self.broadcast(topic, msg, space_twins)
+
 
     async def _task_impl(self) -> None:
         """Override in subclasses for supervised task impl."""