From 9ab7820b5c128511df8a5865314fc897eae80bd9 Mon Sep 17 00:00:00 2001 From: "jake@prefect.io" Date: Fri, 6 Sep 2024 17:46:51 -0400 Subject: [PATCH] go back to helper functions --- src/prefect/context.py | 6 +++--- src/prefect/events/clients.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/prefect/context.py b/src/prefect/context.py index 20b89dd3e41e..c991557bbe71 100644 --- a/src/prefect/context.py +++ b/src/prefect/context.py @@ -39,8 +39,8 @@ from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_client from prefect.client.schemas import FlowRun, TaskRun from prefect.events.clients import ( + araise_for_events_connection_error, raise_for_events_connection_error, - sync_raise_for_events_connection_error, ) from prefect.events.worker import EventsWorker from prefect.exceptions import MissingContextError @@ -215,7 +215,7 @@ def __enter__(self): if self._context_stack == 1: self.client.__enter__() self.client.raise_for_api_version_mismatch() - sync_raise_for_events_connection_error() + raise_for_events_connection_error() return super().__enter__() else: return self @@ -274,7 +274,7 @@ async def __aenter__(self): if self._context_stack == 1: await self.client.__aenter__() await self.client.raise_for_api_version_mismatch() - await raise_for_events_connection_error() + await araise_for_events_connection_error() return super().__enter__() else: return self diff --git a/src/prefect/events/clients.py b/src/prefect/events/clients.py index 0233ebc4fe62..eb81be43a61f 100644 --- a/src/prefect/events/clients.py +++ b/src/prefect/events/clients.py @@ -164,7 +164,7 @@ def _get_socket_url_and_headers(): return events_in_socket_from_api_url(api_url), headers -def sync_raise_for_events_connection_error(): +def raise_for_events_connection_error(): socket_url, headers = _get_socket_url_and_headers() try: @@ -177,7 +177,7 @@ def sync_raise_for_events_connection_error(): ) from e -async def raise_for_events_connection_error(): +async def araise_for_events_connection_error(): socket_url, headers = _get_socket_url_and_headers() try: