chore: cleanup and refactor

Author: Julien Barreau

src/agent_rpc/forestadmin/agent_rpc/agent.py

@@ -1,16 +1,12 @@
 import asyncio
 import json
-from enum import Enum
 from uuid import UUID
 
 from aiohttp import web
 from aiohttp_sse import sse_response
-
-# import grpc
 from forestadmin.agent_rpc.options import RpcOptions
-
-# from forestadmin.agent_rpc.services.datasource import DatasourceService
 from forestadmin.agent_toolkit.agent import Agent
+from forestadmin.agent_toolkit.options import Options
 from forestadmin.datasource_toolkit.interfaces.fields import PrimitiveType
 from forestadmin.datasource_toolkit.utils.schema import SchemaUtils
 from forestadmin.rpc_common.hmac import is_valid_hmac
@@ -30,48 +26,19 @@
 from forestadmin.rpc_common.serializers.schema.schema import SchemaSerializer
 from forestadmin.rpc_common.serializers.utils import CallerSerializer
 
-# from concurrent import futures
-
-
-# from forestadmin.rpc_common.proto import datasource_pb2_grpc
-
-
-class RcpJsonEncoder(json.JSONEncoder):
-    def default(self, o):
-        if isinstance(o, Enum):
-            return o.value
-        if isinstance(o, set):
-            return list(sorted(o, key=lambda x: x.value if isinstance(x, Enum) else str(x)))
-        if isinstance(o, set):
-            return list(sorted(o, key=lambda x: x.value if isinstance(x, Enum) else str(x)))
-
-        try:
-            return super().default(o)
-        except Exception as exc:
-            print(f"error on seriliaze {o}, {type(o)}: {exc}")
-
 
 class RpcAgent(Agent):
-    # TODO: options to add:
-    # * listen addr
     def __init__(self, options: RpcOptions):
-        # self.server = grpc.server(futures.ThreadPoolExecutor(max_workers=1))
-        # self.server = grpc.aio.server()
         self.listen_addr, self.listen_port = options["listen_addr"].rsplit(":", 1)
+        agent_options: Options = {**options}  # type:ignore
+        agent_options["skip_schema_update"] = True
+        agent_options["env_secret"] = "f" * 64
+        agent_options["server_url"] = "http://fake"
+        agent_options["schema_path"] = "./.forestadmin-schema.json"
+        super().__init__(agent_options)
+
         self.app = web.Application(middlewares=[self.hmac_middleware])
-        # self.server.add_insecure_port(options["listen_addr"])
-        options["skip_schema_update"] = True
-        options["env_secret"] = "f" * 64
-        options["server_url"] = "http://fake"
-        # options["auth_secret"] = "f48186505a3c5d62c27743126d6a76c1dd8b3e2d8897de19"
-        options["schema_path"] = "./.forestadmin-schema.json"
-        super().__init__(options)
-
-        self.aes_key = self.options["auth_secret"][:16].encode()
-        self.aes_iv = self.options["auth_secret"][-16:].encode()
-        self._server_stop = False
         self.setup_routes()
-        # signal.signal(signal.SIGUSR1, self.stop_handler)
 
     @web.middleware
     async def hmac_middleware(self, request: web.Request, handler):
@@ -80,11 +47,10 @@ async def hmac_middleware(self, request: web.Request, handler):
             if not is_valid_hmac(
                 self.options["auth_secret"].encode(), body, request.headers.get("X-FOREST-HMAC", "").encode("utf-8")
             ):
-                return web.Response(status=401)
+                return web.Response(status=401, text="Unauthorized from HMAC verification")
         return await handler(request)
 
     def setup_routes(self):
-        # self.app.middlewares.append(self.hmac_middleware)
         self.app.router.add_route("GET", "/sse", self.sse_handler)
         self.app.router.add_route("GET", "/schema", self.schema)
         self.app.router.add_route("POST", "/collection/list", self.collection_list)
@@ -98,11 +64,11 @@ def setup_routes(self):
 
         self.app.router.add_route("POST", "/execute-native-query", self.native_query)
         self.app.router.add_route("POST", "/render-chart", self.render_chart)
-        self.app.router.add_route("GET", "/", lambda _: web.Response(text="OK"))
+        self.app.router.add_route("GET", "/", lambda _: web.Response(text="OK"))  # type: ignore
 
     async def sse_handler(self, request: web.Request) -> web.StreamResponse:
         async with sse_response(request) as resp:
-            while resp.is_connected() and not self._server_stop:
+            while resp.is_connected():
                 await resp.send("", event="heartbeat")
                 await asyncio.sleep(1)
             data = json.dumps({"event": "RpcServerStop"})
@@ -112,44 +78,42 @@ async def sse_handler(self, request: web.Request) -> web.StreamResponse:
     async def schema(self, request):
         await self.customizer.get_datasource()
 
-        return web.Response(text=json.dumps(await SchemaSerializer(await self.customizer.get_datasource()).serialize()))
+        return web.json_response(await SchemaSerializer(await self.customizer.get_datasource()).serialize())
 
     async def collection_list(self, request: web.Request):
         body_params = await request.json()
         ds = await self.customizer.get_datasource()
         collection = ds.get_collection(body_params["collectionName"])
         caller = CallerSerializer.deserialize(body_params["caller"])
-        filter_ = PaginatedFilterSerializer.deserialize(body_params["filter"], collection)
+        filter_ = PaginatedFilterSerializer.deserialize(body_params["filter"], collection)  # type:ignore
         projection = ProjectionSerializer.deserialize(body_params["projection"])
 
         records = await collection.list(caller, filter_, projection)
         records = [RecordSerializer.serialize(record) for record in records]
-        return web.Response(text=aes_encrypt(json.dumps(records), self.aes_key, self.aes_iv))
+        return web.json_response(records)
 
     async def collection_create(self, request: web.Request):
         body_params = await request.text()
-        body_params = aes_decrypt(body_params, self.aes_key, self.aes_iv)
         body_params = json.loads(body_params)
         ds = await self.customizer.get_datasource()
 
         collection = ds.get_collection(body_params["collectionName"])
         caller = CallerSerializer.deserialize(body_params["caller"])
-        data = [RecordSerializer.deserialize(r, collection) for r in body_params["data"]]
+        data = [RecordSerializer.deserialize(r, collection) for r in body_params["data"]]  # type:ignore
 
         records = await collection.create(caller, data)
         records = [RecordSerializer.serialize(record) for record in records]
         return web.json_response(records)
 
     async def collection_update(self, request: web.Request):
         body_params = await request.text()
-        body_params = aes_decrypt(body_params, self.aes_key, self.aes_iv)
         body_params = json.loads(body_params)
 
         ds = await self.customizer.get_datasource()
         collection = ds.get_collection(body_params["collectionName"])
         caller = CallerSerializer.deserialize(body_params["caller"])
-        filter_ = FilterSerializer.deserialize(body_params["filter"], collection)
-        patch = RecordSerializer.deserialize(body_params["patch"], collection)
+        filter_ = FilterSerializer.deserialize(body_params["filter"], collection)  # type:ignore
+        patch = RecordSerializer.deserialize(body_params["patch"], collection)  # type:ignore
 
         await collection.update(caller, filter_, patch)
         return web.Response(text="OK")
@@ -159,7 +123,7 @@ async def collection_delete(self, request: web.Request):
         ds = await self.customizer.get_datasource()
         collection = ds.get_collection(body_params["collectionName"])
         caller = CallerSerializer.deserialize(body_params["caller"])
-        filter_ = FilterSerializer.deserialize(body_params["filter"], collection)
+        filter_ = FilterSerializer.deserialize(body_params["filter"], collection)  # type:ignore
 
         await collection.delete(caller, filter_)
         return web.Response(text="OK")
@@ -169,53 +133,51 @@ async def collection_aggregate(self, request: web.Request):
         ds = await self.customizer.get_datasource()
         collection = ds.get_collection(body_params["collectionName"])
         caller = CallerSerializer.deserialize(body_params["caller"])
-        filter_ = FilterSerializer.deserialize(body_params["filter"], collection)
+        filter_ = FilterSerializer.deserialize(body_params["filter"], collection)  # type:ignore
         aggregation = AggregationSerializer.deserialize(body_params["aggregation"])
 
         records = await collection.aggregate(caller, filter_, aggregation)
-        # records = [RecordSerializer.serialize(record) for record in records]
-        return web.Response(text=aes_encrypt(json.dumps(records), self.aes_key, self.aes_iv))
+        return web.json_response(records)
 
     async def collection_get_form(self, request: web.Request):
         body_params = await request.text()
-        body_params = aes_decrypt(body_params, self.aes_key, self.aes_iv)
         body_params = json.loads(body_params)
 
         ds = await self.customizer.get_datasource()
         collection = ds.get_collection(body_params["collectionName"])
 
-        caller = CallerSerializer.deserialize(body_params["caller"]) if body_params["caller"] else None
+        caller = CallerSerializer.deserialize(body_params["caller"])
         action_name = body_params["actionName"]
-        filter_ = FilterSerializer.deserialize(body_params["filter"], collection) if body_params["filter"] else None
+        if body_params["filter"]:
+            filter_ = FilterSerializer.deserialize(body_params["filter"], collection)  # type:ignore
+        else:
+            filter_ = None
         data = ActionFormValuesSerializer.deserialize(body_params["data"])
         meta = body_params["meta"]
 
         form = await collection.get_form(caller, action_name, data, filter_, meta)
-        return web.Response(
-            text=aes_encrypt(json.dumps(ActionFormSerializer.serialize(form)), self.aes_key, self.aes_iv)
-        )
+        return web.json_response(ActionFormSerializer.serialize(form))
 
     async def collection_execute(self, request: web.Request):
         body_params = await request.text()
-        body_params = aes_decrypt(body_params, self.aes_key, self.aes_iv)
         body_params = json.loads(body_params)
 
         ds = await self.customizer.get_datasource()
         collection = ds.get_collection(body_params["collectionName"])
 
-        caller = CallerSerializer.deserialize(body_params["caller"]) if body_params["caller"] else None
+        caller = CallerSerializer.deserialize(body_params["caller"])
         action_name = body_params["actionName"]
-        filter_ = FilterSerializer.deserialize(body_params["filter"], collection) if body_params["filter"] else None
+        if body_params["filter"]:
+            filter_ = FilterSerializer.deserialize(body_params["filter"], collection)  # type:ignore
+        else:
+            filter_ = None
         data = ActionFormValuesSerializer.deserialize(body_params["data"])
 
         result = await collection.execute(caller, action_name, data, filter_)
-        return web.Response(
-            text=aes_encrypt(json.dumps(ActionResultSerializer.serialize(result)), self.aes_key, self.aes_iv)
-        )
+        return web.json_response(ActionResultSerializer.serialize(result))  # type:ignore
 
     async def collection_render_chart(self, request: web.Request):
         body_params = await request.text()
-        body_params = aes_decrypt(body_params, self.aes_key, self.aes_iv)
         body_params = json.loads(body_params)
 
         ds = await self.customizer.get_datasource()
@@ -244,11 +206,10 @@ async def collection_render_chart(self, request: web.Request):
                 ret.append(value)
 
         result = await collection.render_chart(caller, name, record_id)
-        return web.Response(text=aes_encrypt(json.dumps(result), self.aes_key, self.aes_iv))
+        return web.json_response(result)
 
     async def render_chart(self, request: web.Request):
         body_params = await request.text()
-        body_params = aes_decrypt(body_params, self.aes_key, self.aes_iv)
         body_params = json.loads(body_params)
 
         ds = await self.customizer.get_datasource()
@@ -257,11 +218,10 @@ async def render_chart(self, request: web.Request):
         name = body_params["name"]
 
         result = await ds.render_chart(caller, name)
-        return web.Response(text=aes_encrypt(json.dumps(result), self.aes_key, self.aes_iv))
+        return web.json_response(result)
 
     async def native_query(self, request: web.Request):
         body_params = await request.text()
-        body_params = aes_decrypt(body_params, self.aes_key, self.aes_iv)
         body_params = json.loads(body_params)
 
         ds = await self.customizer.get_datasource()
@@ -270,7 +230,7 @@ async def native_query(self, request: web.Request):
         parameters = body_params["parameters"]
 
         result = await ds.execute_native_query(connection_name, native_query, parameters)
-        return web.Response(text=aes_encrypt(json.dumps(result), self.aes_key, self.aes_iv))
+        return web.json_response(result)
 
     def start(self):
         web.run_app(self.app, host=self.listen_addr, port=int(self.listen_port))