diff --git a/skyvern/forge/sdk/schemas/totp_codes.py b/skyvern/forge/sdk/schemas/totp_codes.py index 392fb34f..8acdddbd 100644 --- a/skyvern/forge/sdk/schemas/totp_codes.py +++ b/skyvern/forge/sdk/schemas/totp_codes.py @@ -1,6 +1,8 @@ from datetime import datetime -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, field_validator + +from skyvern.forge.sdk.utils.sanitization import sanitize_postgres_text class TOTPCodeBase(BaseModel): @@ -19,6 +21,12 @@ class TOTPCodeCreate(TOTPCodeBase): totp_identifier: str content: str + @field_validator("content") + @classmethod + def sanitize_content(cls, value: str) -> str: + """Remove NUL (0x00) bytes from content to avoid PostgreSQL DataError.""" + return sanitize_postgres_text(value) + class TOTPCode(TOTPCodeCreate): totp_code_id: str diff --git a/skyvern/forge/sdk/utils/sanitization.py b/skyvern/forge/sdk/utils/sanitization.py new file mode 100644 index 00000000..9f3e91a2 --- /dev/null +++ b/skyvern/forge/sdk/utils/sanitization.py @@ -0,0 +1,18 @@ +""" +Utility functions for sanitizing content before storing in the database. +""" + + +def sanitize_postgres_text(text: str) -> str: + """ + Sanitize text to be stored in PostgreSQL by removing NUL bytes. + + PostgreSQL text fields cannot contain NUL (0x00) bytes, so we remove them. + + Args: + text: The text to sanitize + + Returns: + The sanitized text without NUL bytes + """ + return text.replace("\0", "")