Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/database/orm_models.py: 97%
551 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 02:04 +0000
1import datetime
2import uuid
3from abc import ABC, abstractmethod
4from collections.abc import Hashable, Iterable
5from pathlib import Path
6from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
8import sqlalchemy as sa
9from sqlalchemy import FetchedValue
10from sqlalchemy.dialects import postgresql
11from sqlalchemy.ext.asyncio import AsyncSession
12from sqlalchemy.ext.hybrid import hybrid_property
13from sqlalchemy.orm import (
14 DeclarativeBase,
15 Mapped,
16 declared_attr,
17 mapped_column,
18 registry,
19 relationship,
20 synonym,
21)
22from sqlalchemy.orm.decl_api import registry as RegistryType
23from sqlalchemy.sql import roles
24from sqlalchemy.sql.functions import coalesce
26import prefect.server.schemas as schemas
27from prefect.server.events.actions import ServerActionTypes
28from prefect.server.events.schemas.automations import (
29 AutomationSort,
30 Firing,
31 ServerTriggerTypes,
32)
33from prefect.server.events.schemas.events import ReceivedEvent
34from prefect.server.schemas.statuses import (
35 DeploymentStatus,
36 WorkerStatus,
37 WorkPoolStatus,
38 WorkQueueStatus,
39)
40from prefect.server.utilities.database import (
41 CAMEL_TO_SNAKE,
42 JSON,
43 UUID,
44 GenerateUUID,
45 Pydantic,
46 Timestamp,
47)
48from prefect.server.utilities.encryption import decrypt_fernet, encrypt_fernet
49from prefect.types._datetime import DateTime, now
50from prefect.utilities.names import generate_slug
52# for 'plain JSON' columns, use the postgresql variant (which comes with an
53# extra operator) and fall back to the generic JSON variant for SQLite
54sa_JSON: postgresql.JSON = postgresql.JSON().with_variant(sa.JSON(), "sqlite")
57class Base(DeclarativeBase):
58 """
59 Base SQLAlchemy model that automatically infers the table name
60 and provides ID, created, and updated columns
61 """
63 registry: ClassVar[RegistryType] = registry(
64 metadata=sa.schema.MetaData(
65 # define naming conventions for our Base class to use
66 # sqlalchemy will use the following templated strings
67 # to generate the names of indices, constraints, and keys
68 #
69 # we offset the table name with two underscores (__) to
70 # help differentiate, for example, between "flow_run.state_type"
71 # and "flow_run_state.type".
72 #
73 # more information on this templating and available
74 # customization can be found here
75 # https://docs.sqlalchemy.org/en/14/core/metadata.html#sqlalchemy.schema.MetaData
76 #
77 # this also allows us to avoid having to specify names explicitly
78 # when using sa.ForeignKey.use_alter = True
79 # https://docs.sqlalchemy.org/en/14/core/constraints.html
80 naming_convention={
81 "ix": "ix_%(table_name)s__%(column_0_N_name)s",
82 "uq": "uq_%(table_name)s__%(column_0_N_name)s",
83 "ck": "ck_%(table_name)s__%(constraint_name)s",
84 "fk": "fk_%(table_name)s__%(column_0_N_name)s__%(referred_table_name)s",
85 "pk": "pk_%(table_name)s",
86 }
87 ),
88 type_annotation_map={
89 uuid.UUID: UUID,
90 DateTime: Timestamp,
91 },
92 )
94 # required in order to access columns with server defaults
95 # or SQL expression defaults, subsequent to a flush, without
96 # triggering an expired load
97 #
98 # this allows us to load attributes with a server default after
99 # an INSERT, for example
100 #
101 # https://docs.sqlalchemy.org/en/14/orm/extensions/asyncio.html#preventing-implicit-io-when-using-asyncsession
102 __mapper_args__: dict[str, Any] = {"eager_defaults": True}
104 def __repr__(self) -> str:
105 return f"{self.__class__.__name__}(id={self.id})"
107 @declared_attr.directive
108 def __tablename__(cls) -> str:
109 """
110 By default, turn the model's camel-case class name
111 into a snake-case table name. Override by providing
112 an explicit `__tablename__` class property.
113 """
114 return CAMEL_TO_SNAKE.sub("_", cls.__name__).lower()
116 id: Mapped[uuid.UUID] = mapped_column(
117 primary_key=True,
118 server_default=GenerateUUID(),
119 default=uuid.uuid4,
120 )
122 created: Mapped[DateTime] = mapped_column(
123 server_default=sa.func.now(), default=lambda: now("UTC")
124 )
126 # onupdate is only called when statements are actually issued
127 # against the database. until COMMIT is issued, this column
128 # will not be updated
129 updated: Mapped[DateTime] = mapped_column(
130 index=True,
131 server_default=sa.func.now(),
132 default=lambda: now("UTC"),
133 onupdate=sa.func.now(),
134 server_onupdate=FetchedValue(),
135 )
138class Flow(Base):
139 """SQLAlchemy mixin of a flow."""
141 name: Mapped[str]
142 tags: Mapped[list[str]] = mapped_column(JSON, server_default="[]", default=list)
143 labels: Mapped[Optional[schemas.core.KeyValueLabels]] = mapped_column(JSON)
145 flow_runs: Mapped[list["FlowRun"]] = relationship(
146 back_populates="flow", lazy="raise"
147 )
148 deployments: Mapped[list["Deployment"]] = relationship(
149 back_populates="flow", lazy="raise"
150 )
152 __table_args__: Any = (
153 sa.UniqueConstraint("name"),
154 sa.Index("ix_flow__created", "created"),
155 sa.Index("trgm_ix_flow_name", "name", postgresql_using="gin").ddl_if(
156 dialect="postgresql"
157 ),
158 )
161class FlowRunState(Base):
162 """SQLAlchemy mixin of a flow run state."""
164 flow_run_id: Mapped[uuid.UUID] = mapped_column(
165 sa.ForeignKey("flow_run.id", ondelete="cascade")
166 )
168 type: Mapped[schemas.states.StateType] = mapped_column(
169 sa.Enum(schemas.states.StateType, name="state_type"), index=True
170 )
171 timestamp: Mapped[DateTime] = mapped_column(
172 server_default=sa.func.now(), default=lambda: now("UTC")
173 )
174 name: Mapped[str] = mapped_column(index=True)
175 message: Mapped[Optional[str]]
176 state_details: Mapped[schemas.states.StateDetails] = mapped_column(
177 Pydantic(schemas.states.StateDetails),
178 server_default="{}",
179 default=schemas.states.StateDetails,
180 )
181 _data: Mapped[Optional[Any]] = mapped_column(JSON, name="data")
183 result_artifact_id: Mapped[Optional[uuid.UUID]] = mapped_column(
184 sa.ForeignKey("artifact.id", ondelete="SET NULL", use_alter=True),
185 index=True,
186 )
188 _result_artifact: Mapped[Optional["Artifact"]] = relationship(
189 lazy="selectin",
190 foreign_keys=[result_artifact_id],
191 primaryjoin="Artifact.id==FlowRunState.result_artifact_id",
192 )
194 @hybrid_property
195 def data(self) -> Optional[Any]:
196 if self._data: 196 ↛ 198line 196 didn't jump to line 198 because the condition on line 196 was never true
197 # ensures backwards compatibility for results stored on state objects
198 return self._data
199 if not self.result_artifact_id:
200 # do not try to load the relationship if there's no artifact id
201 return None
202 if TYPE_CHECKING: 202 ↛ 203line 202 didn't jump to line 203 because the condition on line 202 was never true
203 assert self._result_artifact is not None
204 return self._result_artifact.data
206 flow_run: Mapped["FlowRun"] = relationship(lazy="raise", foreign_keys=[flow_run_id])
208 def as_state(self) -> schemas.states.State:
209 return schemas.states.State.model_validate(self, from_attributes=True)
211 @declared_attr.directive
212 @classmethod
213 def __table_args__(cls) -> Iterable[sa.Index]:
214 return (
215 sa.Index(
216 "uq_flow_run_state__flow_run_id_timestamp_desc",
217 cls.flow_run_id,
218 cls.timestamp.desc(),
219 unique=True,
220 ),
221 )
224class TaskRunState(Base):
225 """SQLAlchemy model of a task run state."""
227 # this column isn't explicitly indexed because it is included in
228 # the unique compound index on (task_run_id, timestamp)
229 task_run_id: Mapped[uuid.UUID] = mapped_column(
230 sa.ForeignKey("task_run.id", ondelete="cascade")
231 )
233 type: Mapped[schemas.states.StateType] = mapped_column(
234 sa.Enum(schemas.states.StateType, name="state_type"), index=True
235 )
236 timestamp: Mapped[DateTime] = mapped_column(
237 server_default=sa.func.now(), default=lambda: now("UTC")
238 )
239 name: Mapped[str] = mapped_column(index=True)
240 message: Mapped[Optional[str]]
241 state_details: Mapped[schemas.states.StateDetails] = mapped_column(
242 Pydantic(schemas.states.StateDetails),
243 server_default="{}",
244 default=schemas.states.StateDetails,
245 )
246 _data: Mapped[Optional[Any]] = mapped_column(JSON, name="data")
248 result_artifact_id: Mapped[Optional[uuid.UUID]] = mapped_column(
249 sa.ForeignKey("artifact.id", ondelete="SET NULL", use_alter=True), index=True
250 )
252 _result_artifact: Mapped[Optional["Artifact"]] = relationship(
253 lazy="selectin",
254 foreign_keys=[result_artifact_id],
255 primaryjoin="Artifact.id==TaskRunState.result_artifact_id",
256 )
258 @hybrid_property
259 def data(self) -> Optional[Any]:
260 if self._data: 260 ↛ 262line 260 didn't jump to line 262 because the condition on line 260 was never true
261 # ensures backwards compatibility for results stored on state objects
262 return self._data
263 if not self.result_artifact_id:
264 # do not try to load the relationship if there's no artifact id
265 return None
266 if TYPE_CHECKING: 266 ↛ 267line 266 didn't jump to line 267 because the condition on line 266 was never true
267 assert self._result_artifact is not None
268 return self._result_artifact.data
270 task_run: Mapped["TaskRun"] = relationship(lazy="raise", foreign_keys=[task_run_id])
272 def as_state(self) -> schemas.states.State:
273 return schemas.states.State.model_validate(self, from_attributes=True)
275 @declared_attr.directive
276 @classmethod
277 def __table_args__(cls) -> Iterable[sa.Index]:
278 return (
279 sa.Index(
280 "uq_task_run_state__task_run_id_timestamp_desc",
281 cls.task_run_id,
282 cls.timestamp.desc(),
283 unique=True,
284 ),
285 )
288class Artifact(Base):
289 """
290 SQLAlchemy model of artifacts.
291 """
293 key: Mapped[Optional[str]] = mapped_column(index=True)
295 task_run_id: Mapped[Optional[uuid.UUID]] = mapped_column(index=True)
297 flow_run_id: Mapped[Optional[uuid.UUID]] = mapped_column(index=True)
299 type: Mapped[Optional[str]]
300 data: Mapped[Optional[Any]] = mapped_column(sa_JSON)
301 description: Mapped[Optional[str]]
303 # Suffixed with underscore as attribute name 'metadata' is reserved for the MetaData instance when using a declarative base class.
304 metadata_: Mapped[Optional[dict[str, str]]] = mapped_column(sa_JSON)
306 @declared_attr.directive
307 @classmethod
308 def __table_args__(cls) -> Iterable[sa.Index]:
309 return (
310 sa.Index(
311 "ix_artifact__key",
312 cls.key,
313 ),
314 sa.Index(
315 "ix_artifact__key_created_desc",
316 cls.key,
317 cls.created.desc(),
318 postgresql_include=[
319 "id",
320 "updated",
321 "type",
322 "task_run_id",
323 "flow_run_id",
324 ],
325 ),
326 )
329class ArtifactCollection(Base):
330 key: Mapped[str]
332 latest_id: Mapped[uuid.UUID]
334 task_run_id: Mapped[Optional[uuid.UUID]]
336 flow_run_id: Mapped[Optional[uuid.UUID]]
338 type: Mapped[Optional[str]]
339 data: Mapped[Optional[Any]] = mapped_column(sa_JSON)
340 description: Mapped[Optional[str]]
341 metadata_: Mapped[Optional[dict[str, str]]] = mapped_column(sa_JSON)
343 __table_args__: Any = (
344 sa.UniqueConstraint("key"),
345 sa.Index(
346 "ix_artifact_collection__key_latest_id",
347 "key",
348 "latest_id",
349 ),
350 )
353class TaskRunStateCache(Base):
354 """
355 SQLAlchemy model of a task run state cache.
356 """
358 cache_key: Mapped[str] = mapped_column()
359 cache_expiration: Mapped[Optional[DateTime]]
360 task_run_state_id: Mapped[uuid.UUID]
362 @declared_attr.directive
363 @classmethod
364 def __table_args__(cls) -> Iterable[sa.Index]:
365 return (
366 sa.Index(
367 "ix_task_run_state_cache__cache_key_created_desc",
368 cls.cache_key,
369 cls.created.desc(),
370 ),
371 )
374class Run(Base):
375 """
376 Common columns and logic for FlowRun and TaskRun models
377 """
379 __abstract__ = True
381 name: Mapped[str] = mapped_column(default=lambda: generate_slug(2), index=True)
382 state_type: Mapped[Optional[schemas.states.StateType]] = mapped_column(
383 sa.Enum(schemas.states.StateType, name="state_type")
384 )
385 state_name: Mapped[Optional[str]]
386 state_timestamp: Mapped[Optional[DateTime]]
387 run_count: Mapped[int] = mapped_column(server_default="0", default=0)
388 expected_start_time: Mapped[Optional[DateTime]]
389 next_scheduled_start_time: Mapped[Optional[DateTime]]
390 start_time: Mapped[Optional[DateTime]]
391 end_time: Mapped[Optional[DateTime]]
392 total_run_time: Mapped[datetime.timedelta] = mapped_column(
393 server_default="0", default=datetime.timedelta(0)
394 )
396 @hybrid_property
397 def estimated_run_time(self) -> datetime.timedelta:
398 """Total run time is incremented in the database whenever a RUNNING
399 state is exited. To give up-to-date estimates, we estimate incremental
400 run time for any runs currently in a RUNNING state."""
401 if self.state_type and self.state_type == schemas.states.StateType.RUNNING:
402 if TYPE_CHECKING: 402 ↛ 403line 402 didn't jump to line 403 because the condition on line 402 was never true
403 assert self.state_timestamp is not None
404 return self.total_run_time + (now("UTC") - self.state_timestamp)
405 else:
406 return self.total_run_time
408 @estimated_run_time.inplace.expression
409 @classmethod
410 def _estimated_run_time_expression(cls) -> sa.Label[datetime.timedelta]:
411 return (
412 sa.select(
413 sa.case(
414 (
415 cls.state_type == schemas.states.StateType.RUNNING,
416 sa.func.interval_add(
417 cls.total_run_time,
418 sa.func.date_diff(sa.func.now(), cls.state_timestamp),
419 ),
420 ),
421 else_=cls.total_run_time,
422 )
423 )
424 # add a correlate statement so this can reuse the `FROM` clause
425 # of any parent query
426 .correlate(cls)
427 .label("estimated_run_time")
428 )
430 @hybrid_property
431 def estimated_start_time_delta(self) -> datetime.timedelta:
432 """The delta to the expected start time (or "lateness") is computed as
433 the difference between the actual start time and expected start time. To
434 give up-to-date estimates, we estimate lateness for any runs that don't
435 have a start time and are not in a final state and were expected to
436 start already."""
437 if (
438 self.start_time
439 and self.expected_start_time is not None
440 and self.start_time > (self.expected_start_time)
441 ):
442 return self.start_time - self.expected_start_time
443 elif (
444 self.start_time is None
445 and self.expected_start_time
446 and self.expected_start_time < now("UTC")
447 and self.state_type not in schemas.states.TERMINAL_STATES
448 ):
449 return now("UTC") - self.expected_start_time
450 else:
451 return datetime.timedelta(0)
453 @estimated_start_time_delta.inplace.expression
454 @classmethod
455 def _estimated_start_time_delta_expression(
456 cls,
457 ) -> sa.SQLColumnExpression[datetime.timedelta]:
458 return sa.case(
459 (
460 cls.start_time > cls.expected_start_time,
461 sa.func.date_diff(cls.start_time, cls.expected_start_time),
462 ),
463 (
464 sa.and_(
465 cls.start_time.is_(None),
466 cls.state_type.not_in(schemas.states.TERMINAL_STATES),
467 cls.expected_start_time < sa.func.now(),
468 ),
469 sa.func.date_diff(sa.func.now(), cls.expected_start_time),
470 ),
471 else_=datetime.timedelta(0),
472 )
475class FlowRun(Run):
476 """SQLAlchemy model of a flow run."""
478 flow_id: Mapped[uuid.UUID] = mapped_column(
479 sa.ForeignKey("flow.id", ondelete="cascade"), index=True
480 )
482 deployment_id: Mapped[Optional[uuid.UUID]] = mapped_column()
483 work_queue_name: Mapped[Optional[str]] = mapped_column(index=True)
484 flow_version: Mapped[Optional[str]] = mapped_column(index=True)
485 deployment_version: Mapped[Optional[str]] = mapped_column(index=True)
486 parameters: Mapped[dict[str, Any]] = mapped_column(
487 JSON, server_default="{}", default=dict
488 )
489 idempotency_key: Mapped[Optional[str]] = mapped_column()
490 context: Mapped[dict[str, Any]] = mapped_column(
491 JSON, server_default="{}", default=dict
492 )
493 empirical_policy: Mapped[schemas.core.FlowRunPolicy] = mapped_column(
494 Pydantic(schemas.core.FlowRunPolicy),
495 server_default="{}",
496 default=schemas.core.FlowRunPolicy,
497 )
498 tags: Mapped[list[str]] = mapped_column(JSON, server_default="[]", default=list)
499 labels: Mapped[Optional[schemas.core.KeyValueLabels]] = mapped_column(JSON)
501 created_by: Mapped[Optional[schemas.core.CreatedBy]] = mapped_column(
502 Pydantic(schemas.core.CreatedBy)
503 )
505 infrastructure_pid: Mapped[Optional[str]]
506 job_variables: Mapped[Optional[dict[str, Any]]] = mapped_column(
507 JSON, server_default="{}", default=dict
508 )
510 infrastructure_document_id: Mapped[Optional[uuid.UUID]] = mapped_column(
511 sa.ForeignKey("block_document.id", ondelete="CASCADE"), index=True
512 )
514 parent_task_run_id: Mapped[Optional[uuid.UUID]] = mapped_column(
515 sa.ForeignKey("task_run.id", ondelete="SET NULL", use_alter=True), index=True
516 )
518 auto_scheduled: Mapped[bool] = mapped_column(server_default="0", default=False)
520 # TODO remove this foreign key for significant delete performance gains
521 state_id: Mapped[Optional[uuid.UUID]] = mapped_column(
522 sa.ForeignKey("flow_run_state.id", ondelete="SET NULL", use_alter=True),
523 index=True,
524 )
526 work_queue_id: Mapped[Optional[uuid.UUID]] = mapped_column(
527 sa.ForeignKey("work_queue.id", ondelete="SET NULL"), index=True
528 )
530 # -------------------------- relationships
532 # current states are eagerly loaded unless otherwise specified
533 _state: Mapped[Optional["FlowRunState"]] = relationship(
534 lazy="selectin",
535 foreign_keys=[state_id],
536 primaryjoin="FlowRunState.id==FlowRun.state_id",
537 )
539 @hybrid_property
540 def state(self) -> Optional[FlowRunState]:
541 return self._state
543 @state.inplace.setter
544 def _set_state(self, value: Optional[FlowRunState]) -> None:
545 # because this is a slightly non-standard SQLAlchemy relationship, we
546 # prefer an explicit setter method to a setter property, because
547 # user expectations about SQLAlchemy attribute assignment might not be
548 # met, namely that an unrelated (from SQLAlchemy's perspective) field of
549 # the provided state is also modified. However, property assignment
550 # still works because the ORM model's __init__ depends on it.
551 return self.set_state(value)
553 def set_state(self, state: Optional[FlowRunState]) -> None:
554 """
555 If a state is assigned to this run, populate its run id.
557 This would normally be handled by the back-populated SQLAlchemy
558 relationship, but because this is a one-to-one pointer to a
559 one-to-many relationship, SQLAlchemy can't figure it out.
560 """
561 if state is not None: 561 ↛ 563line 561 didn't jump to line 563 because the condition on line 561 was always true
562 state.flow_run_id = self.id
563 self._state = state
565 flow: Mapped["Flow"] = relationship(back_populates="flow_runs", lazy="raise")
567 task_runs: Mapped[list["TaskRun"]] = relationship(
568 back_populates="flow_run",
569 lazy="raise",
570 # foreign_keys=lambda: [flow_run_id],
571 primaryjoin="TaskRun.flow_run_id==FlowRun.id",
572 )
574 parent_task_run: Mapped[Optional["TaskRun"]] = relationship(
575 back_populates="subflow_run",
576 lazy="raise",
577 foreign_keys=[parent_task_run_id],
578 )
580 work_queue: Mapped[Optional["WorkQueue"]] = relationship(
581 lazy="selectin", foreign_keys=[work_queue_id]
582 )
584 @declared_attr.directive
585 @classmethod
586 def __table_args__(cls) -> Iterable[sa.Index]:
587 return (
588 sa.Index(
589 "uq_flow_run__flow_id_idempotency_key",
590 cls.flow_id,
591 cls.idempotency_key,
592 unique=True,
593 ),
594 sa.Index(
595 "ix_flow_run__coalesce_start_time_expected_start_time_desc",
596 coalesce(cls.start_time, cls.expected_start_time).desc(),
597 ),
598 sa.Index(
599 "ix_flow_run__coalesce_start_time_expected_start_time_asc",
600 coalesce(cls.start_time, cls.expected_start_time).asc(),
601 ),
602 sa.Index(
603 "ix_flow_run__expected_start_time_desc",
604 cls.expected_start_time.desc(),
605 ),
606 sa.Index(
607 "ix_flow_run__next_scheduled_start_time_asc",
608 cls.next_scheduled_start_time.asc(),
609 ),
610 sa.Index(
611 "ix_flow_run__end_time_desc",
612 cls.end_time.desc(),
613 ),
614 sa.Index(
615 "ix_flow_run__start_time",
616 cls.start_time,
617 ),
618 sa.Index(
619 "ix_flow_run__state_type",
620 cls.state_type,
621 ),
622 sa.Index(
623 "ix_flow_run__state_name",
624 cls.state_name,
625 ),
626 sa.Index(
627 "ix_flow_run__state_timestamp",
628 cls.state_timestamp,
629 ),
630 sa.Index("trgm_ix_flow_run_name", cls.name, postgresql_using="gin").ddl_if(
631 dialect="postgresql"
632 ),
633 sa.Index(
634 # index names are at most 63 characters long.
635 "ix_flow_run__scheduler_deployment_id_auto_scheduled_next_schedu",
636 cls.deployment_id,
637 cls.auto_scheduled,
638 cls.next_scheduled_start_time,
639 postgresql_where=cls.state_type == schemas.states.StateType.SCHEDULED,
640 sqlite_where=cls.state_type == schemas.states.StateType.SCHEDULED,
641 ),
642 )
645_TaskInput = Union[
646 schemas.core.TaskRunResult,
647 schemas.core.FlowRunResult,
648 schemas.core.Parameter,
649 schemas.core.Constant,
650]
651_TaskInputs = dict[str, list[_TaskInput]]
654class TaskRun(Run):
655 """SQLAlchemy model of a task run."""
657 flow_run_id: Mapped[Optional[uuid.UUID]] = mapped_column(
658 sa.ForeignKey("flow_run.id", ondelete="cascade"), index=True
659 )
661 task_key: Mapped[str] = mapped_column()
662 dynamic_key: Mapped[str] = mapped_column()
663 cache_key: Mapped[Optional[str]]
664 cache_expiration: Mapped[Optional[DateTime]]
665 task_version: Mapped[Optional[str]]
666 flow_run_run_count: Mapped[int] = mapped_column(server_default="0", default=0)
667 empirical_policy: Mapped[schemas.core.TaskRunPolicy] = mapped_column(
668 Pydantic(schemas.core.TaskRunPolicy),
669 server_default="{}",
670 default=schemas.core.TaskRunPolicy,
671 )
672 task_inputs: Mapped[_TaskInputs] = mapped_column(
673 Pydantic(_TaskInputs), server_default="{}", default=dict
674 )
675 tags: Mapped[list[str]] = mapped_column(JSON, server_default="[]", default=list)
676 labels: Mapped[Optional[schemas.core.KeyValueLabels]] = mapped_column(JSON)
678 state_id: Mapped[Optional[uuid.UUID]] = mapped_column(
679 UUID,
680 index=True,
681 )
683 # -------------------------- relationships
685 # current states are eagerly loaded unless otherwise specified
686 _state: Mapped[Optional[TaskRunState]] = relationship(
687 lazy="selectin",
688 foreign_keys=[state_id],
689 primaryjoin="TaskRunState.id==TaskRun.state_id",
690 )
692 @hybrid_property
693 def state(self) -> Optional[TaskRunState]:
694 return self._state
696 @state.inplace.setter
697 def _set_state(self, value: Optional[TaskRunState]) -> None:
698 # because this is a slightly non-standard SQLAlchemy relationship, we
699 # prefer an explicit setter method to a setter property, because
700 # user expectations about SQLAlchemy attribute assignment might not be
701 # met, namely that an unrelated (from SQLAlchemy's perspective) field of
702 # the provided state is also modified. However, property assignment
703 # still works because the ORM model's __init__ depends on it.
704 return self.set_state(value)
706 def set_state(self, state: Optional[TaskRunState]) -> None:
707 """
708 If a state is assigned to this run, populate its run id.
710 This would normally be handled by the back-populated SQLAlchemy
711 relationship, but because this is a one-to-one pointer to a
712 one-to-many relationship, SQLAlchemy can't figure it out.
713 """
714 if state is not None:
715 state.task_run_id = self.id
716 self._state = state
718 flow_run: Mapped[Optional["FlowRun"]] = relationship(
719 back_populates="task_runs",
720 lazy="raise",
721 foreign_keys=[flow_run_id],
722 )
724 subflow_run: Mapped["FlowRun"] = relationship(
725 back_populates="parent_task_run",
726 lazy="raise",
727 # foreign_keys=["FlowRun.parent_task_run_id"],
728 primaryjoin="FlowRun.parent_task_run_id==TaskRun.id",
729 uselist=False,
730 )
732 @declared_attr.directive
733 @classmethod
734 def __table_args__(cls) -> Iterable[sa.Index]:
735 return (
736 sa.Index(
737 "uq_task_run__flow_run_id_task_key_dynamic_key",
738 cls.flow_run_id,
739 cls.task_key,
740 cls.dynamic_key,
741 unique=True,
742 ),
743 sa.Index(
744 "ix_task_run__expected_start_time_desc",
745 cls.expected_start_time.desc(),
746 ),
747 sa.Index(
748 "ix_task_run__next_scheduled_start_time_asc",
749 cls.next_scheduled_start_time.asc(),
750 ),
751 sa.Index(
752 "ix_task_run__end_time_desc",
753 cls.end_time.desc(),
754 ),
755 sa.Index(
756 "ix_task_run__start_time",
757 cls.start_time,
758 ),
759 sa.Index(
760 "ix_task_run__state_type",
761 cls.state_type,
762 ),
763 sa.Index(
764 "ix_task_run__state_type_start_time",
765 cls.state_type,
766 cls.start_time,
767 ),
768 sa.Index(
769 "ix_task_run__state_name",
770 cls.state_name,
771 ),
772 sa.Index(
773 "ix_task_run__state_timestamp",
774 cls.state_timestamp,
775 ),
776 sa.Index("trgm_ix_task_run_name", cls.name, postgresql_using="gin").ddl_if(
777 dialect="postgresql"
778 ),
779 )
782class DeploymentSchedule(Base):
783 deployment_id: Mapped[uuid.UUID] = mapped_column(
784 sa.ForeignKey("deployment.id", ondelete="CASCADE"), index=True
785 )
787 schedule: Mapped[schemas.schedules.SCHEDULE_TYPES] = mapped_column(
788 Pydantic(schemas.schedules.SCHEDULE_TYPES)
789 )
790 active: Mapped[bool] = mapped_column(default=True)
791 max_scheduled_runs: Mapped[Optional[int]]
792 parameters: Mapped[dict[str, Any]] = mapped_column(
793 JSON, server_default="{}", default=dict, nullable=False
794 )
795 slug: Mapped[Optional[str]] = mapped_column(sa.String, nullable=True)
797 @declared_attr.directive
798 @classmethod
799 def __table_args__(cls) -> Iterable[sa.Index]:
800 return (
801 sa.Index(
802 "ix_deployment_schedule__deployment_id__slug",
803 cls.deployment_id,
804 cls.slug,
805 unique=True,
806 ),
807 sa.Index(
808 "ix_deployment_schedule__slug",
809 cls.slug,
810 unique=False,
811 ),
812 )
815class Deployment(Base):
816 """SQLAlchemy model of a deployment."""
818 name: Mapped[str]
819 version: Mapped[Optional[str]]
820 description: Mapped[Optional[str]] = mapped_column(sa.Text())
821 work_queue_name: Mapped[Optional[str]] = mapped_column(index=True)
822 infra_overrides: Mapped[dict[str, Any]] = mapped_column(
823 JSON, server_default="{}", default=dict
824 )
825 path: Mapped[Optional[str]]
826 entrypoint: Mapped[Optional[str]]
828 last_polled: Mapped[Optional[DateTime]]
829 status: Mapped[DeploymentStatus] = mapped_column(
830 sa.Enum(DeploymentStatus, name="deployment_status"),
831 default=DeploymentStatus.NOT_READY,
832 server_default="NOT_READY",
833 )
835 @declared_attr
836 def job_variables(self) -> Mapped[dict[str, Any]]:
837 return synonym("infra_overrides")
839 flow_id: Mapped[uuid.UUID] = mapped_column(
840 sa.ForeignKey("flow.id", ondelete="CASCADE"), index=True
841 )
843 work_queue_id: Mapped[Optional[uuid.UUID]] = mapped_column(
844 sa.ForeignKey("work_queue.id", ondelete="SET NULL"), index=True
845 )
846 paused: Mapped[bool] = mapped_column(server_default="0", default=False, index=True)
848 schedules: Mapped[list["DeploymentSchedule"]] = relationship(
849 lazy="selectin", order_by=lambda: DeploymentSchedule.updated.desc()
850 )
852 # deprecated in favor of `concurrency_limit_id` FK
853 _concurrency_limit: Mapped[Optional[int]] = mapped_column(name="concurrency_limit")
854 concurrency_limit_id: Mapped[Optional[uuid.UUID]] = mapped_column(
855 sa.ForeignKey("concurrency_limit_v2.id", ondelete="SET NULL"),
856 )
857 global_concurrency_limit: Mapped[Optional["ConcurrencyLimitV2"]] = relationship(
858 lazy="selectin",
859 )
860 concurrency_options: Mapped[Optional[schemas.core.ConcurrencyOptions]] = (
861 mapped_column(
862 Pydantic(schemas.core.ConcurrencyOptions),
863 server_default=None,
864 nullable=True,
865 default=None,
866 )
867 )
869 tags: Mapped[list[str]] = mapped_column(JSON, server_default="[]", default=list)
870 labels: Mapped[Optional[schemas.core.KeyValueLabels]] = mapped_column(JSON)
871 parameters: Mapped[dict[str, Any]] = mapped_column(
872 JSON, server_default="{}", default=dict
873 )
874 pull_steps: Mapped[Optional[list[dict[str, Any]]]] = mapped_column(
875 JSON, default=list
876 )
877 parameter_openapi_schema: Mapped[Optional[dict[str, Any]]] = mapped_column(
878 JSON, default=dict
879 )
880 enforce_parameter_schema: Mapped[bool] = mapped_column(
881 default=True, server_default="0"
882 )
883 created_by: Mapped[Optional[schemas.core.CreatedBy]] = mapped_column(
884 Pydantic(schemas.core.CreatedBy)
885 )
886 updated_by: Mapped[Optional[schemas.core.UpdatedBy]] = mapped_column(
887 Pydantic(schemas.core.UpdatedBy)
888 )
890 infrastructure_document_id: Mapped[Optional[uuid.UUID]] = mapped_column(
891 sa.ForeignKey("block_document.id", ondelete="CASCADE"), index=False
892 )
894 storage_document_id: Mapped[Optional[uuid.UUID]] = mapped_column(
895 sa.ForeignKey("block_document.id", ondelete="CASCADE"),
896 index=False,
897 )
899 flow: Mapped["Flow"] = relationship(
900 "Flow", back_populates="deployments", lazy="raise"
901 )
903 work_queue: Mapped[Optional["WorkQueue"]] = relationship(
904 lazy="selectin", foreign_keys=[work_queue_id]
905 )
907 __table_args__: Any = (
908 sa.Index(
909 "uq_deployment__flow_id_name",
910 "flow_id",
911 "name",
912 unique=True,
913 ),
914 sa.Index(
915 "ix_deployment__created",
916 "created",
917 ),
918 sa.Index("trgm_ix_deployment_name", "name", postgresql_using="gin").ddl_if(
919 dialect="postgresql"
920 ),
921 )
924class Log(Base):
925 """
926 SQLAlchemy model of a logging statement.
927 """
929 name: Mapped[str]
930 level: Mapped[int] = mapped_column(sa.SmallInteger, index=True)
931 flow_run_id: Mapped[Optional[uuid.UUID]] = mapped_column(index=True)
932 task_run_id: Mapped[Optional[uuid.UUID]] = mapped_column(index=True)
933 message: Mapped[str] = mapped_column(sa.Text)
935 # The client-side timestamp of this logged statement.
936 timestamp: Mapped[DateTime] = mapped_column(index=True)
938 __table_args__: Any = (
939 sa.Index(
940 "ix_log__flow_run_id_timestamp",
941 "flow_run_id",
942 "timestamp",
943 ),
944 )
947class ConcurrencyLimit(Base):
948 tag: Mapped[str]
949 concurrency_limit: Mapped[int]
950 active_slots: Mapped[list[str]] = mapped_column(
951 JSON, server_default="[]", default=list
952 )
954 __table_args__: Any = (sa.Index("uq_concurrency_limit__tag", "tag", unique=True),)
957class ConcurrencyLimitV2(Base):
958 active: Mapped[bool] = mapped_column(default=True)
959 name: Mapped[str]
960 limit: Mapped[int]
961 active_slots: Mapped[int] = mapped_column(default=0)
962 denied_slots: Mapped[int] = mapped_column(default=0)
964 slot_decay_per_second: Mapped[float] = mapped_column(default=0.0)
965 avg_slot_occupancy_seconds: Mapped[float] = mapped_column(default=2.0)
967 __table_args__: Any = (sa.UniqueConstraint("name"),)
970class BlockType(Base):
971 name: Mapped[str]
972 slug: Mapped[str]
973 logo_url: Mapped[Optional[str]]
974 documentation_url: Mapped[Optional[str]]
975 description: Mapped[Optional[str]]
976 code_example: Mapped[Optional[str]]
977 is_protected: Mapped[bool] = mapped_column(server_default="0", default=False)
979 __table_args__: Any = (
980 sa.Index(
981 "uq_block_type__slug",
982 "slug",
983 unique=True,
984 ),
985 sa.Index("trgm_ix_block_type_name", "name", postgresql_using="gin").ddl_if(
986 dialect="postgresql"
987 ),
988 )
991class BlockSchema(Base):
992 checksum: Mapped[str]
993 fields: Mapped[dict[str, Any]] = mapped_column(
994 JSON, server_default="{}", default=dict
995 )
996 capabilities: Mapped[list[str]] = mapped_column(
997 JSON, server_default="[]", default=list
998 )
999 version: Mapped[str] = mapped_column(
1000 server_default=schemas.core.DEFAULT_BLOCK_SCHEMA_VERSION,
1001 )
1003 block_type_id: Mapped[uuid.UUID] = mapped_column(
1004 sa.ForeignKey("block_type.id", ondelete="cascade"), index=True
1005 )
1007 block_type: Mapped["BlockType"] = relationship(lazy="selectin")
1009 __table_args__: Any = (
1010 sa.Index(
1011 "uq_block_schema__checksum_version",
1012 "checksum",
1013 "version",
1014 unique=True,
1015 ),
1016 sa.Index("ix_block_schema__created", "created"),
1017 sa.Index(
1018 "ix_block_schema__capabilities", "capabilities", postgresql_using="gin"
1019 ).ddl_if(dialect="postgresql"),
1020 )
1023class BlockSchemaReference(Base):
1024 name: Mapped[str]
1026 parent_block_schema_id: Mapped[uuid.UUID] = mapped_column(
1027 sa.ForeignKey("block_schema.id", ondelete="cascade")
1028 )
1030 reference_block_schema_id: Mapped[uuid.UUID] = mapped_column(
1031 sa.ForeignKey("block_schema.id", ondelete="cascade")
1032 )
1035class BlockDocument(Base):
1036 name: Mapped[str] = mapped_column(index=True)
1037 data: Mapped[Any] = mapped_column(JSON, server_default="{}", default=dict)
1038 is_anonymous: Mapped[bool] = mapped_column(server_default="0", index=True)
1040 block_type_name: Mapped[Optional[str]]
1042 block_type_id: Mapped[uuid.UUID] = mapped_column(
1043 sa.ForeignKey("block_type.id", ondelete="cascade")
1044 )
1046 block_type: Mapped["BlockType"] = relationship(lazy="selectin")
1048 block_schema_id: Mapped[uuid.UUID] = mapped_column(
1049 sa.ForeignKey("block_schema.id", ondelete="cascade")
1050 )
1052 block_schema: Mapped["BlockSchema"] = relationship(lazy="selectin")
1054 __table_args__: Any = (
1055 sa.Index(
1056 "uq_block__type_id_name",
1057 "block_type_id",
1058 "name",
1059 unique=True,
1060 ),
1061 sa.Index("ix_block_document__block_type_name__name", "block_type_name", "name"),
1062 sa.Index("trgm_ix_block_document_name", "name", postgresql_using="gin").ddl_if(
1063 dialect="postgresql"
1064 ),
1065 )
1067 async def encrypt_data(self, session: AsyncSession, data: dict[str, Any]) -> None:
1068 """
1069 Store encrypted data on the ORM model
1071 Note: will only succeed if the caller has sufficient permission.
1072 """
1073 self.data = await encrypt_fernet(session, data)
1075 async def decrypt_data(self, session: AsyncSession) -> dict[str, Any]:
1076 """
1077 Retrieve decrypted data from the ORM model.
1079 Note: will only succeed if the caller has sufficient permission.
1080 """
1081 return await decrypt_fernet(session, self.data)
1084class BlockDocumentReference(Base):
1085 name: Mapped[str]
1087 parent_block_document_id: Mapped[uuid.UUID] = mapped_column(
1088 sa.ForeignKey("block_document.id", ondelete="cascade"),
1089 )
1091 reference_block_document_id: Mapped[uuid.UUID] = mapped_column(
1092 sa.ForeignKey("block_document.id", ondelete="cascade"),
1093 )
1095 __table_args__: Any = (
1096 sa.CheckConstraint(
1097 "parent_block_document_id != reference_block_document_id",
1098 name="ck_block_document_reference__no_self_reference",
1099 ),
1100 )
1103class Configuration(Base):
1104 key: Mapped[str] = mapped_column(index=True)
1105 value: Mapped[dict[str, Any]] = mapped_column(JSON)
1107 __table_args__: Any = (sa.UniqueConstraint("key"),)
1110class SavedSearch(Base):
1111 """SQLAlchemy model of a saved search."""
1113 name: Mapped[str]
1114 filters: Mapped[list[dict[str, Any]]] = mapped_column(
1115 JSON, server_default="[]", default=list
1116 )
1118 __table_args__: Any = (sa.UniqueConstraint("name"),)
1121class WorkQueue(Base):
1122 """SQLAlchemy model of a work queue"""
1124 name: Mapped[str]
1126 filter: Mapped[Optional[schemas.core.QueueFilter]] = mapped_column(
1127 Pydantic(schemas.core.QueueFilter)
1128 )
1129 description: Mapped[str] = mapped_column(default="", server_default="")
1130 is_paused: Mapped[bool] = mapped_column(server_default="0", default=False)
1131 concurrency_limit: Mapped[Optional[int]]
1132 priority: Mapped[int]
1134 last_polled: Mapped[Optional[DateTime]]
1135 status: Mapped[WorkQueueStatus] = mapped_column(
1136 sa.Enum(WorkQueueStatus, name="work_queue_status"),
1137 default=WorkQueueStatus.NOT_READY,
1138 server_default=WorkQueueStatus.NOT_READY,
1139 )
1141 work_pool_id: Mapped[uuid.UUID] = mapped_column(
1142 sa.ForeignKey("work_pool.id", ondelete="cascade"), index=True
1143 )
1145 work_pool: Mapped["WorkPool"] = relationship(
1146 lazy="selectin", foreign_keys=[work_pool_id]
1147 )
1149 __table_args__: ClassVar[Any] = (
1150 sa.UniqueConstraint("work_pool_id", "name"),
1151 sa.Index("ix_work_queue__work_pool_id_priority", "work_pool_id", "priority"),
1152 sa.Index("trgm_ix_work_queue_name", "name", postgresql_using="gin").ddl_if(
1153 dialect="postgresql"
1154 ),
1155 )
1158class WorkPool(Base):
1159 """SQLAlchemy model of an worker"""
1161 name: Mapped[str]
1162 description: Mapped[Optional[str]]
1163 type: Mapped[str] = mapped_column(index=True)
1164 base_job_template: Mapped[dict[str, Any]] = mapped_column(
1165 JSON, server_default="{}", default={}
1166 )
1167 is_paused: Mapped[bool] = mapped_column(server_default="0", default=False)
1168 default_queue_id: Mapped[Optional[uuid.UUID]] = mapped_column(
1169 UUID,
1170 sa.ForeignKey("work_queue.id", ondelete="RESTRICT", use_alter=True),
1171 nullable=True,
1172 )
1173 concurrency_limit: Mapped[Optional[int]]
1175 status: Mapped[WorkPoolStatus] = mapped_column(
1176 sa.Enum(WorkPoolStatus, name="work_pool_status"),
1177 default=WorkPoolStatus.NOT_READY,
1178 server_default=WorkPoolStatus.NOT_READY,
1179 )
1180 last_transitioned_status_at: Mapped[Optional[DateTime]]
1181 last_status_event_id: Mapped[Optional[uuid.UUID]]
1183 storage_configuration: Mapped[schemas.core.WorkPoolStorageConfiguration] = (
1184 mapped_column(
1185 Pydantic(schemas.core.WorkPoolStorageConfiguration),
1186 server_default="{}",
1187 default=schemas.core.WorkPoolStorageConfiguration,
1188 nullable=False,
1189 )
1190 )
1192 __table_args__: Any = (sa.UniqueConstraint("name"),)
1195class Worker(Base):
1196 """SQLAlchemy model of an worker"""
1198 work_pool_id: Mapped[uuid.UUID] = mapped_column(
1199 sa.ForeignKey("work_pool.id", ondelete="cascade"), index=True
1200 )
1202 name: Mapped[str]
1203 last_heartbeat_time: Mapped[DateTime] = mapped_column(
1204 server_default=sa.func.now(), default=lambda: now("UTC")
1205 )
1206 heartbeat_interval_seconds: Mapped[Optional[int]]
1208 status: Mapped[WorkerStatus] = mapped_column(
1209 sa.Enum(WorkerStatus, name="worker_status"),
1210 default=WorkerStatus.OFFLINE,
1211 server_default=WorkerStatus.OFFLINE,
1212 )
1214 __table_args__: Any = (
1215 sa.UniqueConstraint("work_pool_id", "name"),
1216 sa.Index(
1217 "ix_worker__work_pool_id_last_heartbeat_time",
1218 "work_pool_id",
1219 "last_heartbeat_time",
1220 ),
1221 )
1224class Agent(Base):
1225 """SQLAlchemy model of an agent"""
1227 name: Mapped[str]
1229 work_queue_id: Mapped[uuid.UUID] = mapped_column(
1230 sa.ForeignKey("work_queue.id"), index=True
1231 )
1233 last_activity_time: Mapped[DateTime] = mapped_column(
1234 server_default=sa.func.now(), default=lambda: now("UTC")
1235 )
1237 __table_args__: Any = (sa.UniqueConstraint("name"),)
1240class Variable(Base):
1241 name: Mapped[str]
1242 value: Mapped[Optional[Any]] = mapped_column(JSON)
1243 tags: Mapped[list[str]] = mapped_column(JSON, server_default="[]", default=list)
1245 __table_args__: Any = (sa.UniqueConstraint("name"),)
1248class FlowRunInput(Base):
1249 flow_run_id: Mapped[uuid.UUID] = mapped_column(
1250 sa.ForeignKey("flow_run.id", ondelete="cascade")
1251 )
1253 key: Mapped[str]
1254 value: Mapped[str] = mapped_column(sa.Text())
1255 sender: Mapped[Optional[str]]
1257 __table_args__: Any = (sa.UniqueConstraint("flow_run_id", "key"),)
1260class CsrfToken(Base):
1261 token: Mapped[str]
1262 client: Mapped[str] = mapped_column(unique=True)
1263 expiration: Mapped[DateTime]
1266class Automation(Base):
1267 name: Mapped[str]
1268 description: Mapped[str] = mapped_column(default="")
1270 enabled: Mapped[bool] = mapped_column(server_default="1", default=True)
1271 tags: Mapped[list[str]] = mapped_column(JSON, server_default="[]", default=list)
1273 trigger: Mapped[ServerTriggerTypes] = mapped_column(Pydantic(ServerTriggerTypes))
1275 actions: Mapped[ServerActionTypes] = mapped_column(
1276 Pydantic(list[ServerActionTypes])
1277 )
1278 actions_on_trigger: Mapped[list[ServerActionTypes]] = mapped_column(
1279 Pydantic(list[ServerActionTypes]), server_default="[]", default=list
1280 )
1281 actions_on_resolve: Mapped[list[ServerActionTypes]] = mapped_column(
1282 Pydantic(list[ServerActionTypes]), server_default="[]", default=list
1283 )
1285 related_resources: Mapped[list["AutomationRelatedResource"]] = relationship(
1286 "AutomationRelatedResource", back_populates="automation", lazy="raise"
1287 )
1289 @classmethod
1290 def sort_expression(cls, value: AutomationSort) -> sa.ColumnExpressionArgument[Any]:
1291 """Return an expression used to sort Automations"""
1292 sort_mapping: dict[AutomationSort, sa.ColumnExpressionArgument[Any]] = {
1293 AutomationSort.CREATED_DESC: cls.created.desc(),
1294 AutomationSort.UPDATED_DESC: cls.updated.desc(),
1295 AutomationSort.NAME_ASC: cls.name.asc(),
1296 AutomationSort.NAME_DESC: cls.name.desc(),
1297 }
1298 return sort_mapping[value]
1301class AutomationBucket(Base):
1302 __table_args__: Any = (
1303 sa.Index(
1304 "uq_automation_bucket__automation_id__trigger_id__bucketing_key",
1305 "automation_id",
1306 "trigger_id",
1307 "bucketing_key",
1308 unique=True,
1309 ),
1310 sa.Index(
1311 "ix_automation_bucket__automation_id__end",
1312 "automation_id",
1313 "end",
1314 ),
1315 )
1317 automation_id: Mapped[uuid.UUID] = mapped_column(
1318 sa.ForeignKey("automation.id", ondelete="CASCADE")
1319 )
1321 trigger_id: Mapped[uuid.UUID]
1323 bucketing_key: Mapped[list[str]] = mapped_column(
1324 JSON, server_default="[]", default=list
1325 )
1327 last_event: Mapped[Optional[ReceivedEvent]] = mapped_column(Pydantic(ReceivedEvent))
1329 start: Mapped[DateTime]
1330 end: Mapped[DateTime]
1332 count: Mapped[int]
1334 last_operation: Mapped[Optional[str]]
1336 triggered_at: Mapped[Optional[DateTime]]
1339class AutomationRelatedResource(Base):
1340 __table_args__: Any = (
1341 sa.Index(
1342 "uq_automation_related_resource__automation_id__resource_id",
1343 "automation_id",
1344 "resource_id",
1345 unique=True,
1346 ),
1347 )
1349 automation_id: Mapped[uuid.UUID] = mapped_column(
1350 sa.ForeignKey("automation.id", ondelete="CASCADE")
1351 )
1353 resource_id: Mapped[Optional[str]] = mapped_column(index=True)
1354 automation_owned_by_resource: Mapped[bool] = mapped_column(
1355 default=False, server_default="0"
1356 )
1358 automation: Mapped["Automation"] = relationship(
1359 "Automation", back_populates="related_resources", lazy="raise"
1360 )
1363class CompositeTriggerChildFiring(Base):
1364 __table_args__: Any = (
1365 sa.Index(
1366 "uq_composite_trigger_child_firing__a_id__pt_id__ct__id",
1367 "automation_id",
1368 "parent_trigger_id",
1369 "child_trigger_id",
1370 unique=True,
1371 ),
1372 )
1374 automation_id: Mapped[uuid.UUID] = mapped_column(
1375 sa.ForeignKey("automation.id", ondelete="CASCADE")
1376 )
1378 parent_trigger_id: Mapped[uuid.UUID]
1380 child_trigger_id: Mapped[uuid.UUID]
1381 child_firing_id: Mapped[uuid.UUID]
1382 child_fired_at: Mapped[Optional[DateTime]]
1383 child_firing: Mapped[Firing] = mapped_column(Pydantic(Firing))
1386class AutomationEventFollower(Base):
1387 __table_args__: Any = (
1388 sa.Index(
1389 "uq_follower_for_scope",
1390 "scope",
1391 "follower_event_id",
1392 unique=True,
1393 ),
1394 # allows lookup on (scope, leader_event_id) to use an index-only instead of full table scan
1395 sa.Index(
1396 "ix_ae_follower_scope_leader",
1397 "scope",
1398 "leader_event_id",
1399 ),
1400 )
1401 scope: Mapped[str] = mapped_column(default="", index=True)
1402 leader_event_id: Mapped[uuid.UUID] = mapped_column(index=True)
1403 follower_event_id: Mapped[uuid.UUID]
1404 received: Mapped[DateTime] = mapped_column(index=True)
1405 follower: Mapped[ReceivedEvent] = mapped_column(Pydantic(ReceivedEvent))
1408class Event(Base):
1409 @declared_attr.directive
1410 def __tablename__(cls) -> str:
1411 return "events"
1413 __table_args__: Any = (
1414 sa.Index(
1415 "ix_events__related_resource_ids_gin",
1416 "related_resource_ids",
1417 postgresql_using="gin",
1418 ),
1419 sa.Index("ix_events__occurred", "occurred"),
1420 sa.Index("ix_events__event__id", "event", "id"),
1421 sa.Index(
1422 "ix_events__event_resource_id_occurred",
1423 "event",
1424 "resource_id",
1425 "occurred",
1426 ),
1427 sa.Index("ix_events__occurred_id", "occurred", "id"),
1428 sa.Index("ix_events__event_occurred_id", "event", "occurred", "id"),
1429 sa.Index(
1430 "ix_events__related_gin",
1431 "related",
1432 postgresql_using="gin",
1433 ),
1434 sa.Index(
1435 "ix_events__event_occurred",
1436 "event",
1437 "occurred",
1438 ),
1439 )
1441 occurred: Mapped[DateTime]
1442 event: Mapped[str] = mapped_column(sa.Text())
1443 resource_id: Mapped[str] = mapped_column(sa.Text())
1444 resource: Mapped[dict[str, Any]] = mapped_column(JSON())
1445 related_resource_ids: Mapped[list[str]] = mapped_column(
1446 JSON(), server_default="[]", default=list
1447 )
1448 related: Mapped[list[dict[str, Any]]] = mapped_column(
1449 JSON(), server_default="[]", default=list
1450 )
1451 payload: Mapped[dict[str, Any]] = mapped_column(JSON())
1452 received: Mapped[DateTime]
1453 recorded: Mapped[DateTime]
1454 follows: Mapped[Optional[uuid.UUID]]
1457class EventResource(Base):
1458 @declared_attr.directive
1459 def __tablename__(cls) -> str:
1460 return "event_resources"
1462 __table_args__: Any = (
1463 sa.Index(
1464 "ix_event_resources__resource_id__occurred",
1465 "resource_id",
1466 "occurred",
1467 ),
1468 sa.Index(
1469 "ix_event_resources__occurred",
1470 "occurred",
1471 ),
1472 )
1474 occurred: Mapped[DateTime]
1475 resource_id: Mapped[str] = mapped_column(sa.Text())
1476 resource_role: Mapped[str] = mapped_column(sa.Text())
1477 resource: Mapped[dict[str, Any]] = mapped_column(sa_JSON)
1478 event_id: Mapped[uuid.UUID]
1481# These are temporary until we've migrated all the references to the new,
1482# non-ORM names
1484ORMFlow = Flow
1485ORMFlowRunState = FlowRunState
1486ORMTaskRunState = TaskRunState
1487ORMArtifact = Artifact
1488ORMArtifactCollection = ArtifactCollection
1489ORMTaskRunStateCache = TaskRunStateCache
1490ORMRun = Run
1491ORMFlowRun = FlowRun
1492ORMTaskRun = TaskRun
1493ORMDeploymentSchedule = DeploymentSchedule
1494ORMDeployment = Deployment
1495ORMLog = Log
1496ORMConcurrencyLimit = ConcurrencyLimit
1497ORMConcurrencyLimitV2 = ConcurrencyLimitV2
1498ORMBlockType = BlockType
1499ORMBlockSchema = BlockSchema
1500ORMBlockSchemaReference = BlockSchemaReference
1501ORMBlockDocument = BlockDocument
1502ORMBlockDocumentReference = BlockDocumentReference
1503ORMConfiguration = Configuration
1504ORMSavedSearch = SavedSearch
1505ORMWorkQueue = WorkQueue
1506ORMWorkPool = WorkPool
1507ORMWorker = Worker
1508ORMAgent = Agent
1509ORMVariable = Variable
1510ORMFlowRunInput = FlowRunInput
1511ORMCsrfToken = CsrfToken
1512ORMAutomation = Automation
1513ORMAutomationBucket = AutomationBucket
1514ORMAutomationRelatedResource = AutomationRelatedResource
1515ORMCompositeTriggerChildFiring = CompositeTriggerChildFiring
1516ORMAutomationEventFollower = AutomationEventFollower
1517ORMEvent = Event
1518ORMEventResource = EventResource
1521_UpsertColumns = Iterable[Union[str, "sa.Column[Any]", roles.DDLConstraintColumnRole]]
1524class BaseORMConfiguration(ABC):
1525 """
1526 Abstract base class used to inject database-specific ORM configuration into Prefect.
1528 Modifications to core Prefect REST API data structures can have unintended consequences.
1529 Use with caution.
1530 """
1532 def unique_key(self) -> tuple[Hashable, ...]:
1533 """
1534 Returns a key used to determine whether to instantiate a new DB interface.
1535 """
1536 return (self.__class__, Base.metadata)
1538 @property
1539 @abstractmethod
1540 def versions_dir(self) -> Path:
1541 """Directory containing migrations"""
1542 ...
1544 @property
1545 def deployment_unique_upsert_columns(self) -> _UpsertColumns:
1546 """Unique columns for upserting a Deployment"""
1547 return [Deployment.flow_id, Deployment.name]
1549 @property
1550 def concurrency_limit_unique_upsert_columns(self) -> _UpsertColumns:
1551 """Unique columns for upserting a ConcurrencyLimit"""
1552 return [ConcurrencyLimit.tag]
1554 @property
1555 def flow_run_unique_upsert_columns(self) -> _UpsertColumns:
1556 """Unique columns for upserting a FlowRun"""
1557 return [FlowRun.flow_id, FlowRun.idempotency_key]
1559 @property
1560 def block_type_unique_upsert_columns(self) -> _UpsertColumns:
1561 """Unique columns for upserting a BlockType"""
1562 return [BlockType.slug]
1564 @property
1565 def artifact_collection_unique_upsert_columns(self) -> _UpsertColumns:
1566 """Unique columns for upserting an ArtifactCollection"""
1567 return [ArtifactCollection.key]
1569 @property
1570 def block_schema_unique_upsert_columns(self) -> _UpsertColumns:
1571 """Unique columns for upserting a BlockSchema"""
1572 return [BlockSchema.checksum, BlockSchema.version]
1574 @property
1575 def flow_unique_upsert_columns(self) -> _UpsertColumns:
1576 """Unique columns for upserting a Flow"""
1577 return [Flow.name]
1579 @property
1580 def saved_search_unique_upsert_columns(self) -> _UpsertColumns:
1581 """Unique columns for upserting a SavedSearch"""
1582 return [SavedSearch.name]
1584 @property
1585 def task_run_unique_upsert_columns(self) -> _UpsertColumns:
1586 """Unique columns for upserting a TaskRun"""
1587 return [
1588 TaskRun.flow_run_id,
1589 TaskRun.task_key,
1590 TaskRun.dynamic_key,
1591 ]
1593 @property
1594 def block_document_unique_upsert_columns(self) -> _UpsertColumns:
1595 """Unique columns for upserting a BlockDocument"""
1596 return [BlockDocument.block_type_id, BlockDocument.name]
1599class AsyncPostgresORMConfiguration(BaseORMConfiguration):
1600 """Postgres specific orm configuration"""
1602 @property
1603 def versions_dir(self) -> Path:
1604 """Directory containing migrations"""
1605 import prefect.server.database
1607 return (
1608 Path(prefect.server.database.__file__).parent
1609 / "_migrations"
1610 / "versions"
1611 / "postgresql"
1612 )
1615class AioSqliteORMConfiguration(BaseORMConfiguration):
1616 """SQLite specific orm configuration"""
1618 @property
1619 def versions_dir(self) -> Path:
1620 """Directory containing migrations"""
1621 import prefect.server.database
1623 return (
1624 Path(prefect.server.database.__file__).parent
1625 / "_migrations"
1626 / "versions"
1627 / "sqlite"
1628 )