[Backend] Add City and State targeting for Massive geo-targeting (#4133)

skyvern/forge/sdk/db/utils.py

@@ -55,7 +55,7 @@ from skyvern.forge.sdk.workflow.models.workflow import (
     WorkflowRunStatus,
     WorkflowStatus,
 )
-from skyvern.schemas.runs import ProxyLocation, ScriptRunResponse
+from skyvern.schemas.runs import GeoTarget, ProxyLocation, ProxyLocationInput, ScriptRunResponse
 from skyvern.schemas.scripts import Script, ScriptBlock, ScriptFile
 from skyvern.schemas.workflows import BlockStatus, BlockType
 from skyvern.webeye.actions.actions import (
@@ -85,6 +85,50 @@ from skyvern.webeye.actions.actions import (
 
 LOG = structlog.get_logger()
 
+
+def _deserialize_proxy_location(value: str | None) -> ProxyLocationInput:
+    """
+    Deserialize proxy_location from database storage.
+
+    Handles:
+    - None -> None
+    - ProxyLocation enum string (e.g., "RESIDENTIAL") -> ProxyLocation enum
+    - JSON string (e.g., '{"country": "US", ...}') -> GeoTarget object
+    """
+    if value is None:
+        return None
+
+    result: ProxyLocationInput = None
+
+    # Try to parse as JSON first (for GeoTarget)
+    if value.startswith("{"):
+        try:
+            data = json.loads(value)
+            result = GeoTarget.model_validate(data)
+            LOG.info(
+                "Deserialized proxy_location as GeoTarget",
+                db_value=value,
+                result=str(result),
+            )
+            return result
+        except (json.JSONDecodeError, ValueError):
+            pass
+
+    # Try as ProxyLocation enum
+    try:
+        result = ProxyLocation(value)
+        LOG.info(
+            "Deserialized proxy_location as ProxyLocation enum",
+            db_value=value,
+            result=str(result),
+        )
+        return result
+    except ValueError:
+        # If all else fails, return as-is (shouldn't happen with valid data)
+        LOG.warning("Failed to deserialize proxy_location", db_value=value)
+        return None
+
+
 # Mapping of action types to their corresponding action classes
 ACTION_TYPE_TO_CLASS = {
     ActionType.CLICK: ClickAction,
@@ -142,7 +186,7 @@ def convert_to_task(task_obj: TaskModel, debug_enabled: bool = False, workflow_p
         extracted_information=task_obj.extracted_information,
         failure_reason=task_obj.failure_reason,
         organization_id=task_obj.organization_id,
-        proxy_location=(ProxyLocation(task_obj.proxy_location) if task_obj.proxy_location else None),
+        proxy_location=_deserialize_proxy_location(task_obj.proxy_location),
         extracted_information_schema=task_obj.extracted_information_schema,
         extra_http_headers=task_obj.extra_http_headers,
         workflow_run_id=task_obj.workflow_run_id,
@@ -272,7 +316,7 @@ def convert_to_workflow(workflow_model: WorkflowModel, debug_enabled: bool = Fal
         totp_identifier=workflow_model.totp_identifier,
         persist_browser_session=workflow_model.persist_browser_session,
         model=workflow_model.model,
-        proxy_location=(ProxyLocation(workflow_model.proxy_location) if workflow_model.proxy_location else None),
+        proxy_location=_deserialize_proxy_location(workflow_model.proxy_location),
         max_screenshot_scrolls=workflow_model.max_screenshot_scrolling_times,
         version=workflow_model.version,
         is_saved_task=workflow_model.is_saved_task,
@@ -312,9 +356,7 @@ def convert_to_workflow_run(
         browser_profile_id=workflow_run_model.browser_profile_id,
         status=WorkflowRunStatus[workflow_run_model.status],
         failure_reason=workflow_run_model.failure_reason,
-        proxy_location=(
-            ProxyLocation(workflow_run_model.proxy_location) if workflow_run_model.proxy_location else None
-        ),
+        proxy_location=_deserialize_proxy_location(workflow_run_model.proxy_location),
         webhook_callback_url=workflow_run_model.webhook_callback_url,
         webhook_failure_reason=workflow_run_model.webhook_failure_reason,
         totp_verification_url=workflow_run_model.totp_verification_url,