feat(robot-server): add protocol_kind to the protocols endpoint and p…

…ersistance layer. (#15353)

robot-server/robot_server/persistence/_migrations/v4_to_v5.py

@@ -0,0 +1,52 @@
+"""Migrate the persistence directory from schema 4 to 5.
+
+Summary of changes from schema 4:
+
+- Adds a new "protocol_kind" column to protocols table
+"""
+
+from pathlib import Path
+from contextlib import ExitStack
+import shutil
+from typing import Any
+
+import sqlalchemy
+
+from ..database import sql_engine_ctx
+from ..tables import schema_5
+from .._folder_migrator import Migration
+
+_DB_FILE = "robot_server.db"
+
+
+class Migration4to5(Migration):  # noqa: D101
+    def migrate(self, source_dir: Path, dest_dir: Path) -> None:
+        """Migrate the persistence directory from schema 4 to 5."""
+        # Copy over all existing directories and files to new version
+        for item in source_dir.iterdir():
+            if item.is_dir():
+                shutil.copytree(src=item, dst=dest_dir / item.name)
+            else:
+                shutil.copy(src=item, dst=dest_dir / item.name)
+        dest_db_file = dest_dir / _DB_FILE
+
+        # Append the new column to existing protocols in v4 database
+        with ExitStack() as exit_stack:
+            dest_engine = exit_stack.enter_context(sql_engine_ctx(dest_db_file))
+            schema_5.metadata.create_all(dest_engine)
+
+            def add_column(
+                engine: sqlalchemy.engine.Engine,
+                table_name: str,
+                column: Any,
+            ) -> None:
+                column_type = column.type.compile(engine.dialect)
+                engine.execute(
+                    f"ALTER TABLE {table_name} ADD COLUMN {column.key} {column_type}"
+                )
+
+            add_column(
+                dest_engine,
+                schema_5.protocol_table.name,
+                schema_5.protocol_table.c.protocol_kind,
+            )

robot-server/robot_server/persistence/persistence_directory.py

@@ -11,7 +11,7 @@
 from anyio import Path as AsyncPath, to_thread
 
 from ._folder_migrator import MigrationOrchestrator
-from ._migrations import up_to_3, v3_to_v4
+from ._migrations import up_to_3, v3_to_v4, v4_to_v5
 
 
 _TEMP_PERSISTENCE_DIR_PREFIX: Final = "opentrons-robot-server-"
@@ -51,6 +51,7 @@ async def prepare_active_subdirectory(prepared_root: Path) -> Path:
         migrations=[
             up_to_3.MigrationUpTo3(subdirectory="3"),
             v3_to_v4.Migration3to4(subdirectory="4"),
+            v4_to_v5.Migration4to5(subdirectory="5"),
         ],
         temp_file_prefix="temp-",
     )

robot-server/robot_server/persistence/tables/__init__.py

@@ -1,7 +1,7 @@
 """SQL database schemas."""
 
 # Re-export the latest schema.
-from .schema_4 import (
+from .schema_5 import (
     metadata,
     protocol_table,
     analysis_table,

robot-server/robot_server/persistence/tables/schema_5.py

@@ -0,0 +1,138 @@
+"""v5 of our SQLite schema."""
+
+import sqlalchemy
+
+from robot_server.persistence._utc_datetime import UTCDateTime
+
+metadata = sqlalchemy.MetaData()
+
+protocol_table = sqlalchemy.Table(
+    "protocol",
+    metadata,
+    sqlalchemy.Column(
+        "id",
+        sqlalchemy.String,
+        primary_key=True,
+    ),
+    sqlalchemy.Column(
+        "created_at",
+        UTCDateTime,
+        nullable=False,
+    ),
+    sqlalchemy.Column("protocol_key", sqlalchemy.String, nullable=True),
+    sqlalchemy.Column("protocol_kind", sqlalchemy.String, nullable=True),
+)
+
+analysis_table = sqlalchemy.Table(
+    "analysis",
+    metadata,
+    sqlalchemy.Column(
+        "id",
+        sqlalchemy.String,
+        primary_key=True,
+    ),
+    sqlalchemy.Column(
+        "protocol_id",
+        sqlalchemy.String,
+        sqlalchemy.ForeignKey("protocol.id"),
+        index=True,
+        nullable=False,
+    ),
+    sqlalchemy.Column(
+        "analyzer_version",
+        sqlalchemy.String,
+        nullable=False,
+    ),
+    sqlalchemy.Column(
+        "completed_analysis",
+        # Stores a JSON string. See CompletedAnalysisStore.
+        sqlalchemy.String,
+        nullable=False,
+    ),
+    # column added in schema v4
+    sqlalchemy.Column(
+        "run_time_parameter_values_and_defaults",
+        sqlalchemy.String,
+        nullable=True,
+    ),
+)
+
+run_table = sqlalchemy.Table(
+    "run",
+    metadata,
+    sqlalchemy.Column(
+        "id",
+        sqlalchemy.String,
+        primary_key=True,
+    ),
+    sqlalchemy.Column(
+        "created_at",
+        UTCDateTime,
+        nullable=False,
+    ),
+    sqlalchemy.Column(
+        "protocol_id",
+        sqlalchemy.String,
+        sqlalchemy.ForeignKey("protocol.id"),
+        nullable=True,
+    ),
+    # column added in schema v1
+    sqlalchemy.Column(
+        "state_summary",
+        sqlalchemy.String,
+        nullable=True,
+    ),
+    # column added in schema v1
+    sqlalchemy.Column("engine_status", sqlalchemy.String, nullable=True),
+    # column added in schema v1
+    sqlalchemy.Column("_updated_at", UTCDateTime, nullable=True),
+    # column added in schema v4
+    sqlalchemy.Column(
+        "run_time_parameters",
+        # Stores a JSON string. See RunStore.
+        sqlalchemy.String,
+        nullable=True,
+    ),
+)
+
+action_table = sqlalchemy.Table(
+    "action",
+    metadata,
+    sqlalchemy.Column(
+        "id",
+        sqlalchemy.String,
+        primary_key=True,
+    ),
+    sqlalchemy.Column("created_at", UTCDateTime, nullable=False),
+    sqlalchemy.Column("action_type", sqlalchemy.String, nullable=False),
+    sqlalchemy.Column(
+        "run_id",
+        sqlalchemy.String,
+        sqlalchemy.ForeignKey("run.id"),
+        nullable=False,
+    ),
+)
+
+run_command_table = sqlalchemy.Table(
+    "run_command",
+    metadata,
+    sqlalchemy.Column("row_id", sqlalchemy.Integer, primary_key=True),
+    sqlalchemy.Column(
+        "run_id", sqlalchemy.String, sqlalchemy.ForeignKey("run.id"), nullable=False
+    ),
+    sqlalchemy.Column("index_in_run", sqlalchemy.Integer, nullable=False),
+    sqlalchemy.Column("command_id", sqlalchemy.String, nullable=False),
+    sqlalchemy.Column("command", sqlalchemy.String, nullable=False),
+    sqlalchemy.Index(
+        "ix_run_run_id_command_id",  # An arbitrary name for the index.
+        "run_id",
+        "command_id",
+        unique=True,
+    ),
+    sqlalchemy.Index(
+        "ix_run_run_id_index_in_run",  # An arbitrary name for the index.
+        "run_id",
+        "index_in_run",
+        unique=True,
+    ),
+)

robot-server/robot_server/protocols/protocol_models.py

@@ -1,7 +1,9 @@
 """Protocol file models."""
+
 from datetime import datetime
 from pydantic import BaseModel, Extra, Field
 from typing import Any, List, Optional
+from enum import Enum
 
 from opentrons.protocol_reader import (
     ProtocolType as ProtocolType,
@@ -14,6 +16,21 @@
 from .analysis_models import AnalysisSummary
 
 
+class ProtocolKind(str, Enum):
+    """Kind of protocol, standard or quick-transfer."""
+
+    STANDARD = "standard"
+    QUICK_TRANSFER = "quick-transfer"
+
+    @staticmethod
+    def from_string(name: Optional[str]) -> Optional["ProtocolKind"]:
+        """Get the ProtocolKind from a string."""
+        for item in ProtocolKind:
+            if name == item.value:
+                return item
+        return None
+
+
 class ProtocolFile(BaseModel):
     """A file in a protocol."""
 
@@ -109,3 +126,10 @@ class Protocol(ResourceModel):
             " See `POST /protocols`."
         ),
     )
+
+    protocolKind: Optional[ProtocolKind] = Field(
+        ...,
+        description="The kind of protocol (standard or quick-transfer)."
+        "The client provides this field when the protocol is uploaded."
+        " See `POST /protocols`.",
+    )

robot-server/robot_server/protocols/protocol_store.py

@@ -1,4 +1,5 @@
 """Store and retrieve information about uploaded protocols."""
+
 from __future__ import annotations
 
 from dataclasses import dataclass
@@ -35,6 +36,7 @@ class ProtocolResource:
     created_at: datetime
     source: ProtocolSource
     protocol_key: Optional[str]
+    protocol_kind: Optional[str]
 
 
 @dataclass(frozen=True)
@@ -166,6 +168,7 @@ def insert(self, resource: ProtocolResource) -> None:
                 protocol_id=resource.protocol_id,
                 created_at=resource.created_at,
                 protocol_key=resource.protocol_key,
+                protocol_kind=resource.protocol_kind,
             )
         )
         self._sources_by_id[resource.protocol_id] = resource.source
@@ -183,6 +186,7 @@ def get(self, protocol_id: str) -> ProtocolResource:
             protocol_id=sql_resource.protocol_id,
             created_at=sql_resource.created_at,
             protocol_key=sql_resource.protocol_key,
+            protocol_kind=sql_resource.protocol_kind,
             source=self._sources_by_id[sql_resource.protocol_id],
         )
 
@@ -198,6 +202,7 @@ def get_all(self) -> List[ProtocolResource]:
                 protocol_id=r.protocol_id,
                 created_at=r.created_at,
                 protocol_key=r.protocol_key,
+                protocol_kind=r.protocol_kind,
                 source=self._sources_by_id[r.protocol_id],
             )
             for r in all_sql_resources
@@ -471,6 +476,7 @@ class _DBProtocolResource:
     protocol_id: str
     created_at: datetime
     protocol_key: Optional[str]
+    protocol_kind: Optional[str]
 
 
 def _convert_sql_row_to_dataclass(
@@ -479,16 +485,21 @@ def _convert_sql_row_to_dataclass(
     protocol_id = sql_row.id
     protocol_key = sql_row.protocol_key
     created_at = sql_row.created_at
+    protocol_kind = sql_row.protocol_kind
 
     assert isinstance(protocol_id, str), f"Protocol ID {protocol_id} not a string"
     assert protocol_key is None or isinstance(
         protocol_key, str
     ), f"Protocol Key {protocol_key} not a string or None"
+    assert protocol_kind is None or isinstance(
+        protocol_kind, str
+    ), f"Protocol Kind {protocol_kind} not a string or None"
 
     return _DBProtocolResource(
         protocol_id=protocol_id,
         created_at=created_at,
         protocol_key=protocol_key,
+        protocol_kind=protocol_kind,
     )
 
 
@@ -499,4 +510,5 @@ def _convert_dataclass_to_sql_values(
         "id": resource.protocol_id,
         "created_at": resource.created_at,
         "protocol_key": resource.protocol_key,
+        "protocol_kind": resource.protocol_kind,
     }