From 8a37a26bc05952358dfa4bd3131ede1379b1cd01 Mon Sep 17 00:00:00 2001 From: lawyzheng Date: Thu, 6 Feb 2025 02:43:08 +0800 Subject: [PATCH] change proxy location db type --- ...ebf093461132_change_proxy_location_enum.py | 198 ++++++++++++++++++ skyvern/forge/sdk/db/models.py | 9 +- 2 files changed, 202 insertions(+), 5 deletions(-) create mode 100644 alembic/versions/2025_02_05_1839-ebf093461132_change_proxy_location_enum.py diff --git a/alembic/versions/2025_02_05_1839-ebf093461132_change_proxy_location_enum.py b/alembic/versions/2025_02_05_1839-ebf093461132_change_proxy_location_enum.py new file mode 100644 index 0000000000..efbac35106 --- /dev/null +++ b/alembic/versions/2025_02_05_1839-ebf093461132_change_proxy_location_enum.py @@ -0,0 +1,198 @@ +"""change proxy location enum + +Revision ID: ebf093461132 +Revises: 5dd8928389c5 +Create Date: 2025-02-05 18:39:29.253202+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "ebf093461132" +down_revision: Union[str, None] = "5dd8928389c5" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "observer_cruises", + "proxy_location", + existing_type=postgresql.ENUM( + "US_CA", + "US_NY", + "US_TX", + "US_FL", + "US_WA", + "RESIDENTIAL", + "NONE", + "RESIDENTIAL_IE", + "RESIDENTIAL_GB", + "RESIDENTIAL_IN", + "RESIDENTIAL_JP", + "RESIDENTIAL_DE", + name="proxylocation", + ), + type_=sa.String(), + existing_nullable=True, + ) + op.alter_column( + "tasks", + "proxy_location", + existing_type=postgresql.ENUM( + "US_CA", + "US_NY", + "US_TX", + "US_FL", + "US_WA", + "RESIDENTIAL", + "NONE", + "RESIDENTIAL_IE", + "RESIDENTIAL_GB", + "RESIDENTIAL_IN", + "RESIDENTIAL_JP", + "RESIDENTIAL_DE", + name="proxylocation", + ), + type_=sa.String(), + existing_nullable=True, + ) + op.alter_column( + "workflow_runs", + "proxy_location", + existing_type=postgresql.ENUM( + "US_CA", + "US_NY", + "US_TX", + "US_FL", + "US_WA", + "RESIDENTIAL", + "NONE", + "RESIDENTIAL_IE", + "RESIDENTIAL_GB", + "RESIDENTIAL_IN", + "RESIDENTIAL_JP", + "RESIDENTIAL_DE", + name="proxylocation", + ), + type_=sa.String(), + existing_nullable=True, + ) + op.alter_column( + "workflows", + "proxy_location", + existing_type=postgresql.ENUM( + "US_CA", + "US_NY", + "US_TX", + "US_FL", + "US_WA", + "RESIDENTIAL", + "NONE", + "RESIDENTIAL_IE", + "RESIDENTIAL_GB", + "RESIDENTIAL_IN", + "RESIDENTIAL_JP", + "RESIDENTIAL_DE", + name="proxylocation", + ), + type_=sa.String(), + existing_nullable=True, + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.alter_column( + "workflows", + "proxy_location", + existing_type=sa.String(), + type_=postgresql.ENUM( + "US_CA", + "US_NY", + "US_TX", + "US_FL", + "US_WA", + "RESIDENTIAL", + "NONE", + "RESIDENTIAL_IE", + "RESIDENTIAL_GB", + "RESIDENTIAL_IN", + "RESIDENTIAL_JP", + "RESIDENTIAL_DE", + name="proxylocation", + ), + existing_nullable=True, + ) + op.alter_column( + "workflow_runs", + "proxy_location", + existing_type=sa.String(), + type_=postgresql.ENUM( + "US_CA", + "US_NY", + "US_TX", + "US_FL", + "US_WA", + "RESIDENTIAL", + "NONE", + "RESIDENTIAL_IE", + "RESIDENTIAL_GB", + "RESIDENTIAL_IN", + "RESIDENTIAL_JP", + "RESIDENTIAL_DE", + name="proxylocation", + ), + existing_nullable=True, + ) + op.alter_column( + "tasks", + "proxy_location", + existing_type=sa.String(), + type_=postgresql.ENUM( + "US_CA", + "US_NY", + "US_TX", + "US_FL", + "US_WA", + "RESIDENTIAL", + "NONE", + "RESIDENTIAL_IE", + "RESIDENTIAL_GB", + "RESIDENTIAL_IN", + "RESIDENTIAL_JP", + "RESIDENTIAL_DE", + name="proxylocation", + ), + existing_nullable=True, + ) + op.alter_column( + "observer_cruises", + "proxy_location", + existing_type=sa.String(), + type_=postgresql.ENUM( + "US_CA", + "US_NY", + "US_TX", + "US_FL", + "US_WA", + "RESIDENTIAL", + "NONE", + "RESIDENTIAL_IE", + "RESIDENTIAL_GB", + "RESIDENTIAL_IN", + "RESIDENTIAL_JP", + "RESIDENTIAL_DE", + name="proxylocation", + ), + existing_nullable=True, + ) + # ### end Alembic commands ### diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 6c60cf9d9c..ffbb267af4 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -44,7 +44,6 @@ generate_workflow_run_id, ) from skyvern.forge.sdk.schemas.observers import ObserverThoughtType -from skyvern.forge.sdk.schemas.tasks import ProxyLocation class Base(AsyncAttrs, DeclarativeBase): @@ -70,7 +69,7 @@ class TaskModel(Base): navigation_payload = Column(JSON) extracted_information = Column(JSON) failure_reason = Column(String) - proxy_location = Column(Enum(ProxyLocation)) + proxy_location = Column(String) extracted_information_schema = Column(JSON) workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), index=True) order = Column(Integer, nullable=True) @@ -204,7 +203,7 @@ class WorkflowModel(Base): title = Column(String, nullable=False) description = Column(String, nullable=True) workflow_definition = Column(JSON, nullable=False) - proxy_location = Column(Enum(ProxyLocation)) + proxy_location = Column(String) webhook_callback_url = Column(String) totp_verification_url = Column(String) totp_identifier = Column(String) @@ -236,7 +235,7 @@ class WorkflowRunModel(Base): organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=False, index=True) status = Column(String, nullable=False) failure_reason = Column(String) - proxy_location = Column(Enum(ProxyLocation)) + proxy_location = Column(String) webhook_callback_url = Column(String) totp_verification_url = Column(String) totp_identifier = Column(String) @@ -564,7 +563,7 @@ class ObserverCruiseModel(Base): webhook_callback_url = Column(String, nullable=True) totp_verification_url = Column(String, nullable=True) totp_identifier = Column(String, nullable=True) - proxy_location = Column(Enum(ProxyLocation), nullable=True) + proxy_location = 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)