Merge pull request #126 from iKonoTelecomunicaciones/125-disable-bot

feat: ✨ Enable/disable menuflow bots

menuflow/api/client.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from json import JSONDecodeError
 from logging import Logger, getLogger
 from typing import Dict, Optional
@@ -196,6 +197,28 @@ async def reload_client_flow(request: web.Request) -> web.Response:
     return resp.ok(client.to_dict())
 
 
+@routes.patch("/client/{mxid}/{action}")
+async def enable_disable_client(request: web.Request) -> web.Response:
+    mxid = request.match_info["mxid"]
+    action = request.match_info["action"]
+    client: MenuClient = await MenuClient.get(mxid)
+    if not client:
+        return resp.client_not_found(mxid)
+
+    if action == "enable":
+        client.enabled = True
+        await client.start()
+    elif action == "disable":
+        client.enabled = False
+        asyncio.create_task(client.leave_rooms(), name=f"{mxid}-leave_rooms")
+        await client.stop()
+    else:
+        return resp.bad_request("Invalid action provided")
+
+    await client.update()
+    return resp.ok({"detail": {"message": f"Client {action}d successfully"}})
+
+
 @routes.get("/client/email_servers")
 async def get_id_email_servers(request: web.Request) -> web.Response:
     name_email_servers = [name for name in FlowUtils.email_servers_by_id.keys()]

menuflow/db/client.py

@@ -24,6 +24,7 @@ class Client(SyncStore):
     filter_id: FilterID
 
     autojoin: bool
+    enabled: bool
     flow: int | None = None
 
     @classmethod
@@ -32,7 +33,9 @@ def _from_row(cls, row: Record | None) -> Client | None:
             return None
         return cls(**row)
 
-    _columns = "id, homeserver, access_token, device_id, next_batch, filter_id, autojoin, flow"
+    _columns = (
+        "id, homeserver, access_token, device_id, next_batch, filter_id, autojoin, enabled, flow"
+    )
 
     @property
     def _values(self):
@@ -44,12 +47,13 @@ def _values(self):
             self.next_batch,
             self.filter_id,
             self.autojoin,
+            self.enabled,
             self.flow,
         )
 
     @classmethod
     async def all(cls) -> list[Client]:
-        rows = await cls.db.fetch(f"SELECT {cls._columns} FROM client")
+        rows = await cls.db.fetch(f"SELECT {cls._columns} FROM client WHERE enabled IS TRUE")
         return [cls._from_row(row) for row in rows]
 
     @classmethod
@@ -59,14 +63,14 @@ async def get(cls, id: str) -> Client | None:
 
     @classmethod
     async def get_by_flow_id(cls, flow_id: int) -> list[Client]:
-        q = f"SELECT {cls._columns} FROM client WHERE flow=$1"
+        q = f"SELECT {cls._columns} FROM client WHERE flow=$1 AND enabled IS TRUE"
         rows = await cls.db.fetch(q, flow_id)
         return [cls._from_row(row) for row in rows]
 
     async def insert(self) -> None:
         q = (
             "INSERT INTO client (id, homeserver, access_token, device_id, next_batch, filter_id, "
-            "autojoin, flow) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)"
+            "autojoin, enabled, flow) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)"
         )
         await self.db.execute(q, *self._values)
 
@@ -80,7 +84,7 @@ async def get_next_batch(self) -> SyncToken:
     async def update(self) -> None:
         q = (
             "UPDATE client SET homeserver=$2, access_token=$3, device_id=$4, next_batch=$5, "
-            "filter_id=$6, autojoin=$7, flow=$8 WHERE id=$1"
+            "filter_id=$6, autojoin=$7, enabled=$8, flow=$9 WHERE id=$1"
         )
         await self.db.execute(q, *self._values)
 

menuflow/db/migrations.py

@@ -86,3 +86,8 @@ async def upgrade_v4(conn: Connection) -> None:
     await conn.execute(
         "ALTER TABLE client ADD CONSTRAINT FK_flow_client FOREIGN KEY (flow) references flow (id)"
     )
+
+
+@upgrade_table.register(description="Add enable column to client table")
+async def upgrade_v5(conn: Connection) -> None:
+    await conn.execute("ALTER TABLE client ADD COLUMN enabled BOOLEAN NOT NULL DEFAULT TRUE")

menuflow/menu.py

@@ -52,6 +52,7 @@ def __init__(
         next_batch: SyncToken = "",
         filter_id: FilterID = "",
         autojoin: bool = True,
+        enabled: bool = True,
         flow: int | None = None,
     ) -> None:
         super().__init__(
@@ -62,6 +63,7 @@ def __init__(
             next_batch=next_batch,
             filter_id=filter_id,
             autojoin=bool(autojoin),
+            enabled=bool(enabled),
             flow=flow,
         )
         self._postinited = False
@@ -239,6 +241,11 @@ async def delete(self) -> None:
             pass
         await super().delete()
 
+    async def leave_rooms(self) -> None:
+        rooms = await self.matrix_handler.get_joined_rooms()
+        for room_id in rooms:
+            await self.matrix_handler.leave_room(room_id)
+
     @classmethod
     async def all(cls) -> AsyncGenerator[MenuClient, None]:
         users = await super().all()