feat: add poc flow for async requests

pipelines/flows.py

@@ -5,5 +5,6 @@
 from pipelines.deteccao_alagamento_cameras import *  # noqa
 from pipelines.exemplo import *  # noqa
 from pipelines.lgpd import *  # noqa
+from pipelines.sandbox import *  # noqa
 from pipelines.stress import *  # noqa
 from pipelines.templates import *  # noqa

pipelines/sandbox/__init__.py

@@ -0,0 +1,2 @@
+# -*- coding: utf-8 -*-
+from pipelines.sandbox.async_requests.flows import *  # noqa

pipelines/sandbox/async_requests/flows.py

@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+from prefect import Flow, Parameter
+from prefect.run_configs import KubernetesRun
+from prefect.storage import GCS
+from prefeitura_rio.pipelines_utils.state_handlers import handler_inject_bd_credentials
+
+from pipelines.constants import constants
+from pipelines.sandbox.async_requests.tasks import get_random_cat_pictures
+
+with Flow(
+    name="Sandbox: Test async requests",
+    state_handlers=[handler_inject_bd_credentials],
+) as sandbox__async_requests_flow:
+    # Parameters
+    batch_size = Parameter("batch_size", default=50)
+    n_times = Parameter("n_times", default=500)
+
+    get_random_cat_pictures(n_times=n_times, batch_size=batch_size)
+
+sandbox__async_requests_flow.storage = GCS(constants.GCS_FLOWS_BUCKET.value)
+sandbox__async_requests_flow.run_config = KubernetesRun(
+    image=constants.DOCKER_IMAGE.value,
+    labels=["escritoriodedados"],
+)

pipelines/sandbox/async_requests/tasks.py

@@ -0,0 +1,27 @@
+# -*- coding: utf-8 -*-
+import asyncio
+
+from httpx import AsyncClient, AsyncHTTPTransport
+from prefect import task
+from prefeitura_rio.pipelines_utils.logging import log
+
+
+@task
+def get_random_cat_pictures(n_times: int, batch_size: int):
+    async def get_single_cat_picture():
+        transport = AsyncHTTPTransport(retries=3)
+        async with AsyncClient(transport=transport, timeout=None) as client:
+            response = await client.get("https://cataas.com/cat")
+            return response.headers["Content-Type"]
+
+    async def main():
+        log("Starting to fetch cat pictures")
+        awaitables = [get_single_cat_picture() for _ in range(n_times)]
+        awaitables = [
+            awaitables[i : i + batch_size] for i in range(0, len(awaitables), batch_size)  # noqa
+        ]
+        for awaitables_batch in awaitables:
+            responses = await asyncio.gather(*awaitables_batch)
+            log(f"Fetched {len(responses)} cat pictures")
+
+    asyncio.run(main())

pyproject.toml

@@ -24,6 +24,7 @@ pillow = "^10.1.0"
 shapely = "^2.0.2"
 h3 = "^3.7.6"
 google-cloud-asset = "^3.24.1"
+httpx = "^0.27.0"
 
 [tool.poetry.group.dev]
 optional = true