From 68e0e73020c6790f57868838fe4302d63fc65f9a Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Thu, 6 Feb 2025 17:35:59 +0800
Subject: [PATCH] get proxy_location enum back in the tasks table (#1736)

---
 ..._get_proxy_location_enum_back_for_tasks.py | 77 +++++++++++++++++++
 skyvern/forge/sdk/db/models.py                |  3 +-
 2 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 alembic/versions/2025_02_06_0923-8a981a8df4bb_get_proxy_location_enum_back_for_tasks.py

diff --git a/alembic/versions/2025_02_06_0923-8a981a8df4bb_get_proxy_location_enum_back_for_tasks.py b/alembic/versions/2025_02_06_0923-8a981a8df4bb_get_proxy_location_enum_back_for_tasks.py
new file mode 100644
index 00000000..775b36bb
--- /dev/null
+++ b/alembic/versions/2025_02_06_0923-8a981a8df4bb_get_proxy_location_enum_back_for_tasks.py
@@ -0,0 +1,77 @@
+"""get proxy_location enum back for tasks
+
+Revision ID: 8a981a8df4bb
+Revises: ebf093461132
+Create Date: 2025-02-06 09:23:24.003503+00:00
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "8a981a8df4bb"
+down_revision: Union[str, None] = "ebf093461132"
+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(
+        "tasks",
+        "proxy_location",
+        existing_type=sa.VARCHAR(),
+        type_=sa.Enum(
+            "US_CA",
+            "US_NY",
+            "US_TX",
+            "US_FL",
+            "US_WA",
+            "RESIDENTIAL",
+            "RESIDENTIAL_ES",
+            "RESIDENTIAL_IE",
+            "RESIDENTIAL_GB",
+            "RESIDENTIAL_IN",
+            "RESIDENTIAL_JP",
+            "RESIDENTIAL_FR",
+            "RESIDENTIAL_DE",
+            "NONE",
+            name="proxylocation",
+        ),
+        postgresql_using="proxy_location::proxylocation",
+        existing_nullable=True,
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column(
+        "tasks",
+        "proxy_location",
+        existing_type=sa.Enum(
+            "US_CA",
+            "US_NY",
+            "US_TX",
+            "US_FL",
+            "US_WA",
+            "RESIDENTIAL",
+            "RESIDENTIAL_ES",
+            "RESIDENTIAL_IE",
+            "RESIDENTIAL_GB",
+            "RESIDENTIAL_IN",
+            "RESIDENTIAL_JP",
+            "RESIDENTIAL_FR",
+            "RESIDENTIAL_DE",
+            "NONE",
+            name="proxylocation",
+        ),
+        type_=sa.VARCHAR(),
+        postgresql_using="proxy_location::text",
+        existing_nullable=True,
+    )
+    # ### end Alembic commands ###
diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py
index ffbb267a..da93d2ee 100644
--- a/skyvern/forge/sdk/db/models.py
+++ b/skyvern/forge/sdk/db/models.py
@@ -44,6 +44,7 @@ from skyvern.forge.sdk.db.id import (
     generate_workflow_run_id,
 )
 from skyvern.forge.sdk.schemas.observers import ObserverThoughtType
+from skyvern.forge.sdk.schemas.tasks import ProxyLocation
 
 
 class Base(AsyncAttrs, DeclarativeBase):
@@ -69,7 +70,7 @@ class TaskModel(Base):
     navigation_payload = Column(JSON)
     extracted_information = Column(JSON)
     failure_reason = Column(String)
-    proxy_location = Column(String)
+    proxy_location = Column(Enum(ProxyLocation))
     extracted_information_schema = Column(JSON)
     workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), index=True)
     order = Column(Integer, nullable=True)