[CZID-9248] update loader signatures (#166)

chanzuckerberg · Jan 31, 2024 · 139b53f · 139b53f
1 parent 321d8d3
commit 139b53f
Show file tree

Hide file tree

Showing 21 changed files with 201 additions and 148 deletions.
diff --git a/workflows/api/loader/loader.py b/workflows/api/loader/loader.py
@@ -13,11 +13,11 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 from entity_interface import create_entities
 
-from plugins.plugin_types import EventBus, EntityInputLoader, EntityOutputLoader, WorkflowSucceededMessage
+from plugins.plugin_types import EventBus, InputLoader, OutputLoader, WorkflowSucceededMessage
 from database.models import WorkflowRun
 from manifest.manifest import EntityInputArgument, Manifest, RawInputArgument
 
-T = TypeVar("T", bound=Type[EntityInputLoader] | Type[EntityOutputLoader])
+T = TypeVar("T", bound=Type[InputLoader] | Type[OutputLoader])
 
 
 def load_loader_plugins(input_or_output: Literal["input", "output"], cls: T) -> Dict[str, List[Tuple[Version, T]]]:
@@ -35,10 +35,10 @@ def load_loader_plugins(input_or_output: Literal["input", "output"], cls: T) ->
     return loaders
 
 
-input_loaders = load_loader_plugins("input", EntityInputLoader)
+input_loaders = load_loader_plugins("input", InputLoader)
 
 
-def resolve_entity_input_loaders(workflow_manifest: Manifest) -> List[type[EntityInputLoader]]:
+def resolve_entity_input_loaders(workflow_manifest: Manifest) -> List[type[InputLoader]]:
     """
     Given a manifest, resolve input loaders
     """
@@ -55,11 +55,11 @@ def resolve_entity_input_loaders(workflow_manifest: Manifest) -> List[type[Entit
     return resolved_loaders
 
 
-output_loaders = load_loader_plugins("output", EntityOutputLoader)
+output_loaders = load_loader_plugins("output", OutputLoader)
 
 
 # TODO: DRY with above but make the types work poperly
-def resolve_entity_output_loaders(workflow_manifest: Manifest) -> List[type[EntityOutputLoader]]:
+def resolve_entity_output_loaders(workflow_manifest: Manifest) -> List[type[OutputLoader]]:
     """
     Given a manifest, resolve output loaders
     """

diff --git a/workflows/api/schema.graphql b/workflows/api/schema.graphql
@@ -180,6 +180,9 @@ enum WorkflowCountColumns {
   producing_run_id
   owner_user_id
   collection_id
+  created_at
+  updated_at
+  deleted_at
 }
 
 input WorkflowCreateInput {
@@ -293,6 +296,9 @@ enum WorkflowRunCountColumns {
   producing_run_id
   owner_user_id
   collection_id
+  created_at
+  updated_at
+  deleted_at
 }
 
 input WorkflowRunCreateInput {
@@ -360,6 +366,9 @@ enum WorkflowRunEntityInputCountColumns {
   producing_run_id
   owner_user_id
   collection_id
+  created_at
+  updated_at
+  deleted_at
 }
 
 input WorkflowRunEntityInputCreateInput {
@@ -494,6 +503,9 @@ enum WorkflowRunStepCountColumns {
   producing_run_id
   owner_user_id
   collection_id
+  created_at
+  updated_at
+  deleted_at
 }
 
 input WorkflowRunStepCreateInput {
@@ -673,6 +685,9 @@ enum WorkflowVersionCountColumns {
   producing_run_id
   owner_user_id
   collection_id
+  created_at
+  updated_at
+  deleted_at
 }
 
 input WorkflowVersionCreateInput {

diff --git a/workflows/api/schema.json b/workflows/api/schema.json
@@ -2948,6 +2948,15 @@
             },
             {
               "name": "collection_id"
+            },
+            {
+              "name": "created_at"
+            },
+            {
+              "name": "updated_at"
+            },
+            {
+              "name": "deleted_at"
             }
           ],
           "fields": null,
@@ -3325,6 +3334,15 @@
             },
             {
               "name": "collection_id"
+            },
+            {
+              "name": "created_at"
+            },
+            {
+              "name": "updated_at"
+            },
+            {
+              "name": "deleted_at"
             }
           ],
           "fields": null,
@@ -3773,6 +3791,15 @@
             },
             {
               "name": "collection_id"
+            },
+            {
+              "name": "created_at"
+            },
+            {
+              "name": "updated_at"
+            },
+            {
+              "name": "deleted_at"
             }
           ],
           "fields": null,
@@ -4200,6 +4227,15 @@
             },
             {
               "name": "collection_id"
+            },
+            {
+              "name": "created_at"
+            },
+            {
+              "name": "updated_at"
+            },
+            {
+              "name": "deleted_at"
             }
           ],
           "fields": null,
@@ -4449,6 +4485,15 @@
             },
             {
               "name": "collection_id"
+            },
+            {
+              "name": "created_at"
+            },
+            {
+              "name": "updated_at"
+            },
+            {
+              "name": "deleted_at"
             }
           ],
           "fields": null,

diff --git a/workflows/api/types/workflow.py b/workflows/api/types/workflow.py
@@ -202,6 +202,9 @@ class WorkflowCountColumns(enum.Enum):
     producing_run_id = "producing_run_id"
     owner_user_id = "owner_user_id"
     collection_id = "collection_id"
+    created_at = "created_at"
+    updated_at = "updated_at"
+    deleted_at = "deleted_at"
 
 
 """

diff --git a/workflows/api/types/workflow_run.py b/workflows/api/types/workflow_run.py
@@ -296,6 +296,9 @@ class WorkflowRunCountColumns(enum.Enum):
     producing_run_id = "producing_run_id"
     owner_user_id = "owner_user_id"
     collection_id = "collection_id"
+    created_at = "created_at"
+    updated_at = "updated_at"
+    deleted_at = "deleted_at"
 
 
 """

diff --git a/workflows/api/types/workflow_run_entity_input.py b/workflows/api/types/workflow_run_entity_input.py
@@ -173,6 +173,9 @@ class WorkflowRunEntityInputCountColumns(enum.Enum):
     producing_run_id = "producing_run_id"
     owner_user_id = "owner_user_id"
     collection_id = "collection_id"
+    created_at = "created_at"
+    updated_at = "updated_at"
+    deleted_at = "deleted_at"
 
 
 """

diff --git a/workflows/api/types/workflow_run_step.py b/workflows/api/types/workflow_run_step.py
@@ -180,6 +180,9 @@ class WorkflowRunStepCountColumns(enum.Enum):
     producing_run_id = "producing_run_id"
     owner_user_id = "owner_user_id"
     collection_id = "collection_id"
+    created_at = "created_at"
+    updated_at = "updated_at"
+    deleted_at = "deleted_at"
 
 
 """

diff --git a/workflows/api/types/workflow_version.py b/workflows/api/types/workflow_version.py
@@ -224,6 +224,9 @@ class WorkflowVersionCountColumns(enum.Enum):
     producing_run_id = "producing_run_id"
     owner_user_id = "owner_user_id"
     collection_id = "collection_id"
+    created_at = "created_at"
+    updated_at = "updated_at"
+    deleted_at = "deleted_at"
 
 
 """

diff --git a/workflows/cli/gql_schema.py b/workflows/cli/gql_schema.py
@@ -38,13 +38,16 @@ class WorkflowCountColumns(sgqlc.types.Enum):
     __schema__ = gql_schema
     __choices__ = (
         "collection_id",
+        "created_at",
         "default_version",
+        "deleted_at",
         "entity_id",
         "id",
         "minimum_supported_version",
         "name",
         "owner_user_id",
         "producing_run_id",
+        "updated_at",
         "versions",
     )
 
@@ -53,6 +56,8 @@ class WorkflowRunCountColumns(sgqlc.types.Enum):
     __schema__ = gql_schema
     __choices__ = (
         "collection_id",
+        "created_at",
+        "deleted_at",
         "deprecated_by",
         "ended_at",
         "entity_id",
@@ -66,6 +71,7 @@ class WorkflowRunCountColumns(sgqlc.types.Enum):
         "started_at",
         "status",
         "steps",
+        "updated_at",
         "workflow_runner_inputs_json",
         "workflow_version",
     )
@@ -75,12 +81,15 @@ class WorkflowRunEntityInputCountColumns(sgqlc.types.Enum):
     __schema__ = gql_schema
     __choices__ = (
         "collection_id",
+        "created_at",
+        "deleted_at",
         "entity_id",
         "field_name",
         "id",
         "input_entity_id",
         "owner_user_id",
         "producing_run_id",
+        "updated_at",
         "workflow_run",
     )
 
@@ -94,13 +103,16 @@ class WorkflowRunStepCountColumns(sgqlc.types.Enum):
     __schema__ = gql_schema
     __choices__ = (
         "collection_id",
+        "created_at",
+        "deleted_at",
         "ended_at",
         "entity_id",
         "id",
         "owner_user_id",
         "producing_run_id",
         "started_at",
         "status",
+        "updated_at",
         "workflow_run",
     )
 
@@ -114,13 +126,16 @@ class WorkflowVersionCountColumns(sgqlc.types.Enum):
     __schema__ = gql_schema
     __choices__ = (
         "collection_id",
+        "created_at",
+        "deleted_at",
         "entity_id",
         "graph_json",
         "id",
         "manifest",
         "owner_user_id",
         "producing_run_id",
         "runs",
+        "updated_at",
         "version",
         "workflow",
         "workflow_uri",

diff --git a/workflows/database_migrations/versions/20240130_195510_.py b/workflows/database_migrations/versions/20240130_195510_.py
@@ -0,0 +1,33 @@
+"""
+
+Revision ID: 20240130_195510
+Revises: 20240128_222905
+Create Date: 2024-01-31 03:55:11.432738
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "20240130_195510"
+down_revision = "20240128_222905"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "entity", sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False)
+    )
+    op.add_column("entity", sa.Column("updated_at", sa.DateTime(timezone=True), nullable=True))
+    op.add_column("entity", sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("entity", "deleted_at")
+    op.drop_column("entity", "updated_at")
+    op.drop_column("entity", "created_at")
+    # ### end Alembic commands ###
diff --git a/workflows/manifest/manifest.py b/workflows/manifest/manifest.py
@@ -261,7 +261,7 @@ def _unique_input_names(self):  # type: ignore
 
     @model_validator(mode="after")
     def _validate_references(self):  # type: ignore
-        errors: list[InvalidInputReference] = []
+        errors = []
         inputs = set(self.entity_inputs.keys()) | set(self.raw_inputs.keys())
         for input_loader in self.input_loaders:
             errors += list(input_loader._check_references("input", input_loader.name, inputs))

diff --git a/workflows/plugins/entity_input_loaders/fasta/loader_fasta.py b/workflows/plugins/entity_input_loaders/fasta/loader_fasta.py
diff --git a/workflows/plugins/entity_input_loaders/fasta/setup.py b/workflows/plugins/entity_input_loaders/fasta/setup.py