From 9cf474396144d73c8b697d6307e1e86fcc4e684c Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Wed, 11 Dec 2024 00:57:55 -0800
Subject: [PATCH] share validate_url function (#1371)

---
 skyvern/forge/sdk/core/validators.py          | 19 ++++++++++++++++-
 skyvern/forge/sdk/schemas/tasks.py            | 21 ++++---------------
 skyvern/forge/sdk/workflow/models/workflow.py | 21 +++----------------
 3 files changed, 25 insertions(+), 36 deletions(-)

diff --git a/skyvern/forge/sdk/core/validators.py b/skyvern/forge/sdk/core/validators.py
index 054455a3..c4e687d3 100644
--- a/skyvern/forge/sdk/core/validators.py
+++ b/skyvern/forge/sdk/core/validators.py
@@ -1,10 +1,11 @@
 import ipaddress
 from urllib.parse import urlparse
 
+from fastapi import status
 from pydantic import HttpUrl, ValidationError
 
 from skyvern.config import settings
-from skyvern.exceptions import InvalidUrl
+from skyvern.exceptions import BlockedHost, InvalidUrl, SkyvernHTTPException
 
 
 def prepend_scheme_and_validate_url(url: str) -> str:
@@ -40,3 +41,19 @@ def is_blocked_host(host: str) -> bool:
         return False
     except Exception:
         return False
+
+
+def validate_url(url: str) -> str | None:
+    try:
+        url = prepend_scheme_and_validate_url(url=url)
+        v = HttpUrl(url=url)
+    except Exception as e:
+        raise SkyvernHTTPException(message=str(e), status_code=status.HTTP_400_BAD_REQUEST)
+
+    if not v.host:
+        return None
+    host = v.host
+    blocked = is_blocked_host(host)
+    if blocked:
+        raise BlockedHost(host=host)
+    return str(v)
diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py
index a71e2169..abe90b13 100644
--- a/skyvern/forge/sdk/schemas/tasks.py
+++ b/skyvern/forge/sdk/schemas/tasks.py
@@ -4,11 +4,10 @@ from datetime import datetime
 from enum import StrEnum
 from typing import Any
 
-from fastapi import status
-from pydantic import BaseModel, Field, HttpUrl, field_validator
+from pydantic import BaseModel, Field, field_validator
 
-from skyvern.exceptions import BlockedHost, InvalidTaskStatusTransition, SkyvernHTTPException, TaskAlreadyCanceled
-from skyvern.forge.sdk.core.validators import is_blocked_host, prepend_scheme_and_validate_url
+from skyvern.exceptions import InvalidTaskStatusTransition, TaskAlreadyCanceled
+from skyvern.forge.sdk.core.validators import validate_url
 from skyvern.forge.sdk.db.enums import TaskType
 
 
@@ -118,19 +117,7 @@ class TaskRequest(TaskBase):
         if url is None:
             return None
 
-        try:
-            url = prepend_scheme_and_validate_url(url=url)
-            v = HttpUrl(url=url)
-        except Exception as e:
-            raise SkyvernHTTPException(message=str(e), status_code=status.HTTP_400_BAD_REQUEST)
-
-        if not v.host:
-            return None
-        host = v.host
-        blocked = is_blocked_host(host)
-        if blocked:
-            raise BlockedHost(host=host)
-        return str(v)
+        return validate_url(url)
 
 
 class TaskStatus(StrEnum):
diff --git a/skyvern/forge/sdk/workflow/models/workflow.py b/skyvern/forge/sdk/workflow/models/workflow.py
index bf563390..184de1ef 100644
--- a/skyvern/forge/sdk/workflow/models/workflow.py
+++ b/skyvern/forge/sdk/workflow/models/workflow.py
@@ -2,11 +2,9 @@ from datetime import datetime
 from enum import StrEnum
 from typing import Any, List
 
-from fastapi import status
-from pydantic import BaseModel, HttpUrl, field_validator
+from pydantic import BaseModel, field_validator
 
-from skyvern.exceptions import BlockedHost, SkyvernHTTPException
-from skyvern.forge.sdk.core.validators import is_blocked_host, prepend_scheme_and_validate_url
+from skyvern.forge.sdk.core.validators import validate_url
 from skyvern.forge.sdk.schemas.tasks import ProxyLocation
 from skyvern.forge.sdk.workflow.exceptions import WorkflowDefinitionHasDuplicateBlockLabels
 from skyvern.forge.sdk.workflow.models.block import BlockTypeVar
@@ -25,20 +23,7 @@ class WorkflowRequestBody(BaseModel):
     def validate_urls(cls, url: str | None) -> str | None:
         if url is None:
             return None
-
-        try:
-            url = prepend_scheme_and_validate_url(url=url)
-            v = HttpUrl(url=url)
-        except Exception as e:
-            raise SkyvernHTTPException(message=str(e), status_code=status.HTTP_400_BAD_REQUEST)
-
-        if not v.host:
-            return None
-        host = v.host
-        blocked = is_blocked_host(host)
-        if blocked:
-            raise BlockedHost(host=host)
-        return str(v)
+        return validate_url(url)
 
 
 class RunWorkflowResponse(BaseModel):