APP-5362 : feat: integrate uvloop for improved asyncio performance

application_sdk/activities/common/utils.py

@@ -3,12 +3,13 @@
 from functools import wraps
 from typing import Any, Awaitable, Callable, Optional, TypeVar, cast
 
+import uvloop
 from temporalio import activity
 
 from application_sdk.common.logger_adaptors import get_logger
 
 logger = get_logger(__name__)
-
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 F = TypeVar("F", bound=Callable[..., Awaitable[Any]])
 

application_sdk/application/fastapi/routers/server.py

@@ -17,13 +17,14 @@
 import uuid
 
 import psutil
+import uvloop
 from fastapi import APIRouter, status
 from fastapi.responses import JSONResponse
 
 from application_sdk.common.logger_adaptors import get_logger
 
 logger = get_logger(__name__)
-
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 router = APIRouter(
     prefix="/server",
     tags=["server"],

application_sdk/clients/sql.py

@@ -11,6 +11,7 @@
 from enum import Enum
 from typing import Any, Dict, List
 
+import uvloop
 from sqlalchemy import create_engine, text
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
 from temporalio import activity
@@ -19,6 +20,7 @@
 from application_sdk.common.logger_adaptors import get_logger
 
 activity.logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 class SQLConstants(Enum):

application_sdk/decorators/__init__.py

@@ -5,13 +5,14 @@
 from typing import Any, Callable, Dict, Iterator, Optional, Union
 
 import pandas as pd
+import uvloop
 
 from application_sdk.activities import ActivitiesState
 from application_sdk.common.logger_adaptors import get_logger
 from application_sdk.inputs import Input
 
 logger = get_logger(__name__)
-
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 executor = ThreadPoolExecutor()
 

application_sdk/handlers/sql.py

@@ -4,6 +4,7 @@
 from typing import Any, Dict, List, Optional, Set, Tuple
 
 import pandas as pd
+import uvloop
 
 from application_sdk.application.fastapi.models import MetadataType
 from application_sdk.clients.sql import SQLClient
@@ -13,6 +14,7 @@
 from application_sdk.inputs.sql_query import SQLQueryInput
 
 logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 class SQLConstants(Enum):

application_sdk/inputs/sql_query.py

@@ -3,6 +3,7 @@
 from typing import Any, Dict, Iterator, Optional, Union
 
 import pandas as pd
+import uvloop
 from sqlalchemy import text
 from sqlalchemy.engine import Engine
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
@@ -15,6 +16,7 @@
 from application_sdk.inputs import Input
 
 logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 def _get_sql_query(

application_sdk/outputs/parquet.py

@@ -6,12 +6,14 @@
 import orjson
 import pyarrow as pa
 import pyarrow.parquet as pq
+import uvloop
 from temporalio import activity
 
 from application_sdk.common.logger_adaptors import get_logger
 from application_sdk.outputs.objectstore import ObjectStoreOutput
 
 activity.logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 class ChunkedObjectStoreWriterInterface(ABC):

application_sdk/test_utils/workflow_monitoring.py

@@ -7,11 +7,13 @@
 import time
 from typing import Any, Optional
 
+import uvloop
 from temporalio.client import WorkflowExecutionStatus, WorkflowHandle
 
 from application_sdk.clients.temporal import TemporalClient
 
 logger = logging.getLogger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 async def monitor_workflow_execution_and_write_status(

application_sdk/worker.py

@@ -8,12 +8,14 @@
 import threading
 from typing import Any, List, Sequence
 
+import uvloop
 from temporalio.types import CallableType
 
 from application_sdk.clients.temporal import TemporalClient
 from application_sdk.common.logger_adaptors import get_logger
 
 logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 class Worker:

application_sdk/workflows/metadata_extraction/sql.py

@@ -8,6 +8,7 @@
 from datetime import timedelta
 from typing import Any, Callable, Coroutine, Dict, List, Sequence, Type
 
+import uvloop
 from temporalio import workflow
 from temporalio.common import RetryPolicy
 
@@ -21,6 +22,7 @@
 from application_sdk.workflows.metadata_extraction import MetadataExtractionWorkflow
 
 workflow.logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 @workflow.defn

application_sdk/workflows/query_extraction/sql.py

@@ -8,6 +8,7 @@
 from datetime import timedelta
 from typing import Any, Callable, Coroutine, Dict, List, Sequence, Type
 
+import uvloop
 from temporalio import workflow
 from temporalio.common import RetryPolicy
 
@@ -20,6 +21,7 @@
 from application_sdk.workflows.query_extraction import QueryExtractionWorkflow
 
 logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 @workflow.defn

examples/application_custom_fastapi.py

@@ -2,12 +2,15 @@
 import uuid
 from typing import Any, Dict
 
+import uvloop
 from fastapi import APIRouter
 
 from application_sdk.application.fastapi import FastAPIApplication, HttpWorkflowTrigger
 from application_sdk.handlers import HandlerInterface
 from application_sdk.workflows import WorkflowInterface
 
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+
 
 class CustomHandler(HandlerInterface):
     async def load(self, **kwargs: Any) -> None:

examples/application_fastapi.py

@@ -1,10 +1,14 @@
 import asyncio
 from typing import Any, Dict, List
 
+import uvloop
+
 from application_sdk.application.fastapi import FastAPIApplication, HttpWorkflowTrigger
 from application_sdk.handlers.sql import SQLHandler
 from application_sdk.workflows import WorkflowInterface
 
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+
 
 class SampleSQLHandler(SQLHandler):
     async def prepare(self, credentials: Dict[str, Any], **kwargs) -> None:

examples/application_hello_world.py

@@ -1,6 +1,7 @@
 import asyncio
 from typing import Any, Callable, Dict, Sequence
 
+import uvloop
 from temporalio import workflow
 
 from application_sdk.activities import ActivitiesInterface
@@ -12,6 +13,7 @@
 APPLICATION_NAME = "hello-world"
 
 logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 @workflow.defn

examples/application_sql.py

@@ -31,6 +31,8 @@
 from typing import Any, Dict
 from urllib.parse import quote_plus
 
+import uvloop
+
 from application_sdk.activities.metadata_extraction.sql import (
     SQLMetadataExtractionActivities,
 )
@@ -46,6 +48,7 @@
 APPLICATION_NAME = "postgres"
 
 logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 class PostgreSQLClient(AsyncSQLClient):

examples/application_sql_miner.py

@@ -22,6 +22,8 @@
 from typing import Any, Dict
 from urllib.parse import quote_plus
 
+import uvloop
+
 from application_sdk.activities.query_extraction.sql import SQLQueryExtractionActivities
 from application_sdk.clients.sql import SQLClient
 from application_sdk.clients.temporal import TemporalClient
@@ -31,7 +33,7 @@
 from application_sdk.workflows.query_extraction.sql import SQLQueryExtractionWorkflow
 
 logger = get_logger(__name__)
-
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 APPLICATION_NAME = "snowflake"
 
 

examples/application_sql_with_custom_transformer.py

@@ -32,6 +32,7 @@
 from typing import Any, Dict
 from urllib.parse import quote_plus
 
+import uvloop
 from pyatlan.model.assets import Database
 
 from application_sdk.activities.metadata_extraction.sql import (
@@ -52,6 +53,7 @@
 DATABASE_DIALECT = "postgresql"
 
 logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 class PostgreSQLClient(AsyncSQLClient):

examples/application_subscriber.py

@@ -2,6 +2,7 @@
 from datetime import timedelta
 from typing import Any, Callable, Dict, List
 
+import uvloop
 from temporalio import activity, workflow
 
 from application_sdk.activities import ActivitiesInterface
@@ -21,6 +22,7 @@
 from application_sdk.workflows import WorkflowInterface
 
 logger = get_logger(__name__)
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 class SampleActivities(ActivitiesInterface):

examples/run_examples.py

@@ -4,6 +4,8 @@
 
 import asyncio
 
+import uvloop
+
 from application_sdk.clients.temporal import TemporalClient
 from application_sdk.test_utils.workflow_monitoring import run_and_monitor_workflow
 from examples.application_hello_world import application_hello_world
@@ -13,6 +15,8 @@
     application_sql_with_custom_transformer,
 )
 
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+
 
 async def main():
     temporal_client = TemporalClient()