prefect-client 3.1.13__py3-none-any.whl → 3.1.15__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
prefect/blocks/core.py CHANGED
@@ -669,7 +669,7 @@ class Block(BaseModel, ABC):
669
669
  module_str = ".".join(qualified_name.split(".")[:-1])
670
670
  origin = cls.__pydantic_generic_metadata__.get("origin") or cls
671
671
  class_name = origin.__name__
672
- block_variable_name = f'{cls.get_block_type_slug().replace("-", "_")}_block'
672
+ block_variable_name = f"{cls.get_block_type_slug().replace('-', '_')}_block"
673
673
 
674
674
  return dedent(
675
675
  f"""\
@@ -570,6 +570,10 @@ class MattermostWebhook(AbstractAppriseNotificationBlock):
570
570
  description="The hostname of your Mattermost server.",
571
571
  examples=["Mattermost.example.com"],
572
572
  )
573
+ secure: bool = Field(
574
+ default=False,
575
+ description="Whether to use secure https connection.",
576
+ )
573
577
 
574
578
  token: SecretStr = Field(
575
579
  default=...,
@@ -621,6 +625,7 @@ class MattermostWebhook(AbstractAppriseNotificationBlock):
621
625
  channels=self.channels,
622
626
  include_image=self.include_image,
623
627
  port=self.port,
628
+ secure=self.secure,
624
629
  ).url() # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] incomplete type hints in apprise
625
630
  )
626
631
  self._start_apprise_client(url)
prefect/cache_policies.py CHANGED
@@ -80,10 +80,12 @@ class CachePolicy:
80
80
  raise NotImplementedError
81
81
 
82
82
  def __sub__(self, other: str) -> "CachePolicy":
83
+ "No-op for all policies except Inputs and Compound"
84
+
85
+ # for interface compatibility
83
86
  if not isinstance(other, str): # type: ignore[reportUnnecessaryIsInstance]
84
87
  raise TypeError("Can only subtract strings from key policies.")
85
- new = Inputs(exclude=[other])
86
- return CompoundCachePolicy(policies=[self, new])
88
+ return self
87
89
 
88
90
  def __add__(self, other: "CachePolicy") -> "CachePolicy":
89
91
  # adding _None is a no-op
@@ -157,6 +159,23 @@ class CompoundCachePolicy(CachePolicy):
157
159
 
158
160
  policies: list[CachePolicy] = field(default_factory=list)
159
161
 
162
+ def __post_init__(self) -> None:
163
+ # flatten any CompoundCachePolicies
164
+ self.policies = [
165
+ policy
166
+ for p in self.policies
167
+ for policy in (p.policies if isinstance(p, CompoundCachePolicy) else [p])
168
+ ]
169
+
170
+ # deduplicate any Inputs policies
171
+ inputs_policies = [p for p in self.policies if isinstance(p, Inputs)]
172
+ self.policies = [p for p in self.policies if not isinstance(p, Inputs)]
173
+ if inputs_policies:
174
+ all_excludes: set[str] = set()
175
+ for inputs_policy in inputs_policies:
176
+ all_excludes.update(inputs_policy.exclude)
177
+ self.policies.append(Inputs(exclude=sorted(all_excludes)))
178
+
160
179
  def compute_key(
161
180
  self,
162
181
  task_ctx: TaskRunContext,
@@ -178,6 +197,35 @@ class CompoundCachePolicy(CachePolicy):
178
197
  return None
179
198
  return hash_objects(*keys, raise_on_failure=True)
180
199
 
200
+ def __add__(self, other: "CachePolicy") -> "CachePolicy":
201
+ # Call the superclass add method to handle validation
202
+ super().__add__(other)
203
+
204
+ if isinstance(other, CompoundCachePolicy):
205
+ policies = [*self.policies, *other.policies]
206
+ else:
207
+ policies = [*self.policies, other]
208
+
209
+ return CompoundCachePolicy(
210
+ policies=policies,
211
+ key_storage=self.key_storage or other.key_storage,
212
+ isolation_level=self.isolation_level or other.isolation_level,
213
+ lock_manager=self.lock_manager or other.lock_manager,
214
+ )
215
+
216
+ def __sub__(self, other: str) -> "CachePolicy":
217
+ if not isinstance(other, str): # type: ignore[reportUnnecessaryIsInstance]
218
+ raise TypeError("Can only subtract strings from key policies.")
219
+
220
+ inputs_policies = [p for p in self.policies if isinstance(p, Inputs)]
221
+
222
+ if inputs_policies:
223
+ new = Inputs(exclude=[other])
224
+ return CompoundCachePolicy(policies=[*self.policies, new])
225
+ else:
226
+ # no dependency on inputs already
227
+ return self
228
+
181
229
 
182
230
  @dataclass
183
231
  class _None(CachePolicy):
@@ -45,6 +45,8 @@ class AutomationClient(BaseClient):
45
45
  return Automation.model_validate_list(response.json())
46
46
 
47
47
  def find_automation(self, id_or_name: "str | UUID") -> "Automation | None":
48
+ from uuid import UUID
49
+
48
50
  if isinstance(id_or_name, str):
49
51
  name = id_or_name
50
52
  try:
@@ -202,6 +204,8 @@ class AutomationAsyncClient(BaseAsyncClient):
202
204
  return Automation.model_validate_list(response.json())
203
205
 
204
206
  async def find_automation(self, id_or_name: "str | UUID") -> "Automation | None":
207
+ from uuid import UUID
208
+
205
209
  if isinstance(id_or_name, str):
206
210
  name = id_or_name
207
211
  try:
@@ -1031,8 +1031,8 @@ class DeploymentAsyncClient(BaseAsyncClient):
1031
1031
  deployment_ids: list["UUID"],
1032
1032
  scheduled_before: "datetime.datetime | None" = None,
1033
1033
  limit: int | None = None,
1034
- ) -> list["FlowRunResponse"]:
1035
- from prefect.client.schemas.responses import FlowRunResponse
1034
+ ) -> list["FlowRun"]:
1035
+ from prefect.client.schemas.objects import FlowRun
1036
1036
 
1037
1037
  body: dict[str, Any] = dict(deployment_ids=[str(id) for id in deployment_ids])
1038
1038
  if scheduled_before:
@@ -1046,7 +1046,7 @@ class DeploymentAsyncClient(BaseAsyncClient):
1046
1046
  json=body,
1047
1047
  )
1048
1048
 
1049
- return FlowRunResponse.model_validate_list(response.json())
1049
+ return FlowRun.model_validate_list(response.json())
1050
1050
 
1051
1051
  async def create_flow_run_from_deployment(
1052
1052
  self,
@@ -1,5 +1,5 @@
1
1
  from copy import deepcopy
2
- from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
2
+ from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
3
3
  from uuid import UUID, uuid4
4
4
 
5
5
  import jsonschema
@@ -93,6 +93,16 @@ class DeploymentScheduleCreate(ActionBaseModel):
93
93
  description="The maximum number of scheduled runs for the schedule.",
94
94
  )
95
95
 
96
+ @field_validator("active", mode="wrap")
97
+ @classmethod
98
+ def validate_active(cls, v: Any, handler: Callable[[Any], Any]) -> bool:
99
+ try:
100
+ return handler(v)
101
+ except Exception:
102
+ raise ValueError(
103
+ f"active must be able to be parsed as a boolean, got {v!r} of type {type(v)}"
104
+ )
105
+
96
106
  @field_validator("max_scheduled_runs")
97
107
  @classmethod
98
108
  def validate_max_scheduled_runs(cls, v: Optional[int]) -> Optional[int]:
@@ -145,7 +145,7 @@ class CronSchedule(PrefectBaseModel):
145
145
 
146
146
  @field_validator("timezone")
147
147
  @classmethod
148
- def valid_timezone(cls, v: Optional[str]) -> str:
148
+ def valid_timezone(cls, v: Optional[str]) -> Optional[str]:
149
149
  return default_timezone(v)
150
150
 
151
151
  @field_validator("cron")
@@ -276,7 +276,7 @@ class RRuleSchedule(PrefectBaseModel):
276
276
  kwargs.update(
277
277
  until=until.replace(tzinfo=timezone),
278
278
  )
279
- return rrule.replace(**kwargs)
279
+ return rrule.replace(**kwargs) # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] missing type hints
280
280
 
281
281
  # update rrules
282
282
  localized_rrules: list[dateutil.rrule.rrule] = []
@@ -286,7 +286,7 @@ class RRuleSchedule(PrefectBaseModel):
286
286
  kwargs: dict[str, Any] = dict(dtstart=dtstart.replace(tzinfo=timezone))
287
287
  if until := _rrule_dt(rr, "_until"):
288
288
  kwargs.update(until=until.replace(tzinfo=timezone))
289
- localized_rrules.append(rr.replace(**kwargs))
289
+ localized_rrules.append(rr.replace(**kwargs)) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] missing type hints
290
290
  setattr(rrule, "_rrule", localized_rrules)
291
291
 
292
292
  # update exrules
@@ -297,7 +297,7 @@ class RRuleSchedule(PrefectBaseModel):
297
297
  kwargs = dict(dtstart=dtstart.replace(tzinfo=timezone))
298
298
  if until := _rrule_dt(exr, "_until"):
299
299
  kwargs.update(until=until.replace(tzinfo=timezone))
300
- localized_exrules.append(exr.replace(**kwargs))
300
+ localized_exrules.append(exr.replace(**kwargs)) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] missing type hints
301
301
  setattr(rrule, "_exrule", localized_exrules)
302
302
 
303
303
  # update rdates
prefect/context.py CHANGED
@@ -337,10 +337,15 @@ class EngineContext(RunContext):
337
337
  flow: The flow instance associated with the run
338
338
  flow_run: The API metadata for the flow run
339
339
  task_runner: The task runner instance being used for the flow run
340
- task_run_futures: A list of futures for task runs submitted within this flow run
341
- task_run_states: A list of states for task runs created within this flow run
342
340
  task_run_results: A mapping of result ids to task run states for this flow run
343
- flow_run_states: A list of states for flow runs created within this flow run
341
+ log_prints: Whether to log print statements from the flow run
342
+ parameters: The parameters passed to the flow run
343
+ detached: Flag indicating if context has been serialized and sent to remote infrastructure
344
+ result_store: The result store used to persist results
345
+ persist_result: Whether to persist the flow run result
346
+ task_run_dynamic_keys: Counter for task calls allowing unique keys
347
+ observed_flow_pauses: Counter for flow pauses
348
+ events: Events worker to emit events
344
349
  """
345
350
 
346
351
  flow: Optional["Flow[Any, Any]"] = None
@@ -373,7 +378,7 @@ class EngineContext(RunContext):
373
378
  __var__: ClassVar[ContextVar[Self]] = ContextVar("flow_run")
374
379
 
375
380
  def serialize(self: Self, include_secrets: bool = True) -> dict[str, Any]:
376
- return self.model_dump(
381
+ serialized = self.model_dump(
377
382
  include={
378
383
  "flow_run",
379
384
  "flow",
@@ -381,13 +386,18 @@ class EngineContext(RunContext):
381
386
  "log_prints",
382
387
  "start_time",
383
388
  "input_keyset",
384
- "result_store",
385
389
  "persist_result",
386
390
  },
387
391
  exclude_unset=True,
388
- serialize_as_any=True,
389
392
  context={"include_secrets": include_secrets},
390
393
  )
394
+ if self.result_store:
395
+ serialized["result_store"] = self.result_store.model_dump(
396
+ serialize_as_any=True,
397
+ exclude_unset=True,
398
+ context={"include_secrets": include_secrets},
399
+ )
400
+ return serialized
391
401
 
392
402
 
393
403
  FlowRunContext = EngineContext # for backwards compatibility
@@ -10,12 +10,11 @@ from __future__ import annotations
10
10
  import os
11
11
  from copy import deepcopy
12
12
  from pathlib import Path
13
- from typing import Any, Dict, List, Optional, cast
13
+ from typing import Any, Dict, List, Optional
14
14
 
15
15
  import yaml
16
16
  from ruamel.yaml import YAML
17
17
 
18
- from prefect.client.schemas.actions import DeploymentScheduleCreate
19
18
  from prefect.client.schemas.objects import ConcurrencyLimitStrategy
20
19
  from prefect.client.schemas.schedules import IntervalSchedule
21
20
  from prefect.utilities._git import get_git_branch, get_git_remote_origin_url
@@ -207,8 +206,8 @@ def initialize_project(
207
206
 
208
207
 
209
208
  def _format_deployment_for_saving_to_prefect_file(
210
- deployment: Dict,
211
- ) -> Dict:
209
+ deployment: dict[str, Any],
210
+ ) -> dict[str, Any]:
212
211
  """
213
212
  Formats a deployment into a templated deploy config for saving to prefect.yaml.
214
213
 
@@ -227,10 +226,8 @@ def _format_deployment_for_saving_to_prefect_file(
227
226
  deployment.pop("flow_name", None)
228
227
 
229
228
  if deployment.get("schedules"):
230
- schedules = []
231
- for deployment_schedule in cast(
232
- List[DeploymentScheduleCreate], deployment["schedules"]
233
- ):
229
+ schedules: list[dict[str, Any]] = []
230
+ for deployment_schedule in deployment["schedules"]:
234
231
  if isinstance(deployment_schedule.schedule, IntervalSchedule):
235
232
  schedule_config = _interval_schedule_to_dict(
236
233
  deployment_schedule.schedule
@@ -257,7 +254,7 @@ def _format_deployment_for_saving_to_prefect_file(
257
254
  return deployment
258
255
 
259
256
 
260
- def _interval_schedule_to_dict(schedule: IntervalSchedule) -> Dict:
257
+ def _interval_schedule_to_dict(schedule: IntervalSchedule) -> dict[str, Any]:
261
258
  """
262
259
  Converts an IntervalSchedule to a dictionary.
263
260
 
@@ -265,7 +262,7 @@ def _interval_schedule_to_dict(schedule: IntervalSchedule) -> Dict:
265
262
  - schedule (IntervalSchedule): the schedule to convert
266
263
 
267
264
  Returns:
268
- - Dict: the schedule as a dictionary
265
+ - dict[str, Any]: the schedule as a dictionary
269
266
  """
270
267
  schedule_config = schedule.model_dump()
271
268
  schedule_config["interval"] = schedule_config["interval"].total_seconds()
@@ -275,12 +272,12 @@ def _interval_schedule_to_dict(schedule: IntervalSchedule) -> Dict:
275
272
 
276
273
 
277
274
  def _save_deployment_to_prefect_file(
278
- deployment: Dict,
279
- build_steps: Optional[List[Dict]] = None,
280
- push_steps: Optional[List[Dict]] = None,
281
- pull_steps: Optional[List[Dict]] = None,
282
- triggers: Optional[List[Dict]] = None,
283
- sla: Optional[list[dict]] = None,
275
+ deployment: dict[str, Any],
276
+ build_steps: list[dict[str, Any]] | None = None,
277
+ push_steps: list[dict[str, Any]] | None = None,
278
+ pull_steps: list[dict[str, Any]] | None = None,
279
+ triggers: list[dict[str, Any]] | None = None,
280
+ sla: list[dict[str, Any]] | None = None,
284
281
  prefect_file: Path = Path("prefect.yaml"),
285
282
  ):
286
283
  """
@@ -49,6 +49,7 @@ from rich.progress import Progress, SpinnerColumn, TextColumn, track
49
49
  from rich.table import Table
50
50
 
51
51
  from prefect._experimental.sla.objects import SlaTypes
52
+ from prefect._internal.compatibility.async_dispatch import async_dispatch
52
53
  from prefect._internal.concurrency.api import create_call, from_async
53
54
  from prefect._internal.schemas.validators import (
54
55
  reconcile_paused_deployment,
@@ -82,7 +83,7 @@ from prefect.settings import (
82
83
  )
83
84
  from prefect.types import ListOfNonEmptyStrings
84
85
  from prefect.types.entrypoint import EntrypointType
85
- from prefect.utilities.asyncutils import sync_compatible
86
+ from prefect.utilities.asyncutils import run_coro_as_sync, sync_compatible
86
87
  from prefect.utilities.callables import ParameterSchema, parameter_schema
87
88
  from prefect.utilities.collections import get_from_dict, isiterable
88
89
  from prefect.utilities.dockerutils import (
@@ -717,8 +718,7 @@ class RunnerDeployment(BaseModel):
717
718
  return deployment
718
719
 
719
720
  @classmethod
720
- @sync_compatible
721
- async def from_storage(
721
+ async def afrom_storage(
722
722
  cls,
723
723
  storage: RunnerStorage,
724
724
  entrypoint: str,
@@ -742,7 +742,7 @@ class RunnerDeployment(BaseModel):
742
742
  work_queue_name: Optional[str] = None,
743
743
  job_variables: Optional[dict[str, Any]] = None,
744
744
  _sla: Optional[Union[SlaTypes, list[SlaTypes]]] = None, # experimental
745
- ):
745
+ ) -> "RunnerDeployment":
746
746
  """
747
747
  Create a RunnerDeployment from a flow located at a given entrypoint and stored in a
748
748
  local storage location.
@@ -831,6 +831,119 @@ class RunnerDeployment(BaseModel):
831
831
 
832
832
  return deployment
833
833
 
834
+ @classmethod
835
+ @async_dispatch(afrom_storage)
836
+ def from_storage(
837
+ cls,
838
+ storage: RunnerStorage,
839
+ entrypoint: str,
840
+ name: str,
841
+ flow_name: Optional[str] = None,
842
+ interval: Optional[
843
+ Union[Iterable[Union[int, float, timedelta]], int, float, timedelta]
844
+ ] = None,
845
+ cron: Optional[Union[Iterable[str], str]] = None,
846
+ rrule: Optional[Union[Iterable[str], str]] = None,
847
+ paused: Optional[bool] = None,
848
+ schedules: Optional["FlexibleScheduleList"] = None,
849
+ concurrency_limit: Optional[Union[int, ConcurrencyLimitConfig, None]] = None,
850
+ parameters: Optional[dict[str, Any]] = None,
851
+ triggers: Optional[List[Union[DeploymentTriggerTypes, TriggerTypes]]] = None,
852
+ description: Optional[str] = None,
853
+ tags: Optional[List[str]] = None,
854
+ version: Optional[str] = None,
855
+ enforce_parameter_schema: bool = True,
856
+ work_pool_name: Optional[str] = None,
857
+ work_queue_name: Optional[str] = None,
858
+ job_variables: Optional[dict[str, Any]] = None,
859
+ _sla: Optional[Union[SlaTypes, list[SlaTypes]]] = None, # experimental
860
+ ) -> "RunnerDeployment":
861
+ """
862
+ Create a RunnerDeployment from a flow located at a given entrypoint and stored in a
863
+ local storage location.
864
+
865
+ Args:
866
+ entrypoint: The path to a file containing a flow and the name of the flow function in
867
+ the format `./path/to/file.py:flow_func_name`.
868
+ name: A name for the deployment
869
+ flow_name: The name of the flow to deploy
870
+ storage: A storage object to use for retrieving flow code. If not provided, a
871
+ URL must be provided.
872
+ interval: An interval on which to execute the current flow. Accepts either a number
873
+ or a timedelta object. If a number is given, it will be interpreted as seconds.
874
+ cron: A cron schedule of when to execute runs of this flow.
875
+ rrule: An rrule schedule of when to execute runs of this flow.
876
+ triggers: A list of triggers that should kick of a run of this flow.
877
+ parameters: A dictionary of default parameter values to pass to runs of this flow.
878
+ description: A description for the created deployment. Defaults to the flow's
879
+ description if not provided.
880
+ tags: A list of tags to associate with the created deployment for organizational
881
+ purposes.
882
+ version: A version for the created deployment. Defaults to the flow's version.
883
+ enforce_parameter_schema: Whether or not the Prefect API should enforce the
884
+ parameter schema for this deployment.
885
+ work_pool_name: The name of the work pool to use for this deployment.
886
+ work_queue_name: The name of the work queue to use for this deployment's scheduled runs.
887
+ If not provided the default work queue for the work pool will be used.
888
+ job_variables: Settings used to override the values specified default base job template
889
+ of the chosen work pool. Refer to the base job template of the chosen work pool for
890
+ available settings.
891
+ _sla: (Experimental) SLA configuration for the deployment. May be removed or modified at any time. Currently only supported on Prefect Cloud.
892
+ """
893
+ from prefect.flows import load_flow_from_entrypoint
894
+
895
+ constructed_schedules = cls._construct_deployment_schedules(
896
+ interval=interval,
897
+ cron=cron,
898
+ rrule=rrule,
899
+ schedules=schedules,
900
+ )
901
+
902
+ if isinstance(concurrency_limit, ConcurrencyLimitConfig):
903
+ concurrency_options = {
904
+ "collision_strategy": concurrency_limit.collision_strategy
905
+ }
906
+ concurrency_limit = concurrency_limit.limit
907
+ else:
908
+ concurrency_options = None
909
+
910
+ job_variables = job_variables or {}
911
+
912
+ with tempfile.TemporaryDirectory() as tmpdir:
913
+ storage.set_base_path(Path(tmpdir))
914
+ run_coro_as_sync(storage.pull_code())
915
+
916
+ full_entrypoint = str(storage.destination / entrypoint)
917
+ flow = load_flow_from_entrypoint(full_entrypoint)
918
+
919
+ deployment = cls(
920
+ name=Path(name).stem,
921
+ flow_name=flow_name or flow.name,
922
+ schedules=constructed_schedules,
923
+ concurrency_limit=concurrency_limit,
924
+ concurrency_options=concurrency_options,
925
+ paused=paused,
926
+ tags=tags or [],
927
+ triggers=triggers or [],
928
+ parameters=parameters or {},
929
+ description=description,
930
+ version=version,
931
+ entrypoint=entrypoint,
932
+ enforce_parameter_schema=enforce_parameter_schema,
933
+ storage=storage,
934
+ work_pool_name=work_pool_name,
935
+ work_queue_name=work_queue_name,
936
+ job_variables=job_variables,
937
+ )
938
+ deployment._sla = _sla
939
+ deployment._path = str(storage.destination).replace(
940
+ tmpdir, "$STORAGE_BASE_PATH"
941
+ )
942
+
943
+ cls._set_defaults_from_flow(deployment, flow)
944
+
945
+ return deployment
946
+
834
947
 
835
948
  @sync_compatible
836
949
  async def deploy(
prefect/events/clients.py CHANGED
@@ -1,6 +1,7 @@
1
1
  import abc
2
2
  import asyncio
3
3
  import os
4
+ import ssl
4
5
  from types import TracebackType
5
6
  from typing import (
6
7
  TYPE_CHECKING,
@@ -37,6 +38,7 @@ from prefect.events import Event
37
38
  from prefect.logging import get_logger
38
39
  from prefect.settings import (
39
40
  PREFECT_API_KEY,
41
+ PREFECT_API_TLS_INSECURE_SKIP_VERIFY,
40
42
  PREFECT_API_URL,
41
43
  PREFECT_CLOUD_API_URL,
42
44
  PREFECT_DEBUG_MODE,
@@ -120,6 +122,13 @@ class WebsocketProxyConnect(Connect):
120
122
  self._host = host
121
123
  self._port = port
122
124
 
125
+ if PREFECT_API_TLS_INSECURE_SKIP_VERIFY:
126
+ # Create an unverified context for insecure connections
127
+ ctx = ssl.create_default_context()
128
+ ctx.check_hostname = False
129
+ ctx.verify_mode = ssl.CERT_NONE
130
+ self._kwargs.setdefault("ssl", ctx)
131
+
123
132
  async def _proxy_connect(self: Self) -> WebSocketClientProtocol:
124
133
  if self._proxy:
125
134
  sock = await self._proxy.connect(
@@ -348,16 +357,26 @@ class PrefectEventsClient(EventsClient):
348
357
  await self._connect.__aexit__(exc_type, exc_val, exc_tb)
349
358
  return await super().__aexit__(exc_type, exc_val, exc_tb)
350
359
 
360
+ def _log_debug(self, message: str, *args, **kwargs) -> None:
361
+ message = f"EventsClient(id={id(self)}): " + message
362
+ logger.debug(message, *args, **kwargs)
363
+
351
364
  async def _reconnect(self) -> None:
365
+ logger.debug("Reconnecting websocket connection.")
366
+
352
367
  if self._websocket:
353
368
  self._websocket = None
354
369
  await self._connect.__aexit__(None, None, None)
370
+ logger.debug("Cleared existing websocket connection.")
355
371
 
356
372
  try:
373
+ logger.debug("Opening websocket connection.")
357
374
  self._websocket = await self._connect.__aenter__()
358
375
  # make sure we have actually connected
376
+ logger.debug("Pinging to ensure websocket connected.")
359
377
  pong = await self._websocket.ping()
360
378
  await pong
379
+ logger.debug("Pong received. Websocket connected.")
361
380
  except Exception as e:
362
381
  # The client is frequently run in a background thread
363
382
  # so we log an additional warning to ensure
@@ -375,11 +394,13 @@ class PrefectEventsClient(EventsClient):
375
394
  raise
376
395
 
377
396
  events_to_resend = self._unconfirmed_events
397
+ logger.debug("Resending %s unconfirmed events.", len(events_to_resend))
378
398
  # Clear the unconfirmed events here, because they are going back through emit
379
399
  # and will be added again through the normal checkpointing process
380
400
  self._unconfirmed_events = []
381
401
  for event in events_to_resend:
382
402
  await self.emit(event)
403
+ logger.debug("Finished resending unconfirmed events.")
383
404
 
384
405
  async def _checkpoint(self, event: Event) -> None:
385
406
  assert self._websocket
@@ -387,11 +408,20 @@ class PrefectEventsClient(EventsClient):
387
408
  self._unconfirmed_events.append(event)
388
409
 
389
410
  unconfirmed_count = len(self._unconfirmed_events)
411
+
412
+ logger.debug(
413
+ "Added event id=%s to unconfirmed events list. "
414
+ "There are now %s unconfirmed events.",
415
+ event.id,
416
+ unconfirmed_count,
417
+ )
390
418
  if unconfirmed_count < self._checkpoint_every:
391
419
  return
392
420
 
421
+ logger.debug("Pinging to checkpoint unconfirmed events.")
393
422
  pong = await self._websocket.ping()
394
423
  await pong
424
+ self._log_debug("Pong received. Events checkpointed.")
395
425
 
396
426
  # once the pong returns, we know for sure that we've sent all the messages
397
427
  # we had enqueued prior to that. There could be more that came in after, so
@@ -401,7 +431,9 @@ class PrefectEventsClient(EventsClient):
401
431
  EVENT_WEBSOCKET_CHECKPOINTS.labels(self.client_name).inc()
402
432
 
403
433
  async def _emit(self, event: Event) -> None:
434
+ self._log_debug("Emitting event id=%s.", event.id)
404
435
  for i in range(self._reconnection_attempts + 1):
436
+ self._log_debug("Emit reconnection attempt %s.", i)
405
437
  try:
406
438
  # If we're here and the websocket is None, then we've had a failure in a
407
439
  # previous reconnection attempt.
@@ -410,14 +442,18 @@ class PrefectEventsClient(EventsClient):
410
442
  # from a ConnectionClosed, so reconnect now, resending any unconfirmed
411
443
  # events before we send this one.
412
444
  if not self._websocket or i > 0:
445
+ self._log_debug("Attempting websocket reconnection.")
413
446
  await self._reconnect()
414
447
  assert self._websocket
415
448
 
449
+ self._log_debug("Sending event id=%s.", event.id)
416
450
  await self._websocket.send(event.model_dump_json())
451
+ self._log_debug("Checkpointing event id=%s.", event.id)
417
452
  await self._checkpoint(event)
418
453
 
419
454
  return
420
455
  except ConnectionClosed:
456
+ self._log_debug("Got ConnectionClosed error.")
421
457
  if i == self._reconnection_attempts:
422
458
  # this was our final chance, raise the most recent error
423
459
  raise
@@ -426,6 +462,9 @@ class PrefectEventsClient(EventsClient):
426
462
  # let the first two attempts happen quickly in case this is just
427
463
  # a standard load balancer timeout, but after that, just take a
428
464
  # beat to let things come back around.
465
+ logger.debug(
466
+ "Sleeping for 1 second before next reconnection attempt."
467
+ )
429
468
  await asyncio.sleep(1)
430
469
 
431
470