mirror of
https://github.com/Skyvern-AI/skyvern.git
synced 2026-07-09 16:09:13 +00:00
Add workflow run tag event data layer (#7204)
This commit is contained in:
Aaron Perez
GitHub
parent
bee2b836ce
commit
78fa30ed11
5 changed files with 927 additions and 113 deletions
|
|
@ -0,0 +1,133 @@
|
|||
"""add workflow run tag events table
|
||||
|
||||
Revision ID: bec06d149264
|
||||
Revises: 5b618ec6dce6
|
||||
Create Date: 2026-07-08T17:54:31.712270+00:00
|
||||
|
||||
"""
|
||||
|
||||
from typing import Sequence, Union
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "bec06d149264"
|
||||
down_revision: Union[str, None] = "5b618ec6dce6"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"workflow_run_tag_events",
|
||||
sa.Column("tag_event_id", sa.String(), nullable=False),
|
||||
sa.Column("workflow_run_id", sa.String(), nullable=False),
|
||||
sa.Column("organization_id", sa.String(), nullable=False),
|
||||
sa.Column("key", sa.String(), nullable=True),
|
||||
sa.Column("value", sa.String(), nullable=True),
|
||||
sa.Column("event_type", sa.String(), nullable=False),
|
||||
sa.Column("set_at", sa.DateTime(), nullable=False),
|
||||
sa.Column("set_by", sa.String(), nullable=False),
|
||||
sa.Column("source", sa.String(), nullable=False),
|
||||
sa.Column("caller_type", sa.String(), nullable=True),
|
||||
sa.Column("superseded_at", sa.DateTime(), nullable=True),
|
||||
sa.Column("inherited_from_tag_event_id", sa.String(), nullable=True),
|
||||
sa.Column("created_at", sa.DateTime(), nullable=False),
|
||||
sa.Column("modified_at", sa.DateTime(), nullable=False),
|
||||
sa.CheckConstraint("event_type IN ('set', 'delete')", name="ck_workflow_run_tag_events_event_type"),
|
||||
sa.CheckConstraint(
|
||||
"source IN ('manual', 'bulk_apply', 'backfill', 'inherited', 'import', 'system')",
|
||||
name="ck_workflow_run_tag_events_source",
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
"caller_type IS NULL OR caller_type IN ('user', 'api_key', 'system')",
|
||||
name="ck_workflow_run_tag_events_caller_type",
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
"event_type != 'set' OR value IS NOT NULL",
|
||||
name="ck_workflow_run_tag_events_set_has_value",
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
["workflow_run_id"],
|
||||
["workflow_runs.workflow_run_id"],
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
["organization_id"],
|
||||
["organizations.organization_id"],
|
||||
),
|
||||
sa.PrimaryKeyConstraint("tag_event_id"),
|
||||
)
|
||||
op.create_index(
|
||||
"workflow_run_tag_events_org_wr_set_at_idx",
|
||||
"workflow_run_tag_events",
|
||||
["organization_id", "workflow_run_id", "set_at"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"workflow_run_tag_events_org_key_value_active_idx",
|
||||
"workflow_run_tag_events",
|
||||
["organization_id", "key", "value"],
|
||||
unique=False,
|
||||
postgresql_include=["workflow_run_id"],
|
||||
postgresql_where=sa.text("superseded_at IS NULL AND event_type = 'set'"),
|
||||
)
|
||||
op.create_index(
|
||||
"workflow_run_tag_events_org_value_active_idx",
|
||||
"workflow_run_tag_events",
|
||||
["organization_id", "value"],
|
||||
unique=False,
|
||||
postgresql_include=["workflow_run_id"],
|
||||
postgresql_where=sa.text("superseded_at IS NULL AND event_type = 'set'"),
|
||||
)
|
||||
op.create_index(
|
||||
"workflow_run_tag_events_org_set_at_idx",
|
||||
"workflow_run_tag_events",
|
||||
["organization_id", "set_at"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"workflow_run_tag_events_active_grouped_unique",
|
||||
"workflow_run_tag_events",
|
||||
["organization_id", "workflow_run_id", "key"],
|
||||
unique=True,
|
||||
postgresql_where=sa.text("superseded_at IS NULL AND event_type = 'set' AND key IS NOT NULL"),
|
||||
sqlite_where=sa.text("superseded_at IS NULL AND event_type = 'set' AND key IS NOT NULL"),
|
||||
)
|
||||
op.create_index(
|
||||
"workflow_run_tag_events_active_label_unique",
|
||||
"workflow_run_tag_events",
|
||||
["organization_id", "workflow_run_id", "value"],
|
||||
unique=True,
|
||||
postgresql_where=sa.text("superseded_at IS NULL AND event_type = 'set' AND key IS NULL"),
|
||||
sqlite_where=sa.text("superseded_at IS NULL AND event_type = 'set' AND key IS NULL"),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index(
|
||||
"workflow_run_tag_events_active_label_unique",
|
||||
table_name="workflow_run_tag_events",
|
||||
postgresql_where=sa.text("superseded_at IS NULL AND event_type = 'set' AND key IS NULL"),
|
||||
)
|
||||
op.drop_index(
|
||||
"workflow_run_tag_events_active_grouped_unique",
|
||||
table_name="workflow_run_tag_events",
|
||||
postgresql_where=sa.text("superseded_at IS NULL AND event_type = 'set' AND key IS NOT NULL"),
|
||||
)
|
||||
op.drop_index("workflow_run_tag_events_org_set_at_idx", table_name="workflow_run_tag_events")
|
||||
op.drop_index(
|
||||
"workflow_run_tag_events_org_value_active_idx",
|
||||
table_name="workflow_run_tag_events",
|
||||
postgresql_include=["workflow_run_id"],
|
||||
postgresql_where=sa.text("superseded_at IS NULL AND event_type = 'set'"),
|
||||
)
|
||||
op.drop_index(
|
||||
"workflow_run_tag_events_org_key_value_active_idx",
|
||||
table_name="workflow_run_tag_events",
|
||||
postgresql_include=["workflow_run_id"],
|
||||
postgresql_where=sa.text("superseded_at IS NULL AND event_type = 'set'"),
|
||||
)
|
||||
op.drop_index("workflow_run_tag_events_org_wr_set_at_idx", table_name="workflow_run_tag_events")
|
||||
op.drop_table("workflow_run_tag_events")
|
||||
|
|
@ -79,6 +79,7 @@ WORKFLOW_RUN_CREDENTIAL_SELECTION_PREFIX = "wrcs"
|
|||
SCRIPT_FALLBACK_EPISODE_PREFIX = "sfe"
|
||||
WORKFLOW_SCHEDULE_PREFIX = "wfs"
|
||||
TAG_EVENT_PREFIX = "tge"
|
||||
RUN_TAG_EVENT_PREFIX = "rtge"
|
||||
TAG_KEY_PREFIX = "tkey"
|
||||
TAG_VALUE_PREFIX = "tval"
|
||||
|
||||
|
|
@ -263,6 +264,11 @@ def generate_tag_event_id() -> str:
|
|||
return f"{TAG_EVENT_PREFIX}_{int_id}"
|
||||
|
||||
|
||||
def generate_run_tag_event_id() -> str:
|
||||
int_id = generate_id()
|
||||
return f"{RUN_TAG_EVENT_PREFIX}_{int_id}"
|
||||
|
||||
|
||||
def generate_tag_key_id() -> str:
|
||||
int_id = generate_id()
|
||||
return f"{TAG_KEY_PREFIX}_{int_id}"
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ from skyvern.forge.sdk.db.id import (
|
|||
generate_organization_bitwarden_collection_id,
|
||||
generate_output_parameter_id,
|
||||
generate_persistent_browser_session_id,
|
||||
generate_run_tag_event_id,
|
||||
generate_script_block_id,
|
||||
generate_script_fallback_episode_id,
|
||||
generate_script_file_id,
|
||||
|
|
@ -417,6 +418,78 @@ class WorkflowTagEventModel(Base):
|
|||
deleted_at = Column(DateTime, nullable=True)
|
||||
|
||||
|
||||
class WorkflowRunTagEventModel(Base):
|
||||
__tablename__ = "workflow_run_tag_events"
|
||||
__table_args__ = (
|
||||
Index("workflow_run_tag_events_org_wr_set_at_idx", "organization_id", "workflow_run_id", "set_at"),
|
||||
Index(
|
||||
"workflow_run_tag_events_org_key_value_active_idx",
|
||||
"organization_id",
|
||||
"key",
|
||||
"value",
|
||||
postgresql_include=["workflow_run_id"],
|
||||
postgresql_where=text("superseded_at IS NULL AND event_type = 'set'"),
|
||||
),
|
||||
Index(
|
||||
"workflow_run_tag_events_org_value_active_idx",
|
||||
"organization_id",
|
||||
"value",
|
||||
postgresql_include=["workflow_run_id"],
|
||||
postgresql_where=text("superseded_at IS NULL AND event_type = 'set'"),
|
||||
),
|
||||
Index("workflow_run_tag_events_org_set_at_idx", "organization_id", "set_at"),
|
||||
Index(
|
||||
"workflow_run_tag_events_active_grouped_unique",
|
||||
"organization_id",
|
||||
"workflow_run_id",
|
||||
"key",
|
||||
unique=True,
|
||||
postgresql_where=text("superseded_at IS NULL AND event_type = 'set' AND key IS NOT NULL"),
|
||||
sqlite_where=text("superseded_at IS NULL AND event_type = 'set' AND key IS NOT NULL"),
|
||||
),
|
||||
Index(
|
||||
"workflow_run_tag_events_active_label_unique",
|
||||
"organization_id",
|
||||
"workflow_run_id",
|
||||
"value",
|
||||
unique=True,
|
||||
postgresql_where=text("superseded_at IS NULL AND event_type = 'set' AND key IS NULL"),
|
||||
sqlite_where=text("superseded_at IS NULL AND event_type = 'set' AND key IS NULL"),
|
||||
),
|
||||
CheckConstraint("event_type IN ('set', 'delete')", name="ck_workflow_run_tag_events_event_type"),
|
||||
CheckConstraint(
|
||||
"source IN ('manual', 'bulk_apply', 'backfill', 'inherited', 'import', 'system')",
|
||||
name="ck_workflow_run_tag_events_source",
|
||||
),
|
||||
CheckConstraint(
|
||||
"caller_type IS NULL OR caller_type IN ('user', 'api_key', 'system')",
|
||||
name="ck_workflow_run_tag_events_caller_type",
|
||||
),
|
||||
CheckConstraint("event_type != 'set' OR value IS NOT NULL", name="ck_workflow_run_tag_events_set_has_value"),
|
||||
)
|
||||
|
||||
tag_event_id = Column(String, primary_key=True, default=generate_run_tag_event_id)
|
||||
workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), nullable=False)
|
||||
organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=False)
|
||||
key = Column(String, nullable=True)
|
||||
value = Column(String, nullable=True)
|
||||
event_type = Column(String, nullable=False)
|
||||
set_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
|
||||
set_by = Column(String, nullable=False)
|
||||
source = Column(String, nullable=False)
|
||||
caller_type = Column(String, nullable=True)
|
||||
superseded_at = Column(DateTime, nullable=True)
|
||||
inherited_from_tag_event_id = Column(String, nullable=True)
|
||||
|
||||
created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
|
||||
modified_at = Column(
|
||||
DateTime,
|
||||
default=datetime.datetime.utcnow,
|
||||
onupdate=datetime.datetime.utcnow,
|
||||
nullable=False,
|
||||
)
|
||||
|
||||
|
||||
class TagKeyModel(Base):
|
||||
"""Org-scoped registry of tag keys and their descriptions. Auto-registered
|
||||
on first use; partial UNIQUE on (org, key) WHERE deleted_at IS NULL races
|
||||
|
|
|
|||
|
|
@ -12,7 +12,14 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
|
||||
from skyvern.forge.sdk.db._error_handling import db_operation, register_passthrough_exception
|
||||
from skyvern.forge.sdk.db.base_repository import BaseRepository
|
||||
from skyvern.forge.sdk.db.models import TagKeyModel, TagValueModel, WorkflowModel, WorkflowTagEventModel
|
||||
from skyvern.forge.sdk.db.models import (
|
||||
TagKeyModel,
|
||||
TagValueModel,
|
||||
WorkflowModel,
|
||||
WorkflowRunModel,
|
||||
WorkflowRunTagEventModel,
|
||||
WorkflowTagEventModel,
|
||||
)
|
||||
from skyvern.forge.sdk.workflow.models.tags import TagEventType, TagWriteContext
|
||||
from skyvern.forge.sdk.workflow.models.validators import RUN_METADATA_MAX_KEYS, random_tag_color
|
||||
|
||||
|
|
@ -20,6 +27,7 @@ LOG = structlog.get_logger()
|
|||
|
||||
# Aliased to RUN_METADATA_MAX_KEYS so tag and run_metadata caps stay coupled.
|
||||
MAX_TAGS_PER_WORKFLOW = RUN_METADATA_MAX_KEYS
|
||||
MAX_TAGS_PER_RUN = RUN_METADATA_MAX_KEYS
|
||||
|
||||
|
||||
class TagCountLimitExceeded(ValueError):
|
||||
|
|
@ -31,10 +39,15 @@ class TagValueRenameCollision(ValueError):
|
|||
already exists active org-wide. v1 rejects rather than merging."""
|
||||
|
||||
|
||||
class RunTagWorkflowRunMismatch(ValueError):
|
||||
"""Raised when a run-tag write targets a workflow run outside the supplied org."""
|
||||
|
||||
|
||||
# Cap breaches and rename collisions are user input, not infra failures — log as
|
||||
# BusinessLogicError (WARN), not UnexpectedError (ERROR).
|
||||
register_passthrough_exception(TagCountLimitExceeded)
|
||||
register_passthrough_exception(TagValueRenameCollision)
|
||||
register_passthrough_exception(RunTagWorkflowRunMismatch)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
|
@ -78,6 +91,65 @@ class TagsRepository(BaseRepository):
|
|||
(keys), standalone labels via ``label_sets``/``label_deletes``; set wins over delete.
|
||||
``colors`` maps a grouped tag's key to a palette color for the value being set;
|
||||
a SET key absent from ``colors`` keeps its existing color or gets a random one."""
|
||||
return await self._apply_tag_changes(
|
||||
event_model=WorkflowTagEventModel,
|
||||
entity_id_name="workflow_permanent_id",
|
||||
entity_id=workflow_permanent_id,
|
||||
entity_label="workflow",
|
||||
organization_id=organization_id,
|
||||
sets=sets,
|
||||
deletes=deletes,
|
||||
context=context,
|
||||
label_sets=label_sets,
|
||||
label_deletes=label_deletes,
|
||||
colors=colors,
|
||||
max_tags=MAX_TAGS_PER_WORKFLOW,
|
||||
)
|
||||
|
||||
@db_operation("apply_run_tag_changes")
|
||||
async def apply_run_tag_changes(
|
||||
self,
|
||||
workflow_run_id: str,
|
||||
organization_id: str,
|
||||
sets: dict[str, str],
|
||||
deletes: set[str],
|
||||
context: TagWriteContext,
|
||||
label_sets: list[str] | None = None,
|
||||
label_deletes: list[str] | None = None,
|
||||
colors: dict[str, str] | None = None,
|
||||
) -> list[TagChange]:
|
||||
"""Atomically apply SET/DELETE tag events to a workflow run."""
|
||||
return await self._apply_tag_changes(
|
||||
event_model=WorkflowRunTagEventModel,
|
||||
entity_id_name="workflow_run_id",
|
||||
entity_id=workflow_run_id,
|
||||
entity_label="workflow run",
|
||||
organization_id=organization_id,
|
||||
sets=sets,
|
||||
deletes=deletes,
|
||||
context=context,
|
||||
label_sets=label_sets,
|
||||
label_deletes=label_deletes,
|
||||
colors=colors,
|
||||
max_tags=MAX_TAGS_PER_RUN,
|
||||
)
|
||||
|
||||
async def _apply_tag_changes(
|
||||
self,
|
||||
*,
|
||||
event_model: type[WorkflowTagEventModel] | type[WorkflowRunTagEventModel],
|
||||
entity_id_name: str,
|
||||
entity_id: str,
|
||||
entity_label: str,
|
||||
organization_id: str,
|
||||
sets: dict[str, str],
|
||||
deletes: set[str],
|
||||
context: TagWriteContext,
|
||||
label_sets: list[str] | None,
|
||||
label_deletes: list[str] | None,
|
||||
colors: dict[str, str] | None,
|
||||
max_tags: int,
|
||||
) -> list[TagChange]:
|
||||
label_sets = label_sets or []
|
||||
label_deletes = label_deletes or []
|
||||
colors = colors or {}
|
||||
|
|
@ -90,9 +162,18 @@ class TagsRepository(BaseRepository):
|
|||
now = datetime.now(timezone.utc)
|
||||
|
||||
async with self.Session() as session:
|
||||
active_rows = await self._get_active_set_rows(
|
||||
if event_model is WorkflowRunTagEventModel:
|
||||
await self._validate_workflow_run_org(
|
||||
session,
|
||||
workflow_run_id=entity_id,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
|
||||
active_rows = await self._get_active_set_rows_for_entity(
|
||||
session,
|
||||
workflow_permanent_id=workflow_permanent_id,
|
||||
event_model=event_model,
|
||||
entity_id_name=entity_id_name,
|
||||
entity_id=entity_id,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
# Grouped tags keyed by their group; standalone labels keyed by value.
|
||||
|
|
@ -104,26 +185,27 @@ class TagsRepository(BaseRepository):
|
|||
projected_grouped = (set(grouped_current.keys()) | set(sets.keys())) - effective_deletes
|
||||
projected_labels = (set(label_current.keys()) | label_set_values) - effective_label_deletes
|
||||
projected_total = len(projected_grouped) + len(projected_labels)
|
||||
if projected_total > MAX_TAGS_PER_WORKFLOW:
|
||||
if projected_total > max_tags:
|
||||
raise TagCountLimitExceeded(
|
||||
f"workflow {workflow_permanent_id} would have {projected_total} tags; "
|
||||
f"max is {MAX_TAGS_PER_WORKFLOW}"
|
||||
f"{entity_label} {entity_id} would have {projected_total} tags; max is {max_tags}"
|
||||
)
|
||||
|
||||
changes: list[TagChange] = []
|
||||
|
||||
def _add_event(*, key: str | None, value: str | None, event_type: TagEventType) -> None:
|
||||
session.add(
|
||||
WorkflowTagEventModel(
|
||||
workflow_permanent_id=workflow_permanent_id,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=value,
|
||||
event_type=event_type.value,
|
||||
set_at=now,
|
||||
set_by=context.caller_id,
|
||||
source=context.source.value,
|
||||
caller_type=context.caller_type.value if context.caller_type else None,
|
||||
event_model(
|
||||
**{
|
||||
entity_id_name: entity_id,
|
||||
"organization_id": organization_id,
|
||||
"key": key,
|
||||
"value": value,
|
||||
"event_type": event_type.value,
|
||||
"set_at": now,
|
||||
"set_by": context.caller_id,
|
||||
"source": context.source.value,
|
||||
"caller_type": context.caller_type.value if context.caller_type else None,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -219,6 +301,26 @@ class TagsRepository(BaseRepository):
|
|||
await session.commit()
|
||||
return changes
|
||||
|
||||
async def _validate_workflow_run_org(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
*,
|
||||
workflow_run_id: str,
|
||||
organization_id: str,
|
||||
) -> None:
|
||||
row = (
|
||||
await session.execute(
|
||||
select(WorkflowRunModel.workflow_run_id)
|
||||
.where(WorkflowRunModel.workflow_run_id == workflow_run_id)
|
||||
.where(WorkflowRunModel.organization_id == organization_id)
|
||||
.limit(1)
|
||||
)
|
||||
).first()
|
||||
if row is None:
|
||||
raise RunTagWorkflowRunMismatch(
|
||||
f"workflow run {workflow_run_id} does not exist in organization {organization_id}"
|
||||
)
|
||||
|
||||
async def _register_tag_value_colors(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
|
|
@ -270,6 +372,31 @@ class TagsRepository(BaseRepository):
|
|||
)
|
||||
await session.execute(stmt)
|
||||
|
||||
async def _get_active_set_rows_for_entity(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
*,
|
||||
event_model: type[WorkflowTagEventModel] | type[WorkflowRunTagEventModel],
|
||||
entity_id_name: str,
|
||||
entity_id: str,
|
||||
organization_id: str,
|
||||
) -> list[WorkflowTagEventModel] | list[WorkflowRunTagEventModel]:
|
||||
"""Active SET event rows for a workflow or workflow run, including
|
||||
grouped tags and standalone labels."""
|
||||
filters = [
|
||||
event_model.organization_id == organization_id,
|
||||
getattr(event_model, entity_id_name) == entity_id,
|
||||
event_model.superseded_at.is_(None),
|
||||
event_model.event_type == TagEventType.SET.value,
|
||||
]
|
||||
deleted_at = getattr(event_model, "deleted_at", None)
|
||||
if deleted_at is not None:
|
||||
filters.append(deleted_at.is_(None))
|
||||
|
||||
stmt = select(event_model).where(and_(*filters))
|
||||
result = await session.execute(stmt)
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def _get_active_set_rows(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
|
|
@ -277,19 +404,30 @@ class TagsRepository(BaseRepository):
|
|||
workflow_permanent_id: str,
|
||||
organization_id: str,
|
||||
) -> list[WorkflowTagEventModel]:
|
||||
"""Active (non-superseded, non-deleted) SET event rows for a workflow —
|
||||
both grouped tags and standalone labels."""
|
||||
stmt = select(WorkflowTagEventModel).where(
|
||||
and_(
|
||||
WorkflowTagEventModel.organization_id == organization_id,
|
||||
WorkflowTagEventModel.workflow_permanent_id == workflow_permanent_id,
|
||||
WorkflowTagEventModel.superseded_at.is_(None),
|
||||
WorkflowTagEventModel.event_type == TagEventType.SET.value,
|
||||
WorkflowTagEventModel.deleted_at.is_(None),
|
||||
)
|
||||
rows = await self._get_active_set_rows_for_entity(
|
||||
session,
|
||||
event_model=WorkflowTagEventModel,
|
||||
entity_id_name="workflow_permanent_id",
|
||||
entity_id=workflow_permanent_id,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
return list(result.scalars().all())
|
||||
return list(rows)
|
||||
|
||||
async def _get_active_set_rows_for_run(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
*,
|
||||
workflow_run_id: str,
|
||||
organization_id: str,
|
||||
) -> list[WorkflowRunTagEventModel]:
|
||||
rows = await self._get_active_set_rows_for_entity(
|
||||
session,
|
||||
event_model=WorkflowRunTagEventModel,
|
||||
entity_id_name="workflow_run_id",
|
||||
entity_id=workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
return list(rows)
|
||||
|
||||
@db_operation("get_active_grouped_tags_for_workflow")
|
||||
async def get_active_grouped_tags_for_workflow(
|
||||
|
|
@ -323,6 +461,35 @@ class TagsRepository(BaseRepository):
|
|||
result[row.key] = row.value
|
||||
return result
|
||||
|
||||
@db_operation("get_active_grouped_tags_for_run")
|
||||
async def get_active_grouped_tags_for_run(
|
||||
self,
|
||||
workflow_run_id: str,
|
||||
organization_id: str,
|
||||
) -> dict[str, str]:
|
||||
"""Current {key: value} map of a workflow run's grouped tags."""
|
||||
async with self.Session() as session:
|
||||
rows = await self._get_active_set_rows_for_run(
|
||||
session,
|
||||
workflow_run_id=workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
result: dict[str, str] = {}
|
||||
for row in rows:
|
||||
if row.key is None:
|
||||
continue
|
||||
if row.value is None:
|
||||
LOG.warning(
|
||||
"active SET run tag row has null value; skipping",
|
||||
tag_event_id=row.tag_event_id,
|
||||
organization_id=organization_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
key=row.key,
|
||||
)
|
||||
continue
|
||||
result[row.key] = row.value
|
||||
return result
|
||||
|
||||
@db_operation("get_active_tag_events_for_workflow")
|
||||
async def get_active_tag_events_for_workflow(
|
||||
self,
|
||||
|
|
@ -338,6 +505,20 @@ class TagsRepository(BaseRepository):
|
|||
organization_id=organization_id,
|
||||
)
|
||||
|
||||
@db_operation("get_active_tag_events_for_run")
|
||||
async def get_active_tag_events_for_run(
|
||||
self,
|
||||
workflow_run_id: str,
|
||||
organization_id: str,
|
||||
) -> list[WorkflowRunTagEventModel]:
|
||||
"""Active SET rows (grouped + standalone) with full run-tag attribution."""
|
||||
async with self.Session() as session:
|
||||
return await self._get_active_set_rows_for_run(
|
||||
session,
|
||||
workflow_run_id=workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
|
||||
@db_operation("get_tag_event_history")
|
||||
async def get_tag_event_history(
|
||||
self,
|
||||
|
|
@ -348,20 +529,63 @@ class TagsRepository(BaseRepository):
|
|||
key: str | None = None,
|
||||
) -> list[WorkflowTagEventModel]:
|
||||
"""Return tag events newest-first for a workflow. Includes DELETE and superseded SET rows."""
|
||||
rows = await self._get_tag_event_history_for_entity(
|
||||
event_model=WorkflowTagEventModel,
|
||||
entity_id_name="workflow_permanent_id",
|
||||
entity_id=workflow_permanent_id,
|
||||
organization_id=organization_id,
|
||||
limit=limit,
|
||||
since=since,
|
||||
key=key,
|
||||
)
|
||||
return list(rows)
|
||||
|
||||
@db_operation("get_run_tag_event_history")
|
||||
async def get_run_tag_event_history(
|
||||
self,
|
||||
workflow_run_id: str,
|
||||
organization_id: str,
|
||||
limit: int = 100,
|
||||
since: datetime | None = None,
|
||||
key: str | None = None,
|
||||
) -> list[WorkflowRunTagEventModel]:
|
||||
"""Return tag events newest-first for a workflow run, with optional key filter."""
|
||||
rows = await self._get_tag_event_history_for_entity(
|
||||
event_model=WorkflowRunTagEventModel,
|
||||
entity_id_name="workflow_run_id",
|
||||
entity_id=workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
limit=limit,
|
||||
since=since,
|
||||
key=key,
|
||||
)
|
||||
return list(rows)
|
||||
|
||||
async def _get_tag_event_history_for_entity(
|
||||
self,
|
||||
*,
|
||||
event_model: type[WorkflowTagEventModel] | type[WorkflowRunTagEventModel],
|
||||
entity_id_name: str,
|
||||
entity_id: str,
|
||||
organization_id: str,
|
||||
limit: int,
|
||||
since: datetime | None,
|
||||
key: str | None,
|
||||
) -> list[WorkflowTagEventModel] | list[WorkflowRunTagEventModel]:
|
||||
async with self.Session() as session:
|
||||
stmt = (
|
||||
select(WorkflowTagEventModel)
|
||||
.where(WorkflowTagEventModel.organization_id == organization_id)
|
||||
.where(WorkflowTagEventModel.workflow_permanent_id == workflow_permanent_id)
|
||||
.where(WorkflowTagEventModel.deleted_at.is_(None))
|
||||
select(event_model)
|
||||
.where(event_model.organization_id == organization_id)
|
||||
.where(getattr(event_model, entity_id_name) == entity_id)
|
||||
)
|
||||
deleted_at = getattr(event_model, "deleted_at", None)
|
||||
if deleted_at is not None:
|
||||
stmt = stmt.where(deleted_at.is_(None))
|
||||
if since is not None:
|
||||
stmt = stmt.where(WorkflowTagEventModel.set_at >= since)
|
||||
stmt = stmt.where(event_model.set_at >= since)
|
||||
if key is not None:
|
||||
stmt = stmt.where(WorkflowTagEventModel.key == key)
|
||||
stmt = stmt.order_by(WorkflowTagEventModel.set_at.desc(), WorkflowTagEventModel.tag_event_id.desc()).limit(
|
||||
limit
|
||||
)
|
||||
stmt = stmt.where(event_model.key == key)
|
||||
stmt = stmt.order_by(event_model.set_at.desc(), event_model.tag_event_id.desc()).limit(limit)
|
||||
result = await session.execute(stmt)
|
||||
return list(result.scalars().all())
|
||||
|
||||
|
|
@ -373,33 +597,67 @@ class TagsRepository(BaseRepository):
|
|||
) -> dict[str, list[tuple[str | None, str]]]:
|
||||
"""Batch read of current SET tags (grouped + standalone) as a list of
|
||||
(key, value) per workflow. The org filter silently drops foreign wpids."""
|
||||
if not workflow_permanent_ids:
|
||||
return await self._get_active_tags_for_entities(
|
||||
event_model=WorkflowTagEventModel,
|
||||
entity_id_name="workflow_permanent_id",
|
||||
entity_ids=workflow_permanent_ids,
|
||||
organization_id=organization_id,
|
||||
log_entity_name="workflow_permanent_id",
|
||||
)
|
||||
|
||||
@db_operation("get_active_tags_for_runs")
|
||||
async def get_active_tags_for_runs(
|
||||
self,
|
||||
workflow_run_ids: list[str],
|
||||
organization_id: str,
|
||||
) -> dict[str, list[tuple[str | None, str]]]:
|
||||
"""Batch read of current SET tags (grouped + standalone) per workflow run."""
|
||||
return await self._get_active_tags_for_entities(
|
||||
event_model=WorkflowRunTagEventModel,
|
||||
entity_id_name="workflow_run_id",
|
||||
entity_ids=workflow_run_ids,
|
||||
organization_id=organization_id,
|
||||
log_entity_name="workflow_run_id",
|
||||
)
|
||||
|
||||
async def _get_active_tags_for_entities(
|
||||
self,
|
||||
*,
|
||||
event_model: type[WorkflowTagEventModel] | type[WorkflowRunTagEventModel],
|
||||
entity_id_name: str,
|
||||
entity_ids: list[str],
|
||||
organization_id: str,
|
||||
log_entity_name: str,
|
||||
) -> dict[str, list[tuple[str | None, str]]]:
|
||||
if not entity_ids:
|
||||
return {}
|
||||
|
||||
filters = [
|
||||
event_model.organization_id == organization_id,
|
||||
getattr(event_model, entity_id_name).in_(entity_ids),
|
||||
event_model.superseded_at.is_(None),
|
||||
event_model.event_type == TagEventType.SET.value,
|
||||
]
|
||||
deleted_at = getattr(event_model, "deleted_at", None)
|
||||
if deleted_at is not None:
|
||||
filters.append(deleted_at.is_(None))
|
||||
|
||||
async with self.Session() as session:
|
||||
stmt = select(WorkflowTagEventModel).where(
|
||||
and_(
|
||||
WorkflowTagEventModel.organization_id == organization_id,
|
||||
WorkflowTagEventModel.workflow_permanent_id.in_(workflow_permanent_ids),
|
||||
WorkflowTagEventModel.superseded_at.is_(None),
|
||||
WorkflowTagEventModel.event_type == TagEventType.SET.value,
|
||||
WorkflowTagEventModel.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
stmt = select(event_model).where(and_(*filters))
|
||||
rows = (await session.execute(stmt)).scalars().all()
|
||||
|
||||
result: dict[str, list[tuple[str | None, str]]] = {}
|
||||
for row in rows:
|
||||
entity_id = getattr(row, entity_id_name)
|
||||
if row.value is None:
|
||||
LOG.warning(
|
||||
"active SET tag row has null value; skipping",
|
||||
tag_event_id=row.tag_event_id,
|
||||
organization_id=organization_id,
|
||||
workflow_permanent_id=row.workflow_permanent_id,
|
||||
key=row.key,
|
||||
**{log_entity_name: entity_id, "key": row.key},
|
||||
)
|
||||
continue
|
||||
result.setdefault(row.workflow_permanent_id, []).append((row.key, row.value))
|
||||
result.setdefault(entity_id, []).append((row.key, row.value))
|
||||
return result
|
||||
|
||||
@db_operation("list_tag_keys")
|
||||
|
|
@ -513,6 +771,25 @@ class TagsRepository(BaseRepository):
|
|||
rows = await session.execute(stmt)
|
||||
return {key: count for key, count in rows.all()}
|
||||
|
||||
@db_operation("count_active_runs_per_key")
|
||||
async def count_active_runs_per_key(self, organization_id: str) -> dict[str, int]:
|
||||
"""Map of tag key -> number of workflow runs carrying it. Standalone labels are excluded."""
|
||||
async with self.Session() as session:
|
||||
stmt = (
|
||||
select(WorkflowRunTagEventModel.key, func.count(WorkflowRunTagEventModel.tag_event_id))
|
||||
.where(
|
||||
and_(
|
||||
WorkflowRunTagEventModel.organization_id == organization_id,
|
||||
WorkflowRunTagEventModel.superseded_at.is_(None),
|
||||
WorkflowRunTagEventModel.event_type == TagEventType.SET.value,
|
||||
WorkflowRunTagEventModel.key.isnot(None),
|
||||
)
|
||||
)
|
||||
.group_by(WorkflowRunTagEventModel.key)
|
||||
)
|
||||
rows = await session.execute(stmt)
|
||||
return {key: count for key, count in rows.all()}
|
||||
|
||||
@staticmethod
|
||||
def _non_deleted_workflow_exists(organization_id: str) -> Exists:
|
||||
"""Correlated EXISTS: the tag event's workflow still has a live version.
|
||||
|
|
@ -554,6 +831,29 @@ class TagsRepository(BaseRepository):
|
|||
rows = await session.execute(stmt)
|
||||
return {(key, value): count for key, value, count in rows.all()}
|
||||
|
||||
@db_operation("count_active_runs_per_value")
|
||||
async def count_active_runs_per_value(self, organization_id: str) -> dict[tuple[str, str], int]:
|
||||
"""Map of grouped ``(key, value)`` -> number of workflow runs carrying it."""
|
||||
async with self.Session() as session:
|
||||
stmt = (
|
||||
select(
|
||||
WorkflowRunTagEventModel.key,
|
||||
WorkflowRunTagEventModel.value,
|
||||
func.count(func.distinct(WorkflowRunTagEventModel.workflow_run_id)),
|
||||
)
|
||||
.where(
|
||||
and_(
|
||||
WorkflowRunTagEventModel.organization_id == organization_id,
|
||||
WorkflowRunTagEventModel.superseded_at.is_(None),
|
||||
WorkflowRunTagEventModel.event_type == TagEventType.SET.value,
|
||||
WorkflowRunTagEventModel.key.isnot(None),
|
||||
)
|
||||
)
|
||||
.group_by(WorkflowRunTagEventModel.key, WorkflowRunTagEventModel.value)
|
||||
)
|
||||
rows = await session.execute(stmt)
|
||||
return {(key, value): count for key, value, count in rows.all()}
|
||||
|
||||
@db_operation("count_active_workflows_for_value")
|
||||
async def count_active_workflows_for_value(self, organization_id: str, key: str, value: str) -> int:
|
||||
"""Number of non-deleted workflows with an active SET on grouped ``(key, value)``.
|
||||
|
|
@ -596,6 +896,60 @@ class TagsRepository(BaseRepository):
|
|||
stmt = stmt.where(TagValueModel.value == value)
|
||||
await session.execute(stmt.values(deleted_at=now))
|
||||
|
||||
async def _get_active_grouped_set_rows(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
*,
|
||||
event_model: type[WorkflowTagEventModel] | type[WorkflowRunTagEventModel],
|
||||
organization_id: str,
|
||||
key: str,
|
||||
value: str | None = None,
|
||||
) -> list[WorkflowTagEventModel] | list[WorkflowRunTagEventModel]:
|
||||
filters = [
|
||||
event_model.organization_id == organization_id,
|
||||
event_model.key == key,
|
||||
event_model.superseded_at.is_(None),
|
||||
event_model.event_type == TagEventType.SET.value,
|
||||
]
|
||||
if value is not None:
|
||||
filters.append(event_model.value == value)
|
||||
deleted_at = getattr(event_model, "deleted_at", None)
|
||||
if deleted_at is not None:
|
||||
filters.append(deleted_at.is_(None))
|
||||
|
||||
result = await session.execute(select(event_model).where(and_(*filters)))
|
||||
return list(result.scalars().all())
|
||||
|
||||
@staticmethod
|
||||
def _supersede_and_add_delete_event(
|
||||
session: AsyncSession,
|
||||
*,
|
||||
event_model: type[WorkflowTagEventModel] | type[WorkflowRunTagEventModel],
|
||||
entity_id_name: str,
|
||||
existing: WorkflowTagEventModel | WorkflowRunTagEventModel,
|
||||
organization_id: str,
|
||||
key: str,
|
||||
value: str | None,
|
||||
context: TagWriteContext,
|
||||
now: datetime,
|
||||
) -> None:
|
||||
existing.superseded_at = now
|
||||
session.add(
|
||||
event_model(
|
||||
**{
|
||||
entity_id_name: getattr(existing, entity_id_name),
|
||||
"organization_id": organization_id,
|
||||
"key": key,
|
||||
"value": value,
|
||||
"event_type": TagEventType.DELETE.value,
|
||||
"set_at": now,
|
||||
"set_by": context.caller_id,
|
||||
"source": context.source.value,
|
||||
"caller_type": context.caller_type.value if context.caller_type else None,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
@db_operation("delete_tag_key")
|
||||
async def delete_tag_key(
|
||||
self,
|
||||
|
|
@ -643,38 +997,42 @@ class TagsRepository(BaseRepository):
|
|||
if key_row is None:
|
||||
return None
|
||||
|
||||
active_sets = (
|
||||
(
|
||||
await session.execute(
|
||||
select(WorkflowTagEventModel).where(
|
||||
and_(
|
||||
WorkflowTagEventModel.organization_id == organization_id,
|
||||
WorkflowTagEventModel.key == key,
|
||||
WorkflowTagEventModel.superseded_at.is_(None),
|
||||
WorkflowTagEventModel.event_type == TagEventType.SET.value,
|
||||
WorkflowTagEventModel.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
active_sets = await self._get_active_grouped_set_rows(
|
||||
session,
|
||||
event_model=WorkflowTagEventModel,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
)
|
||||
active_run_sets = await self._get_active_grouped_set_rows(
|
||||
session,
|
||||
event_model=WorkflowRunTagEventModel,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
)
|
||||
|
||||
for existing in active_sets:
|
||||
existing.superseded_at = now
|
||||
session.add(
|
||||
WorkflowTagEventModel(
|
||||
workflow_permanent_id=existing.workflow_permanent_id,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=None,
|
||||
event_type=TagEventType.DELETE.value,
|
||||
set_at=now,
|
||||
set_by=context.caller_id,
|
||||
source=context.source.value,
|
||||
caller_type=context.caller_type.value if context.caller_type else None,
|
||||
)
|
||||
self._supersede_and_add_delete_event(
|
||||
session,
|
||||
event_model=WorkflowTagEventModel,
|
||||
entity_id_name="workflow_permanent_id",
|
||||
existing=existing,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=None,
|
||||
context=context,
|
||||
now=now,
|
||||
)
|
||||
for existing in active_run_sets:
|
||||
self._supersede_and_add_delete_event(
|
||||
session,
|
||||
event_model=WorkflowRunTagEventModel,
|
||||
entity_id_name="workflow_run_id",
|
||||
existing=existing,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=None,
|
||||
context=context,
|
||||
now=now,
|
||||
)
|
||||
|
||||
key_row.deleted_at = now
|
||||
|
|
@ -716,42 +1074,47 @@ class TagsRepository(BaseRepository):
|
|||
)
|
||||
).scalar_one_or_none()
|
||||
|
||||
active_sets = (
|
||||
(
|
||||
await session.execute(
|
||||
select(WorkflowTagEventModel).where(
|
||||
and_(
|
||||
WorkflowTagEventModel.organization_id == organization_id,
|
||||
WorkflowTagEventModel.key == key,
|
||||
WorkflowTagEventModel.value == value,
|
||||
WorkflowTagEventModel.superseded_at.is_(None),
|
||||
WorkflowTagEventModel.event_type == TagEventType.SET.value,
|
||||
WorkflowTagEventModel.deleted_at.is_(None),
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
active_sets = await self._get_active_grouped_set_rows(
|
||||
session,
|
||||
event_model=WorkflowTagEventModel,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=value,
|
||||
)
|
||||
active_run_sets = await self._get_active_grouped_set_rows(
|
||||
session,
|
||||
event_model=WorkflowRunTagEventModel,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=value,
|
||||
)
|
||||
|
||||
if value_row is None and not active_sets:
|
||||
if value_row is None and not active_sets and not active_run_sets:
|
||||
return None
|
||||
|
||||
for existing in active_sets:
|
||||
existing.superseded_at = now
|
||||
session.add(
|
||||
WorkflowTagEventModel(
|
||||
workflow_permanent_id=existing.workflow_permanent_id,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=value,
|
||||
event_type=TagEventType.DELETE.value,
|
||||
set_at=now,
|
||||
set_by=context.caller_id,
|
||||
source=context.source.value,
|
||||
caller_type=context.caller_type.value if context.caller_type else None,
|
||||
)
|
||||
self._supersede_and_add_delete_event(
|
||||
session,
|
||||
event_model=WorkflowTagEventModel,
|
||||
entity_id_name="workflow_permanent_id",
|
||||
existing=existing,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=value,
|
||||
context=context,
|
||||
now=now,
|
||||
)
|
||||
for existing in active_run_sets:
|
||||
self._supersede_and_add_delete_event(
|
||||
session,
|
||||
event_model=WorkflowRunTagEventModel,
|
||||
entity_id_name="workflow_run_id",
|
||||
existing=existing,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=value,
|
||||
context=context,
|
||||
now=now,
|
||||
)
|
||||
|
||||
await self._soft_delete_tag_value_rows(
|
||||
|
|
@ -825,8 +1188,15 @@ class TagsRepository(BaseRepository):
|
|||
.scalars()
|
||||
.all()
|
||||
)
|
||||
active_run_sets = await self._get_active_grouped_set_rows(
|
||||
session,
|
||||
event_model=WorkflowRunTagEventModel,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=old_value,
|
||||
)
|
||||
|
||||
if old_row is None and not active_sets:
|
||||
if old_row is None and not active_sets and not active_run_sets:
|
||||
return None
|
||||
|
||||
if await self._grouped_value_active(session, organization_id=organization_id, key=key, value=new_value):
|
||||
|
|
@ -836,6 +1206,8 @@ class TagsRepository(BaseRepository):
|
|||
|
||||
for existing in active_sets:
|
||||
existing.superseded_at = now
|
||||
for existing in active_run_sets:
|
||||
existing.superseded_at = now
|
||||
# Flush supersede UPDATEs before the new SET INSERTs so the active-SET
|
||||
# partial UNIQUE on (org, wpid, key) sees a consistent state.
|
||||
await session.flush()
|
||||
|
|
@ -853,6 +1225,20 @@ class TagsRepository(BaseRepository):
|
|||
caller_type=context.caller_type.value if context.caller_type else None,
|
||||
)
|
||||
)
|
||||
for existing in active_run_sets:
|
||||
session.add(
|
||||
WorkflowRunTagEventModel(
|
||||
workflow_run_id=existing.workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
key=key,
|
||||
value=new_value,
|
||||
event_type=TagEventType.SET.value,
|
||||
set_at=now,
|
||||
set_by=context.caller_id,
|
||||
source=context.source.value,
|
||||
caller_type=context.caller_type.value if context.caller_type else None,
|
||||
)
|
||||
)
|
||||
|
||||
carried_color = old_row.color if old_row is not None else random_tag_color()
|
||||
dialect_name = session.bind.dialect.name if session.bind is not None else "postgresql"
|
||||
|
|
@ -887,7 +1273,7 @@ class TagsRepository(BaseRepository):
|
|||
value: str,
|
||||
) -> bool:
|
||||
"""True when grouped label ``(key, value)`` exists active org-wide — either a
|
||||
registered color row or an active SET event in use on some workflow."""
|
||||
registered color row or an active SET event in use on some workflow/run."""
|
||||
registered = (
|
||||
await session.execute(
|
||||
select(TagValueModel.tag_value_id)
|
||||
|
|
@ -920,4 +1306,21 @@ class TagsRepository(BaseRepository):
|
|||
.limit(1)
|
||||
)
|
||||
).first()
|
||||
return in_use is not None
|
||||
if in_use is not None:
|
||||
return True
|
||||
run_in_use = (
|
||||
await session.execute(
|
||||
select(WorkflowRunTagEventModel.tag_event_id)
|
||||
.where(
|
||||
and_(
|
||||
WorkflowRunTagEventModel.organization_id == organization_id,
|
||||
WorkflowRunTagEventModel.key == key,
|
||||
WorkflowRunTagEventModel.value == value,
|
||||
WorkflowRunTagEventModel.superseded_at.is_(None),
|
||||
WorkflowRunTagEventModel.event_type == TagEventType.SET.value,
|
||||
)
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
).first()
|
||||
return run_in_use is not None
|
||||
|
|
|
|||
|
|
@ -16,10 +16,13 @@ from skyvern.forge.sdk.db.models import (
|
|||
TagKeyModel,
|
||||
TagValueModel,
|
||||
WorkflowModel,
|
||||
WorkflowRunModel,
|
||||
WorkflowRunTagEventModel,
|
||||
WorkflowTagEventModel,
|
||||
)
|
||||
from skyvern.forge.sdk.db.repositories.tags import (
|
||||
MAX_TAGS_PER_WORKFLOW,
|
||||
RunTagWorkflowRunMismatch,
|
||||
TagCountLimitExceeded,
|
||||
TagsRepository,
|
||||
TagValueRenameCollision,
|
||||
|
|
@ -34,6 +37,7 @@ from skyvern.forge.sdk.workflow.models.validators import TAG_COLOR_PALETTE
|
|||
|
||||
ORG_ID = "o_test"
|
||||
WPID = "wpid_alpha"
|
||||
WRID = "wr_alpha"
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
|
|
@ -59,6 +63,16 @@ async def _all_events(repo: TagsRepository) -> list[WorkflowTagEventModel]:
|
|||
return list(result.scalars().all())
|
||||
|
||||
|
||||
async def _all_run_events(repo: TagsRepository) -> list[WorkflowRunTagEventModel]:
|
||||
async with repo.Session() as session:
|
||||
result = await session.execute(
|
||||
select(WorkflowRunTagEventModel).order_by(
|
||||
WorkflowRunTagEventModel.set_at, WorkflowRunTagEventModel.tag_event_id
|
||||
)
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
|
||||
|
||||
async def _registered_keys(repo: TagsRepository) -> list[str]:
|
||||
async with repo.Session() as session:
|
||||
result = await session.execute(select(TagKeyModel.key).order_by(TagKeyModel.key))
|
||||
|
|
@ -95,6 +109,20 @@ async def _seed_workflow(repo: TagsRepository, wpid: str, *, org_id: str = ORG_I
|
|||
await session.commit()
|
||||
|
||||
|
||||
async def _seed_workflow_run(repo: TagsRepository, wrid: str, *, org_id: str = ORG_ID) -> None:
|
||||
async with repo.Session() as session:
|
||||
session.add(
|
||||
WorkflowRunModel(
|
||||
workflow_run_id=wrid,
|
||||
workflow_id=f"w_{wrid}",
|
||||
workflow_permanent_id=f"wpid_{wrid}",
|
||||
organization_id=org_id,
|
||||
status="completed",
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_op_when_sets_and_deletes_empty(repo: TagsRepository) -> None:
|
||||
changes = await repo.apply_tag_changes(WPID, ORG_ID, sets={}, deletes=set(), context=_ctx())
|
||||
|
|
@ -431,6 +459,92 @@ async def test_get_active_tags_for_workflows_empty_input(repo: TagsRepository) -
|
|||
assert await repo.get_active_tags_for_workflows([], ORG_ID) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_run_tag_changes_set_supersede_delete_and_history_key_filter(repo: TagsRepository) -> None:
|
||||
await _seed_workflow_run(repo, WRID)
|
||||
|
||||
await repo.apply_run_tag_changes(WRID, ORG_ID, sets={"env": "prod", "team": "core"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes(WRID, ORG_ID, sets={"env": "stg"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes(WRID, ORG_ID, sets={}, deletes={"team"}, context=_ctx())
|
||||
|
||||
assert await repo.get_active_grouped_tags_for_run(WRID, ORG_ID) == {"env": "stg"}
|
||||
|
||||
history = await repo.get_run_tag_event_history(WRID, ORG_ID, key="env")
|
||||
assert [(event.event_type, event.key, event.value) for event in history] == [
|
||||
(TagEventType.SET.value, "env", "stg"),
|
||||
(TagEventType.SET.value, "env", "prod"),
|
||||
]
|
||||
assert history[1].superseded_at is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_run_tag_changes_supports_standalone_labels_and_cap(repo: TagsRepository) -> None:
|
||||
await _seed_workflow_run(repo, WRID)
|
||||
|
||||
grouped = {f"k{i}": "v" for i in range(MAX_TAGS_PER_WORKFLOW - 1)}
|
||||
await repo.apply_run_tag_changes(
|
||||
WRID, ORG_ID, sets=grouped, deletes=set(), context=_ctx(), label_sets=["one_label"]
|
||||
)
|
||||
|
||||
rows = await repo.get_active_tag_events_for_run(WRID, ORG_ID)
|
||||
assert {row.value for row in rows if row.key is None} == {"one_label"}
|
||||
|
||||
with pytest.raises(TagCountLimitExceeded):
|
||||
await repo.apply_run_tag_changes(WRID, ORG_ID, sets={}, deletes=set(), context=_ctx(), label_sets=["over"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_run_tag_changes_rejects_run_from_other_org(repo: TagsRepository) -> None:
|
||||
await _seed_workflow_run(repo, WRID, org_id=ORG_ID)
|
||||
|
||||
with pytest.raises(RunTagWorkflowRunMismatch):
|
||||
await repo.apply_run_tag_changes(WRID, "o_other", sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
|
||||
assert await repo.get_active_grouped_tags_for_run(WRID, "o_other") == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_active_tags_for_runs_batch(repo: TagsRepository) -> None:
|
||||
await _seed_workflow_run(repo, WRID)
|
||||
await _seed_workflow_run(repo, "wr_beta")
|
||||
await _seed_workflow_run(repo, "wr_other", org_id="o_other")
|
||||
|
||||
await repo.apply_run_tag_changes(
|
||||
WRID, ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx(), label_sets=["urgent"]
|
||||
)
|
||||
await repo.apply_run_tag_changes("wr_beta", ORG_ID, sets={"team": "growth"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes(
|
||||
"wr_other", "o_other", sets={"secret": "leak"}, deletes=set(), context=_ctx("other")
|
||||
)
|
||||
|
||||
result = await repo.get_active_tags_for_runs([WRID, "wr_beta", "wr_missing"], ORG_ID)
|
||||
assert set(result[WRID]) == {("env", "prod"), (None, "urgent")}
|
||||
assert result["wr_beta"] == [("team", "growth")]
|
||||
assert "wr_missing" not in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_active_tags_for_runs_empty_input(repo: TagsRepository) -> None:
|
||||
assert await repo.get_active_tags_for_runs([], ORG_ID) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_active_runs_per_key_and_value(repo: TagsRepository) -> None:
|
||||
for wrid in ("wr_a", "wr_b", "wr_c", "wr_d"):
|
||||
await _seed_workflow_run(repo, wrid)
|
||||
|
||||
await repo.apply_run_tag_changes("wr_a", ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes(
|
||||
"wr_b", ORG_ID, sets={"env": "prod", "team": "core"}, deletes=set(), context=_ctx()
|
||||
)
|
||||
await repo.apply_run_tag_changes("wr_c", ORG_ID, sets={"env": "stg"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes("wr_c", ORG_ID, sets={}, deletes={"env"}, context=_ctx())
|
||||
await repo.apply_run_tag_changes("wr_d", ORG_ID, sets={}, deletes=set(), context=_ctx(), label_sets=["standalone"])
|
||||
|
||||
assert await repo.count_active_runs_per_key(ORG_ID) == {"env": 2, "team": 1}
|
||||
assert await repo.count_active_runs_per_value(ORG_ID) == {("env", "prod"): 2, ("team", "core"): 1}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tag_keys_orders_alphabetically(repo: TagsRepository) -> None:
|
||||
await repo.apply_tag_changes(
|
||||
|
|
@ -560,6 +674,29 @@ async def test_delete_tag_key_cascades_across_workflows(repo: TagsRepository) ->
|
|||
assert [row.key for row in await repo.list_tag_keys(ORG_ID)] == ["team"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_tag_key_cascades_across_workflow_runs(repo: TagsRepository) -> None:
|
||||
await _seed_workflow_run(repo, "wr_a")
|
||||
await _seed_workflow_run(repo, "wr_b")
|
||||
|
||||
await repo.apply_tag_changes("wpid_a", ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes("wr_a", ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes("wr_b", ORG_ID, sets={"env": "stg", "team": "core"}, deletes=set(), context=_ctx())
|
||||
|
||||
removed_workflows = await repo.delete_tag_key(ORG_ID, "env", _ctx())
|
||||
|
||||
assert removed_workflows == 1
|
||||
assert await repo.get_active_grouped_tags_for_workflow("wpid_a", ORG_ID) == {}
|
||||
assert await repo.get_active_grouped_tags_for_run("wr_a", ORG_ID) == {}
|
||||
assert await repo.get_active_grouped_tags_for_run("wr_b", ORG_ID) == {"team": "core"}
|
||||
rows = await _all_run_events(repo)
|
||||
run_deletes = [row for row in rows if row.event_type == TagEventType.DELETE.value]
|
||||
assert {(row.workflow_run_id, row.key, row.value) for row in run_deletes} == {
|
||||
("wr_a", "env", None),
|
||||
("wr_b", "env", None),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_tag_key_unknown_returns_none(repo: TagsRepository) -> None:
|
||||
assert await repo.delete_tag_key(ORG_ID, "never_seen", _ctx()) is None
|
||||
|
|
@ -903,6 +1040,22 @@ async def test_rename_tag_value_cascades_and_carries_color(repo: TagsRepository)
|
|||
assert [(v.key, v.value, v.color) for v in await repo.list_tag_values(ORG_ID)] == [("env", "production", "blue")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rename_tag_value_cascades_across_workflow_runs(repo: TagsRepository) -> None:
|
||||
await _seed_workflow_run(repo, "wr_a")
|
||||
await _seed_workflow_run(repo, "wr_b")
|
||||
await repo.apply_run_tag_changes("wr_a", ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes("wr_b", ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
|
||||
result = await repo.rename_tag_value(ORG_ID, "env", "prod", "production", _ctx())
|
||||
|
||||
assert result is not None
|
||||
assert result.renamed_workflow_count == 0
|
||||
assert await repo.get_active_grouped_tags_for_run("wr_a", ORG_ID) == {"env": "production"}
|
||||
assert await repo.get_active_grouped_tags_for_run("wr_b", ORG_ID) == {"env": "production"}
|
||||
assert [(v.key, v.value) for v in await repo.list_tag_values(ORG_ID)] == [("env", "production")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rename_tag_value_preserves_history(repo: TagsRepository) -> None:
|
||||
"""Append-only invariant: the old value's SET event stays in history (superseded),
|
||||
|
|
@ -954,6 +1107,27 @@ async def test_rename_tag_value_collision_detects_in_use_value(repo: TagsReposit
|
|||
await repo.rename_tag_value(ORG_ID, "env", "prod", "stg", _ctx())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rename_tag_value_collision_detects_in_use_run_value(repo: TagsRepository) -> None:
|
||||
"""A target value in use on a workflow run blocks rename even if its color row was soft-deleted."""
|
||||
from sqlalchemy import update as _update
|
||||
|
||||
await _seed_workflow_run(repo, "wr_source")
|
||||
await _seed_workflow_run(repo, "wr_target")
|
||||
await repo.apply_run_tag_changes("wr_source", ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes("wr_target", ORG_ID, sets={"env": "stg"}, deletes=set(), context=_ctx())
|
||||
async with repo.Session() as session:
|
||||
await session.execute(
|
||||
_update(TagValueModel)
|
||||
.where(TagValueModel.key == "env", TagValueModel.value == "stg")
|
||||
.values(deleted_at=datetime.now(timezone.utc))
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
with pytest.raises(TagValueRenameCollision):
|
||||
await repo.rename_tag_value(ORG_ID, "env", "prod", "stg", _ctx())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rename_tag_value_reuses_soft_deleted_target(repo: TagsRepository) -> None:
|
||||
"""A previously soft-deleted target value is NOT a collision — rename succeeds."""
|
||||
|
|
@ -1003,6 +1177,31 @@ async def test_delete_tag_value_cascades_and_counts(repo: TagsRepository) -> Non
|
|||
assert await repo.get_active_grouped_tags_for_workflow("wpid_c", ORG_ID) == {"env": "stg"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_tag_value_cascades_across_workflow_runs(repo: TagsRepository) -> None:
|
||||
for wrid in ("wr_a", "wr_b", "wr_c"):
|
||||
await _seed_workflow_run(repo, wrid)
|
||||
|
||||
await repo.apply_tag_changes("wpid_a", ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes("wr_a", ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes("wr_b", ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
await repo.apply_run_tag_changes("wr_c", ORG_ID, sets={"env": "stg"}, deletes=set(), context=_ctx())
|
||||
|
||||
removed_workflows = await repo.delete_tag_value(ORG_ID, "env", "prod", _ctx())
|
||||
|
||||
assert removed_workflows == 1
|
||||
assert await repo.get_active_grouped_tags_for_workflow("wpid_a", ORG_ID) == {}
|
||||
assert await repo.get_active_grouped_tags_for_run("wr_a", ORG_ID) == {}
|
||||
assert await repo.get_active_grouped_tags_for_run("wr_b", ORG_ID) == {}
|
||||
assert await repo.get_active_grouped_tags_for_run("wr_c", ORG_ID) == {"env": "stg"}
|
||||
rows = await _all_run_events(repo)
|
||||
run_deletes = [row for row in rows if row.event_type == TagEventType.DELETE.value]
|
||||
assert {(row.workflow_run_id, row.key, row.value) for row in run_deletes} == {
|
||||
("wr_a", "env", "prod"),
|
||||
("wr_b", "env", "prod"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_tag_value_writes_delete_event_carrying_value(repo: TagsRepository) -> None:
|
||||
await repo.apply_tag_changes("wpid_a", ORG_ID, sets={"env": "prod"}, deletes=set(), context=_ctx())
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue