prefect-client 2.19.3__py3-none-any.whl → 3.0.0rc1__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.
Files changed (239) hide show
  1. prefect/__init__.py +8 -56
  2. prefect/_internal/compatibility/deprecated.py +6 -115
  3. prefect/_internal/compatibility/experimental.py +4 -79
  4. prefect/_internal/concurrency/api.py +0 -34
  5. prefect/_internal/concurrency/calls.py +0 -6
  6. prefect/_internal/concurrency/cancellation.py +0 -3
  7. prefect/_internal/concurrency/event_loop.py +0 -20
  8. prefect/_internal/concurrency/inspection.py +3 -3
  9. prefect/_internal/concurrency/threads.py +35 -0
  10. prefect/_internal/concurrency/waiters.py +0 -28
  11. prefect/_internal/pydantic/__init__.py +0 -45
  12. prefect/_internal/pydantic/v1_schema.py +21 -22
  13. prefect/_internal/pydantic/v2_schema.py +0 -2
  14. prefect/_internal/pydantic/v2_validated_func.py +18 -23
  15. prefect/_internal/schemas/bases.py +44 -177
  16. prefect/_internal/schemas/fields.py +1 -43
  17. prefect/_internal/schemas/validators.py +60 -158
  18. prefect/artifacts.py +161 -14
  19. prefect/automations.py +39 -4
  20. prefect/blocks/abstract.py +1 -1
  21. prefect/blocks/core.py +268 -148
  22. prefect/blocks/fields.py +2 -57
  23. prefect/blocks/kubernetes.py +8 -12
  24. prefect/blocks/notifications.py +40 -20
  25. prefect/blocks/system.py +22 -11
  26. prefect/blocks/webhook.py +2 -9
  27. prefect/client/base.py +4 -4
  28. prefect/client/cloud.py +8 -13
  29. prefect/client/orchestration.py +347 -341
  30. prefect/client/schemas/actions.py +92 -86
  31. prefect/client/schemas/filters.py +20 -40
  32. prefect/client/schemas/objects.py +147 -145
  33. prefect/client/schemas/responses.py +16 -24
  34. prefect/client/schemas/schedules.py +47 -35
  35. prefect/client/subscriptions.py +2 -2
  36. prefect/client/utilities.py +5 -2
  37. prefect/concurrency/asyncio.py +3 -1
  38. prefect/concurrency/events.py +1 -1
  39. prefect/concurrency/services.py +6 -3
  40. prefect/context.py +195 -27
  41. prefect/deployments/__init__.py +5 -6
  42. prefect/deployments/base.py +7 -5
  43. prefect/deployments/flow_runs.py +185 -0
  44. prefect/deployments/runner.py +50 -45
  45. prefect/deployments/schedules.py +28 -23
  46. prefect/deployments/steps/__init__.py +0 -1
  47. prefect/deployments/steps/core.py +1 -0
  48. prefect/deployments/steps/pull.py +7 -21
  49. prefect/engine.py +12 -2422
  50. prefect/events/actions.py +17 -23
  51. prefect/events/cli/automations.py +19 -6
  52. prefect/events/clients.py +14 -37
  53. prefect/events/filters.py +14 -18
  54. prefect/events/related.py +2 -2
  55. prefect/events/schemas/__init__.py +0 -5
  56. prefect/events/schemas/automations.py +55 -46
  57. prefect/events/schemas/deployment_triggers.py +7 -197
  58. prefect/events/schemas/events.py +34 -65
  59. prefect/events/schemas/labelling.py +10 -14
  60. prefect/events/utilities.py +2 -3
  61. prefect/events/worker.py +2 -3
  62. prefect/filesystems.py +6 -517
  63. prefect/{new_flow_engine.py → flow_engine.py} +313 -72
  64. prefect/flow_runs.py +377 -5
  65. prefect/flows.py +248 -165
  66. prefect/futures.py +186 -345
  67. prefect/infrastructure/__init__.py +0 -27
  68. prefect/infrastructure/provisioners/__init__.py +5 -3
  69. prefect/infrastructure/provisioners/cloud_run.py +11 -6
  70. prefect/infrastructure/provisioners/container_instance.py +11 -7
  71. prefect/infrastructure/provisioners/ecs.py +6 -4
  72. prefect/infrastructure/provisioners/modal.py +8 -5
  73. prefect/input/actions.py +2 -4
  74. prefect/input/run_input.py +5 -7
  75. prefect/logging/formatters.py +0 -2
  76. prefect/logging/handlers.py +3 -11
  77. prefect/logging/loggers.py +2 -2
  78. prefect/manifests.py +2 -1
  79. prefect/records/__init__.py +1 -0
  80. prefect/records/result_store.py +42 -0
  81. prefect/records/store.py +9 -0
  82. prefect/results.py +43 -39
  83. prefect/runner/runner.py +9 -9
  84. prefect/runner/server.py +6 -10
  85. prefect/runner/storage.py +3 -8
  86. prefect/runner/submit.py +2 -2
  87. prefect/runner/utils.py +2 -2
  88. prefect/serializers.py +24 -35
  89. prefect/server/api/collections_data/views/aggregate-worker-metadata.json +5 -14
  90. prefect/settings.py +70 -133
  91. prefect/states.py +17 -47
  92. prefect/task_engine.py +697 -58
  93. prefect/task_runners.py +269 -301
  94. prefect/task_server.py +53 -34
  95. prefect/tasks.py +327 -337
  96. prefect/transactions.py +220 -0
  97. prefect/types/__init__.py +61 -82
  98. prefect/utilities/asyncutils.py +195 -136
  99. prefect/utilities/callables.py +121 -41
  100. prefect/utilities/collections.py +23 -38
  101. prefect/utilities/dispatch.py +11 -3
  102. prefect/utilities/dockerutils.py +4 -0
  103. prefect/utilities/engine.py +140 -20
  104. prefect/utilities/importtools.py +26 -27
  105. prefect/utilities/pydantic.py +128 -38
  106. prefect/utilities/schema_tools/hydration.py +5 -1
  107. prefect/utilities/templating.py +12 -2
  108. prefect/variables.py +78 -61
  109. prefect/workers/__init__.py +0 -1
  110. prefect/workers/base.py +15 -17
  111. prefect/workers/process.py +3 -8
  112. prefect/workers/server.py +2 -2
  113. {prefect_client-2.19.3.dist-info → prefect_client-3.0.0rc1.dist-info}/METADATA +22 -21
  114. prefect_client-3.0.0rc1.dist-info/RECORD +176 -0
  115. prefect/_internal/pydantic/_base_model.py +0 -51
  116. prefect/_internal/pydantic/_compat.py +0 -82
  117. prefect/_internal/pydantic/_flags.py +0 -20
  118. prefect/_internal/pydantic/_types.py +0 -8
  119. prefect/_internal/pydantic/utilities/__init__.py +0 -0
  120. prefect/_internal/pydantic/utilities/config_dict.py +0 -72
  121. prefect/_internal/pydantic/utilities/field_validator.py +0 -150
  122. prefect/_internal/pydantic/utilities/model_construct.py +0 -56
  123. prefect/_internal/pydantic/utilities/model_copy.py +0 -55
  124. prefect/_internal/pydantic/utilities/model_dump.py +0 -136
  125. prefect/_internal/pydantic/utilities/model_dump_json.py +0 -112
  126. prefect/_internal/pydantic/utilities/model_fields.py +0 -50
  127. prefect/_internal/pydantic/utilities/model_fields_set.py +0 -29
  128. prefect/_internal/pydantic/utilities/model_json_schema.py +0 -82
  129. prefect/_internal/pydantic/utilities/model_rebuild.py +0 -80
  130. prefect/_internal/pydantic/utilities/model_validate.py +0 -75
  131. prefect/_internal/pydantic/utilities/model_validate_json.py +0 -68
  132. prefect/_internal/pydantic/utilities/model_validator.py +0 -87
  133. prefect/_internal/pydantic/utilities/type_adapter.py +0 -71
  134. prefect/_vendor/__init__.py +0 -0
  135. prefect/_vendor/fastapi/__init__.py +0 -25
  136. prefect/_vendor/fastapi/applications.py +0 -946
  137. prefect/_vendor/fastapi/background.py +0 -3
  138. prefect/_vendor/fastapi/concurrency.py +0 -44
  139. prefect/_vendor/fastapi/datastructures.py +0 -58
  140. prefect/_vendor/fastapi/dependencies/__init__.py +0 -0
  141. prefect/_vendor/fastapi/dependencies/models.py +0 -64
  142. prefect/_vendor/fastapi/dependencies/utils.py +0 -877
  143. prefect/_vendor/fastapi/encoders.py +0 -177
  144. prefect/_vendor/fastapi/exception_handlers.py +0 -40
  145. prefect/_vendor/fastapi/exceptions.py +0 -46
  146. prefect/_vendor/fastapi/logger.py +0 -3
  147. prefect/_vendor/fastapi/middleware/__init__.py +0 -1
  148. prefect/_vendor/fastapi/middleware/asyncexitstack.py +0 -25
  149. prefect/_vendor/fastapi/middleware/cors.py +0 -3
  150. prefect/_vendor/fastapi/middleware/gzip.py +0 -3
  151. prefect/_vendor/fastapi/middleware/httpsredirect.py +0 -3
  152. prefect/_vendor/fastapi/middleware/trustedhost.py +0 -3
  153. prefect/_vendor/fastapi/middleware/wsgi.py +0 -3
  154. prefect/_vendor/fastapi/openapi/__init__.py +0 -0
  155. prefect/_vendor/fastapi/openapi/constants.py +0 -2
  156. prefect/_vendor/fastapi/openapi/docs.py +0 -203
  157. prefect/_vendor/fastapi/openapi/models.py +0 -480
  158. prefect/_vendor/fastapi/openapi/utils.py +0 -485
  159. prefect/_vendor/fastapi/param_functions.py +0 -340
  160. prefect/_vendor/fastapi/params.py +0 -453
  161. prefect/_vendor/fastapi/requests.py +0 -4
  162. prefect/_vendor/fastapi/responses.py +0 -40
  163. prefect/_vendor/fastapi/routing.py +0 -1331
  164. prefect/_vendor/fastapi/security/__init__.py +0 -15
  165. prefect/_vendor/fastapi/security/api_key.py +0 -98
  166. prefect/_vendor/fastapi/security/base.py +0 -6
  167. prefect/_vendor/fastapi/security/http.py +0 -172
  168. prefect/_vendor/fastapi/security/oauth2.py +0 -227
  169. prefect/_vendor/fastapi/security/open_id_connect_url.py +0 -34
  170. prefect/_vendor/fastapi/security/utils.py +0 -10
  171. prefect/_vendor/fastapi/staticfiles.py +0 -1
  172. prefect/_vendor/fastapi/templating.py +0 -3
  173. prefect/_vendor/fastapi/testclient.py +0 -1
  174. prefect/_vendor/fastapi/types.py +0 -3
  175. prefect/_vendor/fastapi/utils.py +0 -235
  176. prefect/_vendor/fastapi/websockets.py +0 -7
  177. prefect/_vendor/starlette/__init__.py +0 -1
  178. prefect/_vendor/starlette/_compat.py +0 -28
  179. prefect/_vendor/starlette/_exception_handler.py +0 -80
  180. prefect/_vendor/starlette/_utils.py +0 -88
  181. prefect/_vendor/starlette/applications.py +0 -261
  182. prefect/_vendor/starlette/authentication.py +0 -159
  183. prefect/_vendor/starlette/background.py +0 -43
  184. prefect/_vendor/starlette/concurrency.py +0 -59
  185. prefect/_vendor/starlette/config.py +0 -151
  186. prefect/_vendor/starlette/convertors.py +0 -87
  187. prefect/_vendor/starlette/datastructures.py +0 -707
  188. prefect/_vendor/starlette/endpoints.py +0 -130
  189. prefect/_vendor/starlette/exceptions.py +0 -60
  190. prefect/_vendor/starlette/formparsers.py +0 -276
  191. prefect/_vendor/starlette/middleware/__init__.py +0 -17
  192. prefect/_vendor/starlette/middleware/authentication.py +0 -52
  193. prefect/_vendor/starlette/middleware/base.py +0 -220
  194. prefect/_vendor/starlette/middleware/cors.py +0 -176
  195. prefect/_vendor/starlette/middleware/errors.py +0 -265
  196. prefect/_vendor/starlette/middleware/exceptions.py +0 -74
  197. prefect/_vendor/starlette/middleware/gzip.py +0 -113
  198. prefect/_vendor/starlette/middleware/httpsredirect.py +0 -19
  199. prefect/_vendor/starlette/middleware/sessions.py +0 -82
  200. prefect/_vendor/starlette/middleware/trustedhost.py +0 -64
  201. prefect/_vendor/starlette/middleware/wsgi.py +0 -147
  202. prefect/_vendor/starlette/requests.py +0 -328
  203. prefect/_vendor/starlette/responses.py +0 -347
  204. prefect/_vendor/starlette/routing.py +0 -933
  205. prefect/_vendor/starlette/schemas.py +0 -154
  206. prefect/_vendor/starlette/staticfiles.py +0 -248
  207. prefect/_vendor/starlette/status.py +0 -199
  208. prefect/_vendor/starlette/templating.py +0 -231
  209. prefect/_vendor/starlette/testclient.py +0 -804
  210. prefect/_vendor/starlette/types.py +0 -30
  211. prefect/_vendor/starlette/websockets.py +0 -193
  212. prefect/agent.py +0 -698
  213. prefect/deployments/deployments.py +0 -1042
  214. prefect/deprecated/__init__.py +0 -0
  215. prefect/deprecated/data_documents.py +0 -350
  216. prefect/deprecated/packaging/__init__.py +0 -12
  217. prefect/deprecated/packaging/base.py +0 -96
  218. prefect/deprecated/packaging/docker.py +0 -146
  219. prefect/deprecated/packaging/file.py +0 -92
  220. prefect/deprecated/packaging/orion.py +0 -80
  221. prefect/deprecated/packaging/serializers.py +0 -171
  222. prefect/events/instrument.py +0 -135
  223. prefect/infrastructure/base.py +0 -323
  224. prefect/infrastructure/container.py +0 -818
  225. prefect/infrastructure/kubernetes.py +0 -920
  226. prefect/infrastructure/process.py +0 -289
  227. prefect/new_task_engine.py +0 -423
  228. prefect/pydantic/__init__.py +0 -76
  229. prefect/pydantic/main.py +0 -39
  230. prefect/software/__init__.py +0 -2
  231. prefect/software/base.py +0 -50
  232. prefect/software/conda.py +0 -199
  233. prefect/software/pip.py +0 -122
  234. prefect/software/python.py +0 -52
  235. prefect/workers/block.py +0 -218
  236. prefect_client-2.19.3.dist-info/RECORD +0 -292
  237. {prefect_client-2.19.3.dist-info → prefect_client-3.0.0rc1.dist-info}/LICENSE +0 -0
  238. {prefect_client-2.19.3.dist-info → prefect_client-3.0.0rc1.dist-info}/WHEEL +0 -0
  239. {prefect_client-2.19.3.dist-info → prefect_client-3.0.0rc1.dist-info}/top_level.txt +0 -0
prefect/blocks/core.py CHANGED
@@ -4,17 +4,22 @@ import inspect
4
4
  import sys
5
5
  import warnings
6
6
  from abc import ABC
7
+ from functools import partial
7
8
  from textwrap import dedent
8
9
  from typing import (
9
10
  TYPE_CHECKING,
10
11
  Any,
12
+ Callable,
13
+ ClassVar,
11
14
  Dict,
12
15
  FrozenSet,
13
16
  List,
14
17
  Optional,
18
+ Tuple,
15
19
  Type,
16
20
  TypeVar,
17
21
  Union,
22
+ get_origin,
18
23
  )
19
24
  from uuid import UUID, uuid4
20
25
 
@@ -22,19 +27,22 @@ from griffe.dataclasses import Docstring
22
27
  from griffe.docstrings.dataclasses import DocstringSection, DocstringSectionKind
23
28
  from griffe.docstrings.parsers import Parser, parse
24
29
  from packaging.version import InvalidVersion, Version
25
-
26
- from prefect._internal.pydantic import HAS_PYDANTIC_V2
27
-
28
- if HAS_PYDANTIC_V2:
29
- from pydantic.v1 import BaseModel, HttpUrl, SecretBytes, SecretStr, ValidationError
30
- else:
31
- from pydantic import BaseModel, HttpUrl, SecretBytes, SecretStr, ValidationError
32
-
33
- from typing_extensions import ParamSpec, Self, get_args, get_origin
30
+ from pydantic import (
31
+ BaseModel,
32
+ ConfigDict,
33
+ HttpUrl,
34
+ PrivateAttr,
35
+ SecretBytes,
36
+ SecretStr,
37
+ SerializationInfo,
38
+ ValidationError,
39
+ model_serializer,
40
+ )
41
+ from pydantic.json_schema import GenerateJsonSchema
42
+ from typing_extensions import Literal, ParamSpec, Self, get_args
34
43
 
35
44
  import prefect
36
45
  import prefect.exceptions
37
- from prefect.blocks.fields import SecretDict
38
46
  from prefect.client.schemas import (
39
47
  DEFAULT_BLOCK_SCHEMA_VERSION,
40
48
  BlockDocument,
@@ -43,26 +51,27 @@ from prefect.client.schemas import (
43
51
  BlockTypeUpdate,
44
52
  )
45
53
  from prefect.client.utilities import inject_client
46
- from prefect.events.instrument import (
47
- ResourceTuple,
48
- emit_instance_method_called_event,
49
- instrument_instance_method_call,
50
- instrument_method_calls_on_class_instances,
51
- )
54
+ from prefect.events import emit_event
52
55
  from prefect.logging.loggers import disable_logger
56
+ from prefect.types import SecretDict
53
57
  from prefect.utilities.asyncutils import sync_compatible
54
- from prefect.utilities.collections import listrepr, remove_nested_keys
58
+ from prefect.utilities.collections import listrepr, remove_nested_keys, visit_collection
55
59
  from prefect.utilities.dispatch import lookup_type, register_base_type
56
60
  from prefect.utilities.hashing import hash_objects
57
61
  from prefect.utilities.importtools import to_qualified_name
62
+ from prefect.utilities.pydantic import handle_secret_render
58
63
  from prefect.utilities.slugify import slugify
59
64
 
60
65
  if TYPE_CHECKING:
66
+ from pydantic.main import IncEx
67
+
61
68
  from prefect.client.orchestration import PrefectClient
62
69
 
63
70
  R = TypeVar("R")
64
71
  P = ParamSpec("P")
65
72
 
73
+ ResourceTuple = Tuple[Dict[str, Any], List[Dict[str, Any]]]
74
+
66
75
 
67
76
  def block_schema_to_key(schema: BlockSchema) -> str:
68
77
  """
@@ -123,7 +132,9 @@ def _is_subclass(cls, parent_cls) -> bool:
123
132
  return inspect.isclass(cls) and issubclass(cls, parent_cls)
124
133
 
125
134
 
126
- def _collect_secret_fields(name: str, type_: Type, secrets: List[str]) -> None:
135
+ def _collect_secret_fields(
136
+ name: str, type_: Type[BaseModel], secrets: List[str]
137
+ ) -> None:
127
138
  """
128
139
  Recursively collects all secret fields from a given type and adds them to the
129
140
  secrets list, supporting nested Union / BaseModel fields. Also, note, this function
@@ -134,8 +145,8 @@ def _collect_secret_fields(name: str, type_: Type, secrets: List[str]) -> None:
134
145
  _collect_secret_fields(name, union_type, secrets)
135
146
  return
136
147
  elif _is_subclass(type_, BaseModel):
137
- for field in type_.__fields__.values():
138
- _collect_secret_fields(f"{name}.{field.name}", field.type_, secrets)
148
+ for field_name, field in type_.model_fields.items():
149
+ _collect_secret_fields(f"{name}.{field_name}", field.annotation, secrets)
139
150
  return
140
151
 
141
152
  if type_ in (SecretStr, SecretBytes):
@@ -145,7 +156,9 @@ def _collect_secret_fields(name: str, type_: Type, secrets: List[str]) -> None:
145
156
  # field are secret and should be obfuscated.
146
157
  secrets.append(f"{name}.*")
147
158
  elif Block.is_block_class(type_):
148
- secrets.extend(f"{name}.{s}" for s in type_.schema()["secret_fields"])
159
+ secrets.extend(
160
+ f"{name}.{s}" for s in type_.model_json_schema()["secret_fields"]
161
+ )
149
162
 
150
163
 
151
164
  def _should_update_block_type(
@@ -158,8 +171,10 @@ def _should_update_block_type(
158
171
  """
159
172
  fields = BlockTypeUpdate.updatable_fields()
160
173
 
161
- local_block_fields = local_block_type.dict(include=fields, exclude_unset=True)
162
- server_block_fields = server_block_type.dict(include=fields, exclude_unset=True)
174
+ local_block_fields = local_block_type.model_dump(include=fields, exclude_unset=True)
175
+ server_block_fields = server_block_type.model_dump(
176
+ include=fields, exclude_unset=True
177
+ )
163
178
 
164
179
  if local_block_fields.get("description") is not None:
165
180
  local_block_fields["description"] = html.unescape(
@@ -191,8 +206,49 @@ class BlockNotSavedError(RuntimeError):
191
206
  pass
192
207
 
193
208
 
209
+ def schema_extra(schema: Dict[str, Any], model: Type["Block"]):
210
+ """
211
+ Customizes Pydantic's schema generation feature to add blocks related information.
212
+ """
213
+ schema["block_type_slug"] = model.get_block_type_slug()
214
+ # Ensures args and code examples aren't included in the schema
215
+ description = model.get_description()
216
+ if description:
217
+ schema["description"] = description
218
+ else:
219
+ # Prevent the description of the base class from being included in the schema
220
+ schema.pop("description", None)
221
+
222
+ # create a list of secret field names
223
+ # secret fields include both top-level keys and dot-delimited nested secret keys
224
+ # A wildcard (*) means that all fields under a given key are secret.
225
+ # for example: ["x", "y", "z.*", "child.a"]
226
+ # means the top-level keys "x" and "y", all keys under "z", and the key "a" of a block
227
+ # nested under the "child" key are all secret. There is no limit to nesting.
228
+ secrets = schema["secret_fields"] = []
229
+ for name, field in model.model_fields.items():
230
+ _collect_secret_fields(name, field.annotation, secrets)
231
+
232
+ # create block schema references
233
+ refs = schema["block_schema_references"] = {}
234
+ for name, field in model.model_fields.items():
235
+ if Block.is_block_class(field.annotation):
236
+ refs[name] = field.annotation._to_block_schema_reference_dict()
237
+ if get_origin(field.annotation) in [Union, list]:
238
+ for type_ in get_args(field.annotation):
239
+ if Block.is_block_class(type_):
240
+ if isinstance(refs.get(name), list):
241
+ refs[name].append(type_._to_block_schema_reference_dict())
242
+ elif isinstance(refs.get(name), dict):
243
+ refs[name] = [
244
+ refs[name],
245
+ type_._to_block_schema_reference_dict(),
246
+ ]
247
+ else:
248
+ refs[name] = type_._to_block_schema_reference_dict()
249
+
250
+
194
251
  @register_base_type
195
- @instrument_method_calls_on_class_instances
196
252
  class Block(BaseModel, ABC):
197
253
  """
198
254
  A base class for implementing a block that wraps an external service.
@@ -210,56 +266,10 @@ class Block(BaseModel, ABC):
210
266
  initialization.
211
267
  """
212
268
 
213
- class Config:
214
- extra = "allow"
215
-
216
- json_encoders = {SecretDict: lambda v: v.dict()}
217
-
218
- @staticmethod
219
- def schema_extra(schema: Dict[str, Any], model: Type["Block"]):
220
- """
221
- Customizes Pydantic's schema generation feature to add blocks related information.
222
- """
223
- schema["block_type_slug"] = model.get_block_type_slug()
224
- # Ensures args and code examples aren't included in the schema
225
- description = model.get_description()
226
- if description:
227
- schema["description"] = description
228
- else:
229
- # Prevent the description of the base class from being included in the schema
230
- schema.pop("description", None)
231
-
232
- # create a list of secret field names
233
- # secret fields include both top-level keys and dot-delimited nested secret keys
234
- # A wildcard (*) means that all fields under a given key are secret.
235
- # for example: ["x", "y", "z.*", "child.a"]
236
- # means the top-level keys "x" and "y", all keys under "z", and the key "a" of a block
237
- # nested under the "child" key are all secret. There is no limit to nesting.
238
- secrets = schema["secret_fields"] = []
239
- for field in model.__fields__.values():
240
- _collect_secret_fields(field.name, field.type_, secrets)
241
-
242
- # create block schema references
243
- refs = schema["block_schema_references"] = {}
244
- for field in model.__fields__.values():
245
- if Block.is_block_class(field.type_):
246
- refs[field.name] = field.type_._to_block_schema_reference_dict()
247
- if get_origin(field.type_) is Union:
248
- for type_ in get_args(field.type_):
249
- if Block.is_block_class(type_):
250
- if isinstance(refs.get(field.name), list):
251
- refs[field.name].append(
252
- type_._to_block_schema_reference_dict()
253
- )
254
- elif isinstance(refs.get(field.name), dict):
255
- refs[field.name] = [
256
- refs[field.name],
257
- type_._to_block_schema_reference_dict(),
258
- ]
259
- else:
260
- refs[
261
- field.name
262
- ] = type_._to_block_schema_reference_dict()
269
+ model_config = ConfigDict(
270
+ extra="allow",
271
+ json_schema_extra=schema_extra,
272
+ )
263
273
 
264
274
  def __init__(self, *args, **kwargs):
265
275
  super().__init__(*args, **kwargs)
@@ -270,7 +280,7 @@ class Block(BaseModel, ABC):
270
280
 
271
281
  def __repr_args__(self):
272
282
  repr_args = super().__repr_args__()
273
- data_keys = self.schema()["properties"].keys()
283
+ data_keys = self.model_json_schema()["properties"].keys()
274
284
  return [
275
285
  (key, value) for key, value in repr_args if key is None or key in data_keys
276
286
  ]
@@ -284,25 +294,26 @@ class Block(BaseModel, ABC):
284
294
  # Attribute to customize the name of the block type created
285
295
  # when the block is registered with the API. If not set, block
286
296
  # type name will default to the class name.
287
- _block_type_name: Optional[str] = None
288
- _block_type_slug: Optional[str] = None
297
+ _block_type_name: ClassVar[Optional[str]] = None
298
+ _block_type_slug: ClassVar[Optional[str]] = None
289
299
 
290
300
  # Attributes used to set properties on a block type when registered
291
301
  # with the API.
292
- _logo_url: Optional[HttpUrl] = None
293
- _documentation_url: Optional[HttpUrl] = None
294
- _description: Optional[str] = None
295
- _code_example: Optional[str] = None
302
+ _logo_url: ClassVar[Optional[HttpUrl]] = None
303
+ _documentation_url: ClassVar[Optional[HttpUrl]] = None
304
+ _description: ClassVar[Optional[str]] = None
305
+ _code_example: ClassVar[Optional[str]] = None
306
+ _block_type_id: ClassVar[Optional[UUID]] = None
307
+ _block_schema_id: ClassVar[Optional[UUID]] = None
308
+ _block_schema_capabilities: ClassVar[Optional[List[str]]] = None
309
+ _block_schema_version: ClassVar[Optional[str]] = None
296
310
 
297
311
  # -- private instance variables
298
312
  # these are set when blocks are loaded from the API
299
- _block_type_id: Optional[UUID] = None
300
- _block_schema_id: Optional[UUID] = None
301
- _block_schema_capabilities: Optional[List[str]] = None
302
- _block_schema_version: Optional[str] = None
303
- _block_document_id: Optional[UUID] = None
304
- _block_document_name: Optional[str] = None
305
- _is_anonymous: Optional[bool] = None
313
+
314
+ _block_document_id: Optional[UUID] = PrivateAttr(None)
315
+ _block_document_name: Optional[str] = PrivateAttr(None)
316
+ _is_anonymous: Optional[bool] = PrivateAttr(None)
306
317
 
307
318
  # Exclude `save` as it uses the `sync_compatible` decorator and needs to be
308
319
  # decorated directly.
@@ -314,6 +325,31 @@ class Block(BaseModel, ABC):
314
325
  return None # The base class is abstract
315
326
  return block_schema_to_key(cls._to_block_schema())
316
327
 
328
+ @model_serializer(mode="wrap")
329
+ def ser_model(self, handler: Callable, info: SerializationInfo) -> Any:
330
+ jsonable_self = handler(self)
331
+ if (ctx := info.context) and ctx.get("include_secrets") is True:
332
+ jsonable_self.update(
333
+ {
334
+ field_name: visit_collection(
335
+ expr=getattr(self, field_name),
336
+ visit_fn=partial(handle_secret_render, context=ctx),
337
+ return_data=True,
338
+ )
339
+ for field_name in self.model_fields
340
+ }
341
+ )
342
+ if extra_fields := {
343
+ "block_type_slug": self.get_block_type_slug(),
344
+ "_block_document_id": self._block_document_id,
345
+ "_block_document_name": self._block_document_name,
346
+ "_is_anonymous": self._is_anonymous,
347
+ }:
348
+ jsonable_self |= {
349
+ key: value for key, value in extra_fields.items() if value is not None
350
+ }
351
+ return jsonable_self
352
+
317
353
  @classmethod
318
354
  def get_block_type_name(cls):
319
355
  return cls._block_type_name or cls.__name__
@@ -379,7 +415,9 @@ class Block(BaseModel, ABC):
379
415
  str: The calculated checksum prefixed with the hashing algorithm used.
380
416
  """
381
417
  block_schema_fields = (
382
- cls.schema() if block_schema_fields is None else block_schema_fields
418
+ cls.model_json_schema()
419
+ if block_schema_fields is None
420
+ else block_schema_fields
383
421
  )
384
422
  fields_for_checksum = remove_nested_keys(["secret_fields"], block_schema_fields)
385
423
  if fields_for_checksum.get("definitions"):
@@ -404,6 +442,7 @@ class Block(BaseModel, ABC):
404
442
  block_schema_id: Optional[UUID] = None,
405
443
  block_type_id: Optional[UUID] = None,
406
444
  is_anonymous: Optional[bool] = None,
445
+ include_secrets: bool = False,
407
446
  ) -> BlockDocument:
408
447
  """
409
448
  Creates the corresponding block document based on the data stored in a block.
@@ -441,10 +480,14 @@ class Block(BaseModel, ABC):
441
480
  # The keys passed to `include` must NOT be aliases, else some items will be missed
442
481
  # i.e. must do `self.schema_` vs `self.schema` to get a `schema_ = Field(alias="schema")`
443
482
  # reported from https://github.com/PrefectHQ/prefect-dbt/issues/54
444
- data_keys = self.schema(by_alias=False)["properties"].keys()
483
+ data_keys = self.model_json_schema(by_alias=False)["properties"].keys()
445
484
 
446
485
  # `block_document_data`` must return the aliased version for it to show in the UI
447
- block_document_data = self.dict(by_alias=True, include=data_keys)
486
+ block_document_data = self.model_dump(
487
+ by_alias=True,
488
+ include=data_keys,
489
+ context={"include_secrets": include_secrets},
490
+ )
448
491
 
449
492
  # Iterate through and find blocks that already have saved block documents to
450
493
  # create references to those saved block documents.
@@ -484,7 +527,7 @@ class Block(BaseModel, ABC):
484
527
  Returns:
485
528
  BlockSchema: The corresponding block schema.
486
529
  """
487
- fields = cls.schema()
530
+ fields = cls.model_json_schema()
488
531
  return BlockSchema(
489
532
  id=cls._block_schema_id if cls._block_schema_id is not None else uuid4(),
490
533
  checksum=cls._calculate_schema_checksum(),
@@ -539,7 +582,7 @@ class Block(BaseModel, ABC):
539
582
  code example from the class docstring if an override is not provided.
540
583
  """
541
584
  code_example = (
542
- dedent(cls._code_example) if cls._code_example is not None else None
585
+ dedent(text=cls._code_example) if cls._code_example is not None else None
543
586
  )
544
587
  # If no code example override has been provided, attempt to find a examples
545
588
  # section or an admonition with the annotation "example" and use that as the
@@ -606,7 +649,7 @@ class Block(BaseModel, ABC):
606
649
  )
607
650
 
608
651
  @classmethod
609
- def _from_block_document(cls, block_document: BlockDocument):
652
+ def _from_block_document(cls, block_document: BlockDocument) -> Self:
610
653
  """
611
654
  Instantiates a block from a given block document. The corresponding block class
612
655
  will be looked up in the block registry based on the corresponding block schema
@@ -634,9 +677,7 @@ class Block(BaseModel, ABC):
634
677
  else cls.get_block_class_from_schema(block_document.block_schema)
635
678
  )
636
679
 
637
- block_cls = instrument_method_calls_on_class_instances(block_cls)
638
-
639
- block = block_cls.parse_obj(block_document.data)
680
+ block = block_cls.model_validate(block_document.data)
640
681
  block._block_document_id = block_document.id
641
682
  block.__class__._block_schema_id = block_document.block_schema_id
642
683
  block.__class__._block_type_id = block_document.block_type_id
@@ -646,13 +687,11 @@ class Block(BaseModel, ABC):
646
687
  block_document.block_document_references
647
688
  )
648
689
 
649
- # Due to the way blocks are loaded we can't directly instrument the
650
- # `load` method and have the data be about the block document. Instead
651
- # this will emit a proxy event for the load method so that block
652
- # document data can be included instead of the event being about an
653
- # 'anonymous' block.
654
-
655
- emit_instance_method_called_event(block, "load", successful=True)
690
+ resources: Optional[ResourceTuple] = block._event_method_called_resources()
691
+ if resources:
692
+ kind = block._event_kind()
693
+ resource, related = resources
694
+ emit_event(event=f"{kind}.loaded", resource=resource, related=related)
656
695
 
657
696
  return block
658
697
 
@@ -731,7 +770,7 @@ class Block(BaseModel, ABC):
731
770
  async def _get_block_document(
732
771
  cls,
733
772
  name: str,
734
- client: "PrefectClient" = None,
773
+ client: Optional["PrefectClient"] = None,
735
774
  ):
736
775
  if cls.__name__ == "Block":
737
776
  block_type_slug, block_document_name = name.split("/", 1)
@@ -758,7 +797,7 @@ class Block(BaseModel, ABC):
758
797
  cls,
759
798
  name: str,
760
799
  validate: bool = True,
761
- client: "PrefectClient" = None,
800
+ client: Optional["PrefectClient"] = None,
762
801
  ):
763
802
  """
764
803
  Retrieves data from the block document with the given name for the block type
@@ -845,7 +884,7 @@ class Block(BaseModel, ABC):
845
884
  missing_block_data = {field: None for field in missing_fields}
846
885
  warnings.warn(
847
886
  f"Could not fully load {block_document_name!r} of block type"
848
- f" {cls._block_type_slug!r} - this is likely because one or more"
887
+ f" {cls.get_block_type_slug()!r} - this is likely because one or more"
849
888
  " required fields were added to the schema for"
850
889
  f" {cls.__name__!r} that did not exist on the class when this block"
851
890
  " was last saved. Please specify values for new field(s):"
@@ -853,10 +892,10 @@ class Block(BaseModel, ABC):
853
892
  f' `{cls.__name__}.save("{block_document_name}", overwrite=True)`,'
854
893
  " and load this block again before attempting to use it."
855
894
  )
856
- return cls.construct(**block_document.data, **missing_block_data)
895
+ return cls.model_construct(**block_document.data, **missing_block_data)
857
896
  raise RuntimeError(
858
897
  f"Unable to load {block_document_name!r} of block type"
859
- f" {cls._block_type_slug!r} due to failed validation. To load without"
898
+ f" {cls.get_block_type_slug()!r} due to failed validation. To load without"
860
899
  " validation, try loading again with `validate=False`."
861
900
  ) from e
862
901
 
@@ -864,6 +903,18 @@ class Block(BaseModel, ABC):
864
903
  def is_block_class(block) -> bool:
865
904
  return _is_subclass(block, Block)
866
905
 
906
+ @staticmethod
907
+ def annotation_refers_to_block_class(annotation: Any) -> bool:
908
+ if Block.is_block_class(annotation):
909
+ return True
910
+
911
+ if get_origin(annotation) is Union:
912
+ for annotation in get_args(annotation):
913
+ if Block.is_block_class(annotation):
914
+ return True
915
+
916
+ return False
917
+
867
918
  @classmethod
868
919
  @sync_compatible
869
920
  @inject_client
@@ -888,19 +939,21 @@ class Block(BaseModel, ABC):
888
939
  "subclass and not on a Block interface class directly."
889
940
  )
890
941
 
891
- for field in cls.__fields__.values():
892
- if Block.is_block_class(field.type_):
893
- await field.type_.register_type_and_schema(client=client)
894
- if get_origin(field.type_) is Union:
895
- for type_ in get_args(field.type_):
896
- if Block.is_block_class(type_):
897
- await type_.register_type_and_schema(client=client)
942
+ for field in cls.model_fields.values():
943
+ if Block.is_block_class(field.annotation):
944
+ await field.annotation.register_type_and_schema(client=client)
945
+ if get_origin(field.annotation) is Union:
946
+ for annotation in get_args(field.annotation):
947
+ if Block.is_block_class(annotation):
948
+ await annotation.register_type_and_schema(client=client)
898
949
 
899
950
  try:
900
951
  block_type = await client.read_block_type_by_slug(
901
952
  slug=cls.get_block_type_slug()
902
953
  )
954
+
903
955
  cls._block_type_id = block_type.id
956
+
904
957
  local_block_type = cls._to_block_type()
905
958
  if _should_update_block_type(
906
959
  local_block_type=local_block_type, server_block_type=block_type
@@ -930,7 +983,7 @@ class Block(BaseModel, ABC):
930
983
  name: Optional[str] = None,
931
984
  is_anonymous: bool = False,
932
985
  overwrite: bool = False,
933
- client: "PrefectClient" = None,
986
+ client: Optional["PrefectClient"] = None,
934
987
  ):
935
988
  """
936
989
  Saves the values of a block as a block document with an option to save as an
@@ -978,7 +1031,9 @@ class Block(BaseModel, ABC):
978
1031
  block_document_id = existing_block_document.id
979
1032
  await client.update_block_document(
980
1033
  block_document_id=block_document_id,
981
- block_document=self._to_block_document(name=name),
1034
+ block_document=self._to_block_document(
1035
+ name=name, include_secrets=True
1036
+ ),
982
1037
  )
983
1038
  block_document = await client.read_block_document(
984
1039
  block_document_id=block_document_id
@@ -996,12 +1051,11 @@ class Block(BaseModel, ABC):
996
1051
  return self._block_document_id
997
1052
 
998
1053
  @sync_compatible
999
- @instrument_instance_method_call
1000
1054
  async def save(
1001
1055
  self,
1002
1056
  name: Optional[str] = None,
1003
1057
  overwrite: bool = False,
1004
- client: "PrefectClient" = None,
1058
+ client: Optional["PrefectClient"] = None,
1005
1059
  ):
1006
1060
  """
1007
1061
  Saves the values of a block as a block document.
@@ -1023,25 +1077,12 @@ class Block(BaseModel, ABC):
1023
1077
  async def delete(
1024
1078
  cls,
1025
1079
  name: str,
1026
- client: "PrefectClient" = None,
1080
+ client: Optional["PrefectClient"] = None,
1027
1081
  ):
1028
1082
  block_document, block_document_name = await cls._get_block_document(name)
1029
1083
 
1030
1084
  await client.delete_block_document(block_document.id)
1031
1085
 
1032
- def _iter(self, *, include=None, exclude=None, **kwargs):
1033
- # Injects the `block_type_slug` into serialized payloads for dispatch
1034
- for key_value in super()._iter(include=include, exclude=exclude, **kwargs):
1035
- yield key_value
1036
-
1037
- # Respect inclusion and exclusion still
1038
- if include and "block_type_slug" not in include:
1039
- return
1040
- if exclude and "block_type_slug" in exclude:
1041
- return
1042
-
1043
- yield "block_type_slug", self.get_block_type_slug()
1044
-
1045
1086
  def __new__(cls: Type[Self], **kwargs) -> Self:
1046
1087
  """
1047
1088
  Create an instance of the Block subclass type if a `block_type_slug` is
@@ -1050,21 +1091,9 @@ class Block(BaseModel, ABC):
1050
1091
  block_type_slug = kwargs.pop("block_type_slug", None)
1051
1092
  if block_type_slug:
1052
1093
  subcls = lookup_type(cls, dispatch_key=block_type_slug)
1053
- m = super().__new__(subcls)
1054
- # NOTE: This is a workaround for an obscure issue where copied models were
1055
- # missing attributes. This pattern is from Pydantic's
1056
- # `BaseModel._copy_and_set_values`.
1057
- # The issue this fixes could not be reproduced in unit tests that
1058
- # directly targeted dispatch handling and was only observed when
1059
- # copying then saving infrastructure blocks on deployment models.
1060
- object.__setattr__(m, "__dict__", kwargs)
1061
- object.__setattr__(m, "__fields_set__", set(kwargs.keys()))
1062
- return m
1094
+ return super().__new__(subcls)
1063
1095
  else:
1064
- m = super().__new__(cls)
1065
- object.__setattr__(m, "__dict__", kwargs)
1066
- object.__setattr__(m, "__fields_set__", set(kwargs.keys()))
1067
- return m
1096
+ return super().__new__(cls)
1068
1097
 
1069
1098
  def get_block_placeholder(self) -> str:
1070
1099
  """
@@ -1087,3 +1116,94 @@ class Block(BaseModel, ABC):
1087
1116
  )
1088
1117
 
1089
1118
  return f"prefect.blocks.{self.get_block_type_slug()}.{block_document_name}"
1119
+
1120
+ @classmethod
1121
+ def model_json_schema(
1122
+ cls,
1123
+ by_alias: bool = True,
1124
+ ref_template: str = "#/definitions/{model}",
1125
+ schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema,
1126
+ mode: Literal["validation", "serialization"] = "validation",
1127
+ ) -> Dict[str, Any]:
1128
+ """TODO: stop overriding this method - use GenerateSchema in ConfigDict instead?"""
1129
+ schema = super().model_json_schema(
1130
+ by_alias, ref_template, schema_generator, mode
1131
+ )
1132
+
1133
+ # ensure backwards compatibility by copying $defs into definitions
1134
+ if "$defs" in schema:
1135
+ schema["definitions"] = schema.pop("$defs")
1136
+
1137
+ # we aren't expecting these additional fields in the schema
1138
+ if "additionalProperties" in schema:
1139
+ schema.pop("additionalProperties")
1140
+
1141
+ for _, definition in schema.get("definitions", {}).items():
1142
+ if "additionalProperties" in definition:
1143
+ definition.pop("additionalProperties")
1144
+
1145
+ return schema
1146
+
1147
+ @classmethod
1148
+ def model_validate(
1149
+ cls: type[Self],
1150
+ obj: Any,
1151
+ *,
1152
+ strict: Optional[bool] = None,
1153
+ from_attributes: Optional[bool] = None,
1154
+ context: Optional[Dict[str, Any]] = None,
1155
+ ) -> Self:
1156
+ if isinstance(obj, dict):
1157
+ extra_serializer_fields = {
1158
+ "_block_document_id",
1159
+ "_block_document_name",
1160
+ "_is_anonymous",
1161
+ }.intersection(obj.keys())
1162
+ for field in extra_serializer_fields:
1163
+ obj.pop(field, None)
1164
+
1165
+ return super().model_validate(
1166
+ obj, strict=strict, from_attributes=from_attributes, context=context
1167
+ )
1168
+
1169
+ def model_dump(
1170
+ self,
1171
+ *,
1172
+ mode: Union[Literal["json", "python"], str] = "python",
1173
+ include: "IncEx" = None,
1174
+ exclude: "IncEx" = None,
1175
+ context: Optional[Dict[str, Any]] = None,
1176
+ by_alias: bool = False,
1177
+ exclude_unset: bool = False,
1178
+ exclude_defaults: bool = False,
1179
+ exclude_none: bool = False,
1180
+ round_trip: bool = False,
1181
+ warnings: Union[bool, Literal["none", "warn", "error"]] = True,
1182
+ serialize_as_any: bool = False,
1183
+ ) -> Dict[str, Any]:
1184
+ d = super().model_dump(
1185
+ mode=mode,
1186
+ include=include,
1187
+ exclude=exclude,
1188
+ context=context,
1189
+ by_alias=by_alias,
1190
+ exclude_unset=exclude_unset,
1191
+ exclude_defaults=exclude_defaults,
1192
+ exclude_none=exclude_none,
1193
+ round_trip=round_trip,
1194
+ warnings=warnings,
1195
+ serialize_as_any=serialize_as_any,
1196
+ )
1197
+
1198
+ extra_serializer_fields = {
1199
+ "block_type_slug",
1200
+ "_block_document_id",
1201
+ "_block_document_name",
1202
+ "_is_anonymous",
1203
+ }.intersection(d.keys())
1204
+
1205
+ for field in extra_serializer_fields:
1206
+ if (include and field not in include) or (exclude and field in exclude):
1207
+ d.pop(field)
1208
+
1209
+ return d