share validate_url function (#1371)

This commit is contained in:
Shuchang Zheng 2024-12-11 00:57:55 -08:00 committed by GitHub
parent 4241cfe3f2
commit 9cf4743961
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 25 additions and 36 deletions

View file

@ -1,10 +1,11 @@
import ipaddress import ipaddress
from urllib.parse import urlparse from urllib.parse import urlparse
from fastapi import status
from pydantic import HttpUrl, ValidationError from pydantic import HttpUrl, ValidationError
from skyvern.config import settings 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: def prepend_scheme_and_validate_url(url: str) -> str:
@ -40,3 +41,19 @@ def is_blocked_host(host: str) -> bool:
return False return False
except Exception: except Exception:
return False 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)

View file

@ -4,11 +4,10 @@ from datetime import datetime
from enum import StrEnum from enum import StrEnum
from typing import Any from typing import Any
from fastapi import status from pydantic import BaseModel, Field, field_validator
from pydantic import BaseModel, Field, HttpUrl, field_validator
from skyvern.exceptions import BlockedHost, InvalidTaskStatusTransition, SkyvernHTTPException, TaskAlreadyCanceled from skyvern.exceptions import InvalidTaskStatusTransition, TaskAlreadyCanceled
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.db.enums import TaskType from skyvern.forge.sdk.db.enums import TaskType
@ -118,19 +117,7 @@ class TaskRequest(TaskBase):
if url is None: if url is None:
return None return None
try: return validate_url(url)
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)
class TaskStatus(StrEnum): class TaskStatus(StrEnum):

View file

@ -2,11 +2,9 @@ from datetime import datetime
from enum import StrEnum from enum import StrEnum
from typing import Any, List from typing import Any, List
from fastapi import status from pydantic import BaseModel, field_validator
from pydantic import BaseModel, HttpUrl, field_validator
from skyvern.exceptions import BlockedHost, SkyvernHTTPException from skyvern.forge.sdk.core.validators import validate_url
from skyvern.forge.sdk.core.validators import is_blocked_host, prepend_scheme_and_validate_url
from skyvern.forge.sdk.schemas.tasks import ProxyLocation from skyvern.forge.sdk.schemas.tasks import ProxyLocation
from skyvern.forge.sdk.workflow.exceptions import WorkflowDefinitionHasDuplicateBlockLabels from skyvern.forge.sdk.workflow.exceptions import WorkflowDefinitionHasDuplicateBlockLabels
from skyvern.forge.sdk.workflow.models.block import BlockTypeVar 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: def validate_urls(cls, url: str | None) -> str | None:
if url is None: if url is None:
return None return None
return validate_url(url)
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)
class RunWorkflowResponse(BaseModel): class RunWorkflowResponse(BaseModel):