Skip to content

Commit

Permalink
change proxy location db type
Browse files Browse the repository at this point in the history
  • Loading branch information
LawyZheng committed Feb 5, 2025
1 parent cc449dc commit 396e767
Show file tree
Hide file tree
Showing 2 changed files with 202 additions and 5 deletions.
Original file line number Diff line number Diff line change
@@ -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 ###
9 changes: 4 additions & 5 deletions skyvern/forge/sdk/db/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 396e767

Please sign in to comment.