Update abstra-lib

abstra_internals/contracts_generated.py

@@ -1,7 +1,7 @@
 """
 This file is auto generated. Do not edit this file manually.
 Source code: [monorepo]/contracts/generator.py
-Last generated: 2024-04-03 19:32:51.808772
+Last generated: 2024-04-05 13:39:45.761560
 """
 
 import typing
@@ -365,8 +365,6 @@ def from_dict(data: typing.Dict[str, typing.Any]) -> "AbstraLibApiStageRun":
 
 AbstraLibApiStageRunCardContent = AbstraLibApiStageCardContents
 
-AbstraLibApiStageRunCardStage = str
-
 
 @dataclass
 class AbstraLibApiStageRunCard:
@@ -376,7 +374,6 @@ class AbstraLibApiStageRunCard:
     assignee: typing.Optional[AbstraLibApiStageRunCardAssignee]
     status: AbstraLibApiStageRunCardStatus
     content: AbstraLibApiStageRunCardContent
-    stage: typing.Optional[AbstraLibApiStageRunCardStage]
 
     def to_dict(self) -> typing.Dict[str, typing.Any]:
         return {
@@ -386,7 +383,6 @@ def to_dict(self) -> typing.Dict[str, typing.Any]:
             "assignee": self.assignee,
             "status": self.status,
             "content": self.content,
-            "stage": self.stage,
         }
 
     @staticmethod
@@ -400,7 +396,6 @@ def from_dict(data: typing.Dict[str, typing.Any]) -> "AbstraLibApiStageRunCard":
             content=[
                 AbstraLibApiStageCardContent.from_dict(item) for item in data["content"]
             ],
-            stage=str(data.get("stage")),
         )
 
 
@@ -467,24 +462,28 @@ def from_dict(data: typing.Dict[str, typing.Any]) -> "AbstraLibApiStage":
 
 AbstraLibApiKanbanColumnStageRunCards = AbstraLibApiStageRunCards
 
-AbstraLibApiKanbanColumnTotalCount = int
+AbstraLibApiStagesItem = AbstraLibApiStage
+
+AbstraLibApiStages = typing.List[AbstraLibApiStagesItem]
+
+AbstraLibApiKanbanColumnStages = AbstraLibApiStages
 
-AbstraLibApiKanbanColumnLoading = bool
+AbstraLibApiKanbanColumnTotalCount = int
 
 
 @dataclass
 class AbstraLibApiKanbanColumn:
     selected_stage: AbstraLibApiKanbanColumnSelectedStage
     stage_run_cards: AbstraLibApiKanbanColumnStageRunCards
+    stages: AbstraLibApiKanbanColumnStages
     total_count: AbstraLibApiKanbanColumnTotalCount
-    loading: typing.Optional[AbstraLibApiKanbanColumnLoading]
 
     def to_dict(self) -> typing.Dict[str, typing.Any]:
         return {
             "selected_stage": self.selected_stage,
             "stage_run_cards": self.stage_run_cards,
+            "stages": self.stages,
             "total_count": self.total_count,
-            "loading": self.loading,
         }
 
     @staticmethod
@@ -495,8 +494,8 @@ def from_dict(data: typing.Dict[str, typing.Any]) -> "AbstraLibApiKanbanColumn":
                 AbstraLibApiStageRunCard.from_dict(item)
                 for item in data["stage_run_cards"]
             ],
+            stages=[AbstraLibApiStage.from_dict(item) for item in data["stages"]],
             total_count=int(data["total_count"]),
-            loading=bool(data.get("loading")),
         )
 
 
@@ -506,14 +505,26 @@ def from_dict(data: typing.Dict[str, typing.Any]) -> "AbstraLibApiKanbanColumn":
 
 AbstraLibApiKanbanDataColumns = AbstraLibApiKanbanColumns
 
+AbstraLibApiKanbanDataNextStageOptions = AbstraLibApiStages
+
+AbstraLibApiKanbanDataNotFoundStagesItem = str
+
+AbstraLibApiKanbanDataNotFoundStages = typing.List[
+    AbstraLibApiKanbanDataNotFoundStagesItem
+]
+
 
 @dataclass
 class AbstraLibApiKanbanData:
     columns: AbstraLibApiKanbanDataColumns
+    next_stage_options: AbstraLibApiKanbanDataNextStageOptions
+    not_found_stages: AbstraLibApiKanbanDataNotFoundStages
 
     def to_dict(self) -> typing.Dict[str, typing.Any]:
         return {
             "columns": self.columns,
+            "next_stage_options": self.next_stage_options,
+            "not_found_stages": self.not_found_stages,
         }
 
     @staticmethod
@@ -522,51 +533,17 @@ def from_dict(data: typing.Dict[str, typing.Any]) -> "AbstraLibApiKanbanData":
             columns=[
                 AbstraLibApiKanbanColumn.from_dict(item) for item in data["columns"]
             ],
-        )
-
-
-AbstraLibApiThreadsStageRunCards = AbstraLibApiStageRunCards
-
-AbstraLibApiThreadsNotFoundStagesItem = str
-
-AbstraLibApiThreadsNotFoundStages = typing.List[AbstraLibApiThreadsNotFoundStagesItem]
-
-AbstraLibApiThreadsTotalCount = int
-
-
-@dataclass
-class AbstraLibApiThreads:
-    stage_run_cards: AbstraLibApiThreadsStageRunCards
-    not_found_stages: AbstraLibApiThreadsNotFoundStages
-    total_count: AbstraLibApiThreadsTotalCount
-
-    def to_dict(self) -> typing.Dict[str, typing.Any]:
-        return {
-            "stage_run_cards": self.stage_run_cards,
-            "not_found_stages": self.not_found_stages,
-            "total_count": self.total_count,
-        }
-
-    @staticmethod
-    def from_dict(data: typing.Dict[str, typing.Any]) -> "AbstraLibApiThreads":
-        return AbstraLibApiThreads(
-            stage_run_cards=[
-                AbstraLibApiStageRunCard.from_dict(item)
-                for item in data["stage_run_cards"]
+            next_stage_options=[
+                AbstraLibApiStage.from_dict(item) for item in data["next_stage_options"]
             ],
             not_found_stages=[str(item) for item in data["not_found_stages"]],
-            total_count=int(data["total_count"]),
         )
 
 
 AbstraLibApiVisualizationsItem = AbstraLibApiVisualization
 
 AbstraLibApiVisualizations = typing.List[AbstraLibApiVisualizationsItem]
 
-AbstraLibApiStagesItem = AbstraLibApiStage
-
-AbstraLibApiStages = typing.List[AbstraLibApiStagesItem]
-
 AbstraLibApiEditorLintersRuleName = str
 
 AbstraLibApiEditorLintersRuleLabel = str

abstra_internals/interface/sdk/workflows/__init__.py

@@ -22,7 +22,7 @@ def match_thread(filter: dict):
         raise Exception("match_thread may only be used inside a hook")
     stage_id = execution.stage.id
     stage_runs = execution.stage_run_manager.find(
-        filter=GetStageRunByQueryFilter(stage=stage_id, data=filter, status=["waiting"])
+        filter=GetStageRunByQueryFilter(stage=stage_id, data=filter, status="waiting")
     )
     if len(stage_runs) == 0:
         raise Exception(f"No thread found for {filter}")

abstra_internals/repositories/stage_run.py

@@ -9,6 +9,7 @@
 from pydantic import ConfigDict, Field
 from pydantic.dataclasses import dataclass
 
+from ..utils import serialize
 from ..utils.datetime import from_utc_iso_string
 from ..utils.environment import SIDECAR_HEADERS, SIDECAR_URL
 from ..utils.validate import validate_json
@@ -58,9 +59,8 @@ class GetStageRunByQueryFilter:
     stage: Optional[Union[str, List[str]]] = None
     assignee: Optional[str] = None
     parent_id: Optional[str] = None
-    status: Optional[List[str]] = None
+    status: Optional[str] = None
     data: Optional[Dict] = None
-    search: Optional[str] = None
 
     @staticmethod
     def from_dict(data: dict) -> "GetStageRunByQueryFilter":
@@ -70,7 +70,6 @@ def from_dict(data: dict) -> "GetStageRunByQueryFilter":
             parent_id=data.get("parent_id", None),
             status=data.get("status", None),
             data=data.get("data", None),
-            search=data.get("search", None),
         )
 
     def to_dict(self) -> dict:
@@ -170,9 +169,7 @@ def find_ancestors(self, id: str) -> List[StageRun]:
 
     @abstractmethod
     def find_leaves(
-        self,
-        filter: GetStageRunByQueryFilter,
-        pagination: Pagination,
+        self, filter: GetStageRunByQueryFilter, pagination: Pagination
     ) -> PaginatedListResponse:
         raise NotImplementedError()
 
@@ -231,39 +228,32 @@ def find(self, filter: GetStageRunByQueryFilter) -> List[StageRun]:
         results: List[StageRun] = []
 
         for stage_run in self._stage_runs:
-            if filter.stage and stage_run.stage not in filter.stage:
-                continue
+            if filter.stage:
+                if isinstance(filter.stage, list):
+                    if stage_run.stage not in filter.stage:
+                        continue
+                else:
+                    if stage_run.stage != filter.stage:
+                        continue
 
             if filter.data and not self._compare_data(stage_run.data, filter.data):
                 continue
 
-            if filter.status and stage_run.status not in filter.status:
+            if filter.status and stage_run.status != filter.status:
                 continue
 
             if filter.parent_id and stage_run.parent_id != filter.parent_id:
                 continue
 
-            if filter.search and not self._compare_search(stage_run, filter.search):
-                continue
-
             results.append(stage_run)
 
         return results
 
     def _compare_data(self, data: dict, filter: dict) -> bool:
-        for key, value in filter.items():
-            if key not in data or data[key] not in value:
-                return False
-        return True
-
-    def _compare_search(self, stage_run: StageRun, search: str) -> bool:
-        # TODO: filter by stage title and stage type
-        status_match = search.lower() in stage_run.status.lower()
-        data_match = False
-        for key, value in stage_run.data.items():
-            if search.lower() in value.lower():
-                data_match = True
-        return status_match or data_match
+        filtered_data = {k: v for k, v in data.items() if k in filter}
+        return serialize(filtered_data, sort_keys=True) == serialize(
+            filter, sort_keys=True
+        )
 
     def find_ancestors(self, id: str) -> List[StageRun]:
         ancestors = []
@@ -275,9 +265,7 @@ def find_ancestors(self, id: str) -> List[StageRun]:
         return ancestors
 
     def find_leaves(
-        self,
-        filter: GetStageRunByQueryFilter,
-        pagination: Pagination,
+        self, filter: GetStageRunByQueryFilter, pagination: Pagination
     ) -> PaginatedListResponse:
         filter_matches = self.find(filter)
         parent_ids = set(u.parent_id for u in self._stage_runs)
@@ -295,6 +283,7 @@ def find_leaves(
         total_count = len(leaves)
         if pagination.offset >= total_count:
             return PaginatedListResponse(stage_runs=[], total_count=total_count)
+
         return PaginatedListResponse(
             stage_runs=leaves[pagination.offset : pagination.limit],
             total_count=total_count,
@@ -468,9 +457,7 @@ def change_state(
         return True
 
     def find_leaves(
-        self,
-        filter: GetStageRunByQueryFilter,
-        pagination: Pagination,
+        self, filter: GetStageRunByQueryFilter, pagination: Pagination
     ) -> PaginatedListResponse:
         r = self._request(
             "GET",